From d9bf016dbcafb3c7c3fac4d081c5e63bed6516c5 Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Thu, 9 Mar 2017 20:34:14 -0800
Subject: [PATCH] Remove runInLoop

Summary: Existing runInLoop implementation was generally not safe (it did not capture the KeepAlive token). runInLoop() is mostly used in library code though, and often times it's ok not capture the KeepAlive token. It's better to just have the caller make the decision, rather then keep runInLoop() and make it always capture the KeepAlive token.

Reviewed By: yfeldblum

Differential Revision: D4684025

fbshipit-source-id: 65907bbe9c774e2a7b92580d0c0387d346b495d5
---
 folly/fibers/EventBaseLoopController-inl.h | 25 ++++++++++++++++++++--
 folly/io/async/VirtualEventBase.cpp        |  1 -
 folly/io/async/VirtualEventBase.h          | 14 ++++--------
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/folly/fibers/EventBaseLoopController-inl.h b/folly/fibers/EventBaseLoopController-inl.h
index 025341ce..4c01ed00 100644
--- a/folly/fibers/EventBaseLoopController-inl.h
+++ b/folly/fibers/EventBaseLoopController-inl.h
@@ -52,8 +52,8 @@ inline void EventBaseLoopControllerT<EventBaseT>::setFiberManager(
   fm_ = fm;
 }
 
-template <typename EventBaseT>
-inline void EventBaseLoopControllerT<EventBaseT>::schedule() {
+template <>
+inline void EventBaseLoopControllerT<folly::EventBase>::schedule() {
   if (eventBase_ == nullptr) {
     // In this case we need to postpone scheduling.
     awaitingScheduling_ = true;
@@ -64,6 +64,22 @@ inline void EventBaseLoopControllerT<EventBaseT>::schedule() {
   }
 }
 
+template <>
+inline void EventBaseLoopControllerT<folly::VirtualEventBase>::schedule() {
+  if (eventBase_ == nullptr) {
+    // In this case we need to postpone scheduling.
+    awaitingScheduling_ = true;
+  } else {
+    // Schedule it to run in current iteration.
+
+    if (!eventBaseKeepAlive_) {
+      eventBaseKeepAlive_ = eventBase_->getKeepAliveToken();
+    }
+    eventBase_->getEventBase().runInLoop(&callback_, true);
+    awaitingScheduling_ = false;
+  }
+}
+
 template <typename EventBaseT>
 inline void EventBaseLoopControllerT<EventBaseT>::cancel() {
   callback_.cancelLoopCallback();
@@ -72,6 +88,11 @@ inline void EventBaseLoopControllerT<EventBaseT>::cancel() {
 template <typename EventBaseT>
 inline void EventBaseLoopControllerT<EventBaseT>::runLoop() {
   if (!eventBaseKeepAlive_) {
+    // runLoop can be called twice if both schedule() and scheduleThreadSafe()
+    // were called.
+    if (!fm_->hasTasks()) {
+      return;
+    }
     eventBaseKeepAlive_ = eventBase_->getKeepAliveToken();
   }
   if (loopRunner_) {
diff --git a/folly/io/async/VirtualEventBase.cpp b/folly/io/async/VirtualEventBase.cpp
index 688faf5e..46c187b8 100644
--- a/folly/io/async/VirtualEventBase.cpp
+++ b/folly/io/async/VirtualEventBase.cpp
@@ -19,7 +19,6 @@ namespace folly {
 
 VirtualEventBase::VirtualEventBase(EventBase& evb) : evb_(evb) {
   evbLoopKeepAlive_ = evb_.getKeepAliveToken();
-  loopKeepAlive_ = getKeepAliveToken();
 }
 
 std::future<void> VirtualEventBase::destroy() {
diff --git a/folly/io/async/VirtualEventBase.h b/folly/io/async/VirtualEventBase.h
index c67d2cb2..a1194d82 100644
--- a/folly/io/async/VirtualEventBase.h
+++ b/folly/io/async/VirtualEventBase.h
@@ -61,14 +61,6 @@ class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
    */
   void runOnDestruction(EventBase::LoopCallback* callback);
 
-  /**
-   * @see EventBase::runInLoop
-   */
-  template <typename F>
-  void runInLoop(F&& f, bool thisIteration = false) {
-    evb_.runInLoop(std::forward<F>(f), thisIteration);
-  }
-
   /**
    * VirtualEventBase destructor blocks until all tasks scheduled through its
    * runInEventBaseThread are complete.
@@ -129,6 +121,8 @@ class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
    * KeepAlive handle can be released from EventBase loop only.
    */
   KeepAlive getKeepAliveToken() override {
+    DCHECK(loopKeepAliveCount_ + loopKeepAliveCountAtomic_.load() > 0);
+
     if (evb_.inRunningEventBaseThread()) {
       ++loopKeepAliveCount_;
     } else {
@@ -163,11 +157,11 @@ class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
 
   EventBase& evb_;
 
-  ssize_t loopKeepAliveCount_{0};
+  ssize_t loopKeepAliveCount_{1};
   std::atomic<ssize_t> loopKeepAliveCountAtomic_{0};
   std::promise<void> destroyPromise_;
   std::future<void> destroyFuture_{destroyPromise_.get_future()};
-  KeepAlive loopKeepAlive_;
+  KeepAlive loopKeepAlive_{makeKeepAlive()};
 
   KeepAlive evbLoopKeepAlive_;
 
-- 
2.34.1