Add a default timeout parameter to HHWheelTimer.

[folly.git] / folly / io / async / HHWheelTimer.cpp

diff --git a/folly/io/async/HHWheelTimer.cpp b/folly/io/async/HHWheelTimer.cpp
index 587a271003c09a6d90d6adb93596583dea3b2ded..1fd5e6d25d81b041bf691c154b610b63664bd456 100644 (file)
--- a/folly/io/async/HHWheelTimer.cpp
+++ b/folly/io/async/HHWheelTimer.cpp
@@ -1,4 +1,6 @@
 /*
+ * Copyright 2015 Facebook, Inc.
+ *
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements. See the NOTICE file
  * distributed with this work for additional information
@@ -54,9 +56,9 @@ void HHWheelTimer::Callback::setScheduled(HHWheelTimer* wheel,
 
   wheel_ = wheel;
 
-  if (wheel_->count_  == 0) {
-    wheel_->now_ = std::chrono::duration_cast<milliseconds>(
-      std::chrono::steady_clock::now().time_since_epoch());
+  // Only update the now_ time if we're not in a timeout expired callback
+  if (wheel_->count_  == 0 && !wheel_->processingCallbacksGuard_) {
+    wheel_->now_ = getCurTime();
   }
 
   expiration_ = wheel_->now_ + timeout;
@@ -74,13 +76,17 @@ void HHWheelTimer::Callback::cancelTimeoutImpl() {
 }
 
 HHWheelTimer::HHWheelTimer(folly::EventBase* eventBase,
-                           std::chrono::milliseconds intervalMS)
-  : AsyncTimeout(eventBase)
+                           std::chrono::milliseconds intervalMS,
+                           AsyncTimeout::InternalEnum internal,
+                           std::chrono::milliseconds defaultTimeoutMS)
+  : AsyncTimeout(eventBase, internal)
   , interval_(intervalMS)
+  , defaultTimeout_(defaultTimeoutMS)
   , nextTick_(1)
   , count_(0)
   , catchupEveryN_(DEFAULT_CATCHUP_EVERY_N)
   , expirationsSinceCatchup_(0)
+  , processingCallbacksGuard_(false)
 {
 }
 
@@ -124,7 +130,7 @@ void HHWheelTimer::scheduleTimeout(Callback* callback,
 
   callback->context_ = RequestContext::saveContext();
 
-  if (count_ == 0) {
+  if (count_ == 0 && !processingCallbacksGuard_) {
     this->AsyncTimeout::scheduleTimeout(interval_.count());
   }
 
@@ -133,6 +139,12 @@ void HHWheelTimer::scheduleTimeout(Callback* callback,
   count_++;
 }
 
+void HHWheelTimer::scheduleTimeout(Callback* callback) {
+  CHECK(std::chrono::milliseconds(-1) != defaultTimeout_)
+      << "Default timeout was not initialized";
+  scheduleTimeout(callback, defaultTimeout_);
+}
+
 bool HHWheelTimer::cascadeTimers(int bucket, int tick) {
   CallbackList cbs;
   cbs.swap(buckets_[bucket][tick]);
@@ -150,6 +162,13 @@ void HHWheelTimer::timeoutExpired() noexcept {
   // If destroy() is called inside timeoutExpired(), delay actual destruction
   // until timeoutExpired() returns
   DestructorGuard dg(this);
+  // If scheduleTimeout is called from a callback in this function, it may
+  // cause inconsistencies in the state of this object. As such, we need
+  // to treat these calls slightly differently.
+  processingCallbacksGuard_ = true;
+  auto reEntryGuard = folly::makeGuard([&] {
+    processingCallbacksGuard_ = false;
+  });
 
   // timeoutExpired() can only be invoked directly from the event base loop.
   // It should never be invoked recursively.
@@ -193,4 +212,36 @@ void HHWheelTimer::timeoutExpired() noexcept {
   }
 }
 
+size_t HHWheelTimer::cancelAll() {
+  decltype(buckets_) buckets;
+
+// Work around std::swap() bug in libc++
+//
+// http://llvm.org/bugs/show_bug.cgi?id=22106
+#if FOLLY_USE_LIBCPP
+  for (size_t i = 0; i < WHEEL_BUCKETS; ++i) {
+    for (size_t ii = 0; ii < WHEEL_SIZE; ++ii) {
+      std::swap(buckets_[i][ii], buckets[i][ii]);
+    }
+  }
+#else
+  std::swap(buckets, buckets_);
+#endif
+
+  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