From 09a84b282cc6d0d5498877eb16e50bc7f5c9869f Mon Sep 17 00:00:00 2001
From: Dave Watson <davejwatson@fb.com>
Date: Wed, 10 Aug 2016 15:41:29 -0700
Subject: [PATCH] Fix scheduling bug

Summary:
Noticed this while debugging other timerwheel changes.  Scheduling new events in callbacks may result in us *running them right away* if they land in a bucket we are currently processing.

Instead, round up all the timeouts, then run them separately.  This means you can never schedule a timeout that will run immediately (0ticks), but that's probably fine.

Reviewed By: yfeldblum

Differential Revision: D3679413

fbshipit-source-id: 7e18632f23ea33c072c2718e30b378641895b094
---
 folly/io/async/HHWheelTimer.cpp          | 34 +++++++++-----
 folly/io/async/HHWheelTimer.h            |  1 +
 folly/io/async/test/HHWheelTimerTest.cpp | 57 ++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 11 deletions(-)

diff --git a/folly/io/async/HHWheelTimer.cpp b/folly/io/async/HHWheelTimer.cpp
index 31213ba5..9adcf2d4 100644
--- a/folly/io/async/HHWheelTimer.cpp
+++ b/folly/io/async/HHWheelTimer.cpp
@@ -97,6 +97,12 @@ HHWheelTimer::~HHWheelTimer() {
       *processingCallbacksGuard_ = true;
     }
   });
+  while (!timeouts.empty()) {
+    auto* cb = &timeouts.front();
+    timeouts.pop_front();
+    cb->cancelTimeout();
+    cb->callbackCanceled();
+  }
   cancelAll();
 }
 
@@ -198,17 +204,23 @@ void HHWheelTimer::timeoutExpired() noexcept {
     while (!cbs->empty()) {
       auto* cb = &cbs->front();
       cbs->pop_front();
-      count_--;
-      cb->wheel_ = nullptr;
-      cb->expiration_ = milliseconds(0);
-      RequestContextScopeGuard rctx(cb->context_);
-      cb->timeoutExpired();
-      if (isDestroyed) {
-        // The HHWheelTimer itself has been destroyed. The other callbacks
-        // will have been cancelled from the destructor. Bail before causing
-        // damage.
-        return;
-      }
+      timeouts.push_back(*cb);
+    }
+  }
+
+  while (!timeouts.empty()) {
+    auto* cb = &timeouts.front();
+    timeouts.pop_front();
+    count_--;
+    cb->wheel_ = nullptr;
+    cb->expiration_ = milliseconds(0);
+    RequestContextScopeGuard rctx(cb->context_);
+    cb->timeoutExpired();
+    if (isDestroyed) {
+      // The HHWheelTimer itself has been destroyed. The other callbacks
+      // will have been cancelled from the destructor. Bail before causing
+      // damage.
+      return;
     }
   }
   if (count_ > 0) {
diff --git a/folly/io/async/HHWheelTimer.h b/folly/io/async/HHWheelTimer.h
index 33300828..c8602336 100644
--- a/folly/io/async/HHWheelTimer.h
+++ b/folly/io/async/HHWheelTimer.h
@@ -297,6 +297,7 @@ class HHWheelTimer : private folly::AsyncTimeout,
   std::chrono::milliseconds now_;
 
   bool* processingCallbacksGuard_;
+  CallbackList timeouts; // Timeouts queued to run
 };
 
 } // folly
diff --git a/folly/io/async/test/HHWheelTimerTest.cpp b/folly/io/async/test/HHWheelTimerTest.cpp
index bcbf72d8..0ee82125 100644
--- a/folly/io/async/test/HHWheelTimerTest.cpp
+++ b/folly/io/async/test/HHWheelTimerTest.cpp
@@ -282,6 +282,63 @@ TEST_F(HHWheelTimerTest, SlowFast) {
   T_CHECK_TIMEOUT(start, t2.timestamps[0], milliseconds(5), milliseconds(1));
 }
 
+TEST_F(HHWheelTimerTest, ReschedTest) {
+  StackWheelTimer t(&eventBase, milliseconds(1));
+
+  TestTimeout t1;
+  TestTimeout t2;
+
+  ASSERT_EQ(t.count(), 0);
+
+  t.scheduleTimeout(&t1, milliseconds(128));
+  TimePoint start2;
+  t1.fn = [&]() {
+    t.scheduleTimeout(&t2, milliseconds(255)); // WHEEL_SIZE - 1
+    start2.reset();
+    ASSERT_EQ(t.count(), 1);
+  };
+
+  ASSERT_EQ(t.count(), 1);
+
+  TimePoint start;
+  eventBase.loop();
+  TimePoint end;
+
+  ASSERT_EQ(t1.timestamps.size(), 1);
+  ASSERT_EQ(t2.timestamps.size(), 1);
+  ASSERT_EQ(t.count(), 0);
+
+  T_CHECK_TIMEOUT(start, t1.timestamps[0], milliseconds(128), milliseconds(1));
+  T_CHECK_TIMEOUT(start2, t2.timestamps[0], milliseconds(255), milliseconds(1));
+}
+
+TEST_F(HHWheelTimerTest, DeleteWheelInTimeout) {
+  auto t = HHWheelTimer::newTimer(&eventBase, milliseconds(1));
+
+  TestTimeout t1;
+  TestTimeout t2;
+  TestTimeout t3;
+
+  ASSERT_EQ(t->count(), 0);
+
+  t->scheduleTimeout(&t1, milliseconds(128));
+  t->scheduleTimeout(&t2, milliseconds(128));
+  t->scheduleTimeout(&t3, milliseconds(128));
+  t1.fn = [&]() { t2.cancelTimeout(); };
+  t3.fn = [&]() { t.reset(); };
+
+  ASSERT_EQ(t->count(), 3);
+
+  TimePoint start;
+  eventBase.loop();
+  TimePoint end;
+
+  ASSERT_EQ(t1.timestamps.size(), 1);
+  ASSERT_EQ(t2.timestamps.size(), 0);
+
+  T_CHECK_TIMEOUT(start, t1.timestamps[0], milliseconds(128), milliseconds(1));
+}
+
 /*
  * Test scheduling a mix of timers with default timeout and variable timeout.
  */
-- 
2.34.1