Move keepalive-acquire code into Executor::keepAliveAcquire overrides

Summary:
[Folly] Move keepalive-acquire code into `Executor::keepAliveAcquire` overrides.

This makes the API more symmetric and allows for more flexibility. Also serves as preparation for generalizing the keepalive-token interface.

Reviewed By: andriigrynenko

Differential Revision: D5974080

fbshipit-source-id: 26209e49a0f5834ba229d4bbfc9272c2e4ffb3fd

folly/Executor.cpp

@@ -27,8 +27,13 @@ void Executor::addWithPriority(Func, int8_t /* priority */) {
       "addWithPriority() is not implemented for this Executor");
 }
 
+void Executor::keepAliveAcquire() {
+  LOG(FATAL) << __func__ << "() should not be called for folly::Executor types "
+             << "which do not override getKeepAliveToken()";
+}
+
 void Executor::keepAliveRelease() {
-  LOG(FATAL) << "keepAliveRelease() should not be called for folly::Executors "
-             << "which do not implement getKeepAliveToken()";
+  LOG(FATAL) << __func__ << "() should not be called for folly::Executor types "
+             << "which do not override getKeepAliveToken()";
 }
 }

folly/Executor.h

@@ -90,6 +90,7 @@ class Executor {
   }
 
  protected:
+  virtual void keepAliveAcquire();
   virtual void keepAliveRelease();
 
   KeepAlive makeKeepAlive() {

folly/io/async/EventBase.h

@@ -615,11 +615,7 @@ class EventBase : private boost::noncopyable,
   /// destroyed. loop() will return to its original behavior only when all
   /// loop keep-alives are released.
   KeepAlive getKeepAliveToken() override {
-    if (inRunningEventBaseThread()) {
-      loopKeepAliveCount_++;
-    } else {
-      loopKeepAliveCountAtomic_.fetch_add(1, std::memory_order_relaxed);
-    }
+    keepAliveAcquire();
     return makeKeepAlive();
   }
 
@@ -649,6 +645,14 @@ class EventBase : private boost::noncopyable,
   folly::VirtualEventBase& getVirtualEventBase();
 
  protected:
+  void keepAliveAcquire() override {
+    if (inRunningEventBaseThread()) {
+      loopKeepAliveCount_++;
+    } else {
+      loopKeepAliveCountAtomic_.fetch_add(1, std::memory_order_relaxed);
+    }
+  }
+
   void keepAliveRelease() override {
     if (inRunningEventBaseThread()) {
       loopKeepAliveCount_--;

folly/io/async/VirtualEventBase.h

@@ -120,6 +120,16 @@ class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
    * Returns you a handle which prevents VirtualEventBase from being destroyed.
    */
   KeepAlive getKeepAliveToken() override {
+    keepAliveAcquire();
+    return makeKeepAlive();
+  }
+
+  bool inRunningEventBaseThread() const {
+    return evb_.inRunningEventBaseThread();
+  }
+
+ protected:
+  void keepAliveAcquire() override {
     DCHECK(loopKeepAliveCount_ + loopKeepAliveCountAtomic_.load() > 0);
 
     if (evb_.inRunningEventBaseThread()) {
@@ -127,14 +137,8 @@ class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
     } else {
       ++loopKeepAliveCountAtomic_;
     }
-    return makeKeepAlive();
   }
 
-  bool inRunningEventBaseThread() const {
-    return evb_.inRunningEventBaseThread();
-  }
-
- protected:
   void keepAliveRelease() override {
     if (!getEventBase().inRunningEventBaseThread()) {
       return getEventBase().add([=] { keepAliveRelease(); });