Template fibers baton timed-wait over the deadline types

Summary:
[Folly] Template fibers baton timed-wait over the deadline types, rather than over the timeout types and rather than fixing a single timeout type.

Templating it over the deadline types permits specifying the clock and permits passing a single deadline through unchanged through multiple time-delayed layers of code without skew, and templating it rather than fixing it over a single timeout type permits caller-specified fine-grained or coarse-grained timings.

Currently, the `EventBase` timers do not parameterize over the deadline types, so the fiber version of timed-wait must for now convert to a fixed timeout type.

Reviewed By: andriigrynenko

Differential Revision: D15448177

fbshipit-source-id: aa3fdbffdcb83cbfd1571780ba431a1004beb347

folly/fibers/Baton-inl.h

@@ -71,40 +71,72 @@ void Baton::waitFiber(FiberManager& fm, F&& mainContextFunc) {
   fm.activeFiber_->preempt(Fiber::AWAITING);
 }
 
+template <typename Clock, typename Duration>
+bool Baton::timedWaitThread(
+    const std::chrono::time_point<Clock, Duration>& deadline) {
+  auto waiter = waiter_.load();
+
+  if (LIKELY(
+          waiter == NO_WAITER &&
+          waiter_.compare_exchange_strong(waiter, THREAD_WAITING))) {
+    do {
+      auto* futex = &futex_.futex;
+      const auto wait_rv = folly::detail::futexWaitUntil(
+          futex, uint32_t(THREAD_WAITING), deadline);
+      if (wait_rv == folly::detail::FutexResult::TIMEDOUT) {
+        return false;
+      }
+      waiter = waiter_.load(std::memory_order_relaxed);
+    } while (waiter == THREAD_WAITING);
+  }
+
+  if (LIKELY(waiter == POSTED)) {
+    return true;
+  }
+
+  // Handle errors
+  if (waiter == TIMEOUT) {
+    throw std::logic_error("Thread baton can't have timeout status");
+  }
+  if (waiter == THREAD_WAITING) {
+    throw std::logic_error("Other thread is already waiting on this baton");
+  }
+  throw std::logic_error("Other waiter is already waiting on this baton");
+}
+
 template <typename Rep, typename Period, typename F>
 bool Baton::try_wait_for(
     const std::chrono::duration<Rep, Period>& timeout,
     F&& mainContextFunc) {
+  const auto deadline = std::chrono::steady_clock::now() + timeout;
+  return try_wait_until(deadline, static_cast<F&&>(mainContextFunc));
+}
+
+template <typename Clock, typename Duration, typename F>
+bool Baton::try_wait_until(
+    const std::chrono::time_point<Clock, Duration>& deadline,
+    F&& mainContextFunc) {
   auto fm = FiberManager::getFiberManagerUnsafe();
-  auto timeoutMs =
-      std::chrono::duration_cast<std::chrono::milliseconds>(timeout);
 
   if (!fm || !fm->activeFiber_) {
     mainContextFunc();
-    return timedWaitThread(timeoutMs);
+    return timedWaitThread(deadline);
   }
 
+  assert(Clock::is_steady); // fiber timer assumes deadlines are steady
+
   auto timeoutFunc = [this]() mutable { this->postHelper(TIMEOUT); };
   TimeoutHandler handler;
   handler.timeoutFunc_ = std::ref(timeoutFunc);
 
+  // TODO: have timer support arbitrary clocks
+  const auto now = Clock::now();
+  const auto timeoutMs = std::chrono::duration_cast<std::chrono::milliseconds>(
+      FOLLY_LIKELY(now <= deadline) ? deadline - now : Duration{});
   fm->loopController_->timer().scheduleTimeout(&handler, timeoutMs);
   waitFiber(*fm, static_cast<F&&>(mainContextFunc));
 
   return waiter_ == POSTED;
 }
-
-template <typename Clock, typename Duration, typename F>
-bool Baton::try_wait_until(
-    const std::chrono::time_point<Clock, Duration>& deadline,
-    F&& mainContextFunc) {
-  auto now = Clock::now();
-
-  if (LIKELY(now <= deadline)) {
-    return try_wait_for(deadline - now, static_cast<F&&>(mainContextFunc));
-  } else {
-    return try_wait_for(Duration{}, static_cast<F&&>(mainContextFunc));
-  }
-}
 } // namespace fibers
 } // namespace folly

folly/fibers/Baton.cpp

@@ -24,7 +24,6 @@
 namespace folly {
 namespace fibers {
 
-using folly::detail::futexWaitUntil;
 using folly::detail::futexWake;
 
 void Baton::setWaiter(Waiter& waiter) {
@@ -85,38 +84,6 @@ void Baton::waitThread() {
   throw std::logic_error("Other waiter is already waiting on this baton");
 }
 
-bool Baton::timedWaitThread(std::chrono::milliseconds timeout) {
-  auto waiter = waiter_.load();
-
-  if (LIKELY(
-          waiter == NO_WAITER &&
-          waiter_.compare_exchange_strong(waiter, THREAD_WAITING))) {
-    auto deadline = std::chrono::steady_clock::now() + timeout;
-    do {
-      auto* futex = &futex_.futex;
-      const auto wait_rv =
-          futexWaitUntil(futex, uint32_t(THREAD_WAITING), deadline);
-      if (wait_rv == folly::detail::FutexResult::TIMEDOUT) {
-        return false;
-      }
-      waiter = waiter_.load(std::memory_order_relaxed);
-    } while (waiter == THREAD_WAITING);
-  }
-
-  if (LIKELY(waiter == POSTED)) {
-    return true;
-  }
-
-  // Handle errors
-  if (waiter == TIMEOUT) {
-    throw std::logic_error("Thread baton can't have timeout status");
-  }
-  if (waiter == THREAD_WAITING) {
-    throw std::logic_error("Other thread is already waiting on this baton");
-  }
-  throw std::logic_error("Other waiter is already waiting on this baton");
-}
-
 void Baton::post() {
   postHelper(POSTED);
 }

folly/fibers/Baton.h

@@ -245,7 +245,9 @@ class Baton {
   template <typename F>
   inline void waitFiber(FiberManager& fm, F&& mainContextFunc);
 
-  bool timedWaitThread(std::chrono::milliseconds timeout);
+  template <typename Clock, typename Duration>
+  bool timedWaitThread(
+      const std::chrono::time_point<Clock, Duration>& deadline);
 
   static constexpr intptr_t NO_WAITER = 0;
   static constexpr intptr_t POSTED = -1;