Use loopKeepAlive() mechanism in FiberManager

Reviewed By: yfeldblum

Differential Revision: D4086486

fbshipit-source-id: bd0cca8dd2c9b74d5c30e4cd095191c1d1ecab79

folly/fibers/EventBaseLoopController-inl.h

@@ -24,6 +24,7 @@ inline EventBaseLoopController::EventBaseLoopController()
 
 inline EventBaseLoopController::~EventBaseLoopController() {
   callback_.cancelLoopCallback();
+  eventBaseKeepAlive_.reset();
 }
 
 inline void EventBaseLoopController::attachEventBase(
@@ -62,10 +63,16 @@ inline void EventBaseLoopController::cancel() {
 }
 
 inline void EventBaseLoopController::runLoop() {
+  if (!eventBaseKeepAlive_) {
+    eventBaseKeepAlive_ = eventBase_->loopKeepAlive();
+  }
   if (loopRunner_) {
-    loopRunner_->run([&] { fm_->loopUntilNoReady(); });
+    loopRunner_->run([&] { fm_->loopUntilNoReadyImpl(); });
   } else {
-    fm_->loopUntilNoReady();
+    fm_->loopUntilNoReadyImpl();
+  }
+  if (!fm_->hasTasks()) {
+    eventBaseKeepAlive_.reset();
   }
 }
 

folly/fibers/EventBaseLoopController.h

@@ -91,6 +91,7 @@ class EventBaseLoopController : public LoopController {
 
   bool awaitingScheduling_{false};
   folly::EventBase* eventBase_{nullptr};
+  folly::EventBase::LoopKeepAlive eventBaseKeepAlive_;
   ControllerCallback callback_;
   DestructionCallback destructionCallback_;
   FiberManager* fm_{nullptr};
@@ -103,7 +104,7 @@ class EventBaseLoopController : public LoopController {
   void setFiberManager(FiberManager* fm) override;
   void schedule() override;
   void cancel() override;
-  void runLoop();
+  void runLoop() override;
   void scheduleThreadSafe(std::function<bool()> func) override;
   void timedSchedule(std::function<void()> func, TimePoint time) override;
 

folly/fibers/FiberManagerInternal-inl.h

@@ -185,7 +185,11 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
   }
 }
 
-inline bool FiberManager::loopUntilNoReady() {
+inline void FiberManager::loopUntilNoReady() {
+  return loopController_->runLoop();
+}
+
+inline void FiberManager::loopUntilNoReadyImpl() {
 #ifndef _WIN32
   if (UNLIKELY(!alternateSignalStackRegistered_)) {
     registerAlternateSignalStack();
@@ -243,8 +247,6 @@ inline bool FiberManager::loopUntilNoReady() {
     }
   }
   readyFibers_.splice(readyFibers_.end(), yieldedFibers_);
-
-  return fibersActive_ > 0;
 }
 
 // We need this to be in a struct, not inlined in addTask, because clang crashes

folly/fibers/FiberManagerInternal.h

@@ -155,10 +155,13 @@ class FiberManager : public ::folly::Executor {
 
   /**
    * Keeps running ready tasks until the list of ready tasks is empty.
-   *
-   * @return True if there are any waiting tasks remaining.
    */
-  bool loopUntilNoReady();
+  void loopUntilNoReady();
+
+  /**
+   * This should only be called by a LoopController.
+   */
+  void loopUntilNoReadyImpl();
 
   /**
    * @return true if there are outstanding tasks.

folly/fibers/FiberManagerMap.cpp

@@ -170,11 +170,6 @@ void EventBaseOnDestructionCallback::runLoopCallback() noexcept {
   DCHECK(fm.get() != nullptr);
   ThreadLocalCache::erase(evb_);
 
-  while (fm->hasTasks()) {
-    fm->loopUntilNoReady();
-    evb_.loopOnce();
-  }
-
   delete this;
 }
 

folly/fibers/LoopController.h

@@ -41,6 +41,12 @@ class LoopController {
    */
   virtual void schedule() = 0;
 
+  /**
+   * Run FiberManager loopUntilNoReadyImpl(). May have additional logic specific
+   * to a LoopController.
+   */
+  virtual void runLoop() = 0;
+
   /**
    * Same as schedule(), but safe to call from any thread.
    * Runs func and only schedules if func returned true.

folly/fibers/SimpleLoopController.h

@@ -54,7 +54,8 @@ class SimpleLoopController : public LoopController {
 
       if (scheduled_) {
         scheduled_ = false;
-        waiting = fm_->loopUntilNoReady();
+        runLoop();
+        waiting = fm_->hasTasks();
       }
     }
   }
@@ -70,6 +71,10 @@ class SimpleLoopController : public LoopController {
     return remoteScheduleCalled_;
   }
 
+  void runLoop() override {
+    fm_->loopUntilNoReadyImpl();
+  }
+
   void schedule() override {
     scheduled_ = true;
   }