Always manage lifetime of a SerialExecutor using KeepAlive tokens

Summary: SerialExecutor is always wrapping another Executor, so we shouldn't force users to coordinate their lifetimes.

Reviewed By: yfeldblum

Differential Revision: D7856146

fbshipit-source-id: ac7caaa0f181406dfb6b59d36ae4efe6d1343590

folly/executors/SerialExecutor.cpp

@@ -72,18 +72,54 @@ class SerialExecutor::TaskQueueImpl {
   std::queue<Func> queue_;
 };
 
-SerialExecutor::SerialExecutor(std::shared_ptr<folly::Executor> parent)
+SerialExecutor::SerialExecutor(KeepAlive<Executor> parent)
     : parent_(std::move(parent)),
       taskQueueImpl_(std::make_shared<TaskQueueImpl>()) {}
 
+SerialExecutor::~SerialExecutor() {
+  DCHECK(!keepAliveCounter_);
+}
+
+Executor::KeepAlive<SerialExecutor> SerialExecutor::create(
+    KeepAlive<Executor> parent) {
+  return makeKeepAlive<SerialExecutor>(new SerialExecutor(std::move(parent)));
+}
+
+SerialExecutor::UniquePtr SerialExecutor::createUnique(
+    std::shared_ptr<Executor> parent) {
+  auto executor = new SerialExecutor(getKeepAliveToken(parent.get()));
+  return {executor, Deleter{std::move(parent)}};
+}
+
+bool SerialExecutor::keepAliveAcquire() {
+  auto keepAliveCounter =
+      keepAliveCounter_.fetch_add(1, std::memory_order_relaxed);
+  DCHECK(keepAliveCounter > 0);
+  return true;
+}
+
+void SerialExecutor::keepAliveRelease() {
+  auto keepAliveCounter = --keepAliveCounter_;
+  DCHECK(keepAliveCounter >= 0);
+  if (!keepAliveCounter) {
+    delete this;
+  }
+}
+
 void SerialExecutor::add(Func func) {
   taskQueueImpl_->add(std::move(func));
-  parent_->add([impl = taskQueueImpl_] { impl->run(); });
+  parent_->add([impl = taskQueueImpl_, keepAlive = getKeepAliveToken(this)] {
+    impl->run();
+  });
 }
 
 void SerialExecutor::addWithPriority(Func func, int8_t priority) {
   taskQueueImpl_->add(std::move(func));
-  parent_->addWithPriority([impl = taskQueueImpl_] { impl->run(); }, priority);
+  parent_->addWithPriority(
+      [impl = taskQueueImpl_, keepAlive = getKeepAliveToken(this)] {
+        impl->run();
+      },
+      priority);
 }
 
 } // namespace folly

folly/executors/SerialExecutor.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <memory>
 
 #include <folly/executors/GlobalExecutor.h>
