Revert D53292113: Improve management of EventBase keepalives

Differential Revision:
D53292113

Original commit changeset: c0817cad9866

Original Phabricator Diff: D53292113

fbshipit-source-id: b695869f67602c64d00c4e5e34eaa856dab482f4

folly/io/async/EventBase.cpp

@@ -623,8 +623,14 @@ void EventBase::loopMainCleanup() {
   loopThread_.store({}, std::memory_order_release);
 }
 
-size_t EventBase::loopKeepAliveCount() {
-  return loopKeepAliveCount_.load(std::memory_order_relaxed);
+ssize_t EventBase::loopKeepAliveCount() {
+  if (loopKeepAliveCountAtomic_.load(std::memory_order_relaxed)) {
+    loopKeepAliveCount_ +=
+        loopKeepAliveCountAtomic_.exchange(0, std::memory_order_relaxed);
+  }
+  DCHECK_GE(loopKeepAliveCount_, 0);
+
+  return loopKeepAliveCount_;
 }
 
 void EventBase::applyLoopKeepAlive() {
@@ -698,7 +704,11 @@ void EventBase::terminateLoopSoon() {
   // In this case, it won't wake up and notice that stop_ is set until it
   // receives another event.  Send an empty frame to the notification queue
   // so that the event loop will wake up even if there are no other events.
-  queue_->putMessage([] {});
+  try {
+    queue_->putMessage([] {});
+  } catch (...) {
+    // putMessage() can only fail when the queue is draining in ~EventBase.
+  }
 }
 
 void EventBase::runInLoop(

folly/io/async/EventBase.h

@@ -922,8 +922,8 @@ class EventBase : public TimeoutManager,
 
   /// Implements the DrivableExecutor interface
   void drive() override {
-    loopKeepAliveCount_.fetch_add(1, std::memory_order_relaxed);
-    SCOPE_EXIT { loopKeepAliveCount_.fetch_sub(1, std::memory_order_relaxed); };
+    ++loopKeepAliveCount_;
+    SCOPE_EXIT { --loopKeepAliveCount_; };
     loopOnce();
   }
 
@@ -962,17 +962,19 @@ class EventBase : public TimeoutManager,
 
  protected:
   bool keepAliveAcquire() noexcept override {
-    loopKeepAliveCount_.fetch_add(1, std::memory_order_relaxed);
+    if (inRunningEventBaseThread()) {
+      loopKeepAliveCount_++;
+    } else {
+      loopKeepAliveCountAtomic_.fetch_add(1, std::memory_order_relaxed);
+    }
     return true;
   }
 
   void keepAliveRelease() noexcept override {
-    auto oldCount = loopKeepAliveCount_.fetch_sub(1, std::memory_order_acq_rel);
-    DCHECK_GT(oldCount, 0);
-    if (oldCount == 1 && !inRunningEventBaseThread()) {
-      // Unblock the loop so it can exit.
-      runInEventBaseThreadAlwaysEnqueue([] {});
+    if (!inRunningEventBaseThread()) {
+      return add([this] { loopKeepAliveCount_--; });
     }
+    loopKeepAliveCount_--;
   }
 
  private:
@@ -981,9 +983,10 @@ class EventBase : public TimeoutManager,
 
   folly::VirtualEventBase* tryGetVirtualEventBase();
 
-  size_t loopKeepAliveCount();
   void applyLoopKeepAlive();
 
+  ssize_t loopKeepAliveCount();
+
   /*
    * Helper function that tells us whether we have already handled
    * some event/timeout/callback in this loop iteration.
@@ -1036,7 +1039,8 @@ class EventBase : public TimeoutManager,
   // A notification queue for runInEventBaseThread() to use
   // to send function requests to the EventBase thread.
   std::unique_ptr<EventBaseAtomicNotificationQueue<Func, FuncRunner>> queue_;
-  std::atomic<size_t> loopKeepAliveCount_{0};
+  ssize_t loopKeepAliveCount_{0};
+  std::atomic<ssize_t> loopKeepAliveCountAtomic_{0};
   bool loopKeepAliveActive_{false};
 
   // limit for latency in microseconds (0 disables)

folly/io/async/VirtualEventBase.h

@@ -31,7 +31,7 @@ namespace folly {
  * VirtualEventBase implements a light-weight view onto existing EventBase.
  *
  * Multiple VirtualEventBases can be backed by a single EventBase. Similarly
- * to EventBase, VirtualEventBase implements KeepAlive functionality,
+ * to EventBase, VirtualEventBase implements loopKeepAlive() functionality,
  * which allows callbacks holding KeepAlive token to keep EventBase looping
  * until they are complete.
  *