Add keepAlive() mechanism

[folly.git] / folly / io / async / EventBase.h

diff --git a/folly/io/async/EventBase.h b/folly/io/async/EventBase.h
index cc23fed266b4799fd750a0f81420f53f4896bb18..c318d609ef6c3ada5e04865cf7724ece4a168df7 100644 (file)
--- a/folly/io/async/EventBase.h
+++ b/folly/io/async/EventBase.h
@@ -600,33 +600,17 @@ class EventBase : private boost::noncopyable,
     loopOnce();
   }
 
-  struct LoopKeepAliveDeleter {
-    void operator()(EventBase* evb) {
-      DCHECK(evb->isInEventBaseThread());
-      evb->loopKeepAliveCount_--;
-    }
-  };
-  using LoopKeepAlive = std::unique_ptr<EventBase, LoopKeepAliveDeleter>;
-
   /// 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. Loop holder is safe to release only from
   /// EventBase thread.
-  ///
-  /// May return no op LoopKeepAlive if loopForever() is already running.
-  LoopKeepAlive loopKeepAlive() {
-    DCHECK(isInEventBaseThread());
-    loopKeepAliveCount_++;
-    return LoopKeepAlive(this);
-  }
-
-  // Thread-safe version of loopKeepAlive()
-  LoopKeepAlive loopKeepAliveAtomic() {
+  KeepAlive getKeepAliveToken() override {
     if (inRunningEventBaseThread()) {
-      return loopKeepAlive();
+      loopKeepAliveCount_++;
+    } else {
+      loopKeepAliveCountAtomic_.fetch_add(1, std::memory_order_relaxed);
     }
-    loopKeepAliveCountAtomic_.fetch_add(1, std::memory_order_relaxed);
-    return LoopKeepAlive(this);
+    return makeKeepAlive();
   }
 
   // TimeoutManager
@@ -645,6 +629,12 @@ class EventBase : private boost::noncopyable,
     return isInEventBaseThread();
   }
 
+ protected:
+  void keepAliveRelease() override {
+    DCHECK(isInEventBaseThread());
+    loopKeepAliveCount_--;
+  }
+
  private:
   void applyLoopKeepAlive();