Improve management of EventBase keepalives

Summary:
The `EventBase` keepalive count is split across two counters, one non-atomic for in-loop operations, and one atomic for external operations (which mostly come from futures).
This requires `keepAliveRelease()` to always wake up the loop, because that's the only place where it can be determined whether the counts add up to 0. The cost of doing this is very large, and more than offsets the atomic operation (which is expected to be uncontended if the majority of operations are in-loop). Furthermore, `inRunningEventBaseThread()` is not cheap either.

With this diff we just use a single atomic count, so that the loop can be awoken only when the last keepalive is released, which in the common case never happens as the evbs are run with `loopForever()` in IO executors.

Reviewed By: Gownta

Differential Revision: D53292113

fbshipit-source-id: c0817cad98669f7538477299b7cfd180fd137243

folly/io/async/EventBase.cpp

@@ -623,14 +623,8 @@ void EventBase::loopMainCleanup() {
   loopThread_.store({}, std::memory_order_release);
 }
 
-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_;
+size_t EventBase::loopKeepAliveCount() {
+  return loopKeepAliveCount_.load(std::memory_order_relaxed);
 }
 
 void EventBase::applyLoopKeepAlive() {
@@ -704,11 +698,7 @@ 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.
-  try {
-    queue_->putMessage([] {});
-  } catch (...) {
-    // putMessage() can only fail when the queue is draining in ~EventBase.
-  }
+  queue_->putMessage([] {});
 }
 
 void EventBase::runInLoop(

folly/io/async/EventBase.h

@@ -922,8 +922,8 @@ class EventBase : public TimeoutManager,
 
   /// Implements the DrivableExecutor interface
   void drive() override {
-    ++loopKeepAliveCount_;
-    SCOPE_EXIT { --loopKeepAliveCount_; };
+    loopKeepAliveCount_.fetch_add(1, std::memory_order_relaxed);
+    SCOPE_EXIT { loopKeepAliveCount_.fetch_sub(1, std::memory_order_relaxed); };
     loopOnce();
   }
 
@@ -962,19 +962,17 @@ class EventBase : public TimeoutManager,
 
  protected:
   bool keepAliveAcquire() noexcept override {
-    if (inRunningEventBaseThread()) {
-      loopKeepAliveCount_++;
-    } else {
-      loopKeepAliveCountAtomic_.fetch_add(1, std::memory_order_relaxed);
-    }
+    loopKeepAliveCount_.fetch_add(1, std::memory_order_relaxed);
     return true;
   }
 
   void keepAliveRelease() noexcept override {
-    if (!inRunningEventBaseThread()) {
-      return add([this] { loopKeepAliveCount_--; });
+    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([] {});
     }
-    loopKeepAliveCount_--;
   }
 
  private:
@@ -983,10 +981,9 @@ 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.
@@ -1039,8 +1036,7 @@ 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_;
-  ssize_t loopKeepAliveCount_{0};
-  std::atomic<ssize_t> loopKeepAliveCountAtomic_{0};
+  std::atomic<size_t> loopKeepAliveCount_{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 loopKeepAlive() functionality,
+ * to EventBase, VirtualEventBase implements KeepAlive functionality,
  * which allows callbacks holding KeepAlive token to keep EventBase looping
  * until they are complete.
  *