@@ -48,14 +49,35 @@ namespace folly {
 
 class SerialExecutor : public SequencedExecutor {
  public:
-  ~SerialExecutor() override = default;
+  ~SerialExecutor() override;
   SerialExecutor(SerialExecutor const&) = delete;
   SerialExecutor& operator=(SerialExecutor const&) = delete;
   SerialExecutor(SerialExecutor&&) = default;
   SerialExecutor& operator=(SerialExecutor&&) = default;
 
-  explicit SerialExecutor(
-      std::shared_ptr<folly::Executor> parent = folly::getCPUExecutor());
+  static KeepAlive<SerialExecutor> create(
+      KeepAlive<Executor> parent = getKeepAliveToken(getCPUExecutor().get()));
+
+  class Deleter {
+   public:
+    Deleter() {}
+
+    void operator()(SerialExecutor* executor) {
+      executor->keepAliveRelease();
+    }
+
+   private:
+    friend class SerialExecutor;
+    explicit Deleter(std::shared_ptr<Executor> parent)
+        : parent_(std::move(parent)) {}
+
+    std::shared_ptr<Executor> parent_;
+  };
+
+  using UniquePtr = std::unique_ptr<SerialExecutor, Deleter>;
+  [[deprecated("Replaced by create")]]
+  static UniquePtr createUnique(
+      std::shared_ptr<Executor> parent = getCPUExecutor());
 
   /**
    * Add one task for execution in the parent executor
@@ -77,11 +99,20 @@ class SerialExecutor : public SequencedExecutor {
     return parent_->getNumPriorities();
   }
 
+ protected:
+  bool keepAliveAcquire() override;
+
+  void keepAliveRelease() override;
+
  private:
+  explicit SerialExecutor(KeepAlive<Executor> parent);
+
   class TaskQueueImpl;
 
-  std::shared_ptr<folly::Executor> parent_;
+  KeepAlive<Executor> parent_;
   std::shared_ptr<TaskQueueImpl> taskQueueImpl_;
+
+  std::atomic<ssize_t> keepAliveCounter_{1};
 };
 
 } // namespace folly

folly/executors/test/SequencedExecutorTest.cpp

@@ -26,7 +26,9 @@ namespace folly {
 
 bool isSequencedExecutor(folly::Executor& executor) {
   // Add can be called from different threads, but it should be sequenced.
-  SerialExecutor producer(std::make_shared<CPUThreadPoolExecutor>(4));
+  auto cpuExecutor = std::make_shared<CPUThreadPoolExecutor>(4);
+  auto producer =
+      SerialExecutor::create(Executor::getKeepAliveToken(cpuExecutor.get()));
 
   std::atomic<size_t> nextCallIndex{0};
   std::atomic<bool> result{true};
@@ -36,7 +38,7 @@ bool isSequencedExecutor(folly::Executor& executor) {
 
   constexpr size_t kNumCalls = 10000;
   for (size_t callIndex = 0; callIndex < kNumCalls; ++callIndex) {
-    producer.add([&result, &executor, &nextCallIndex, callIndex, joinPromise] {
+    producer->add([&result, &executor, &nextCallIndex, callIndex, joinPromise] {
       executor.add([&result, &nextCallIndex, callIndex, joinPromise] {
         if (nextCallIndex != callIndex) {
           result = false;
@@ -69,8 +71,10 @@ TEST(SequencedExecutor, CPUThreadPoolExecutor) {
 }
 
 TEST(SequencedExecutor, SerialCPUThreadPoolExecutor) {
-  SerialExecutor executor(std::make_shared<CPUThreadPoolExecutor>(4));
-  testExecutor(executor);
+  auto cpuExecutor = std::make_shared<CPUThreadPoolExecutor>(4);
+  auto executor =
+      SerialExecutor::create(Executor::getKeepAliveToken(cpuExecutor.get()));
+  testExecutor(*executor);
 }
 
 TEST(SequencedExecutor, EventBase) {

folly/executors/test/SerialExecutorTest.cpp

@@ -32,13 +32,14 @@ void burnMs(uint64_t ms) {
 } // namespace
 
 void SimpleTest(std::shared_ptr<folly::Executor> const& parent) {
-  SerialExecutor executor(parent);
+  auto executor =
+      SerialExecutor::create(folly::getKeepAliveToken(parent.get()));
 
   std::vector<int> values;
   std::vector<int> expected;
 
   for (int i = 0; i < 20; ++i) {
-    executor.add([i, &values] {
+    executor->add([i, &values] {
       // make this extra vulnerable to concurrent execution
       values.push_back(0);
       burnMs(10);
@@ -49,7 +50,7 @@ void SimpleTest(std::shared_ptr<folly::Executor> const& parent) {
 
   // wait until last task has executed
   folly::Baton<> finished_baton;
-  executor.add([&finished_baton] { finished_baton.post(); });
+  executor->add([&finished_baton] { finished_baton.post(); });
   finished_baton.wait();
 
   EXPECT_EQ(expected, values);
@@ -67,7 +68,8 @@ TEST(SerialExecutor, SimpleInline) {
 // destroy the SerialExecutor
 TEST(SerialExecutor, Afterlife) {
   auto cpu_executor = std::make_shared<folly::CPUThreadPoolExecutor>(4);
-  auto executor = std::make_unique<SerialExecutor>(cpu_executor);
+  auto executor =
+      SerialExecutor::create(folly::getKeepAliveToken(cpu_executor.get()));
 
   // block executor until we call start_baton.post()
   folly::Baton<> start_baton;
@@ -102,7 +104,8 @@ TEST(SerialExecutor, Afterlife) {
 }
 
 void RecursiveAddTest(std::shared_ptr<folly::Executor> const& parent) {
-  SerialExecutor executor(parent);
+  auto executor =
+      SerialExecutor::create(folly::getKeepAliveToken(parent.get()));
 
   folly::Baton<> finished_baton;
 
@@ -116,7 +119,7 @@ void RecursiveAddTest(std::shared_ptr<folly::Executor> const& parent) {
       values.push_back(0);
       burnMs(10);
       values.back() = i;
-      executor.add(lambda);
+      executor->add(lambda);
     } else if (i < 12) {
       // Below we will post this lambda three times to the executor. When
       // executed, the lambda will re-post itself during the first ten
@@ -128,9 +131,9 @@ void RecursiveAddTest(std::shared_ptr<folly::Executor> const& parent) {
     ++i;
   };
 
-  executor.add(lambda);
-  executor.add(lambda);
-  executor.add(lambda);
+  executor->add(lambda);
+  executor->add(lambda);
+  executor->add(lambda);
 
   // wait until last task has executed
   finished_baton.wait();
@@ -146,9 +149,9 @@ TEST(SerialExecutor, RecursiveAddInline) {
 }
 
 TEST(SerialExecutor, ExecutionThrows) {
-  SerialExecutor executor(std::make_shared<folly::InlineExecutor>());
+  auto executor = SerialExecutor::create();
 
   // an empty Func will throw std::bad_function_call when invoked,
   // but SerialExecutor should catch that exception
-  executor.add(folly::Func{});
+  executor->add(folly::Func{});
 }