Dedup spinning in Baton and SaturatingSemaphore

Summary:
[Folly] Dedup spinning in `Baton` and `SaturatingSemaphore`.

And add spin-yield to spin-only `SaturatingSemaphore`, as is done in `Baton`.

Reviewed By: djwatson

Differential Revision: D6726786

fbshipit-source-id: 428efb950577ab1d500caffe56bf2a8cacf66c34

folly/Makefile.am

@@ -446,6 +446,7 @@ nobase_follyinclude_HEADERS = \
 	synchronization/WaitOptions.h \
 	synchronization/detail/AtomicUtils.h \
 	synchronization/detail/Sleeper.h \
+	synchronization/detail/Spin.h \
 	system/MemoryMapping.h \
 	system/Shell.h \
 	system/ThreadId.h \

folly/synchronization/Baton.h

@@ -27,6 +27,7 @@
 #include <folly/detail/MemoryIdler.h>
 #include <folly/portability/Asm.h>
 #include <folly/synchronization/WaitOptions.h>
+#include <folly/synchronization/detail/Spin.h>
 
 namespace folly {
 
@@ -258,62 +259,17 @@ class Baton {
     TIMED_OUT = 4,
   };
 
-  // Spin for "some time"
-  //
-  // @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
-  template <typename Clock, typename Duration>
-  bool spinWaitForEarlyDelivery(
-      const std::chrono::time_point<Clock, Duration>& deadline,
-      const WaitOptions& opt) noexcept {
-    auto tbegin = Clock::now();
-    while (true) {
-      if (try_wait()) {
-        return true;
-      }
-
-      auto const tnow = Clock::now();
-      if (tnow >= deadline) {
-        return false;
-      }
-
-      // Backward time discontinuity in Clock? revise pre_block starting point
-      tbegin = std::min(tbegin, tnow);
-      auto const dur = std::chrono::duration_cast<Duration>(tnow - tbegin);
-      if (dur >= opt.spin_max()) {
-        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
-      // its other hyperthreads for a dozen cycles or so
-      asm_volatile_pause();
-    }
-  }
-
   template <typename Clock, typename Duration>
   FOLLY_NOINLINE bool tryWaitSlow(
       const std::chrono::time_point<Clock, Duration>& deadline,
       const WaitOptions& opt) noexcept {
-    auto const forever = std::chrono::time_point<Clock, Duration>::max();
-
-    if (spinWaitForEarlyDelivery(deadline, opt)) {
-      assert(state_.load(std::memory_order_acquire) == EARLY_DELIVERY);
+    if (detail::spin_pause_until(deadline, opt, [=] { return ready(); })) {
+      assert(ready());
       return true;
     }
 
     if (!MayBlock) {
-      while (true) {
-        if (try_wait()) {
-          return true;
-        }
-        if (deadline != forever && Clock::now() >= deadline) {
-          return false;
-        }
-        std::this_thread::yield();
-      }
+      return detail::spin_yield_until(deadline, [=] { return ready(); });
     }
 
     // guess we have to block :(
@@ -333,7 +289,7 @@ class Baton {
 
       // Awoken by the deadline passing.
       if (rv == detail::FutexResult::TIMEDOUT) {
-        assert(deadline != forever);
+        assert(deadline != (std::chrono::time_point<Clock, Duration>::max()));
         state_.store(TIMED_OUT, std::memory_order_release);
         return false;
       }

folly/synchronization/SaturatingSemaphore.h

@@ -21,6 +21,7 @@
 #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>
 
@@ -279,36 +280,36 @@ template <typename Clock, typename Duration>
 FOLLY_NOINLINE bool SaturatingSemaphore<MayBlock, Atom>::tryWaitSlow(
     const std::chrono::time_point<Clock, Duration>& deadline,
     const WaitOptions& opt) noexcept {
-  auto tbegin = Clock::now();
-  while (true) {
-    auto before = state_.load(std::memory_order_acquire);
+  if (detail::spin_pause_until(deadline, opt, [=] { return ready(); })) {
+    return true;
+  }
+
+  if (!MayBlock) {
+    return detail::spin_yield_until(deadline, [=] { return ready(); });
+  }
+
+  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) {
       return true;
     }
-    if (Clock::now() >= deadline) {
+  }
+
+  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 (MayBlock) {
-      auto tnow = Clock::now();
-      if (tnow < tbegin) {
-        // backward time discontinuity in Clock, revise spin_max starting point
-        tbegin = tnow;
-      }
-      auto dur = std::chrono::duration_cast<Duration>(tnow - tbegin);
-      if (dur >= opt.spin_max()) {
-        if (before == NOTREADY) {
-          if (!state_.compare_exchange_strong(
-                  before,
-                  BLOCKED,
-                  std::memory_order_relaxed,
-                  std::memory_order_relaxed)) {
-            continue;
-          }
-        }
-        detail::MemoryIdler::futexWaitUntil(state_, BLOCKED, deadline);
-      }
+
+    if (ready()) {
+      return true;
     }
-    asm_volatile_pause();
   }
 }
 

folly/synchronization/detail/Spin.h

+/*
+ * Copyright 2018-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 <algorithm>
+#include <chrono>
+#include <thread>
+
+#include <folly/portability/Asm.h>
+#include <folly/synchronization/WaitOptions.h>
+
+namespace folly {
+namespace detail {
+
+template <typename Clock, typename Duration, typename F>
+bool spin_pause_until(
+    std::chrono::time_point<Clock, Duration> const& deadline,
+    WaitOptions const& opt,
+    F f) {
+  auto tbegin = Clock::now();
+  while (true) {
+    if (f()) {
+      return true;
+    }
+
+    auto const tnow = Clock::now();
+    if (tnow >= deadline) {
+      return false;
+    }
+
+    //  Backward time discontinuity in Clock? revise pre_block starting point
+    tbegin = std::min(tbegin, tnow);
+    if (tnow >= tbegin + opt.spin_max()) {
+      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 its other
+    //  hyperthreads for a dozen cycles or so.
+    asm_volatile_pause();
+  }
+}
+
+template <typename Clock, typename Duration, typename F>
+bool spin_yield_until(
+    std::chrono::time_point<Clock, Duration> const& deadline,
+    F f) {
+  while (true) {
+    if (f()) {
+      return true;
+    }
+
+    auto const max = std::chrono::time_point<Clock, Duration>::max();
+    if (deadline != max && Clock::now() >= deadline) {
+      return false;
+    }
+
+    std::this_thread::yield();
+  }
+}
+
+} // namespace detail
+} // namespace folly