include/boost/capy/ex/io_awaitable_promise_base.hpp

91.1% Lines (1002/1100) 100.0% List of functions (208/208)

[ ] prev/next file n p prev/next uncovered line

f(x) Functions (208)

Function Calls Lines Blocks

Line	Hits	Source Code
1		//
2		// Copyright (c) 2025 Vinnie Falco (vinnie.falco@gmail.com)
3		//
4		// Distributed under the Boost Software License, Version 1.0. (See accompanying
5		// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
6		//
7		// Official repository: https://github.com/cppalliance/capy
8		//
9
10		#ifndef BOOST_CAPY_EX_IO_AWAITABLE_PROMISE_BASE_HPP
11		#define BOOST_CAPY_EX_IO_AWAITABLE_PROMISE_BASE_HPP
12
13		#include <boost/capy/detail/config.hpp>
14		#include <boost/capy/ex/frame_allocator.hpp>
15		#include <boost/capy/ex/io_env.hpp>
16		#include <boost/capy/ex/recycling_memory_resource.hpp>
17		#include <boost/capy/ex/this_coro.hpp>
18
19		#include <coroutine>
20		#include <cstddef>
21		#include <cstring>
22		#include <memory_resource>
23		#include <stop_token>
24		#include <type_traits>
25
26		namespace boost {
27		namespace capy {
28
29		/** CRTP mixin that adds I/O awaitable support to a promise type.
30
31		Inherit from this class to enable these capabilities in your coroutine:
32
33		1. Frame allocation — The mixin provides `operator new/delete` that
34		use the thread-local frame allocator set by `run_async`.
35
36		2. Environment storage — The mixin stores a pointer to the `io_env`
37		containing the executor, stop token, and allocator for this coroutine.
38
39		3. Environment access — Coroutine code can retrieve the environment
40		via `co_await this_coro::environment`, or individual fields via
41		`co_await this_coro::executor`, `co_await this_coro::stop_token`,
42		and `co_await this_coro::frame_allocator`.
43
44		@tparam Derived The derived promise type (CRTP pattern).
45
46		@par Basic Usage
47
48		For coroutines that need to access their execution environment:
49
50		@code
51		struct my_task
52		{
53		struct promise_type : io_awaitable_promise_base<promise_type>
54		{
55		my_task get_return_object();
56		std::suspend_always initial_suspend() noexcept;
57		std::suspend_always final_suspend() noexcept;
58		void return_void();
59		void unhandled_exception();
60		};
61
62		// ... awaitable interface ...
63		};
64
65		my_task example()
66		{
67		auto env = co_await this_coro::environment;
68		// Access env->executor, env->stop_token, env->frame_allocator
69
70		// Or use fine-grained accessors:
71		auto ex = co_await this_coro::executor;
72		auto token = co_await this_coro::stop_token;
73		auto* alloc = co_await this_coro::frame_allocator;
74		}
75		@endcode
76
77		@par Custom Awaitable Transformation
78
79		If your promise needs to transform awaitables (e.g., for affinity or
80		logging), override `transform_awaitable` instead of `await_transform`:
81
82		@code
83		struct promise_type : io_awaitable_promise_base<promise_type>
84		{
85		template<typename A>
86		auto transform_awaitable(A&& a)
87		{
88		// Your custom transformation logic
89		return std::forward<A>(a);
90		}
91		};
92		@endcode
93
94		The mixin's `await_transform` intercepts @ref this_coro::environment_tag
95		and the fine-grained tag types (@ref this_coro::executor_tag,
96		@ref this_coro::stop_token_tag, @ref this_coro::frame_allocator_tag),
97		then delegates all other awaitables to your `transform_awaitable`.
98
99		@par Making Your Coroutine an IoAwaitable
100
101		The mixin handles the "inside the coroutine" part—accessing the
102		environment. To receive the environment when your coroutine is awaited
103		(satisfying @ref IoAwaitable), implement the `await_suspend` overload
104		on your coroutine return type:
105
106		@code
107		struct my_task
108		{
109		struct promise_type : io_awaitable_promise_base<promise_type> { ... };
110
111		std::coroutine_handle<promise_type> h_;
112
113		// IoAwaitable await_suspend receives and stores the environment
114		std::coroutine_handle<> await_suspend(std::coroutine_handle<> cont, io_env const* env)
115		{
116		h_.promise().set_environment(env);
117		// ... rest of suspend logic ...
118		}
119		};
120		@endcode
121
122		@par Thread Safety
123		The environment is stored during `await_suspend` and read during
124		`co_await this_coro::environment`. These occur on the same logical
125		thread of execution, so no synchronization is required.
126
127		@see this_coro::environment, this_coro::executor,
128		this_coro::stop_token, this_coro::frame_allocator
129		@see io_env
130		@see IoAwaitable
131		*/
132		template<typename Derived>
133		class io_awaitable_promise_base
134		{
135		io_env const* env_ = nullptr;
136		mutable std::coroutine_handle<> cont_{std::noop_coroutine()};
137
138		public:
139		/** Allocate a coroutine frame.
140
141		Uses the thread-local frame allocator set by run_async.
142		Falls back to default memory resource if not set.
143		Stores the allocator pointer at the end of each frame for
144		correct deallocation even when TLS changes. Uses memcpy
145		to avoid alignment requirements on the trailing pointer.
146		Bypasses virtual dispatch for the recycling allocator.
147		*/
148	5015x	static void* operator new(std::size_t size)
149		{
150	5015x	static auto* const rmr = get_recycling_memory_resource();
151
152	5015x	auto* mr = get_current_frame_allocator();
153	5015x	if(!mr)
154	2795x	mr = std::pmr::get_default_resource();
155
156	5015x	auto total = size + sizeof(std::pmr::memory_resource*);
157		void* raw;
158	5015x	if(mr == rmr)
159		raw = static_cast<recycling_memory_resource*>(mr)
160	1857x	->allocate_fast(total, alignof(std::max_align_t));
161		else
162	3158x	raw = mr->allocate(total, alignof(std::max_align_t));
163	5015x	std::memcpy(static_cast<char*>(raw) + size, &mr, sizeof(mr));
164	5015x	return raw;
165		}
166
167		/** Deallocate a coroutine frame.
168
169		Reads the allocator pointer stored at the end of the frame
170		to ensure correct deallocation regardless of current TLS.
171		Bypasses virtual dispatch for the recycling allocator.
172		*/
173	5015x	static void operator delete(void* ptr, std::size_t size) noexcept
174		{
175	5015x	static auto* const rmr = get_recycling_memory_resource();
176
177		std::pmr::memory_resource* mr;
178	5015x	std::memcpy(&mr, static_cast<char*>(ptr) + size, sizeof(mr));
179	5015x	auto total = size + sizeof(std::pmr::memory_resource*);
180	5015x	if(mr == rmr)
181		static_cast<recycling_memory_resource*>(mr)
182	1857x	->deallocate_fast(ptr, total, alignof(std::max_align_t));
183		else
184	3158x	mr->deallocate(ptr, total, alignof(std::max_align_t));
185	5015x	}
186
187	5015x	~io_awaitable_promise_base()
188		{
189		// Abnormal teardown: destroy orphaned continuation
190	5015x	if(cont_ != std::noop_coroutine())
191	126x	cont_.destroy();
192	5015x	}
193
194		//----------------------------------------------------------
195		// Continuation support
196		//----------------------------------------------------------
197
198		/** Store the continuation to resume on completion.
199
200		Call this from your coroutine type's `await_suspend` overload
201		to set up the completion path. The `final_suspend` awaiter
202		returns this handle via unconditional symmetric transfer.
203
204		@param cont The continuation to resume on completion.
205		*/
206	4934x	void set_continuation(std::coroutine_handle<> cont) noexcept
207		{
208	4934x	cont_ = cont;
209	4934x	}
210
211		/** Return and consume the stored continuation handle.
212
213		Resets the stored handle to `noop_coroutine()` so the
214		destructor will not double-destroy it.
215
216		@return The continuation for symmetric transfer.
217		*/
218	4866x	std::coroutine_handle<> continuation() const noexcept
219		{
220	4866x	return std::exchange(cont_, std::noop_coroutine());
221		}
222
223		//----------------------------------------------------------
224		// Environment support
225		//----------------------------------------------------------
226
227		/** Store a pointer to the execution environment.
228
229		Call this from your coroutine type's `await_suspend`
230		overload to make the environment available via
231		`co_await this_coro::environment`. The pointed-to
232		`io_env` must outlive this coroutine.
233
234		@param env The environment to store.
235		*/
236	5012x	void set_environment(io_env const* env) noexcept
237		{
238	5012x	env_ = env;
239	5012x	}
240
241		/** Return the stored execution environment.
242
243		@return The environment.
244		*/
245	16481x	io_env const* environment() const noexcept
246		{
247	16481x	BOOST_CAPY_ASSERT(env_);
248	16481x	return env_;
249		}
250
251		/** Transform an awaitable before co_await.
252
253		Override this in your derived promise type to customize how
254		awaitables are transformed. The default implementation passes
255		the awaitable through unchanged.
256
257		@param a The awaitable expression from `co_await a`.
258
259		@return The transformed awaitable.
260		*/
261		template<typename A>
262		decltype(auto) transform_awaitable(A&& a)
263		{
264		return std::forward<A>(a);
265		}
266
267		/** Intercept co_await expressions.
268
269		This function handles @ref this_coro::environment_tag and
270		the fine-grained tags (@ref this_coro::executor_tag,
271		@ref this_coro::stop_token_tag, @ref this_coro::frame_allocator_tag)
272		specially, returning an awaiter that yields the stored value.
273		All other awaitables are delegated to @ref transform_awaitable.
274
275		@param t The awaited expression.
276
277		@return An awaiter for the expression.
278		*/
279		template<typename T>
280	9191x	auto await_transform(T&& t)
281		{
282		using Tag = std::decay_t<T>;
283
284		if constexpr (std::is_same_v<Tag, this_coro::environment_tag>)
285		{
286	18x	BOOST_CAPY_ASSERT(env_);
287		struct awaiter
288		{
289		io_env const* env_;
290	16x	bool await_ready() const noexcept { return true; }
291	2x	void await_suspend(std::coroutine_handle<>) const noexcept { }
292	15x	io_env const* await_resume() const noexcept { return env_; }
293		};
294	18x	return awaiter{env_};
295		}
296		else if constexpr (std::is_same_v<Tag, this_coro::executor_tag>)
297		{
298	3x	BOOST_CAPY_ASSERT(env_);
299		struct awaiter
300		{
301		executor_ref executor_;
302	2x	bool await_ready() const noexcept { return true; }
303		void await_suspend(std::coroutine_handle<>) const noexcept { }
304	2x	executor_ref await_resume() const noexcept { return executor_; }
305		};
306	3x	return awaiter{env_->executor};
307		}
308		else if constexpr (std::is_same_v<Tag, this_coro::stop_token_tag>)
309		{
310	15x	BOOST_CAPY_ASSERT(env_);
311		struct awaiter
312		{
313		std::stop_token token_;
314	14x	bool await_ready() const noexcept { return true; }
315	✗	void await_suspend(std::coroutine_handle<>) const noexcept { }
316	14x	std::stop_token await_resume() const noexcept { return token_; }
317		};
318	15x	return awaiter{env_->stop_token};
319		}
320		else if constexpr (std::is_same_v<Tag, this_coro::frame_allocator_tag>)
321		{
322	8x	BOOST_CAPY_ASSERT(env_);
323		struct awaiter
324		{
325		std::pmr::memory_resource* frame_allocator_;
326	6x	bool await_ready() const noexcept { return true; }
327	✗	void await_suspend(std::coroutine_handle<>) const noexcept { }
328	7x	std::pmr::memory_resource* await_resume() const noexcept { return frame_allocator_; }
329		};
330	8x	return awaiter{env_->frame_allocator};
331		}
332		else
333		{
334	6992x	return static_cast<Derived*>(this)->transform_awaitable(
335	9147x	std::forward<T>(t));
336		}
337		}
338		};
339
340		} // namespace capy
341		} // namespace boost
342
343		#endif
344