From: Hans Fugal <fugalh@fb.com>
Date: Fri, 26 Dec 2014 23:18:08 +0000 (-0800)
Subject: HHWheelTimer::cancelAll
X-Git-Tag: v0.22.0~68
X-Git-Url: http://plrg.eecs.uci.edu/git/?a=commitdiff_plain;h=062bc87d14e1135bff6ee43edc5f3a40d9ee5943;p=folly.git

HHWheelTimer::cancelAll

Summary: Provide a way to massacre outstanding timers.

Test Plan: new test, and using in wangle (different diff)

Reviewed By: davejwatson@fb.com

Subscribers: jsedgwick, exa, folly-diffs@

FB internal diff: D1753866

Tasks: 4548494

Signature: t1:1753866:1419361846:58ff6ca4d01b0d546495b032b454c8bf0fdb0277
---

diff --git a/folly/io/async/HHWheelTimer.cpp b/folly/io/async/HHWheelTimer.cpp
index 5eec0ba4..d7622248 100644
--- a/folly/io/async/HHWheelTimer.cpp
+++ b/folly/io/async/HHWheelTimer.cpp
@@ -203,4 +203,23 @@ void HHWheelTimer::timeoutExpired() noexcept {
   }
 }
 
+size_t HHWheelTimer::cancelAll() {
+  decltype(buckets_) buckets;
+  std::swap(buckets, buckets_);
+  size_t count = 0;
+
+  for (auto& tick : buckets) {
+    for (auto& bucket : tick) {
+      while (!bucket.empty()) {
+        auto& cb = bucket.front();
+        cb.cancelTimeout();
+        cb.callbackCanceled();
+        count++;
+      }
+    }
+  }
+
+  return count;
+}
+
 } // folly
diff --git a/folly/io/async/HHWheelTimer.h b/folly/io/async/HHWheelTimer.h
index 7c6db8a0..b15de3e1 100644
--- a/folly/io/async/HHWheelTimer.h
+++ b/folly/io/async/HHWheelTimer.h
@@ -76,6 +76,13 @@ class HHWheelTimer : private folly::AsyncTimeout,
      */
     virtual void timeoutExpired() noexcept = 0;
 
+    /// This callback was canceled. The default implementation is to just
+    /// proxy to `timeoutExpired` but if you care about the difference between
+    /// the timeout finishing or being canceled you can override this.
+    virtual void callbackCanceled() noexcept {
+      timeoutExpired();
+    }
+
     /**
      * Cancel the timeout, if it is running.
      *
@@ -152,10 +159,18 @@ class HHWheelTimer : private folly::AsyncTimeout,
    * Destroy the HHWheelTimer.
    *
    * A HHWheelTimer should only be destroyed when there are no more
-   * callbacks pending in the set.
+   * callbacks pending in the set. (If it helps you may use cancelAll() to
+   * cancel all pending timeouts explicitly before calling this.)
    */
   virtual void destroy();
 
+  /**
+   * Cancel all outstanding timeouts
+   *
+   * @returns the number of timeouts that were cancelled.
+   */
+  size_t cancelAll();
+
   /**
    * Get the tick interval for this HHWheelTimer.
    *
diff --git a/folly/io/async/test/HHWheelTimerTest.cpp b/folly/io/async/test/HHWheelTimerTest.cpp
index 3193926f..e2b5c096 100644
--- a/folly/io/async/test/HHWheelTimerTest.cpp
+++ b/folly/io/async/test/HHWheelTimerTest.cpp
@@ -35,14 +35,23 @@ class TestTimeout : public HHWheelTimer::Callback {
   TestTimeout(HHWheelTimer* t, milliseconds timeout) {
     t->scheduleTimeout(this, timeout);
   }
-  virtual void timeoutExpired() noexcept {
+
+  void timeoutExpired() noexcept override {
     timestamps.push_back(TimePoint());
     if (fn) {
       fn();
     }
   }
 
+  void callbackCanceled() noexcept override {
+    canceledTimestamps.push_back(TimePoint());
+    if (fn) {
+      fn();
+    }
+  }
+
   std::deque<TimePoint> timestamps;
+  std::deque<TimePoint> canceledTimestamps;
   std::function<void()> fn;
 };
 
@@ -392,7 +401,15 @@ TEST_F(HHWheelTimerTest, lambda) {
 // at the console to confirm logging)
 TEST_F(HHWheelTimerTest, lambdaThrows) {
   StackWheelTimer t(&eventBase, milliseconds(1));
-  t.scheduleTimeoutFn([&]{ throw std::runtime_error("foo"); },
+  t.scheduleTimeoutFn([&]{ throw std::runtime_error("expected"); },
                       milliseconds(1));
   eventBase.loop();
 }
+
+TEST_F(HHWheelTimerTest, cancelAll) {
+  StackWheelTimer t(&eventBase);
+  TestTimeout tt;
+  t.scheduleTimeout(&tt, std::chrono::minutes(1));
+  EXPECT_EQ(1, t.cancelAll());
+  EXPECT_EQ(1, tt.canceledTimestamps.size());
+}