Fix KeepAlive management for eager variants of addTask*

Summary:
LoopController::runLoop is responsible for managing Executor/EventBase keep-alive handle.
This diff introduces a similar level of indirection (LoopController::addTaskEager) to solve the same problem for eager variants of addTask*.

Reviewed By: yfeldblum

Differential Revision: D20613287

fbshipit-source-id: 22388ca3b31664c5fedfd509f2d9bab3d86aab47

folly/fibers/EventBaseLoopController-inl.h

@@ -85,6 +85,20 @@ inline void EventBaseLoopController::runLoop() {
   }
 }
 
+inline void EventBaseLoopController::runEagerFiber(Fiber* fiber) {
+  if (!eventBaseKeepAlive_) {
+    eventBaseKeepAlive_ = getKeepAliveToken(eventBase_);
+  }
+  if (loopRunner_) {
+    loopRunner_->run([&] { fm_->runEagerFiberImpl(fiber); });
+  } else {
+    fm_->runEagerFiberImpl(fiber);
+  }
+  if (!fm_->hasTasks()) {
+    eventBaseKeepAlive_.reset();
+  }
+}
+
 inline void EventBaseLoopController::scheduleThreadSafe() {
   /* The only way we could end up here is if
      1) Fiber thread creates a fiber that awaits (which means we must

folly/fibers/EventBaseLoopController.h

@@ -75,6 +75,7 @@ class EventBaseLoopController : public ExecutorBasedLoopController {
   void setFiberManager(FiberManager* fm) override;
   void schedule() override;
   void runLoop() override;
+  void runEagerFiber(Fiber*) override;
   void scheduleThreadSafe() override;
   HHWheelTimer& timer() override;
 

folly/fibers/ExecutorLoopController-inl.h

@@ -52,6 +52,16 @@ inline void ExecutorLoopController::runLoop() {
   }
 }
 
+inline void ExecutorLoopController::runEagerFiber(Fiber* fiber) {
+  if (!executorKeepAlive_) {
+    executorKeepAlive_ = getKeepAliveToken(executor_);
+  }
+  fm_->runEagerFiberImpl(fiber);
+  if (!fm_->hasTasks()) {
+    executorKeepAlive_.reset();
+  }
+}
+
 inline void ExecutorLoopController::scheduleThreadSafe() {
   executor_->add(
       [this, executorKeepAlive = getKeepAliveToken(executor_)]() mutable {

folly/fibers/ExecutorLoopController.h

@@ -92,6 +92,7 @@ class ExecutorLoopController : public fibers::ExecutorBasedLoopController {
   void setFiberManager(fibers::FiberManager* fm) override;
   void schedule() override;
   void runLoop() override;
+  void runEagerFiber(Fiber*) override;
   void scheduleThreadSafe() override;
   HHWheelTimer& timer() override;
 

folly/fibers/FiberManagerInternal-inl.h

@@ -281,6 +281,10 @@ inline void FiberManager::loopUntilNoReadyImpl() {
 }
 
 inline void FiberManager::runEagerFiber(Fiber* fiber) {
+  loopController_->runEagerFiber(fiber);
+}
+
+inline void FiberManager::runEagerFiberImpl(Fiber* fiber) {
   runInMainContext([&] {
     auto prevCurrentFiber = std::exchange(currentFiber_, fiber);
     SCOPE_EXIT {

folly/fibers/FiberManagerInternal.h

@@ -201,6 +201,11 @@ class FiberManager : public ::folly::Executor {
    */
   void loopUntilNoReadyImpl();
 
+  /**
+   * This should only be called by a LoopController.
+   */
+  void runEagerFiberImpl(Fiber*);
+
   /**
    * This should only be called by a LoopController.
    */

folly/fibers/LoopController.h

@@ -24,6 +24,7 @@
 namespace folly {
 namespace fibers {
 
+class Fiber;
 class FiberManager;
 
 class LoopController {
@@ -50,6 +51,12 @@ class LoopController {
    */
   virtual void runLoop() = 0;
 
+  /**
+   * Run FiberManager runEagerFiberImpl(fiber). May have additional logic
+   * specific to a LoopController.
+   */
+  virtual void runEagerFiber(Fiber*) = 0;
+
   /**
    * Same as schedule(), but safe to call from any thread.
    */

folly/fibers/SimpleLoopController.h

@@ -84,6 +84,10 @@ class SimpleLoopController : public LoopController {
     } while (remoteLoopRun_ < remoteScheduleCalled_);
   }
 
+  void runEagerFiber(Fiber* fiber) override {
+    fm_->runEagerFiberImpl(fiber);
+  }
+
   void schedule() override {
     scheduled_ = true;
   }

folly/fibers/test/FibersTest.cpp

@@ -2585,3 +2585,16 @@ TEST(FiberManager, addTaskRemoteFutureTry) {
           .getVia(&evb)
           .value());
 }
+
+TEST(FiberManager, addTaskEagerKeepAlive) {
+  auto f = [&] {
+    folly::EventBase evb;
+    return getFiberManager(evb).addTaskEagerFuture([&] {
+      folly::futures::sleep(std::chrono::milliseconds{100}).get();
+      return 42;
+    });
+  }();
+
+  EXPECT_TRUE(f.isReady());
+  EXPECT_EQ(42, std::move(f).get());
+}