Fix pool resizer to correctly handle VirtualEventBase destruction

Summary: It's possible that we still have a recurring pool resizer task scheduled after virtual event base is destroyed. This changes makes sure that we don't try to reschedule it if VirtualEventBase is already destroyed.

Reviewed By: yfeldblum

Differential Revision: D21875241

fbshipit-source-id: c4622c5d103c6291386055c153ab4aeebe7795a2

folly/fibers/Baton-inl.h

@@ -134,7 +134,7 @@ bool Baton::try_wait_until(
   const auto now = Clock::now();
   const auto timeoutMs = std::chrono::duration_cast<std::chrono::milliseconds>(
       FOLLY_LIKELY(now <= deadline) ? deadline - now : Duration{});
-  fm->loopController_->timer().scheduleTimeout(&handler, timeoutMs);
+  fm->loopController_->timer()->scheduleTimeout(&handler, timeoutMs);
   waitFiber(*fm, static_cast<F&&>(mainContextFunc));
 
   return waiter_ == POSTED;

folly/fibers/Baton.cpp

@@ -132,7 +132,7 @@ void Baton::TimeoutHandler::scheduleTimeout(std::chrono::milliseconds timeout) {
   assert(timeoutFunc_ != nullptr);
 
   if (timeout.count() > 0) {
-    fiberManager_->loopController_->timer().scheduleTimeout(this, timeout);
+    fiberManager_->loopController_->timer()->scheduleTimeout(this, timeout);
   }
 }
 

folly/fibers/EventBaseLoopController-inl.h

@@ -33,12 +33,17 @@ inline void EventBaseLoopController::attachEventBase(EventBase& eventBase) {
 
 inline void EventBaseLoopController::attachEventBase(
     VirtualEventBase& eventBase) {
-  if (eventBase_ != nullptr) {
+  if (eventBaseAttached_.exchange(true)) {
     LOG(ERROR) << "Attempt to reattach EventBase to LoopController";
+    return;
   }
 
   eventBase_ = &eventBase;
-  eventBaseAttached_ = true;
+
+  CancellationSource source;
+  eventBaseShutdownToken_ = source.getToken();
+  eventBase_->runOnDestruction(
+      [source = std::move(source)] { source.requestCancellation(); });
 
   if (awaitingScheduling_) {
     schedule();
@@ -119,10 +124,14 @@ inline void EventBaseLoopController::scheduleThreadSafe() {
       });
 }
 
-inline HHWheelTimer& EventBaseLoopController::timer() {
+inline HHWheelTimer* EventBaseLoopController::timer() {
   assert(eventBaseAttached_);
 
-  return eventBase_->timer();
+  if (UNLIKELY(eventBaseShutdownToken_.isCancellationRequested())) {
+    return nullptr;
+  }
+
+  return &eventBase_->timer();
 }
 } // namespace fibers
 } // namespace folly

folly/fibers/EventBaseLoopController.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <folly/CancellationToken.h>
 #include <folly/fibers/ExecutorBasedLoopController.h>
 #include <folly/fibers/FiberManagerInternal.h>
 #include <folly/io/async/VirtualEventBase.h>
@@ -62,6 +63,8 @@ class EventBaseLoopController : public ExecutorBasedLoopController {
     EventBaseLoopController& controller_;
   };
 
+  folly::CancellationToken eventBaseShutdownToken_;
+
   bool awaitingScheduling_{false};
   VirtualEventBase* eventBase_{nullptr};
   Executor::KeepAlive<VirtualEventBase> eventBaseKeepAlive_;
@@ -77,7 +80,7 @@ class EventBaseLoopController : public ExecutorBasedLoopController {
   void runLoop() override;
   void runEagerFiber(Fiber*) override;
   void scheduleThreadSafe() override;
-  HHWheelTimer& timer() override;
+  HHWheelTimer* timer() override;
 
   friend class FiberManager;
 };

folly/fibers/ExecutorLoopController-inl.h

@@ -71,8 +71,8 @@ inline void ExecutorLoopController::scheduleThreadSafe() {
       });
 }
 
-inline HHWheelTimer& ExecutorLoopController::timer() {
-  return *timer_;
+inline HHWheelTimer* ExecutorLoopController::timer() {
+  return timer_.get();
 }
 
 } // namespace fibers

folly/fibers/ExecutorLoopController.h

@@ -94,7 +94,7 @@ class ExecutorLoopController : public fibers::ExecutorBasedLoopController {
   void runLoop() override;
   void runEagerFiber(Fiber*) override;
   void scheduleThreadSafe() override;
-  HHWheelTimer& timer() override;
+  HHWheelTimer* timer() override;
 
   friend class fibers::FiberManager;
 };

folly/fibers/FiberManager.cpp

@@ -199,10 +199,12 @@ void FiberManager::doFibersPoolResizing() {
 
 void FiberManager::FibersPoolResizer::run() {
   fiberManager_.doFibersPoolResizing();
-  fiberManager_.loopController_->timer().scheduleTimeout(
+  if (auto timer = fiberManager_.loopController_->timer()) {
+    timer->scheduleTimeout(
         this,
         std::chrono::milliseconds(
             fiberManager_.options_.fibersPoolResizePeriodMs));
+  }
 }
 
 #ifdef FOLLY_SANITIZE_ADDRESS

folly/fibers/LoopController.h

@@ -64,8 +64,10 @@ class LoopController {
 
   /**
    * Used by FiberManager to schedule some function to be run at some time.
+   * May return null, but only if called outside of runLoop() call (e.g. if
+   * Executor backing the timer is already destroyed).
    */
-  virtual HHWheelTimer& timer() = 0;
+  virtual HHWheelTimer* timer() = 0;
 };
 } // namespace fibers
 } // namespace folly

folly/fibers/SimpleLoopController.h

@@ -92,8 +92,8 @@ class SimpleLoopController : public LoopController {
     scheduled_ = true;
   }
 
-  HHWheelTimer& timer() override {
-    return *timer_;
+  HHWheelTimer* timer() override {
+    return timer_.get();
   }
 
   bool isInLoopThread() const {

folly/fibers/test/FibersTest.cpp

@@ -1478,6 +1478,10 @@ TEST(FiberManager, resizePeriodically) {
   evb.loopOnce();
   EXPECT_EQ(5, manager.fibersAllocated());
   EXPECT_EQ(5, manager.fibersPoolSize());
+
+  // Sleep again before destruction to force the case where the
+  // resize timer fires during destruction of the EventBase.
+  std::this_thread::sleep_for(std::chrono::milliseconds(400));
 }
 
 TEST(FiberManager, batonWaitTimeoutHandler) {