Cancel timeouts about to run in the cancelAll

Summary:
Currently if `cancelAll` is called from inside the `timeoutExpired` of one of the callbacks, it will not cancel timeouts that we're about to run (they were extracted from the buckets already).

This diff fixes that behavior by also canceling timeouts in `timeoutsToRunNow_` list (note, we were already doing that in the destructor).

Reviewed By: yfeldblum

Differential Revision: D13531908

fbshipit-source-id: f05ba31f2ac845851c1560d2ebdf41aa995b2deb

folly/io/async/HHWheelTimer.cpp

@@ -99,12 +99,6 @@ HHWheelTimer::~HHWheelTimer() {
       *processingCallbacksGuard_ = true;
     }
   });
-  while (!timeouts.empty()) {
-    auto* cb = &timeouts.front();
-    timeouts.pop_front();
-    cb->cancelTimeout();
-    cb->callbackCanceled();
-  }
   cancelAll();
 }
 
@@ -214,7 +208,7 @@ void HHWheelTimer::timeoutExpired() noexcept {
     while (!cbs->empty()) {
       auto* cb = &cbs->front();
       cbs->pop_front();
-      timeouts.push_back(*cb);
+      timeoutsToRunNow_.push_back(*cb);
     }
 
     if (idx == 0) {
@@ -226,9 +220,9 @@ void HHWheelTimer::timeoutExpired() noexcept {
     }
   }
 
-  while (!timeouts.empty()) {
-    auto* cb = &timeouts.front();
-    timeouts.pop_front();
+  while (!timeoutsToRunNow_.empty()) {
+    auto* cb = &timeoutsToRunNow_.front();
+    timeoutsToRunNow_.pop_front();
     count_--;
     cb->wheel_ = nullptr;
     cb->expiration_ = {};
@@ -266,13 +260,13 @@ size_t HHWheelTimer::cancelAll() {
     }
 
     for (size_t i = 0; i < countBuckets; ++i) {
-      auto& bucket = buckets[i];
-      while (!bucket.empty()) {
-        auto& cb = bucket.front();
-        cb.cancelTimeout();
-        cb.callbackCanceled();
-      }
+      cancelTimeoutsFromList(buckets[i]);
     }
+    // Swap the list to prevent potential recursion if cancelAll is called by
+    // one of the callbacks.
+    CallbackList timeoutsToRunNow;
+    timeoutsToRunNow.swap(timeoutsToRunNow_);
+    count += cancelTimeoutsFromList(timeoutsToRunNow);
   }
 
   return count;
@@ -304,6 +298,17 @@ void HHWheelTimer::scheduleNextTimeout() {
   }
 }
 
+size_t HHWheelTimer::cancelTimeoutsFromList(CallbackList& timeouts) {
+  size_t count = 0;
+  while (!timeouts.empty()) {
+    ++count;
+    auto& cb = timeouts.front();
+    cb.cancelTimeout();
+    cb.callbackCanceled();
+  }
+  return count;
+}
+
 int64_t HHWheelTimer::calcNextTick() {
   auto intervals = (getCurTime() - startTime_) / interval_;
   // Slow eventbases will have skew between the actual time and the

folly/io/async/HHWheelTimer.h

@@ -305,8 +305,13 @@ class HHWheelTimer : private folly::AsyncTimeout,
 
   void scheduleNextTimeout();
 
+  size_t cancelTimeoutsFromList(CallbackList& timeouts);
+
   bool* processingCallbacksGuard_;
-  CallbackList timeouts; // Timeouts queued to run
+  // Timeouts that we're about to run. They're already extracted from their
+  // corresponding buckets, so we need this list for the `cancelAll` to be able
+  // to cancel them.
+  CallbackList timeoutsToRunNow_;
 
   std::chrono::steady_clock::time_point getCurTime() {
     return std::chrono::steady_clock::now();

folly/io/async/test/HHWheelTimerTest.cpp

@@ -399,10 +399,19 @@ TEST_F(HHWheelTimerTest, lambdaThrows) {
 
 TEST_F(HHWheelTimerTest, cancelAll) {
   StackWheelTimer t(&eventBase, milliseconds(1));
-  TestTimeout tt;
-  t.scheduleTimeout(&tt, std::chrono::minutes(1));
-  EXPECT_EQ(1, t.cancelAll());
-  EXPECT_EQ(1, tt.canceledTimestamps.size());
+  TestTimeout t1;
+  TestTimeout t2;
+  t.scheduleTimeout(&t1, std::chrono::milliseconds(1));
+  t.scheduleTimeout(&t2, std::chrono::milliseconds(1));
+  size_t canceled = 0;
+  t1.fn = [&] { canceled += t.cancelAll(); };
+  t2.fn = [&] { canceled += t.cancelAll(); };
+  // Sleep 20ms to ensure both timeouts will fire in a single event (in case
+  // they ended up in different slots)
+  ::usleep(20000);
+  eventBase.loop();
+  EXPECT_EQ(1, t1.canceledTimestamps.size() + t2.canceledTimestamps.size());
+  EXPECT_EQ(1, canceled);
 }
 
 TEST_F(HHWheelTimerTest, IntrusivePtr) {