Workaround for TSAN symmetric transfer bug

Summary: TSAN breaks symmetric transfer behavior in compiler. Work around it by rescheduling work on executor when possible.

Reviewed By: yfeldblum

Differential Revision: D23391328

fbshipit-source-id: 7b25a0c89664ca14114f8ae0dc650ff9bb786fe1

folly/experimental/coro/AsyncGenerator.h

@@ -23,6 +23,7 @@
 #include <folly/experimental/coro/Utils.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
 #include <folly/experimental/coro/WithCancellation.h>
+#include <folly/experimental/coro/detail/Helpers.h>
 #include <folly/experimental/coro/detail/Malloc.h>
 #include <folly/experimental/coro/detail/ManualLifetime.h>
 
@@ -49,7 +50,8 @@ class AsyncGeneratorPromise {
     }
     std::experimental::coroutine_handle<> await_suspend(
         std::experimental::coroutine_handle<AsyncGeneratorPromise> h) noexcept {
-      return h.promise().continuation_;
+      return symmetricTransferMaybeReschedule(
+          h.promise().continuation_, h.promise().clearContext());
     }
     void await_resume() noexcept {}
   };
@@ -77,7 +79,6 @@ class AsyncGeneratorPromise {
 
   YieldAwaiter final_suspend() noexcept {
     DCHECK(!hasValue_);
-    clearContext();
     return {};
   }
 
@@ -86,8 +87,7 @@ class AsyncGeneratorPromise {
     DCHECK(!hasValue_);
     value_.construct(static_cast<Reference&&>(value));
     hasValue_ = true;
-    clearContext();
-    return YieldAwaiter{};
+    return {};
   }
 
   // In the case where 'Reference' is not actually a reference-type we
@@ -106,7 +106,6 @@ class AsyncGeneratorPromise {
     DCHECK(!hasValue_);
     value_.construct(static_cast<U&&>(value));
     hasValue_ = true;
-    clearContext();
     return {};
   }
 
@@ -177,10 +176,12 @@ class AsyncGeneratorPromise {
   }
 
  private:
-  void clearContext() noexcept {
-    executor_ = {};
+  folly::Executor::KeepAlive<> clearContext() noexcept {
+    auto executor = std::exchange(executor_, {});
     cancelToken_ = {};
     hasCancelTokenOverride_ = false;
+
+    return executor;
   }
 
   std::experimental::coroutine_handle<> continuation_;

folly/experimental/coro/Task.h

@@ -35,6 +35,7 @@
 #include <folly/experimental/coro/Utils.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
 #include <folly/experimental/coro/WithCancellation.h>
+#include <folly/experimental/coro/detail/Helpers.h>
 #include <folly/experimental/coro/detail/InlineTask.h>
 #include <folly/experimental/coro/detail/Malloc.h>
 #include <folly/experimental/coro/detail/Traits.h>
@@ -86,7 +87,8 @@ class TaskPromiseBase {
     std::experimental::coroutine_handle<> await_suspend(
         std::experimental::coroutine_handle<Promise> coro) noexcept {
       TaskPromiseBase& promise = coro.promise();
-      return promise.continuation_;
+      return symmetricTransferMaybeReschedule(
+          promise.continuation_, promise.executor_);
     }
 
     [[noreturn]] void await_resume() noexcept {
@@ -597,10 +599,10 @@ class FOLLY_NODISCARD Task {
       return false;
     }
 
-    handle_t await_suspend(
+    auto await_suspend(
         std::experimental::coroutine_handle<> continuation) noexcept {
       coro_.promise().continuation_ = continuation;
-      return coro_;
+      return symmetricTransferMaybeReschedule(coro_, coro_.promise().executor_);
     }
 
     T await_resume() {

folly/experimental/coro/detail/Helpers.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <folly/Executor.h>
+#include <folly/io/async/Request.h>
 
 namespace folly {
 namespace coro {
@@ -39,6 +40,26 @@ class UnsafeResumeInlineSemiAwaitable {
  private:
   Awaitable awaitable_;
 };
+
+template <typename Promise>
+FOLLY_ALWAYS_INLINE folly::conditional_t<
+    kIsSanitizeThread,
+    std::experimental::coroutine_handle<>,
+    std::experimental::coroutine_handle<Promise>>
+symmetricTransferMaybeReschedule(
+    std::experimental::coroutine_handle<Promise> ch,
+    const Executor::KeepAlive<>& ex) {
+  if constexpr (kIsSanitizeThread) {
+    copy(ex).add([ch, rctx = RequestContext::saveContext()](
+                     Executor::KeepAlive<>&&) mutable {
+      RequestContextScopeGuard guard(std::move(rctx));
+      ch.resume();
+    });
+    return std::experimental::noop_coroutine();
+  } else {
+    return ch;
+  }
+}
 } // namespace detail
 } // namespace coro
 } // namespace folly

folly/experimental/coro/test/AsyncGeneratorTest.cpp

@@ -473,4 +473,22 @@ TEST_F(AsyncGeneratorTest, BlockingWaitOnThrowingFinalNextDoesNotDeadlock) {
   }
 }
 
+folly::coro::AsyncGenerator<int> range(int from, int to) {
+  for (int i = from; i < to; ++i) {
+    co_yield i;
+  }
+}
+
+TEST_F(AsyncGeneratorTest, SymmetricTransfer) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    int max = 100000;
+    auto g = range(1, max + 1);
+    long long sum = 0;
+    while (auto result = co_await g.next()) {
+      sum += *result;
+    }
+    EXPECT_EQ(((long long)max + 1) * max / 2, sum);
+  }());
+}
+
 #endif

folly/experimental/coro/test/CoroTest.cpp

@@ -63,7 +63,7 @@ TEST_F(CoroTest, Basic) {
 
   EXPECT_FALSE(future.isReady());
 
-  executor.drive();
+  executor.drain();
 
   EXPECT_TRUE(future.isReady());
   EXPECT_EQ(42, std::move(future).get());
@@ -75,8 +75,7 @@ TEST_F(CoroTest, BasicSemiFuture) {
 
   EXPECT_FALSE(future.isReady());
 
-  executor.drive();
-  executor.drive();
+  executor.drain();
 
   EXPECT_TRUE(future.isReady());
   EXPECT_EQ(42, std::move(future).get());
@@ -103,7 +102,7 @@ TEST_F(CoroTest, Basic2) {
 
   EXPECT_FALSE(future.isReady());
 
-  executor.drive();
+  executor.drain();
 
   EXPECT_TRUE(future.isReady());
 }
@@ -193,7 +192,7 @@ TEST_F(CoroTest, FutureThrow) {
 
   EXPECT_FALSE(future.isReady());
 
-  executor.drive();
+  executor.drain();
 
   EXPECT_TRUE(future.isReady());
   EXPECT_THROW(std::move(future).get(), std::runtime_error);
@@ -211,9 +210,9 @@ coro::Task<int> taskRecursion(int depth) {
 
 TEST_F(CoroTest, LargeStack) {
   ScopedEventBaseThread evbThread;
-  auto task = taskRecursion(5000).scheduleOn(evbThread.getEventBase());
+  auto task = taskRecursion(50000).scheduleOn(evbThread.getEventBase());
 
-  EXPECT_EQ(5000, coro::blockingWait(std::move(task)));
+  EXPECT_EQ(50000, coro::blockingWait(std::move(task)));
 }
 
 coro::Task<void> taskThreadNested(std::thread::id threadId) {
@@ -416,12 +415,12 @@ TEST_F(CoroTest, Baton) {
 
   EXPECT_FALSE(future.isReady());
 
-  executor.run();
+  executor.drain();
 
   EXPECT_FALSE(future.isReady());
 
   baton.post();
-  executor.run();
+  executor.drain();
 
   EXPECT_TRUE(future.isReady());
   EXPECT_EQ(42, std::move(future).get());
@@ -450,12 +449,12 @@ TEST_F(CoroTest, co_invoke) {
       })
           .scheduleOn(&executor)
           .start();
-  executor.run();
+  executor.drain();
   EXPECT_FALSE(coroFuture.isReady());
 
   p.setValue(folly::unit);
 
-  executor.run();
+  executor.drain();
   EXPECT_TRUE(coroFuture.isReady());
 }