fix blocking wait executor crash on destruction

Summary: It is possible for blockingWait to crash if just after it leaves it's drive() method in the destructor, from another thread something releases the keep alive. If then the destructor continues it will clean up, but that other thread will then be manipulating freed data.

Reviewed By: andriigrynenko

Differential Revision: D19860616

fbshipit-source-id: 25df16b5e900f283bbc0a264233770b6d3667ba3

folly/experimental/coro/BlockingWait.h

@@ -322,12 +322,22 @@ class BlockingWaitExecutor final : public folly::DrivableExecutor {
   }
 
   void keepAliveRelease() override {
-    auto keepAliveCount =
-        keepAliveCount_.fetch_sub(1, std::memory_order_acq_rel);
-    DCHECK(keepAliveCount > 0);
-    if (keepAliveCount == 1) {
-      add([] {});
-    }
+    auto keepAliveCount = keepAliveCount_.load(std::memory_order_relaxed);
+    do {
+      DCHECK(keepAliveCount > 0);
+      if (keepAliveCount == 1) {
+        add([this] {
+          // the final count *must* be released from this executor or else if we
+          // are mid-destructor we have a data race
+          keepAliveCount_.fetch_sub(1, std::memory_order_relaxed);
+        });
+        return;
+      }
+    } while (!keepAliveCount_.compare_exchange_weak(
+        keepAliveCount,
+        keepAliveCount - 1,
+        std::memory_order_release,
+        std::memory_order_relaxed));
   }
 
   folly::Synchronized<std::vector<Func>> queue_;

folly/experimental/coro/test/BlockingWaitTest.cpp

@@ -301,4 +301,32 @@ TEST_F(BlockingWaitTest, DrivableExecutor) {
       &executor);
 }
 
+TEST_F(BlockingWaitTest, ReleaseExecutorFromAnotherThread) {
+  auto fn = []() {
+    auto c1 = folly::makePromiseContract<folly::Executor::KeepAlive<>>();
+    auto c2 = folly::makePromiseContract<folly::Unit>();
+    std::thread t{[&] {
+      auto e = std::move(c1.second).get();
+      c2.first.setValue(folly::Unit{});
+      std::this_thread::sleep_for(std::chrono::microseconds(1));
+      e = {};
+    }};
+    folly::ManualExecutor executor;
+    folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
+      folly::Executor::KeepAlive<> taskExecutor =
+          co_await folly::coro::co_current_executor;
+      c1.first.setValue(std::move(taskExecutor));
+      co_await std::move(c2.second);
+    }());
+    t.join();
+  };
+  std::vector<std::thread> threads;
+  for (int i = 0; i < 100; ++i) {
+    threads.emplace_back(fn);
+  }
+  for (auto& t : threads) {
+    t.join();
+  }
+}
+
 #endif