From 4f2c98ac2540ef83b1c8cb4d40110da33ef5125c Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Wed, 2 Nov 2016 17:39:29 -0700
Subject: [PATCH] Use loopKeepAlive() mechanism in FiberManager

Reviewed By: yfeldblum

Differential Revision: D4086486

fbshipit-source-id: bd0cca8dd2c9b74d5c30e4cd095191c1d1ecab79
---
 folly/fibers/EventBaseLoopController-inl.h | 11 +++++++++--
 folly/fibers/EventBaseLoopController.h     |  3 ++-
 folly/fibers/FiberManagerInternal-inl.h    |  8 +++++---
 folly/fibers/FiberManagerInternal.h        |  9 ++++++---
 folly/fibers/FiberManagerMap.cpp           |  5 -----
 folly/fibers/LoopController.h              |  6 ++++++
 folly/fibers/SimpleLoopController.h        |  7 ++++++-
 7 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/folly/fibers/EventBaseLoopController-inl.h b/folly/fibers/EventBaseLoopController-inl.h
index 0e4c2fd7..89a0ece3 100644
--- a/folly/fibers/EventBaseLoopController-inl.h
+++ b/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();
   }
 }
 
diff --git a/folly/fibers/EventBaseLoopController.h b/folly/fibers/EventBaseLoopController.h
index 017b7522..80dd6667 100644
--- a/folly/fibers/EventBaseLoopController.h
+++ b/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;
 
diff --git a/folly/fibers/FiberManagerInternal-inl.h b/folly/fibers/FiberManagerInternal-inl.h
index 4c8878ba..fa526288 100644
--- a/folly/fibers/FiberManagerInternal-inl.h
+++ b/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
diff --git a/folly/fibers/FiberManagerInternal.h b/folly/fibers/FiberManagerInternal.h
index dc0c3e6c..28c74db6 100644
--- a/folly/fibers/FiberManagerInternal.h
+++ b/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.
diff --git a/folly/fibers/FiberManagerMap.cpp b/folly/fibers/FiberManagerMap.cpp
index 45481371..4a45e665 100644
--- a/folly/fibers/FiberManagerMap.cpp
+++ b/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;
 }
 
diff --git a/folly/fibers/LoopController.h b/folly/fibers/LoopController.h
index 917a6557..19cace4e 100644
--- a/folly/fibers/LoopController.h
+++ b/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.
diff --git a/folly/fibers/SimpleLoopController.h b/folly/fibers/SimpleLoopController.h
index 8c728bd6..7928424e 100644
--- a/folly/fibers/SimpleLoopController.h
+++ b/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;
   }
-- 
2.34.1