[k2] Add support for multiple sync/async stack switches in backtrace algorithm

runtime-light/coroutine/async-stack-structs.h

@@ -0,0 +1,102 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#pragma once
+
+/**
+ * This header defines the data structures used to represent a coroutine asynchronous stack.
+ *
+ * Overview:
+ * The asynchronous stack is used to manage the execution state of coroutines, allowing for
+ * efficient context switching and stack management.
+ *
+ * Diagram: Normal and Asynchronous Stacks
+ *
+ *  Base Pointer (%rbp)              async_stack_root
+ *         |                                 |
+ *         V                                 V
+ *    stack_frame                    async_stack_frame
+ *         |                                 |
+ *         V                                 V
+ *    stack_frame                    async_stack_frame
+ *        ...                               ...
+ *         |                                 |
+ *         V                                 V
+ *    stack_frame                    async_stack_frame
+ *         |                                 |
+ *         V                                 V
+ *
+ * In the diagram above, the left side represents a typical call stack with stack frames linked by
+ * the base pointer (%rbp). The right side illustrates an asynchronous stack where `async_stack_frame`
+ * structures are linked by `async_stack_root`.
+ *
+ * Diagram: Backtrace Mechanism
+ *
+ *  Base Pointer (%rbp)
+ *         |
+ *         V
+ *    stack_frame
+ *         |
+ *         V
+ *    stack_frame (stop_sync_frame)  <- async_stack_root
+ *                                             |
+ *                                             V
+ *                                     async_stack_frame (top_async_frame)
+ *                                             |
+ *                                             V
+ *                                     async_stack_frame
+ *                                            ...
+ *                                             |
+ *                                             V
+ *                                     async_stack_frame
+ *
+ * The backtrace mechanism involves traversing the stack frames to capture the call stack.
+ * The `stop_sync_frame` serves as a marker where the transition to the asynchronous stack occurs,
+ * allowing the backtrace to continue through the `async_stack_frame` structures.
+ */
+
+#define STACK_RETURN_ADDRESS __builtin_return_address(0)
+
+#define STACK_FRAME_ADDRESS __builtin_frame_address(0)
+
+namespace kphp::coro {
+
+struct stack_frame {
+  stack_frame* caller_stack_frame{};
+  void* return_address{};
+};
+
+struct async_stack_root;
+
+struct async_stack_frame {
+  async_stack_frame* caller_async_stack_frame{};
+  async_stack_root* async_stack_root{};
+  void* return_address{};
+};
+
+struct async_stack_root {
+  async_stack_frame* top_async_stack_frame{};
+  stack_frame* stop_sync_stack_frame{};
+  async_stack_root* next_async_stack_root{};
+};
+
+/**
+ * The async_stack_element class facilitates working with asynchronous traces in templated code.
+ * This allows for uniform handling of any coroutines in places where async frames are pushed or popped.
+ */
+struct async_stack_element {
+  async_stack_frame& get_async_stack_frame() noexcept {
+    return m_async_stack_frame;
+  }
+
+private:
+  async_stack_frame m_async_stack_frame;
+};
+
+/**
+ * The async_stack_element class facilitates working with asynchronous traces in templated code.
+ * This allows for uniform handling of any coroutines in places where async frames are pushed or popped.
+ */
+
+} // namespace kphp::coro

runtime-light/coroutine/async-stack.h

@@ -4,107 +4,49 @@
 
 #pragma once
 
+#include "runtime-light/coroutine/async-stack-structs.h"
+
 #include <coroutine>
 #include <utility>
 
