Add support for async stack traces to folly::coro::Task coroutines

Summary:
Initial diff that adds support for tracing through a chain of folly::coro::Task coroutines.

This adds some scaffolding that handles saving/restoring the active AsyncStackFrame
when a Task awaits some awaitable type that does not know about AsyncStackFrame
objects but also allows awaitables to opt-in to AsyncStackFrame awareness by
customising the new `folly::coro::co_withAsyncStack()` CPO.

Currently only `Task` and `TaskWithExecutor` awaiters have customised this CPO.

Also updated the awaiters for `Task` and `TaskWithExecutor` to handle being awaited
from coroutines that are not async-stack aware - in which case they just record a
null parent-frame for the Task.

The Task's `final_suspend()` then either deactivates or pops the frame depending
on whether there was a parent frame recorded.

BUG: This change currently breaks the symmetric-transfer stack-overflow avoidance
when awaiting synchronously-completing coroutines from an AsyncGenerator or from
a BarrierTask (eg. inside collectAll implementations).

Reviewed By: andriigrynenko

Differential Revision: D24428736

fbshipit-source-id: 5722e511ad10d95198ae70a5afe567d83cb06285

folly/experimental/coro/Task.h

@@ -34,6 +34,7 @@
 #include <folly/experimental/coro/Traits.h>
 #include <folly/experimental/coro/Utils.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
+#include <folly/experimental/coro/WithAsyncStack.h>
 #include <folly/experimental/coro/WithCancellation.h>
 #include <folly/experimental/coro/detail/InlineTask.h>
 #include <folly/experimental/coro/detail/Malloc.h>
@@ -41,6 +42,7 @@
 #include <folly/futures/Future.h>
 #include <folly/io/async/Request.h>
 #include <folly/lang/Assume.h>
