From a166df2162cab3d8c29e97947c2afc4e12f553d0 Mon Sep 17 00:00:00 2001
From: Jon Maltiel Swenson <jmswen@fb.com>
Date: Thu, 5 Nov 2015 07:24:16 -0800
Subject: [PATCH] Activate server timeout after write success

Summary: Start server timeout after socket write succeeds in mcrouter.  Add neceessary Fibers logic to enable this behavior.

Reviewed By: pavlo-fb

Differential Revision: D2613344

fb-gh-sync-id: 1bc0fbe8b325a3e91cd010f89104b83ebf183679
---
 folly/experimental/fibers/Baton.cpp           | 30 ++++++++++++++
 folly/experimental/fibers/Baton.h             | 41 +++++++++++++++++++
 folly/experimental/fibers/test/FibersTest.cpp | 32 +++++++++++++++
 3 files changed, 103 insertions(+)

diff --git a/folly/experimental/fibers/Baton.cpp b/folly/experimental/fibers/Baton.cpp
index 263dc7d2..b9ca350c 100644
--- a/folly/experimental/fibers/Baton.cpp
+++ b/folly/experimental/fibers/Baton.cpp
@@ -15,7 +15,10 @@
  */
 #include "Baton.h"
 
+#include <chrono>
+
 #include <folly/detail/MemoryIdler.h>
+#include <folly/experimental/fibers/FiberManager.h>
 
 namespace folly { namespace fibers {
 
@@ -23,6 +26,13 @@ void Baton::wait() {
   wait([](){});
 }
 
+void Baton::wait(TimeoutHandler& timeoutHandler) {
+  timeoutHandler.setBaton(this);
+  timeoutHandler.setFiberManager(FiberManager::getFiberManagerUnsafe());
+  wait();
+  timeoutHandler.cancelTimeout();
+}
+
 bool Baton::timed_wait(TimeoutController::Duration timeout) {
   return timed_wait(timeout, [](){});
 }
@@ -153,4 +163,24 @@ void Baton::reset() {
   waitingFiber_.store(NO_WAITER, std::memory_order_relaxed);;
 }
 
+void Baton::TimeoutHandler::scheduleTimeout(uint32_t timeoutMs) {
+  assert(fiberManager_ != nullptr);
+  assert(baton_ != nullptr);
+  if (timeoutMs > 0) {
+    timeoutPtr_ = fiberManager_->timeoutManager_->registerTimeout(
+      [baton = baton_]() {
+        if (!baton->try_wait()) {
+          baton->postHelper(TIMEOUT);
+        }
+      },
+      std::chrono::milliseconds(timeoutMs));
+  }
+}
+
+void Baton::TimeoutHandler::cancelTimeout() {
+  if (timeoutPtr_) {
+    fiberManager_->timeoutManager_->cancel(timeoutPtr_);
+  }
+}
+
 }}
diff --git a/folly/experimental/fibers/Baton.h b/folly/experimental/fibers/Baton.h
index 49585d9e..e0b09329 100644
--- a/folly/experimental/fibers/Baton.h
+++ b/folly/experimental/fibers/Baton.h
@@ -23,6 +23,7 @@
 namespace folly { namespace fibers {
 
 class Fiber;
+class FiberManager;
 
 /**
  * @class Baton
@@ -32,6 +33,8 @@ class Fiber;
  */
 class Baton {
  public:
+  class TimeoutHandler;
+
   Baton();
 
   ~Baton() {}
@@ -41,6 +44,15 @@ class Baton {
    */
   void wait();
 
+  /**
+   * Put active fiber to sleep indefinitely. However, timeoutHandler may
+   * be used elsewhere on the same thread in order to schedule a wakeup
+   * for the active fiber.  Users of timeoutHandler must be on the same thread
+   * as the active fiber and may only schedule one timeout, which must occur
+   * after the active fiber calls wait.
+   */
+  void wait(TimeoutHandler& timeoutHandler);
+
   /**
    * Puts active fiber to sleep. Returns when post is called.
    *
@@ -98,6 +110,35 @@ class Baton {
    */
   void reset();
 
+  /**
+   * Provides a way to schedule a wakeup for a wait()ing fiber.
+   * A TimeoutHandler must be passed to Baton::wait(TimeoutHandler&)
+   * before timeouts are scheduled/cancelled.  It is only safe to use the
+   * TimeoutHandler on the same thread as the wait()ing fiber.
+   * scheduleTimeout() may only be called once prior to the end of the
+   * associated Baton's life.
+   */
+  class TimeoutHandler {
+   public:
+    void scheduleTimeout(uint32_t timeoutMs);
+    void cancelTimeout();
+
+   private:
+    friend class Baton;
+
+    void setFiberManager(FiberManager* fiberManager) {
+      fiberManager_ = fiberManager;
+    }
+    void setBaton(Baton* baton) {
+      baton_ = baton;
+    }
+
+    FiberManager* fiberManager_{nullptr};
+    Baton* baton_{nullptr};
+
+    intptr_t timeoutPtr_{0};
+  };
+
  private:
   enum {
     /**
diff --git a/folly/experimental/fibers/test/FibersTest.cpp b/folly/experimental/fibers/test/FibersTest.cpp
index ed6156f1..d4ce1270 100644
--- a/folly/experimental/fibers/test/FibersTest.cpp
+++ b/folly/experimental/fibers/test/FibersTest.cpp
@@ -1465,6 +1465,38 @@ TEST(FiberManager, resizePeriodically) {
   EXPECT_EQ(5, manager.fibersPoolSize());
 }
 
+TEST(FiberManager, batonWaitTimeoutHandler) {
+  FiberManager manager(folly::make_unique<EventBaseLoopController>());
+
+  folly::EventBase evb;
+  dynamic_cast<EventBaseLoopController&>(manager.loopController())
+    .attachEventBase(evb);
+
+  size_t fibersRun = 0;
+  Baton baton;
+  Baton::TimeoutHandler timeoutHandler;
+
+  manager.addTask([&]() {
+    baton.wait(timeoutHandler);
+    ++fibersRun;
+  });
+  manager.loopUntilNoReady();
+
+  EXPECT_FALSE(baton.try_wait());
+  EXPECT_EQ(0, fibersRun);
+
+  timeoutHandler.scheduleTimeout(250);
+  std::this_thread::sleep_for(std::chrono::milliseconds(500));
+
+  EXPECT_FALSE(baton.try_wait());
+  EXPECT_EQ(0, fibersRun);
+
+  evb.loopOnce();
+  manager.loopUntilNoReady();
+
+  EXPECT_EQ(1, fibersRun);
+}
+
 static size_t sNumAwaits;
 
 void runBenchmark(size_t numAwaits, size_t toSend) {
-- 
2.34.1