Copyright 2014->2015

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

diff --git a/folly/io/async/HHWheelTimer.cpp b/folly/io/async/HHWheelTimer.cpp
index 587a271003c09a6d90d6adb93596583dea3b2ded..a75dc7340b5db1dd69806b1caa3edaf5eac0258a 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;
@@ -81,6 +83,7 @@ HHWheelTimer::HHWheelTimer(folly::EventBase* eventBase,
   , count_(0)
   , catchupEveryN_(DEFAULT_CATCHUP_EVERY_N)
   , expirationsSinceCatchup_(0)
+  , processingCallbacksGuard_(false)
 {
 }
 
@@ -124,7 +127,7 @@ void HHWheelTimer::scheduleTimeout(Callback* callback,
 
   callback->context_ = RequestContext::saveContext();
 
-  if (count_ == 0) {
+  if (count_ == 0 && !processingCallbacksGuard_) {
     this->AsyncTimeout::scheduleTimeout(interval_.count());
   }
 
@@ -150,6 +153,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 +203,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; i < 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