SaturatingSemaphore.h

Go to the documentation of this file.

/*
 * Copyright 2017-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <folly/Likely.h>
#include <folly/detail/Futex.h>
#include <folly/detail/MemoryIdler.h>
#include <folly/portability/Asm.h>
#include <folly/synchronization/WaitOptions.h>
#include <folly/synchronization/detail/Spin.h>

#include <glog/logging.h>

#include <atomic>

namespace folly {


template <bool MayBlock, template <typename> class Atom = std::atomic>
class SaturatingSemaphore {
  detail::Futex<Atom> state_;

  enum State : uint32_t {
    NOTREADY = 0,
    READY = 1,
    BLOCKED = 2,
  };

 public:
  FOLLY_ALWAYS_INLINE static WaitOptions wait_options() {
    return {};
  }

  constexpr SaturatingSemaphore() noexcept : state_(NOTREADY) {}

  ~SaturatingSemaphore() {}

  FOLLY_ALWAYS_INLINE bool ready() const noexcept {
    return state_.load(std::memory_order_acquire) == READY;
  }

  void reset() noexcept {
    state_.store(NOTREADY, std::memory_order_relaxed);
  }

  FOLLY_ALWAYS_INLINE void post() noexcept {
    if (!MayBlock) {
      state_.store(READY, std::memory_order_release);
    } else {
      postFastWaiterMayBlock();
    }
  }

  FOLLY_ALWAYS_INLINE
  void wait(const WaitOptions& opt = wait_options()) noexcept {
    try_wait_until(std::chrono::steady_clock::time_point::max(), opt);
  }

  FOLLY_ALWAYS_INLINE bool try_wait() noexcept {
    return ready();
  }

  template <typename Clock, typename Duration>
  FOLLY_ALWAYS_INLINE bool try_wait_until(
      const std::chrono::time_point<Clock, Duration>& deadline,
      const WaitOptions& opt = wait_options()) noexcept {
    if (LIKELY(try_wait())) {
      return true;
    }
    return tryWaitSlow(deadline, opt);
  }

  template <class Rep, class Period>
  FOLLY_ALWAYS_INLINE bool try_wait_for(
      const std::chrono::duration<Rep, Period>& duration,
      const WaitOptions& opt = wait_options()) noexcept {
    if (LIKELY(try_wait())) {
      return true;
    }
    auto deadline = std::chrono::steady_clock::now() + duration;
    return tryWaitSlow(deadline, opt);
  }

 private:
  FOLLY_ALWAYS_INLINE void postFastWaiterMayBlock() noexcept {
    uint32_t before = NOTREADY;
    if (LIKELY(state_.compare_exchange_strong(
            before,
            READY,
            std::memory_order_release,
            std::memory_order_relaxed))) {
      return;
    }
    postSlowWaiterMayBlock(before);
  }

  void postSlowWaiterMayBlock(uint32_t before) noexcept; // defined below

  template <typename Clock, typename Duration>
  bool tryWaitSlow(
      const std::chrono::time_point<Clock, Duration>& deadline,
      const WaitOptions& opt) noexcept; // defined below
};


template <bool MayBlock, template <typename> class Atom>
FOLLY_NOINLINE void SaturatingSemaphore<MayBlock, Atom>::postSlowWaiterMayBlock(
    uint32_t before) noexcept {
  while (true) {
    if (before == NOTREADY) {
      if (state_.compare_exchange_strong(
              before,
              READY,
              std::memory_order_release,
              std::memory_order_relaxed)) {
        return;
      }
    }
    if (before == READY) { // Only if multiple posters
      // The reason for not simply returning (without the following
      // steps) is to prevent the following case:
      //
      //  T1:             T2:             T3:
      //  local1.post();  local2.post();  global.wait();
      //  global.post();  global.post();  global.reset();
      //                                  seq_cst fence
      //                                  local1.try_wait() == true;
      //                                  local2.try_wait() == false;
      //
      // This following steps correspond to T2's global.post(), where
      // global is already posted by T1.
      //
      // The following fence and load guarantee that T3 does not miss
      // T2's prior stores, i.e., local2.post() in this example.
      //
      // The following case is prevented:
      //
      // Starting with local2 == NOTREADY and global == READY
      //
      // T2:                              T3:
      // store READY to local2 // post    store NOTREADY to global // reset
      // seq_cst fenc                     seq_cst fence
      // load READY from global // post   load NOTREADY from local2 // try_wait
      //
      std::atomic_thread_fence(std::memory_order_seq_cst);
      before = state_.load(std::memory_order_relaxed);
      if (before == READY) {
        return;
      }
      continue;
    }
    DCHECK_EQ(before, BLOCKED);
    if (state_.compare_exchange_strong(
            before,
            READY,
            std::memory_order_release,
            std::memory_order_relaxed)) {
      detail::futexWake(&state_);
      return;
    }
  }
}

template <bool MayBlock, template <typename> class Atom>
template <typename Clock, typename Duration>
FOLLY_NOINLINE bool SaturatingSemaphore<MayBlock, Atom>::tryWaitSlow(
    const std::chrono::time_point<Clock, Duration>& deadline,
    const WaitOptions& opt) noexcept {
  switch (detail::spin_pause_until(deadline, opt, [=] { return ready(); })) {
    case detail::spin_result::success:
      return true;
    case detail::spin_result::timeout:
      return false;
    case detail::spin_result::advance:
      break;
  }

  if (!MayBlock) {
    switch (detail::spin_yield_until(deadline, [=] { return ready(); })) {
      case detail::spin_result::success:
        return true;
      case detail::spin_result::timeout:
        return false;
      case detail::spin_result::advance:
        break;
    }
  }

  auto before = state_.load(std::memory_order_relaxed);
  while (before == NOTREADY &&
         !state_.compare_exchange_strong(
             before,
             BLOCKED,
             std::memory_order_relaxed,
             std::memory_order_relaxed)) {
    if (before == READY) {
      // TODO: move the acquire to the compare_exchange failure load after C++17
      std::atomic_thread_fence(std::memory_order_acquire);
      return true;
    }
  }

  while (true) {
    auto rv = detail::MemoryIdler::futexWaitUntil(state_, BLOCKED, deadline);
    if (rv == detail::FutexResult::TIMEDOUT) {
      assert(deadline != (std::chrono::time_point<Clock, Duration>::max()));
      return false;
    }

    if (ready()) {
      return true;
    }
  }
}

} // namespace folly