Timed wait operations for spin-only Baton

[folly.git] / folly / synchronization / Baton.h

diff --git a/folly/synchronization/Baton.h b/folly/synchronization/Baton.h
index ff6b2978ad7ff5124ff9c77b5934213dbb8ccc2b..b3b1e09ec3f2070fadd01a09535f612b3fe41d60 100644 (file)
--- a/folly/synchronization/Baton.h
+++ b/folly/synchronization/Baton.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2017 Facebook, Inc.
+ * Copyright 2014-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.
@@ -40,9 +40,9 @@ namespace folly {
 /// This is basically a stripped-down semaphore that supports only a
 /// single call to sem_post and a single call to sem_wait.
 ///
-/// The non-blocking version (Blocking == false) provides more speed
+/// The non-blocking version (MayBlock == false) provides more speed
 /// by using only load acquire and store release operations in the
-/// critical path, at the cost of disallowing blocking and timing out.
+/// critical path, at the cost of disallowing blocking.
 ///
 /// The current posix semaphore sem_t isn't too bad, but this provides
 /// more a bit more speed, inlining, smaller size, a guarantee that
@@ -50,9 +50,7 @@ namespace folly {
 /// DeterministicSchedule.  By having a much more restrictive
 /// lifecycle we can also add a bunch of assertions that can help to
 /// catch race conditions ahead of time.
-template <
-    template <typename> class Atom = std::atomic,
-    bool Blocking = true> // blocking vs spinning
+template <bool MayBlock = true, template <typename> class Atom = std::atomic>
 struct Baton {
   constexpr Baton() noexcept : state_(INIT) {}
 
@@ -79,6 +77,12 @@ struct Baton {
     assert(state_.load(std::memory_order_relaxed) != WAITING);
   }
 
+  FOLLY_ALWAYS_INLINE bool ready() const noexcept {
+    auto s = state_.load(std::memory_order_acquire);
+    assert(s == INIT || s == EARLY_DELIVERY);
+    return LIKELY(s == EARLY_DELIVERY);
+  }
+
   /// Equivalent to destroying the Baton and creating a new one.  It is
   /// a bug to call this while there is a waiting thread, so in practice
   /// the waiter will be the one that resets the baton.
@@ -110,8 +114,8 @@ struct Baton {
   /// destruction or reset()) there can be at most one call to post(),
   /// in the single poster version.  Any thread may call post().
   void post() noexcept {
-    if (!Blocking) {
-      /// Non-blocking version
+    if (!MayBlock) {
+      /// Spin-only version
       ///
       assert([&] {
         auto state = state_.load(std::memory_order_relaxed);
@@ -121,7 +125,7 @@ struct Baton {
       return;
     }
 
-    /// Blocking versions
+    /// May-block versions
     ///
     uint32_t before = state_.load(std::memory_order_acquire);
 
@@ -176,9 +180,7 @@ struct Baton {
   ///
   /// @return       true if baton has been posted, false othewise
   FOLLY_ALWAYS_INLINE bool try_wait() const noexcept {
-    auto s = state_.load(std::memory_order_acquire);
-    assert(s == INIT || s == EARLY_DELIVERY);
-    return LIKELY(s == EARLY_DELIVERY);
+    return ready();
   }
 
   /// Similar to wait, but with a timeout. The thread is unblocked if the
@@ -195,9 +197,6 @@ struct Baton {
   template <typename Rep, typename Period>
   FOLLY_ALWAYS_INLINE bool try_wait_for(
       const std::chrono::duration<Rep, Period>& timeout) noexcept {
-    static_assert(
-        Blocking, "Non-blocking Baton does not support try_wait_for.");
-
     if (try_wait()) {
       return true;
     }
@@ -220,9 +219,6 @@ struct Baton {
   template <typename Clock, typename Duration>
   FOLLY_ALWAYS_INLINE bool try_wait_until(
       const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
-    static_assert(
-        Blocking, "Non-blocking Baton does not support try_wait_until.");
-
     if (try_wait()) {
       return true;
     }
@@ -276,15 +272,16 @@ struct Baton {
   // @return       true if we received an early delivery during the wait,
   //               false otherwise. If the function returns true then
   //               state_ is guaranteed to be EARLY_DELIVERY
-  bool spinWaitForEarlyDelivery() noexcept {
-    static_assert(
-        PreBlockAttempts > 0,
-        "isn't this assert clearer than an uninitialized variable warning?");
+  template <typename Clock, typename Duration>
+  bool spinWaitForEarlyDelivery(
+      const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
     for (int i = 0; i < PreBlockAttempts; ++i) {
       if (try_wait()) {
         return true;
       }
-
+      if (Clock::now() >= deadline) {
+        return false;
+      }
       // The pause instruction is the polite way to spin, but it doesn't
       // actually affect correctness to omit it if we don't have it.
       // Pausing donates the full capabilities of the current core to
@@ -296,12 +293,13 @@ struct Baton {
   }
 
   FOLLY_NOINLINE void waitSlow() noexcept {
-    if (spinWaitForEarlyDelivery()) {
+    auto const deadline = std::chrono::steady_clock::time_point::max();
+    if (spinWaitForEarlyDelivery(deadline)) {
       assert(state_.load(std::memory_order_acquire) == EARLY_DELIVERY);
       return;
     }
 
-    if (!Blocking) {
+    if (!MayBlock) {
       while (!try_wait()) {
         std::this_thread::yield();
       }
@@ -349,11 +347,23 @@ struct Baton {
   template <typename Clock, typename Duration>
   FOLLY_NOINLINE bool tryWaitUntilSlow(
       const std::chrono::time_point<Clock, Duration>& deadline) noexcept {
-    if (spinWaitForEarlyDelivery()) {
+    if (spinWaitForEarlyDelivery(deadline)) {
       assert(state_.load(std::memory_order_acquire) == EARLY_DELIVERY);
       return true;
     }
 
+    if (!MayBlock) {
+      while (true) {
+        if (try_wait()) {
+          return true;
+        }
+        if (Clock::now() >= deadline) {
+          return false;
+        }
+        std::this_thread::yield();
+      }
+    }
+
     // guess we have to block :(
     uint32_t expected = INIT;
     if (!state_.compare_exchange_strong(expected, WAITING)) {