+#include <folly/tracing/AsyncStack.h>
 
 namespace folly {
 namespace coro {
@@ -77,9 +79,13 @@ class TaskPromiseBase {
     bool await_ready() noexcept { return false; }
 
     template <typename Promise>
-    std::experimental::coroutine_handle<> await_suspend(
-        std::experimental::coroutine_handle<Promise> coro) noexcept {
-      return coro.promise().continuation_;
+    FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES
+        std::experimental::coroutine_handle<>
+        await_suspend(
+            std::experimental::coroutine_handle<Promise> coro) noexcept {
+      TaskPromiseBase& promise = coro.promise();
+      folly::popAsyncStackFrameCallee(promise.asyncFrame_);
+      return promise.continuation_;
     }
 
     [[noreturn]] void await_resume() noexcept { folly::assume_unreachable(); }
@@ -105,10 +111,10 @@ class TaskPromiseBase {
 
   template <typename Awaitable>
   auto await_transform(Awaitable&& awaitable) {
-    return folly::coro::co_viaIfAsync(
+    return folly::coro::co_withAsyncStack(folly::coro::co_viaIfAsync(
         executor_.get_alias(),
         folly::coro::co_withCancellation(
-            cancelToken_, static_cast<Awaitable&&>(awaitable)));
+            cancelToken_, static_cast<Awaitable&&>(awaitable))));
   }
 
   auto await_transform(co_current_executor_t) noexcept {
@@ -126,6 +132,8 @@ class TaskPromiseBase {
     }
   }
 
+  folly::AsyncStackFrame& getAsyncFrame() noexcept { return asyncFrame_; }
+
  private:
   template <typename T>
   friend class folly::coro::TaskWithExecutor;
@@ -134,6 +142,7 @@ class TaskPromiseBase {
   friend class folly::coro::Task;
 
   std::experimental::coroutine_handle<> continuation_;
+  folly::AsyncStackFrame asyncFrame_;
   folly::Executor::KeepAlive<> executor_;
   folly::CancellationToken cancelToken_;
   bool hasCancelTokenOverride_ = false;
@@ -341,6 +350,8 @@ class FOLLY_NODISCARD TaskWithExecutor {
    public:
     explicit Awaiter(handle_t coro) noexcept : coro_(coro) {}
 
+    Awaiter(Awaiter&& other) noexcept : coro_(std::exchange(other.coro_, {})) {}
+
     ~Awaiter() {
       if (coro_) {
         coro_.destroy();
@@ -349,8 +360,9 @@ class FOLLY_NODISCARD TaskWithExecutor {
 
     bool await_ready() const { return false; }
 
-    FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES void await_suspend(
-        std::experimental::coroutine_handle<> continuation) noexcept {
+    template <typename Promise>
+    FOLLY_NOINLINE void await_suspend(
+        std::experimental::coroutine_handle<Promise> continuation) noexcept {
       auto& promise = coro_.promise();
       DCHECK(!promise.continuation_);
       DCHECK(promise.executor_);
@@ -364,11 +376,20 @@ class FOLLY_NODISCARD TaskWithExecutor {
           << "If you need to run a task inline in a unit-test, you should use "
           << "coro::blockingWait instead.";
 
+      auto& calleeFrame = promise.getAsyncFrame();
+      calleeFrame.setReturnAddress();
+
+      if constexpr (detail::promiseHasAsyncFrame_v<Promise>) {
+        auto& callerFrame = continuation.promise().getAsyncFrame();
+        calleeFrame.setParentFrame(callerFrame);
+        folly::deactivateAsyncStackFrame(callerFrame);
+      }
+
       promise.continuation_ = continuation;
       promise.executor_->add(
           [coro = coro_, ctx = RequestContext::saveContext()]() mutable {
             RequestContextScopeGuard contextScope{std::move(ctx)};
-            coro.resume();
+            folly::resumeCoroutineWithNewAsyncStackRoot(coro);
           });
     }
 
@@ -391,6 +412,9 @@ class FOLLY_NODISCARD TaskWithExecutor {
    public:
     InlineTryAwaitable(handle_t coro) noexcept : coro_(coro) {}
 
+    InlineTryAwaitable(InlineTryAwaitable&& other) noexcept
+        : coro_(std::exchange(other.coro_, {})) {}
+
     ~InlineTryAwaitable() {
       if (coro_) {
         coro_.destroy();
@@ -399,13 +423,24 @@ class FOLLY_NODISCARD TaskWithExecutor {
 
     bool await_ready() { return false; }
 
-    auto await_suspend(std::experimental::coroutine_handle<> continuation) {
+    template <typename Promise>
+    FOLLY_NOINLINE void await_suspend(
+        std::experimental::coroutine_handle<Promise> continuation) {
       auto& promise = coro_.promise();
       DCHECK(!promise.continuation_);
       DCHECK(promise.executor_);
 
       promise.continuation_ = continuation;
-      return coro_;
+
+      auto& calleeFrame = promise.getAsyncFrame();
+      calleeFrame.setReturnAddress();
+
+      // This awaitable is only ever awaited from a DetachedInlineTask
+      // which is not an async-stack-aware coroutine.
+      // Can't use symmetric-transfer here as we need to register a
+      // new AsyncStackRoot (we can't assume there is one already)
+      // and then ensure it is deregistered when the coroutine suspends.
+      folly::resumeCoroutineWithNewAsyncStackRoot(coro_);
     }
 
     folly::Try<StorageType> await_resume() {
@@ -430,6 +465,12 @@ class FOLLY_NODISCARD TaskWithExecutor {
     return std::move(task);
   }
 
+  friend TaskWithExecutor tag_invoke(
+      cpo_t<co_withAsyncStack>,
+      TaskWithExecutor&& task) noexcept {
+    return std::move(task);
+  }
+
  private:
   friend class Task<T>;
 
@@ -551,10 +592,24 @@ class FOLLY_NODISCARD Task {
 
     bool await_ready() noexcept { return false; }
 
-    handle_t await_suspend(
-        std::experimental::coroutine_handle<> continuation) noexcept {
-      coro_.promise().continuation_ = continuation;
-      return coro_;
+    template <typename Promise>
+    FOLLY_NOINLINE auto await_suspend(
+        std::experimental::coroutine_handle<Promise> continuation) noexcept {
+      auto& promise = coro_.promise();
+
+      promise.continuation_ = continuation;
+
+      auto& calleeFrame = promise.getAsyncFrame();
+      calleeFrame.setReturnAddress();
+
+      if constexpr (detail::promiseHasAsyncFrame_v<Promise>) {
+        auto& callerFrame = continuation.promise().getAsyncFrame();
+        folly::pushAsyncStackFrameCallerCallee(callerFrame, calleeFrame);
+        return coro_;
+      } else {
+        folly::resumeCoroutineWithNewAsyncStackRoot(coro_);
+        return;
+      }
     }
 
     T await_resume() {
@@ -568,6 +623,12 @@ class FOLLY_NODISCARD Task {
     }
 
    private:
+    friend Awaiter tag_invoke(
+        cpo_t<co_withAsyncStack>,
+        Awaiter&& awaiter) noexcept {
+      return std::move(awaiter);
+    }
+
     handle_t coro_;
   };
 

folly/experimental/coro/Traits.h

@@ -205,5 +205,17 @@ struct await_result<Awaitable, std::enable_if_t<is_awaitable_v<Awaitable>>> {
 template <typename Awaitable>
 using await_result_t = typename await_result<Awaitable>::type;
 
+namespace detail {
+
+template <typename Promise, typename = void>
+constexpr bool promiseHasAsyncFrame_v = false;
+
+template <typename Promise>
+constexpr bool promiseHasAsyncFrame_v<
+    Promise,
+    std::void_t<decltype(std::declval<Promise&>().getAsyncFrame())>> = true;
+
+} // namespace detail
+
 } // namespace coro
 } // namespace folly

folly/experimental/coro/WithAsyncStack.h

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <folly/experimental/coro/Traits.h>
+#include <folly/functional/Invoke.h>
+#include <folly/lang/Assume.h>
+#include <folly/lang/CustomizationPoint.h>
+#include <folly/tracing/AsyncStack.h>
+
+#include <cassert>
+#include <type_traits>
+#include <utility>
+
+namespace folly::coro {
+
+namespace detail {
+
+class WithAsyncStackCoroutine {
+ public:
+  class promise_type {
+   public:
+    WithAsyncStackCoroutine get_return_object() noexcept {
+      return WithAsyncStackCoroutine{
+          std::experimental::coroutine_handle<promise_type>::from_promise(
+              *this)};
+    }
+
+    std::experimental::suspend_always initial_suspend() noexcept { return {}; }
+
+    struct FinalAwaiter {
+      bool await_ready() noexcept { return false; }
+      void await_suspend(
+          std::experimental::coroutine_handle<promise_type> h) noexcept {
+        auto& promise = h.promise();
+        folly::resumeCoroutineWithNewAsyncStackRoot(
+            promise.continuation_, *promise.parentFrame_);
+      }
+
+      [[noreturn]] void await_resume() noexcept { folly::assume_unreachable(); }
+    };
+
+    FinalAwaiter final_suspend() noexcept { return {}; }
+
+    void return_void() noexcept {}
+
+    [[noreturn]] void unhandled_exception() noexcept {
+      folly::assume_unreachable();
+    }
+
+   private:
+    friend WithAsyncStackCoroutine;
+
+    std::experimental::coroutine_handle<> continuation_;
+    folly::AsyncStackFrame* parentFrame_ = nullptr;
+  };
+
+  WithAsyncStackCoroutine() noexcept : coro_() {}
+
+  WithAsyncStackCoroutine(WithAsyncStackCoroutine&& other) noexcept
+      : coro_(std::exchange(other.coro_, {})) {}
+
+  ~WithAsyncStackCoroutine() {
+    if (coro_) {
+      coro_.destroy();
+    }
+  }
+
+  WithAsyncStackCoroutine& operator=(WithAsyncStackCoroutine other) noexcept {
+    std::swap(coro_, other.coro_);
+    return *this;
+  }
+
+  static WithAsyncStackCoroutine create() { co_return; }
+
+  template <typename Promise>
+  std::experimental::coroutine_handle<promise_type> getWrapperHandleFor(
+      std::experimental::coroutine_handle<Promise> h) noexcept {
+    auto& promise = coro_.promise();
+    promise.continuation_ = h;
+    promise.parentFrame_ = std::addressof(h.promise().getAsyncFrame());
+    return coro_;
+  }
+
+ private:
+  explicit WithAsyncStackCoroutine(
+      std::experimental::coroutine_handle<promise_type> h) noexcept
+      : coro_(h) {}
+
+  std::experimental::coroutine_handle<promise_type> coro_;
+};
+
+template <typename Awaitable>
+class WithAsyncStackAwaiter {
+ public:
+  explicit WithAsyncStackAwaiter(Awaitable&& awaitable)
+      : awaiter_(folly::coro::get_awaiter(static_cast<Awaitable&&>(awaitable))),
+        coroWrapper_(WithAsyncStackCoroutine::create()) {}
+
+  decltype(auto) await_ready() noexcept(noexcept(awaiter_.await_ready())) {
+    return awaiter_.await_ready();
+  }
+
+  template <typename Promise>
+  FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES decltype(auto) await_suspend(
+      std::experimental::coroutine_handle<Promise> h) {
+    AsyncStackFrame& callerFrame = h.promise().getAsyncFrame();
+    AsyncStackRoot* stackRoot = callerFrame.getStackRoot();
+    assert(stackRoot != nullptr);
+
+    auto wrapperHandle = coroWrapper_.getWrapperHandleFor(h);
+
+    folly::deactivateAsyncStackFrame(callerFrame);
+
+    using await_suspend_result_t =
+        decltype(awaiter_.await_suspend(wrapperHandle));
+
+    try {
+      if constexpr (std::is_same_v<await_suspend_result_t, bool>) {
+        if (!awaiter_.await_suspend(wrapperHandle)) {
+          folly::activateAsyncStackFrame(*stackRoot, callerFrame);
+          return false;
+        }
+        return true;
+      } else {
+        return awaiter_.await_suspend(wrapperHandle);
+      }
+    } catch (...) {
+      folly::activateAsyncStackFrame(*stackRoot, callerFrame);
+      throw;
+    }
+  }
+
+  decltype(auto) await_resume() noexcept(noexcept(awaiter_.await_resume())) {
+    coroWrapper_ = {};
+    return awaiter_.await_resume();
+  }
+
+ private:
+  awaiter_type_t<Awaitable> awaiter_;
+  WithAsyncStackCoroutine coroWrapper_;
+};
+
+template <typename Awaitable>
+class WithAsyncStackAwaitable {
+ public:
+  explicit WithAsyncStackAwaitable(Awaitable&& awaitable)
+      : awaitable_(static_cast<Awaitable&&>(awaitable)) {}
+
+  WithAsyncStackAwaiter<Awaitable&> operator co_await() & {
+    return WithAsyncStackAwaiter<Awaitable&>{awaitable_};
+  }
+
+  WithAsyncStackAwaiter<Awaitable> operator co_await() && {
+    return WithAsyncStackAwaiter<Awaitable>{
+        static_cast<Awaitable&&>(awaitable_)};
+  }
+
+ private:
+  Awaitable awaitable_;
+};
+
+struct WithAsyncStackFunction {
+  // Dispatches to a custom implementation using tag_invoke()
+  template <
+      typename Awaitable,
+      std::enable_if_t<
+          folly::is_tag_invocable_v<WithAsyncStackFunction, Awaitable>,
+          int> = 0>
+  auto operator()(Awaitable&& awaitable) const noexcept(
+      folly::is_nothrow_tag_invocable_v<WithAsyncStackFunction, Awaitable>)
+      -> folly::tag_invoke_result_t<WithAsyncStackFunction, Awaitable> {
+    return folly::tag_invoke(
+        WithAsyncStackFunction{}, static_cast<Awaitable&&>(awaitable));
+  }
+
+  // Fallback implementation. Wraps the awaitable in the
+  // WithAsyncStackAwaitable which just saves/restores the
+  // awaiting coroutine's AsyncStackFrame.
+  template <
+      typename Awaitable,
+      std::enable_if_t<
+          !folly::is_tag_invocable_v<WithAsyncStackFunction, Awaitable>,
+          int> = 0,
+      std::enable_if_t<folly::coro::is_awaitable_v<Awaitable>, int> = 0>
+  WithAsyncStackAwaitable<Awaitable> operator()(Awaitable&& awaitable) const
+      noexcept(std::is_nothrow_move_constructible_v<Awaitable>) {
+    return WithAsyncStackAwaitable<Awaitable>{
+        static_cast<Awaitable&&>(awaitable)};
+  }
+};
+
+} // namespace detail
+
+// Coroutines that support the AsyncStack protocol will apply the
+// co_withAsyncStack() customisation-point to an awaitable inside its
+// await_transform() to ensure that the current coroutine's AsyncStackFrame
+// is saved and later restored when the coroutine resumes.
+//
+// The default implementation is used for awaitables that don't know
+// about the AsyncStackFrame and just wraps the awaitable to ensure
+// that the stack-frame is saved/restored if the coroutine suspends.
+//
+// Awaitables that know about the AsyncStackFrame protocol can customise
+// this CPO by defining an overload of tag_invoke() for this CPO
+// for their type.
+//
+// For example:
+//   class MyAwaitable {
+//     friend MyAwaitable&& tag_invoke(
+//         cpo_t<folly::coro::co_withAsyncStack>, MyAwaitable&& awaitable) {
+//       return std::move(awaitable);
+//     }
+//
+//     ...
+//   };
+//
+// If you customise this CPO then it is your responsibility to ensure that
+// if the awaiting coroutine suspends then before the coroutine is resumed
+// that its original AsyncStackFrame is activated on the current thread.
+// e.g. using folly::activateAsyncStackFrame()
+//
+// The awaiting coroutine's AsyncStackFrame can be obtained from its
+// promise, which is assumed to have a 'AsyncStackFrame& getAsyncFrame()'
+// method that returns a reference to the parent coroutine's async frame.
+
+FOLLY_DEFINE_CPO(detail::WithAsyncStackFunction, co_withAsyncStack)
+
+} // namespace folly::coro

folly/experimental/coro/detail/Helpers.h

@@ -19,6 +19,7 @@
 #include <folly/Executor.h>
 #include <folly/SingletonThreadLocal.h>
 #include <folly/io/async/Request.h>
+#include <folly/tracing/AsyncStack.h>
 
 namespace folly {
 namespace coro {