From 3579ddb46bf91312e7c1d24389ad60fd96fad776 Mon Sep 17 00:00:00 2001
From: Yedidya Feldblum <yfeldblum@fb.com>
Date: Tue, 9 Jan 2018 20:11:16 -0800
Subject: [PATCH] Timed wait operations for spin-only Baton

Summary:
[Folly] Timed wait operations for spin-only `Baton`.

Enables `try_wait_for` and `try_wait_until` for `Baton</* MayBlock = */ false, /*...*/>`.

Reviewed By: nbronson

Differential Revision: D6672153

fbshipit-source-id: 95da07260b21c2b88b8f7bf81cbfcbe5f5099ac0
---
 folly/synchronization/Baton.h                 | 36 +++++---
 folly/synchronization/test/BatonTest.cpp      | 88 +++++++++++++------
 folly/synchronization/test/BatonTestHelpers.h | 12 +--
 3 files changed, 90 insertions(+), 46 deletions(-)

diff --git a/folly/synchronization/Baton.h b/folly/synchronization/Baton.h
index 2b0f7eb5..b3b1e09e 100644
--- a/folly/synchronization/Baton.h
+++ b/folly/synchronization/Baton.h
@@ -42,7 +42,7 @@ namespace folly {
 ///
 /// 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
@@ -197,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(
-        MayBlock, "Non-blocking Baton does not support try_wait_for.");
-
     if (try_wait()) {
       return true;
     }
@@ -222,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(
-        MayBlock, "Non-blocking Baton does not support try_wait_until.");
-
     if (try_wait()) {
       return true;
     }
@@ -278,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
@@ -298,7 +293,8 @@ 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;
     }
