Make keep-alive token a valid pointer for every Executor

Summary: This allows using a keep-alive token in places where Executor* was previously used.

Reviewed By: yfeldblum

Differential Revision: D7751516

fbshipit-source-id: af2cb31191a78306439408dbee78dcd923492e30

folly/Executor.cpp

@@ -27,13 +27,12 @@ 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()";
+bool Executor::keepAliveAcquire() {
+  return false;
 }
 
 void Executor::keepAliveRelease() {
   LOG(FATAL) << __func__ << "() should not be called for folly::Executor types "
-             << "which do not override getKeepAliveToken()";
+             << "which do not override keepAliveAcquire()";
 }
 } // namespace folly

folly/Executor.h

@@ -49,7 +49,7 @@ class Executor {
 
   class KeepAlive {
    public:
-    KeepAlive() {}
+    KeepAlive() : executor_(nullptr, Deleter(true)) {}
 
     void reset() {
       executor_.reset();
@@ -65,31 +65,44 @@ class Executor {
 
    private:
     friend class Executor;
-    explicit KeepAlive(folly::Executor* executor) : executor_(executor) {}
+    KeepAlive(folly::Executor* executor, bool dummy)
+        : executor_(executor, Deleter(dummy)) {}
 
     struct Deleter {
+      explicit Deleter(bool dummy) : dummy_(dummy) {}
       void operator()(folly::Executor* executor) {
+        if (dummy_) {
+          return;
+        }
         executor->keepAliveRelease();
       }
+
+     private:
+      bool dummy_;
     };
     std::unique_ptr<folly::Executor, Deleter> executor_;
   };
 
   /// Returns a keep-alive token which guarantees that Executor will keep
-  /// processing tasks until the token is released.
-  ///
-  /// If executor does not support keep-alive functionality - dummy token will
-  /// be returned.
-  virtual KeepAlive getKeepAliveToken() {
-    return {};
+  /// processing tasks until the token is released (if supported by Executor).
+  /// KeepAlive always contains a valid pointer to an Executor.
+  KeepAlive getKeepAliveToken() {
+    if (keepAliveAcquire()) {
+      return makeKeepAlive();
+    }
+    return KeepAlive{this, true};
   }
 
  protected:
-  virtual void keepAliveAcquire();
+  // Acquire a keep alive token. Should return false if keep-alive mechanism
+  // is not supported.
+  virtual bool keepAliveAcquire();
+  // Release a keep alive token previously acquired by keepAliveAcquire().
+  // Will never be called if keepAliveAcquire() returns false.
   virtual void keepAliveRelease();
 
   KeepAlive makeKeepAlive() {
-    return KeepAlive{this};
+    return KeepAlive{this, false};
   }
 };
 

folly/io/async/EventBase.h

@@ -632,14 +632,6 @@ class EventBase : private boost::noncopyable,
   // Implements the ScheduledExecutor interface
   void scheduleAt(Func&& fn, TimePoint const& timeout) override;
 
-  /// Returns you a handle which make loop() behave like loopForever() until
-  /// destroyed. loop() will return to its original behavior only when all
-  /// loop keep-alives are released.
-  KeepAlive getKeepAliveToken() override {
-    keepAliveAcquire();
-    return makeKeepAlive();
-  }
-
   // TimeoutManager
   void attachTimeoutManager(
       AsyncTimeout* obj,
@@ -669,12 +661,13 @@ class EventBase : private boost::noncopyable,
   EventBase* getEventBase() override;
 
  protected:
-  void keepAliveAcquire() override {
+  bool keepAliveAcquire() override {
     if (inRunningEventBaseThread()) {
       loopKeepAliveCount_++;
     } else {
       loopKeepAliveCountAtomic_.fetch_add(1, std::memory_order_relaxed);
     }
+    return true;
   }
 
   void keepAliveRelease() override {

folly/io/async/VirtualEventBase.h

@@ -116,20 +116,12 @@ class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
     runInEventBaseThread(std::move(f));
   }
 
-  /**
-   * 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 {
+  bool keepAliveAcquire() override {
     DCHECK(loopKeepAliveCount_ + loopKeepAliveCountAtomic_.load() > 0);
 
     if (evb_.inRunningEventBaseThread()) {
@@ -137,6 +129,7 @@ class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
     } else {
       ++loopKeepAliveCountAtomic_;
     }
+    return true;
   }
 
   void keepAliveRelease() override {