Update folly::fibers to preserve thread_local AsyncStackRoot

Summary:
Updates the FiberManager to make sure that it saves/restores the current-thread's
top-most active AsyncStackRoot when the coroutine is suspended/resumed.

This is required for correctness to ensure async stack traces remain functional
in the presence of fiber context-switches.

(Note: this ignores all push blocking failures!)

Reviewed By: yfeldblum

Differential Revision: D21130668

fbshipit-source-id: 4002bae0fc12a61c38836f9f099da4e6cf9183f4

folly/fibers/Fiber.h

@@ -29,6 +29,8 @@
 #include <folly/io/async/Request.h>
 
 namespace folly {
+struct AsyncStackRoot;
+
 namespace fibers {
 
 class Baton;
@@ -112,6 +114,7 @@ class Fiber {
   unsigned char* fiberStackLimit_;
   FiberImpl fiberImpl_; /**< underlying fiber implementation */
   std::shared_ptr<RequestContext> rcontext_; /**< current RequestContext */
+  folly::AsyncStackRoot* asyncRoot_ = nullptr;
   folly::Function<void()> func_; /**< task function */
   bool recordStackUsed_{false};
   bool stackFilledWithMagic_{false};

folly/fibers/FiberManagerInternal-inl.h

@@ -31,6 +31,7 @@
 #include <folly/fibers/Fiber.h>
 #include <folly/fibers/LoopController.h>
 #include <folly/fibers/Promise.h>
+#include <folly/tracing/AsyncStack.h>
 
 namespace folly {
 namespace fibers {
@@ -118,6 +119,10 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
   currentFiber_ = fiber;
   // Note: resetting the context is handled by the loop
   RequestContext::setContext(std::move(fiber->rcontext_));
+
+  (void)folly::exchangeCurrentAsyncStackRoot(
+      std::exchange(fiber->asyncRoot_, nullptr));
+
   if (observer_) {
     observer_->starting(reinterpret_cast<uintptr_t>(fiber));
   }
@@ -144,6 +149,7 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
     }
     currentFiber_ = nullptr;
     fiber->rcontext_ = RequestContext::saveContext();
+    fiber->asyncRoot_ = folly::exchangeCurrentAsyncStackRoot(nullptr);
   } else if (fiber->state_ == Fiber::INVALID) {
     assert(fibersActive_ > 0);
     --fibersActive_;
@@ -166,6 +172,11 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
     }
     currentFiber_ = nullptr;
     fiber->rcontext_ = RequestContext::saveContext();
+    // Async stack roots should have been popped by the time the
+    // func_() call has returned.
+    fiber->asyncRoot_ = folly::exchangeCurrentAsyncStackRoot(nullptr);
+    CHECK(fiber->asyncRoot_ == nullptr);
+
     fiber->localData_.reset();
     fiber->rcontext_.reset();
 
@@ -184,6 +195,7 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
     }
     currentFiber_ = nullptr;
     fiber->rcontext_ = RequestContext::saveContext();
+    fiber->asyncRoot_ = folly::exchangeCurrentAsyncStackRoot(nullptr);
     fiber->state_ = Fiber::READY_TO_RUN;
     yieldedFibers_->push_back(*fiber);
   }
@@ -210,10 +222,17 @@ void FiberManager::runFibersHelper(LoopFunc&& loopFunc) {
   // if the Fibers share the same context
   auto curCtx = RequestContext::saveContext();
 
+  auto* curAsyncRoot = folly::exchangeCurrentAsyncStackRoot(nullptr);
+
   FiberTailQueue yieldedFibers;
   auto prevYieldedFibers = std::exchange(yieldedFibers_, &yieldedFibers);
 
   SCOPE_EXIT {
+    // Restore the previous AsyncStackRoot and make sure that none of
+    // the fibers left any AsyncStackRoot pointers lying around.
+    auto* oldAsyncRoot = folly::exchangeCurrentAsyncStackRoot(curAsyncRoot);
+    CHECK(oldAsyncRoot == nullptr);
+
     yieldedFibers_ = prevYieldedFibers;
     if (observer_) {
       for (auto& yielded : yieldedFibers) {

folly/fibers/test/FibersTest.cpp

@@ -40,6 +40,7 @@
 #include <folly/fibers/WhenN.h>
 #include <folly/io/async/ScopedEventBaseThread.h>
 #include <folly/portability/GTest.h>
+#include <folly/tracing/AsyncStack.h>
 
 using namespace folly::fibers;
 
@@ -2705,3 +2706,60 @@ TEST(FiberManager, addTaskEagerKeepAlive) {
   EXPECT_TRUE(f.isReady());
   EXPECT_EQ(42, std::move(f).get());
 }
+
+TEST(FiberManager, fibersPreserveAsyncStackRoots) {
+  folly::EventBase evb;
+  auto& fm = getFiberManager(evb);
+
+  {
+    folly::detail::ScopedAsyncStackRoot root;
+
+    auto f = [&] {
+      // Should be launched with a no active AsyncStackRoot
+      EXPECT_TRUE(folly::tryGetCurrentAsyncStackRoot() == nullptr);
+
+      folly::detail::ScopedAsyncStackRoot scopedRoot1;
+
+      auto* root1 = folly::tryGetCurrentAsyncStackRoot();
+      EXPECT_TRUE(root1 != nullptr);
+
+      fm.yield();
+
+      EXPECT_EQ(root1, folly::tryGetCurrentAsyncStackRoot());
+
+      {
+        folly::detail::ScopedAsyncStackRoot scopedRoot2;
+
+        auto* root2 = folly::tryGetCurrentAsyncStackRoot();
+
+        folly::AsyncStackFrame frame1;
+
+        folly::AsyncStackFrame frame2;
+        frame2.setParentFrame(frame1);
+
+        scopedRoot2.activateFrame(frame2);
+
+        fm.yield();
+
+        EXPECT_EQ(root2, folly::tryGetCurrentAsyncStackRoot());
+
+        folly::deactivateAsyncStackFrame(frame2);
+      }
+    };
+
+    auto* originalRoot = folly::tryGetCurrentAsyncStackRoot();
+
+    auto task1 = fm.addTaskFuture(f);
+    auto task2 = fm.addTaskFuture(f);
+
+    EXPECT_EQ(originalRoot, folly::tryGetCurrentAsyncStackRoot());
+
+    std::move(task1).getVia(&evb);
+
+    EXPECT_EQ(originalRoot, folly::tryGetCurrentAsyncStackRoot());
+
+    std::move(task2).getVia(&evb);
+
+    EXPECT_EQ(originalRoot, folly::tryGetCurrentAsyncStackRoot());
+  }
+}