Extend TSAN symmetric transfer bug workaround to use a thread-local counter

Summary: Many unit tests ended up already depending on the scheduling order. Relax the logic to only reschedule once we reach a high nesting depth.

Reviewed By: iahs

Differential Revision: D23415806

fbshipit-source-id: dde682293ded0061cf331fa5f46e1f6b55928a3e

folly/experimental/coro/AsyncGenerator.h

@@ -48,7 +48,7 @@ class AsyncGeneratorPromise {
     bool await_ready() noexcept {
       return false;
     }
-    std::experimental::coroutine_handle<> await_suspend(
+    auto await_suspend(
         std::experimental::coroutine_handle<AsyncGeneratorPromise> h) noexcept {
       return symmetricTransferMaybeReschedule(
           h.promise().continuation_, h.promise().clearContext());

folly/experimental/coro/Task.h

@@ -84,7 +84,7 @@ class TaskPromiseBase {
     }
 
     template <typename Promise>
-    std::experimental::coroutine_handle<> await_suspend(
+    auto await_suspend(
         std::experimental::coroutine_handle<Promise> coro) noexcept {
       TaskPromiseBase& promise = coro.promise();
       return symmetricTransferMaybeReschedule(

folly/experimental/coro/detail/Helpers.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <folly/Executor.h>
+#include <folly/SingletonThreadLocal.h>
 #include <folly/io/async/Request.h>
 
 namespace folly {
@@ -41,21 +42,46 @@ class UnsafeResumeInlineSemiAwaitable {
   Awaitable awaitable_;
 };
 
+class NestingCounter {
+ public:
+  struct GuardDestructor {
+    void operator()(NestingCounter* counter) {
+      --counter->counter_;
+    }
+  };
+  using Guard = std::unique_ptr<NestingCounter, GuardDestructor>;
+  Guard guard() {
+    const size_t maxNestingDepth = 128;
+    if (counter_ + 1 == maxNestingDepth) {
+      return nullptr;
+    }
+    ++counter_;
+    return Guard{this};
+  }
+
+ private:
+  size_t counter_{0};
+};
+using NestingCounterSingleton = SingletonThreadLocal<NestingCounter>;
+
 template <typename Promise>
 FOLLY_ALWAYS_INLINE folly::conditional_t<
     kIsSanitizeThread,
-    std::experimental::coroutine_handle<>,
+    void,
     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));
+    if (auto nestingGuard = NestingCounterSingleton::get().guard()) {
       ch.resume();
-    });
-    return std::experimental::noop_coroutine();
+    } else {
+      copy(ex).add([ch, rctx = RequestContext::saveContext()](
+                       Executor::KeepAlive<>&&) mutable {
+        RequestContextScopeGuard guard(std::move(rctx));
+        ch.resume();
+      });
+    }
   } else {
     return ch;
   }