@@ -351,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)) {
diff --git a/folly/synchronization/test/BatonTest.cpp b/folly/synchronization/test/BatonTest.cpp
index fc925d6c..dc44f5ea 100644
--- a/folly/synchronization/test/BatonTest.cpp
+++ b/folly/synchronization/test/BatonTest.cpp
@@ -25,6 +25,8 @@
 using namespace folly;
 using namespace folly::test;
 using folly::detail::EmulatedFutexAtomic;
+using std::chrono::steady_clock;
+using std::chrono::system_clock;
 
 /// Basic test
 
@@ -54,54 +56,88 @@ TEST(Baton, pingpong_nonblocking) {
   run_pingpong_test<false, DeterministicAtomic>(1000);
 }
 
-/// Timed wait tests - Nonblocking Baton does not support try_wait_until()
-
 // Timed wait basic system clock tests
 
-TEST(Baton, timed_wait_basic_system_clock) {
-  run_basic_timed_wait_tests<std::atomic, std::chrono::system_clock>();
-  run_basic_timed_wait_tests<EmulatedFutexAtomic, std::chrono::system_clock>();
-  run_basic_timed_wait_tests<DeterministicAtomic, std::chrono::system_clock>();
+TEST(Baton, timed_wait_basic_system_clock_blocking) {
+  run_basic_timed_wait_tests<true, std::atomic, system_clock>();
+  run_basic_timed_wait_tests<true, EmulatedFutexAtomic, system_clock>();
+  run_basic_timed_wait_tests<true, DeterministicAtomic, system_clock>();
+}
+
+TEST(Baton, timed_wait_basic_system_clock_nonblocking) {
+  run_basic_timed_wait_tests<false, std::atomic, system_clock>();
+  run_basic_timed_wait_tests<false, EmulatedFutexAtomic, system_clock>();
+  run_basic_timed_wait_tests<false, DeterministicAtomic, system_clock>();
 }
 
 // Timed wait timeout system clock tests
 
-TEST(Baton, timed_wait_timeout_system_clock) {
-  run_timed_wait_tmo_tests<std::atomic, std::chrono::system_clock>();
-  run_timed_wait_tmo_tests<EmulatedFutexAtomic, std::chrono::system_clock>();
-  run_timed_wait_tmo_tests<DeterministicAtomic, std::chrono::system_clock>();
+TEST(Baton, timed_wait_timeout_system_clock_blocking) {
+  run_timed_wait_tmo_tests<true, std::atomic, system_clock>();
+  run_timed_wait_tmo_tests<true, EmulatedFutexAtomic, system_clock>();
+  run_timed_wait_tmo_tests<true, DeterministicAtomic, system_clock>();
+}
+
+TEST(Baton, timed_wait_timeout_system_clock_nonblocking) {
+  run_timed_wait_tmo_tests<false, std::atomic, system_clock>();
+  run_timed_wait_tmo_tests<false, EmulatedFutexAtomic, system_clock>();
+  run_timed_wait_tmo_tests<false, DeterministicAtomic, system_clock>();
 }
 
 // Timed wait regular system clock tests
 
-TEST(Baton, timed_wait_system_clock) {
-  run_timed_wait_regular_test<std::atomic, std::chrono::system_clock>();
-  run_timed_wait_regular_test<EmulatedFutexAtomic, std::chrono::system_clock>();
-  run_timed_wait_regular_test<DeterministicAtomic, std::chrono::system_clock>();
+TEST(Baton, timed_wait_system_clock_blocking) {
+  run_timed_wait_regular_test<true, std::atomic, system_clock>();
+  run_timed_wait_regular_test<true, EmulatedFutexAtomic, system_clock>();
+  run_timed_wait_regular_test<true, DeterministicAtomic, system_clock>();
+}
+
+TEST(Baton, timed_wait_system_clock_nonblocking) {
+  run_timed_wait_regular_test<false, std::atomic, system_clock>();
+  run_timed_wait_regular_test<false, EmulatedFutexAtomic, system_clock>();
+  run_timed_wait_regular_test<false, DeterministicAtomic, system_clock>();
 }
 
 // Timed wait basic steady clock tests
 
-TEST(Baton, timed_wait_basic_steady_clock) {
-  run_basic_timed_wait_tests<std::atomic, std::chrono::steady_clock>();
-  run_basic_timed_wait_tests<EmulatedFutexAtomic, std::chrono::steady_clock>();
-  run_basic_timed_wait_tests<DeterministicAtomic, std::chrono::steady_clock>();
+TEST(Baton, timed_wait_basic_steady_clock_blocking) {
+  run_basic_timed_wait_tests<true, std::atomic, steady_clock>();
+  run_basic_timed_wait_tests<true, EmulatedFutexAtomic, steady_clock>();
+  run_basic_timed_wait_tests<true, DeterministicAtomic, steady_clock>();
+}
+
+TEST(Baton, timed_wait_basic_steady_clock_nonblocking) {
+  run_basic_timed_wait_tests<false, std::atomic, steady_clock>();
+  run_basic_timed_wait_tests<false, EmulatedFutexAtomic, steady_clock>();
+  run_basic_timed_wait_tests<false, DeterministicAtomic, steady_clock>();
 }
 
 // Timed wait timeout steady clock tests
 
-TEST(Baton, timed_wait_timeout_steady_clock) {
-  run_timed_wait_tmo_tests<std::atomic, std::chrono::steady_clock>();
-  run_timed_wait_tmo_tests<EmulatedFutexAtomic, std::chrono::steady_clock>();
-  run_timed_wait_tmo_tests<DeterministicAtomic, std::chrono::steady_clock>();
+TEST(Baton, timed_wait_timeout_steady_clock_blocking) {
+  run_timed_wait_tmo_tests<true, std::atomic, steady_clock>();
+  run_timed_wait_tmo_tests<true, EmulatedFutexAtomic, steady_clock>();
+  run_timed_wait_tmo_tests<true, DeterministicAtomic, steady_clock>();
+}
+
+TEST(Baton, timed_wait_timeout_steady_clock_nonblocking) {
+  run_timed_wait_tmo_tests<false, std::atomic, steady_clock>();
+  run_timed_wait_tmo_tests<false, EmulatedFutexAtomic, steady_clock>();
+  run_timed_wait_tmo_tests<false, DeterministicAtomic, steady_clock>();
 }
 
 // Timed wait regular steady clock tests
 
-TEST(Baton, timed_wait_steady_clock) {
-  run_timed_wait_regular_test<std::atomic, std::chrono::steady_clock>();
-  run_timed_wait_regular_test<EmulatedFutexAtomic, std::chrono::steady_clock>();
-  run_timed_wait_regular_test<DeterministicAtomic, std::chrono::steady_clock>();
+TEST(Baton, timed_wait_steady_clock_blocking) {
+  run_timed_wait_regular_test<true, std::atomic, steady_clock>();
+  run_timed_wait_regular_test<true, EmulatedFutexAtomic, steady_clock>();
+  run_timed_wait_regular_test<true, DeterministicAtomic, steady_clock>();
+}
+
+TEST(Baton, timed_wait_steady_clock_nonblocking) {
+  run_timed_wait_regular_test<false, std::atomic, steady_clock>();
+  run_timed_wait_regular_test<false, EmulatedFutexAtomic, steady_clock>();
+  run_timed_wait_regular_test<false, DeterministicAtomic, steady_clock>();
 }
 
 /// Try wait tests
diff --git a/folly/synchronization/test/BatonTestHelpers.h b/folly/synchronization/test/BatonTestHelpers.h
index 61b4e8c3..0884927a 100644
--- a/folly/synchronization/test/BatonTestHelpers.h
+++ b/folly/synchronization/test/BatonTestHelpers.h
@@ -53,17 +53,17 @@ void run_pingpong_test(int numRounds) {
   DSched::join(thr);
 }
 
-template <template <typename> class Atom, typename Clock>
+template <bool MayBlock, template <typename> class Atom, typename Clock>
 void run_basic_timed_wait_tests() {
-  Baton<true, Atom> b;
+  Baton<MayBlock, Atom> b;
   b.post();
   // tests if early delivery works fine
   EXPECT_TRUE(b.try_wait_until(Clock::now()));
 }
 
-template <template <typename> class Atom, typename Clock>
+template <bool MayBlock, template <typename> class Atom, typename Clock>
 void run_timed_wait_tmo_tests() {
-  Baton<true, Atom> b;
+  Baton<MayBlock, Atom> b;
 
   auto thr = DSched::thread([&] {
     bool rv = b.try_wait_until(Clock::now() + std::chrono::milliseconds(1));
@@ -73,9 +73,9 @@ void run_timed_wait_tmo_tests() {
   DSched::join(thr);
 }
 
-template <template <typename> class Atom, typename Clock>
+template <bool MayBlock, template <typename> class Atom, typename Clock>
 void run_timed_wait_regular_test() {
-  Baton<true, Atom> b;
+  Baton<MayBlock, Atom> b;
 
   auto thr = DSched::thread([&] {
     // To wait forever we'd like to use time_point<Clock>::max, but
-- 
2.34.1