From 06b04cb147733a296278c42c5d0364ec3078e1c9 Mon Sep 17 00:00:00 2001
From: Jon Maltiel Swenson <jmswen@fb.com>
Date: Thu, 5 Nov 2015 18:54:38 -0800
Subject: [PATCH] Cancel timeout only if not run

Summary: Do not cancel TimeoutHandler timeout if the timeout has already run.

Reviewed By: spalamarchuk

Differential Revision: D2622280

fb-gh-sync-id: 27d83b44ab225e2859695f4e5f21cef934824a35
---
 folly/experimental/fibers/Baton.cpp           | 26 ++++++++-------
 folly/experimental/fibers/Baton.h             | 14 +++-----
 folly/experimental/fibers/test/FibersTest.cpp | 32 ++++++++++++++++++-
 3 files changed, 50 insertions(+), 22 deletions(-)

diff --git a/folly/experimental/fibers/Baton.cpp b/folly/experimental/fibers/Baton.cpp
index b9ca350c..b5167b15 100644
--- a/folly/experimental/fibers/Baton.cpp
+++ b/folly/experimental/fibers/Baton.cpp
@@ -27,8 +27,14 @@ void Baton::wait() {
 }
 
 void Baton::wait(TimeoutHandler& timeoutHandler) {
-  timeoutHandler.setBaton(this);
-  timeoutHandler.setFiberManager(FiberManager::getFiberManagerUnsafe());
+  auto timeoutFunc = [this, &timeoutHandler] {
+    if (!try_wait()) {
+      postHelper(TIMEOUT);
+    }
+    timeoutHandler.timeoutPtr_ = 0;
+  };
+  timeoutHandler.timeoutFunc_ = std::ref(timeoutFunc);
+  timeoutHandler.fiberManager_ = FiberManager::getFiberManagerUnsafe();
   wait();
   timeoutHandler.cancelTimeout();
 }
@@ -163,17 +169,15 @@ void Baton::reset() {
   waitingFiber_.store(NO_WAITER, std::memory_order_relaxed);;
 }
 
-void Baton::TimeoutHandler::scheduleTimeout(uint32_t timeoutMs) {
+void Baton::TimeoutHandler::scheduleTimeout(
+    TimeoutController::Duration timeout) {
   assert(fiberManager_ != nullptr);
-  assert(baton_ != nullptr);
-  if (timeoutMs > 0) {
+  assert(timeoutFunc_ != nullptr);
+  assert(timeoutPtr_ == 0);
+
+  if (timeout.count() > 0) {
     timeoutPtr_ = fiberManager_->timeoutManager_->registerTimeout(
-      [baton = baton_]() {
-        if (!baton->try_wait()) {
-          baton->postHelper(TIMEOUT);
-        }
-      },
-      std::chrono::milliseconds(timeoutMs));
+        timeoutFunc_, timeout);
   }
 }
 
diff --git a/folly/experimental/fibers/Baton.h b/folly/experimental/fibers/Baton.h
index e0b09329..862820bb 100644
--- a/folly/experimental/fibers/Baton.h
+++ b/folly/experimental/fibers/Baton.h
@@ -113,28 +113,22 @@ class Baton {
   /**
    * 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
+   * before a timeout is scheduled. 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();
+    void scheduleTimeout(TimeoutController::Duration timeoutMs);
 
    private:
     friend class Baton;
 
-    void setFiberManager(FiberManager* fiberManager) {
-      fiberManager_ = fiberManager;
-    }
-    void setBaton(Baton* baton) {
-      baton_ = baton;
-    }
+    void cancelTimeout();
 
+    std::function<void()> timeoutFunc_{nullptr};
     FiberManager* fiberManager_{nullptr};
-    Baton* baton_{nullptr};
 
     intptr_t timeoutPtr_{0};
   };
diff --git a/folly/experimental/fibers/test/FibersTest.cpp b/folly/experimental/fibers/test/FibersTest.cpp
index d4ce1270..641c8db8 100644
--- a/folly/experimental/fibers/test/FibersTest.cpp
+++ b/folly/experimental/fibers/test/FibersTest.cpp
@@ -1485,7 +1485,7 @@ TEST(FiberManager, batonWaitTimeoutHandler) {
   EXPECT_FALSE(baton.try_wait());
   EXPECT_EQ(0, fibersRun);
 
-  timeoutHandler.scheduleTimeout(250);
+  timeoutHandler.scheduleTimeout(std::chrono::milliseconds(250));
   std::this_thread::sleep_for(std::chrono::milliseconds(500));
 
   EXPECT_FALSE(baton.try_wait());
@@ -1497,6 +1497,36 @@ TEST(FiberManager, batonWaitTimeoutHandler) {
   EXPECT_EQ(1, fibersRun);
 }
 
+TEST(FiberManager, batonWaitTimeoutMany) {
+  FiberManager manager(folly::make_unique<EventBaseLoopController>());
+
+  folly::EventBase evb;
+  dynamic_cast<EventBaseLoopController&>(manager.loopController())
+    .attachEventBase(evb);
+
+  constexpr size_t kNumTimeoutTasks = 10000;
+  size_t tasksCount = kNumTimeoutTasks;
+
+  // We add many tasks to hit timeout queue deallocation logic.
+  for (size_t i = 0; i < kNumTimeoutTasks; ++i) {
+    manager.addTask([&]() {
+      Baton baton;
+      Baton::TimeoutHandler timeoutHandler;
+
+      folly::fibers::addTask([&] {
+        timeoutHandler.scheduleTimeout(std::chrono::milliseconds(1000));
+      });
+
+      baton.wait(timeoutHandler);
+      if (--tasksCount == 0) {
+        evb.terminateLoopSoon();
+      }
+    });
+  }
+
+  evb.loopForever();
+}
+
 static size_t sNumAwaits;
 
 void runBenchmark(size_t numAwaits, size_t toSend) {
-- 
2.34.1