-/**
- * This header defines the data structures used to represent a coroutine asynchronous stack.
- *
- * Overview:
- * The asynchronous stack is used to manage the execution state of coroutines, allowing for
- * efficient context switching and stack management.
- *
- * Diagram: Normal and Asynchronous Stacks
- *
- *  Base Pointer (%rbp)              async_stack_root
- *         |                                 |
- *         V                                 V
- *    stack_frame                    async_stack_frame
- *         |                                 |
- *         V                                 V
- *    stack_frame                    async_stack_frame
- *        ...                               ...
- *         |                                 |
- *         V                                 V
- *    stack_frame                    async_stack_frame
- *         |                                 |
- *         V                                 V
- *
- * In the diagram above, the left side represents a typical call stack with stack frames linked by
- * the base pointer (%rbp). The right side illustrates an asynchronous stack where `async_stack_frame`
- * structures are linked by `async_stack_root`.
- *
- * Diagram: Backtrace Mechanism
- *
- *  Base Pointer (%rbp)
- *         |
- *         V
- *    stack_frame
- *         |
- *         V
- *    stack_frame (stop_sync_frame)  <- async_stack_root
- *                                             |
- *                                             V
- *                                     async_stack_frame (top_async_frame)
- *                                             |
- *                                             V
- *                                     async_stack_frame
- *                                            ...
- *                                             |
- *                                             V
- *                                     async_stack_frame
- *
- * The backtrace mechanism involves traversing the stack frames to capture the call stack.
- * The `stop_sync_frame` serves as a marker where the transition to the asynchronous stack occurs,
- * allowing the backtrace to continue through the `async_stack_frame` structures.
- */
-
-#define STACK_RETURN_ADDRESS __builtin_return_address(0)
-
-#define STACK_FRAME_ADDRESS __builtin_frame_address(0)
+#include "runtime-light/coroutine/coroutine-state.h"
+#include "runtime-light/stdlib/diagnostics/logs.h"
 
 namespace kphp::coro {
 
-struct stack_frame {
-  stack_frame* caller_stack_frame{};
-  void* return_address{};
-};
-
-struct async_stack_root;
+struct async_stack_root_wrapper {
+  async_stack_root root;
 
-struct async_stack_frame {
-  async_stack_frame* caller_async_stack_frame{};
-  async_stack_root* async_stack_root{};
-  void* return_address{};
-};
-
-struct async_stack_root {
-  async_stack_frame* top_async_stack_frame{};
-  stack_frame* stop_sync_stack_frame{};
+  async_stack_root_wrapper() = default;
+  async_stack_root_wrapper(const async_stack_root_wrapper& other) = delete;
+  async_stack_root_wrapper& operator=(const async_stack_root_wrapper& other) = delete;
+  async_stack_root_wrapper(async_stack_root_wrapper&& other) = default;
+  async_stack_root_wrapper& operator=(async_stack_root_wrapper&& other) = default;
+  ~async_stack_root_wrapper() {
+    CoroutineInstanceState::get().current_async_stack_root = root.next_async_stack_root;
+  }
 };
 
-/**
- * The `resume` function is responsible for storing the current synchronous stack frame
- * in async_stack_root::stop_sync_frame before resuming the coroutine. This allows
- * capturing one of the stack frames in the synchronous stack trace.
- */
-inline void resume(std::coroutine_handle<> handle, async_stack_root& stack_root) noexcept {
-  auto* previous_stack_frame{std::exchange(stack_root.stop_sync_stack_frame, reinterpret_cast<stack_frame*>(STACK_FRAME_ADDRESS))};
-  handle.resume();
-  stack_root.stop_sync_stack_frame = previous_stack_frame;
+inline void preparation_for_resume(kphp::coro::async_stack_root* root, void* stack_frame_addr,
+                                   CoroutineInstanceState& coroutine_state = CoroutineInstanceState::get()) {
+  root->stop_sync_stack_frame = reinterpret_cast<stack_frame*>(stack_frame_addr);
+  coroutine_state.current_async_stack_root = root;
 }
 
 /**
- * The async_stack_element class facilitates working with asynchronous traces in templated code.
- * This allows for uniform handling of any coroutines in places where async frames are pushed or popped.
+ * The `resume_with_new_root` function is responsible for storing the current synchronous stack frame
+ * in async_stack_root::stop_sync_frame before resuming the coroutine. This allows
+ * capturing one of the stack frames in the synchronous stack trace and jump between parts of the async_stack.
  */
-struct async_stack_element {
-  async_stack_frame& get_async_stack_frame() noexcept {
-    return m_async_stack_frame;
-  }
+[[clang::noinline]] inline void resume_with_new_root(std::coroutine_handle<> handle, async_stack_root* new_async_stack_root) noexcept {
+  kphp::log::assertion(new_async_stack_root != nullptr);
+  auto& coroutine_st{CoroutineInstanceState::get()};
 
-private:
-  async_stack_frame m_async_stack_frame;
-};
+  new_async_stack_root->next_async_stack_root = coroutine_st.current_async_stack_root;
+  new_async_stack_root->stop_sync_stack_frame = reinterpret_cast<stack_frame*>(STACK_FRAME_ADDRESS);
 
+  coroutine_st.current_async_stack_root = new_async_stack_root;
+  handle.resume();
+  new_async_stack_root->stop_sync_stack_frame = nullptr;
+}
 } // namespace kphp::coro

runtime-light/coroutine/await-set.h

@@ -19,21 +19,17 @@ namespace kphp::coro {
 template<typename return_type>
 class await_set {
   std::unique_ptr<detail::await_set::await_broker<return_type>> m_await_broker;
-  kphp::coro::async_stack_root& m_coroutine_stack_root;
 
 public:
   await_set() noexcept
-      : m_await_broker(std::make_unique<detail::await_set::await_broker<return_type>>()),
-        m_coroutine_stack_root(CoroutineInstanceState::get().coroutine_stack_root) {}
+      : m_await_broker(std::make_unique<detail::await_set::await_broker<return_type>>()) {}
 
   await_set(await_set&& other) noexcept
-      : m_await_broker(std::move(other.m_await_broker)),
-        m_coroutine_stack_root(other.m_coroutine_stack_root) {}
+      : m_await_broker(std::move(other.m_await_broker)) {}
 
   await_set& operator=(await_set&& other) noexcept {
     if (this != std::addressof(other)) {
       m_await_broker = std::move(other.m_await_broker);
-      m_coroutine_stack_root = other.m_coroutine_stack_root;
     }
     return *this;
   }
@@ -44,7 +40,7 @@ class await_set {
   template<typename awaitable_type>
   requires kphp::coro::concepts::awaitable<awaitable_type> && std::is_same_v<typename awaitable_traits<awaitable_type>::awaiter_return_type, return_type>
   void push(awaitable_type awaitable) noexcept {
-    m_await_broker->start_task(detail::await_set::make_await_set_task(std::move(awaitable)), m_coroutine_stack_root, STACK_RETURN_ADDRESS);
+    m_await_broker->start_task(detail::await_set::make_await_set_task(std::move(awaitable)));
   }
 
   auto next() noexcept {

runtime-light/coroutine/coroutine-state.h

@@ -6,13 +6,14 @@
 
 #include "common/mixin/not_copyable.h"
 
-#include "runtime-light/coroutine/async-stack.h"
+#include "runtime-light/coroutine/async-stack-structs.h"
 
 struct CoroutineInstanceState final : private vk::not_copyable {
 
   CoroutineInstanceState() noexcept = default;
 
   static CoroutineInstanceState& get() noexcept;
 
-  kphp::coro::async_stack_root coroutine_stack_root;
+  kphp::coro::async_stack_root base_coroutine_stack_root; // used to reduce the number of accesses to thread-local state
+  kphp::coro::async_stack_root* current_async_stack_root{&base_coroutine_stack_root};
 };