[k2] implement async stack (#1320)

Author: astrophysik

runtime-light/coroutine/async-stack.h

@@ -0,0 +1,110 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#pragma once
+
+#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)
+
+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{};
+};
+
+/**
+ * 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;
+}
+
+/**
+ * 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 async_stack_frame_;
+  }
+
+private:
+  async_stack_frame async_stack_frame_;
+};
+
+} // namespace kphp::coro

runtime-light/coroutine/awaitable.h

@@ -8,11 +8,14 @@
 #include <concepts>
 #include <coroutine>
 #include <cstdint>
+#include <memory>
 #include <optional>
 #include <tuple>
 #include <type_traits>
 #include <utility>
 
+#include "runtime-light/coroutine/async-stack.h"
+#include "runtime-light/coroutine/coroutine-state.h"
 #include "runtime-light/coroutine/shared-task.h"
 #include "runtime-light/coroutine/task.h"
 #include "runtime-light/k2-platform/k2-api.h"
@@ -55,9 +58,24 @@ class fork_id_watcher_t {
   }
 };
 
+class async_stack_watcher_t {
+  kphp::coro::async_stack_root* const async_stack_root;
+  kphp::coro::async_stack_frame* const suspended_async_stack_frame;
+
+protected:
+  void await_resume() const noexcept {
+    async_stack_root->top_async_stack_frame = suspended_async_stack_frame;
+  }
+
+public:
+  async_stack_watcher_t() noexcept
+      : async_stack_root(std::addressof(CoroutineInstanceState::get().coroutine_stack_root)),
+        suspended_async_stack_frame(async_stack_root->top_async_stack_frame) {}
+};
+
 } // namespace awaitable_impl_
 
-class wait_for_update_t : public awaitable_impl_::fork_id_watcher_t {
+class wait_for_update_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::async_stack_watcher_t {
   uint64_t stream_d;
   SuspendToken suspend_token;
   awaitable_impl_::state state{awaitable_impl_::state::init};
@@ -96,6 +114,7 @@ class wait_for_update_t : public awaitable_impl_::fork_id_watcher_t {
 
   constexpr void await_resume() noexcept {
     state = awaitable_impl_::state::end;
+    async_stack_watcher_t::await_resume();
     fork_id_watcher_t::await_resume();
   }
 
@@ -111,7 +130,7 @@ class wait_for_update_t : public awaitable_impl_::fork_id_watcher_t {
 
 // ================================================================================================
 
-class wait_for_incoming_stream_t : awaitable_impl_::fork_id_watcher_t {
+class wait_for_incoming_stream_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::async_stack_watcher_t {
   SuspendToken suspend_token{std::noop_coroutine(), WaitEvent::IncomingStream{}};
   awaitable_impl_::state state{awaitable_impl_::state::init};
 
@@ -146,6 +165,7 @@ class wait_for_incoming_stream_t : awaitable_impl_::fork_id_watcher_t {
 
   uint64_t await_resume() noexcept {
     state = awaitable_impl_::state::end;
+    async_stack_watcher_t::await_resume();
     fork_id_watcher_t::await_resume();
     const auto incoming_stream_d{InstanceState::get().take_incoming_stream()};
     kphp::log::assertion(incoming_stream_d != k2::INVALID_PLATFORM_DESCRIPTOR);
@@ -164,7 +184,7 @@ class wait_for_incoming_stream_t : awaitable_impl_::fork_id_watcher_t {
 
 // ================================================================================================
 
-class wait_for_reschedule_t : awaitable_impl_::fork_id_watcher_t {
+class wait_for_reschedule_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::async_stack_watcher_t {
   SuspendToken suspend_token{std::noop_coroutine(), WaitEvent::Rechedule{}};
   awaitable_impl_::state state{awaitable_impl_::state::init};
 
@@ -198,6 +218,7 @@ class wait_for_reschedule_t : awaitable_impl_::fork_id_watcher_t {
 
   constexpr void await_resume() noexcept {
     state = awaitable_impl_::state::end;
+    async_stack_watcher_t::await_resume();
     fork_id_watcher_t::await_resume();
   }
 
@@ -213,7 +234,7 @@ class wait_for_reschedule_t : awaitable_impl_::fork_id_watcher_t {
 
 // ================================================================================================
 
-class wait_for_timer_t : awaitable_impl_::fork_id_watcher_t {
+class wait_for_timer_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::async_stack_watcher_t {
   std::chrono::nanoseconds duration;
   uint64_t timer_d{k2::INVALID_PLATFORM_DESCRIPTOR};
   SuspendToken suspend_token{std::noop_coroutine(), WaitEvent::Rechedule{}};
@@ -259,6 +280,7 @@ class wait_for_timer_t : awaitable_impl_::fork_id_watcher_t {
 
   constexpr void await_resume() noexcept {
     state = awaitable_impl_::state::end;
+    async_stack_watcher_t::await_resume();
     fork_id_watcher_t::await_resume();
   }
 
@@ -275,7 +297,7 @@ class wait_for_timer_t : awaitable_impl_::fork_id_watcher_t {
 // ================================================================================================
 
 template<typename T>
-class start_fork_t : awaitable_impl_::fork_id_watcher_t {
+class start_fork_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::async_stack_watcher_t {
   ForkInstanceState& fork_instance_st{ForkInstanceState::get()};
 
   int64_t fork_id{};
@@ -304,7 +326,8 @@ class start_fork_t : awaitable_impl_::fork_id_watcher_t {
     return fork_awaiter.await_ready();
   }
 
-  std::coroutine_handle<> await_suspend(std::coroutine_handle<> current_coro) noexcept {
+  template<typename promise_t>
+  std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_t> current_coro) noexcept {
     state = awaitable_impl_::state::suspend;
     fork_instance_st.current_id = fork_id;
     if (fork_awaiter.await_suspend(current_coro)) [[unlikely]] {
@@ -318,6 +341,7 @@ class start_fork_t : awaitable_impl_::fork_id_watcher_t {
 
   int64_t await_resume() noexcept {
     state = awaitable_impl_::state::end;
+    async_stack_watcher_t::await_resume();
     fork_id_watcher_t::await_resume();
     fork_awaiter.await_resume();
     return fork_id;
@@ -330,7 +354,7 @@ class start_fork_t : awaitable_impl_::fork_id_watcher_t {
 // The main difference between them is that wait_fork_t can use shared_task::when_ready as an awaiter,
 // whereas wait_fork_result_t can use shared_task::operator co_await.
 template<typename T>
-class wait_fork_t : awaitable_impl_::fork_id_watcher_t {
+class wait_fork_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::async_stack_watcher_t {
   kphp::coro::shared_task<T> fork_task;
   std::remove_cvref_t<decltype(std::declval<kphp::coro::shared_task<T>>().operator co_await())> fork_awaiter;
   awaitable_impl_::state state{awaitable_impl_::state::init};
@@ -364,13 +388,15 @@ class wait_fork_t : awaitable_impl_::fork_id_watcher_t {
     return state == awaitable_impl_::state::ready;
   }
 
-  constexpr bool await_suspend(std::coroutine_handle<> coro) noexcept {
+  template<typename promise_t>
+  constexpr bool await_suspend(std::coroutine_handle<promise_t> coro) noexcept {
     state = awaitable_impl_::state::suspend;
     return fork_awaiter.await_suspend(coro);
   }
 
   await_resume_t await_resume() noexcept {
     state = awaitable_impl_::state::end;
+    async_stack_watcher_t::await_resume();
     fork_id_watcher_t::await_resume();
     if constexpr (std::is_void_v<await_resume_t>) {
       fork_awaiter.await_resume();
@@ -431,7 +457,8 @@ class wait_with_timeout_t {
   // 3. await_suspend returns std::coroutine_handle<>.
   // we must guarantee that 'co_await wait_with_timeout_t{awaitable, timeout}' behaves like 'co_await awaitable' except
   // it may cancel 'co_await awaitable' if the timeout has elapsed.
-  await_suspend_return_t await_suspend(std::coroutine_handle<> coro) noexcept {
+  template<typename promise_t>
+  await_suspend_return_t await_suspend(std::coroutine_handle<promise_t> coro) noexcept {
     // as we don't rely on coroutine scheduler implementation, let's always suspend awaitable first. in case of some smart scheduler
     // it won't have any effect, but it will have an effect if our scheduler is quite simple.
     state = awaitable_impl_::state::suspend;

runtime-light/coroutine/coroutine-state.cpp

@@ -0,0 +1,11 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#include "runtime-light/coroutine/coroutine-state.h"
+
+#include "runtime-light/state/instance-state.h"
+
+CoroutineInstanceState& CoroutineInstanceState::get() noexcept {
+  return InstanceState::get().coroutine_instance_state;
+}

runtime-light/coroutine/coroutine-state.h

@@ -0,0 +1,18 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#pragma once
+
+#include "common/mixin/not_copyable.h"
+
+#include "runtime-light/coroutine/async-stack.h"
+
+struct CoroutineInstanceState final : private vk::not_copyable {
+
+  CoroutineInstanceState() noexcept = default;
+
+  static CoroutineInstanceState& get() noexcept;
+
+  kphp::coro::async_stack_root coroutine_stack_root;
+};

runtime-light/coroutine/coroutine.cmake

@@ -0,0 +1 @@
+prepend(RUNTIME_LIGHT_COROUTINE_SRC coroutine/ coroutine-state.cpp)

runtime-light/coroutine/shared-task.h

@@ -13,6 +13,7 @@
 #include <type_traits>
 #include <utility>
 
+#include "runtime-light/coroutine/async-stack.h"
 #include "runtime-light/k2-platform/k2-api.h"
 #include "runtime-light/utils/logs.h"
 
@@ -27,7 +28,7 @@ struct shared_task_waiter final {
 };
 
 template<typename promise_type>
-struct promise_base {
+struct promise_base : async_stack_element {
   constexpr auto initial_suspend() const noexcept -> std::suspend_always {
     return {};
   }
@@ -50,7 +51,8 @@ struct promise_base {
           // read the m_next pointer before resuming the coroutine
           // since resuming the coroutine may destroy the shared_task_waiter value
           auto* next{waiter->m_next};
-          waiter->m_continuation.resume();
+          auto& async_stack_root{*promise.get_async_stack_frame().async_stack_root};
+          kphp::coro::resume(waiter->m_continuation, async_stack_root);
           waiter = next;
         }
         // return last waiter's coroutine_handle to allow it to potentially be compiled as a tail-call
@@ -96,7 +98,9 @@ struct promise_base {
     // start the coroutine if not yet started
     if (m_waiters == NOT_STARTED_VAL) {
       m_waiters = STARTED_NO_WAITERS_VAL;
-      std::coroutine_handle<promise_type>::from_promise(*static_cast<promise_type*>(this)).resume();
+      const auto& handle{std::coroutine_handle<promise_type>::from_promise(*static_cast<promise_type*>(this))};
+      auto& async_stack_root{*get_async_stack_frame().async_stack_root};
+      kphp::coro::resume(handle, async_stack_root);
     }
     // coroutine already completed, don't suspend
     if (done()) {
@@ -167,6 +171,23 @@ class awaiter_base {
   enum class state : uint8_t { init, suspend, end };
   state m_state{state::init};
 
+  void set_async_top_frame(async_stack_frame& caller_frame, void* return_address) noexcept {
+    /**
+     * shared_task is the top of the stack for calls from it.
+     * Therefore, it's awaiter doesn't store caller_frame, but it save `await_suspend()` return address
+     * */
+    async_stack_frame& callee_frame{m_coro.promise().get_async_stack_frame()};
+
+    callee_frame.return_address = return_address;
+    auto* async_stack_root{caller_frame.async_stack_root};
+    callee_frame.async_stack_root = async_stack_root;
+    async_stack_root->top_async_stack_frame = std::addressof(callee_frame);
+  }
+
+  void reset_async_top_frame(async_stack_frame& caller_frame) noexcept {
+    caller_frame.async_stack_root->top_async_stack_frame = std::addressof(caller_frame);
+  }
+
 protected:
   std::coroutine_handle<promise_type> m_coro;
   shared_task_impl::shared_task_waiter m_waiter{};
@@ -195,10 +216,14 @@ class awaiter_base {
     return m_coro.promise().done();
   }
 
-  auto await_suspend(std::coroutine_handle<> awaiter) noexcept -> bool {
+  template<typename promise_t>
+  [[clang::noinline]] auto await_suspend(std::coroutine_handle<promise_t> awaiter) noexcept -> bool {
+    set_async_top_frame(awaiter.promise().get_async_stack_frame(), STACK_RETURN_ADDRESS);
     m_state = state::suspend;
     m_waiter.m_continuation = awaiter;
-    return m_coro.promise().suspend_awaiter(m_waiter);
+    bool should_be_suspended{m_coro.promise().suspend_awaiter(m_waiter)};
+    reset_async_top_frame(awaiter.promise().get_async_stack_frame());
+    return should_be_suspended;
   }
 
   auto await_resume() noexcept -> void {