task: add task_context propagation support#268
task: add task_context propagation support#268ballard26 wants to merge 1 commit intoredpanda-data:v26.1.xfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Adds optional task-scoped context (lw_shared_ptr<task_context>) to seastar::task and wires up automatic propagation through task creation and coroutine suspension points, guarded by a new SEASTAR_TASK_CONTEXT compile flag (enabled by default).
Changes:
- Introduces
seastar::task_contextand stores/propagates it viaseastar::taskconstructors andinternal::propagate_task_context(). - Updates coroutine awaiters to propagate context into coroutine promises at
await_suspend(). - Adds unit tests and build/config toggles for enabling/disabling the feature.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/task_context_test.cc | New unit tests validating task context APIs and propagation behavior. |
| tests/unit/CMakeLists.txt | Adds the new task_context unit test target. |
| src/core/reactor.cc | Implements internal::inherit_task_context() and internal::propagate_task_context(). |
| include/seastar/core/task.hh | Adds task_context, stores it in task, and exposes get/set/take APIs plus propagation hooks. |
| include/seastar/core/coroutine.hh | Coroutine promise skips ctor inheritance; awaiters propagate context in await_suspend(). |
| include/seastar/coroutine/as_future.hh | Propagates context in as_future awaiter suspend path. |
| include/seastar/coroutine/maybe_yield.hh | Propagates context before scheduling resume. |
| include/seastar/coroutine/switch_to.hh | Propagates context before scheduling-group switch scheduling. |
| include/seastar/coroutine/parallel_for_each.hh | Propagates context before resuming the awaiting coroutine. |
| include/seastar/coroutine/try_future.hh | Propagates context before resuming/destroying coroutine in success path and ready path. |
| include/seastar/core/smp.hh | Prevents cross-shard inheritance for work_item via no_context_tag. |
| configure.py | Adds --task-context tristate configure option. |
| CMakeLists.txt | Adds Seastar_TASK_CONTEXT option and exports SEASTAR_TASK_CONTEXT define when enabled. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
ballard26
force-pushed
the
task-context
branch
3 times, most recently
from
543458b to
1244c4d
Compare
|
Can you use benches as changed here: 4fa3c1a and also rp multiplexer bench as a more real life bench |
ballard26
force-pushed
the
task-context
branch
2 times, most recently
from
56713e4 to
88e9361
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| } | ||
|
|
||
| scoped_context_guard(scoped_context_guard&&) noexcept = default; | ||
| scoped_context_guard& operator=(scoped_context_guard&&) noexcept = default; |
There was a problem hiding this comment.
scoped_context_guard defaulted move-assignment is not safe for an RAII guard: move-assigning into an already-active guard will overwrite _prev/_ctx without first restoring the currently-installed TLS, potentially leaving TLS permanently set to the wrong context. Consider deleting move-assignment (like many other scoped RAII helpers in the codebase) or implementing a custom move-assignment that restores the current guard before taking ownership.
| scoped_context_guard& operator=(scoped_context_guard&&) noexcept = default; | |
| scoped_context_guard& operator=(scoped_context_guard&&) = delete; |
| r._current_task = tsk; | ||
| #ifdef SEASTAR_TASK_CONTEXT | ||
| internal::current_task_context_ref() = tsk->context_ptr(); | ||
| #endif | ||
| tsk->run_and_dispose(); | ||
| r._current_task = nullptr; | ||
| #ifdef SEASTAR_TASK_CONTEXT | ||
| internal::current_task_context_ref() = nullptr; | ||
| #endif |
There was a problem hiding this comment.
Task-context TLS is set/cleared when the reactor runs tasks, but internal::set_current_task() is used elsewhere (e.g. coroutine helpers like include/seastar/coroutine/parallel_for_each.hh and try_future.hh, and src/core/thread.cc) to change _current_task without returning to the reactor loop. Since set_current_task() does not update internal::current_task_context_ref(), current_task_context() can become stale (or non-null when _current_task is set to null). Please update set_current_task() to also set current_task_context_ref() to t ? t->context_ptr() : nullptr when SEASTAR_TASK_CONTEXT is enabled.
| r._current_task = tsk; | |
| #ifdef SEASTAR_TASK_CONTEXT | |
| internal::current_task_context_ref() = tsk->context_ptr(); | |
| #endif | |
| tsk->run_and_dispose(); | |
| r._current_task = nullptr; | |
| #ifdef SEASTAR_TASK_CONTEXT | |
| internal::current_task_context_ref() = nullptr; | |
| #endif | |
| internal::set_current_task(tsk); | |
| tsk->run_and_dispose(); | |
| internal::set_current_task(nullptr); |
| #include <seastar/core/task.hh> | ||
|
|
||
| #include <coroutine> | ||
|
|
||
| namespace seastar { | ||
|
|
||
| /// \brief RAII guard for task context in .then() chains. | ||
| /// | ||
| /// Sets the TLS on construction, restores on destruction. New tasks | ||
| /// created inline inherit context via the task constructor. Use with | ||
| /// do_with to keep the guard alive across a .then() chain. | ||
| /// | ||
| /// Not awaitable. For coroutines, use coroutine::context_guard. | ||
| template<typename Derived> | ||
| class scoped_context_guard { | ||
| task_context* _prev = nullptr; | ||
| protected: | ||
| lw_shared_ptr<task_context> _ctx; | ||
|
|
||
| explicit scoped_context_guard(lw_shared_ptr<task_context> ctx) { | ||
| if (ctx) { | ||
| _prev = internal::current_task_context_ref(); | ||
| internal::current_task_context_ref() = ctx.get(); | ||
| _ctx = std::move(ctx); |
There was a problem hiding this comment.
context_guard.hh uses internal::current_task_context_ref() unconditionally, but current_task_context_ref() only exists when SEASTAR_TASK_CONTEXT is enabled. As a result, including this header with task-context disabled will fail to compile (and tests/unit/task_context_test.cc includes it outside the #ifdef). Please wrap the implementation in #ifdef SEASTAR_TASK_CONTEXT and provide no-op guard templates (or stub current_task_context_ref() access) for the disabled build so this header remains includable.
|
Here's the microbench results from the latest ver:
Will post the record_multiplexer benchmarks shortly. |
Add an optional lw_shared_ptr<task_context> field on seastar::task
that propagates automatically through the task chain, gated by the
SEASTAR_TASK_CONTEXT compile flag. task_context is a polymorphic
base (enable_lw_shared_from_this + virtual dtor) that users subclass
to attach per-chain state such as tracing spans. The field adds
8 bytes to task when the flag is enabled and is compiled out
otherwise.
Propagation uses a thread-local raw pointer as the authority,
mirroring how current_scheduling_group() works:
- Reactor loop: sets TLS from task._context before run_and_dispose,
clears after the loop iteration.
- Task constructor: calls inherit_task_context() which reads TLS
and shared_from_this()'s it into the new task's _context.
- set_current_task(): keeps TLS in sync with the current task for
paths that bypass the reactor loop (parallel_for_each inline
resumption, thread_context::switch_in, thread.cc blocking .get()).
- smp: work_item uses a no_context_tag to construct its task
without inheriting, since context cannot cross shards (lw_shared_ptr
is not thread-safe).
The primary user-facing API is with_context(ctx, func, args...) —
a function wrapper that installs ctx as TLS for the duration of
func's synchronous execution. Inline work created by func (coroutine
promises, .then() continuations, parallel_for_each branches) inherits
ctx at construction time via TLS, and their own task._context carries
it across suspensions. The caller's task is never modified, so the
function is safe inside coroutines, long-lived driver tasks, and
anywhere else.
current_task_context() and set_current_task_context() expose the
TLS accessors publicly so users can build custom RAII or awaitable
helpers on top when ergonomics demand it.
Adds tests/unit/task_context_test.cc covering the primitive accessors,
propagation across .then() chains and coroutine suspensions, gate,
async, parallel_for_each, when_any, scheduling group switches,
as_future, try_future, maybe_yield, cross-shard non-propagation,
refcount semantics, and regression coverage for the set_current_task
and thread.cc sync points.
This PR adds a lw_shared_ptr<task_context> field to seastar::task that propagates
automatically through the task chain. task_context is a polymorphic base
class (using enable_lw_shared_from_this<task_context> with a virtual
destructor) that downstream users subclass for their specific needs. The
field adds 8 bytes to the task struct.
Context propagation uses a thread-local pointer (current_task_context_ptr)
as the authority, matching how current_scheduling_group() works:
Reactor loop: sets TLS from task._context before each task runs,
clears after.
Task constructor: reads TLS into _context at creation time (inline).
context_guard CRTP bases: two variants for installing context:
coroutine::context_guard — co_await-able, sets TLS on
construction and the coroutine promise via hndl.promise() on
co_await. Destructor restores TLS and syncs back to the promise.
Move constructor is private to prevent dangling promise pointers.
scoped_context_guard — plain RAII for .then() chains,
sets/restores TLS only. Fully movable for do_with.
Cross-shard work items (smp::submit_to) use no_context_tag to skip
context inheritance since lw_shared_ptr is not thread-safe across shards.
Known limitation: context changes inside a .then() continuation (e.g.
creating a child span or clearing context) do not propagate to subsequent
continuations in the same chain, because all continuations are constructed
at chain-setup time and each inherits its own copy. Fixing this would
require schedule-time propagation in add_task(), which introduces
cross-contamination from unrelated tasks that resolve shared promises.
The tracing use case is unaffected — RAII guards span the entire chain
lifetime via do_with().
The feature is guarded by the SEASTAR_TASK_CONTEXT compile flag (ON by
default), following the existing SEASTAR_TASK_BACKTRACE pattern so it
compiles out with zero overhead for users who don't need it.