diff --git a/CMakeLists.txt b/CMakeLists.txt index be2c307110c..b84204ca99b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -57,6 +57,10 @@ option (Seastar_TASK_BACKTRACE "Collect backtrace at deferring points." OFF) +option (Seastar_TASK_CONTEXT + "Enable task context propagation for tracing and request-scoped state." + ON) + option (Seastar_DEBUG_ALLOCATIONS "For now just writes 0xab to newly allocated memory" OFF) @@ -967,6 +971,11 @@ if (Seastar_TASK_BACKTRACE) PUBLIC SEASTAR_TASK_BACKTRACE) endif () +if (Seastar_TASK_CONTEXT) + target_compile_definitions (seastar + PUBLIC SEASTAR_TASK_CONTEXT) +endif () + if (Seastar_DEBUG_ALLOCATIONS) target_compile_definitions (seastar PRIVATE SEASTAR_DEBUG_ALLOCATIONS) diff --git a/configure.py b/configure.py index dd02280ec48..dd24c611eb6 100755 --- a/configure.py +++ b/configure.py @@ -118,6 +118,11 @@ def standard_supported(standard, compiler='g++'): name='task-backtrace', dest='task_backtrace', help='Collect backtrace at deferring points') +add_tristate( + arg_parser, + name='task-context', + dest='task_context', + help='Enable task context propagation for tracing and request-scoped state') add_tristate( arg_parser, name='unused-result-error', @@ -207,6 +212,7 @@ def configure_mode(mode): tr(args.io_uring, 'IO_URING', value_when_none=None), tr(args.alloc_failure_injection, 'ALLOC_FAILURE_INJECTION', value_when_none='DEFAULT'), tr(args.task_backtrace, 'TASK_BACKTRACE'), + tr(args.task_context, 'TASK_CONTEXT'), tr(args.alloc_page_size, 'ALLOC_PAGE_SIZE'), tr(args.split_dwarf, 'SPLIT_DWARF'), tr(args.heap_profiling, 'HEAP_PROFILING'), diff --git a/include/seastar/core/smp.hh b/include/seastar/core/smp.hh index a542d316b02..7795ce01e27 100644 --- a/include/seastar/core/smp.hh +++ b/include/seastar/core/smp.hh @@ -209,7 +209,8 @@ class smp_message_queue { size_t _last_rcv_batch = 0; }; struct work_item : public task { - explicit work_item(smp_service_group ssg) : task(current_scheduling_group()), ssg(ssg) {} + explicit work_item(smp_service_group ssg) + : task(current_scheduling_group(), no_context_tag{}), ssg(ssg) {} smp_service_group ssg; virtual ~work_item() {} virtual void fail_with(std::exception_ptr) = 0; diff --git a/include/seastar/core/task.hh b/include/seastar/core/task.hh index 76b1f4c69e7..56e442c9497 100644 --- a/include/seastar/core/task.hh +++ b/include/seastar/core/task.hh @@ -22,12 +22,34 @@ #pragma once #include +#include #include #include namespace seastar { +/// \brief Base class for user-defined task context. +/// +/// Derive from this to attach application-specific context (e.g., +/// tracing) that propagates automatically through the task chain. +/// Use \ref with_context to scope a context to a function call. +/// +/// Context does NOT propagate across shards (smp::submit_to) since +/// lw_shared_ptr is not thread-safe. +/// +/// Guarded by the SEASTAR_TASK_CONTEXT compile flag. +struct task_context : enable_lw_shared_from_this { + virtual ~task_context() = default; +}; + +namespace internal { +#ifdef SEASTAR_TASK_CONTEXT +task_context*& current_task_context_ref() noexcept; +lw_shared_ptr inherit_task_context() noexcept; +#endif +} // namespace internal + class task { protected: scheduling_group _sg; @@ -35,6 +57,9 @@ private: #ifdef SEASTAR_TASK_BACKTRACE shared_backtrace _bt; #endif +#ifdef SEASTAR_TASK_CONTEXT + lw_shared_ptr _context; +#endif protected: // Task destruction is performed by run_and_dispose() via a concrete type, // so no need for a virtual destructor here. Derived classes that implement @@ -46,7 +71,18 @@ protected: return std::exchange(_sg, new_sg); } public: - explicit task(scheduling_group sg = current_scheduling_group()) noexcept : _sg(sg) {} + explicit task(scheduling_group sg = current_scheduling_group()) noexcept + : _sg(sg) +#ifdef SEASTAR_TASK_CONTEXT + , _context(internal::inherit_task_context()) +#endif + {} + + /// Tag type to construct a task without inheriting context. + /// Used by work_item (cross-shard) to avoid a wasted + /// inherit-then-clear cycle. + struct no_context_tag {}; + explicit task(scheduling_group sg, no_context_tag) noexcept : _sg(sg) {} virtual void run_and_dispose() noexcept = 0; /// Returns the next task which is waiting for this task to complete execution, or nullptr. virtual task* waiting_task() noexcept = 0; @@ -58,6 +94,16 @@ public: void make_backtrace() noexcept {} shared_backtrace get_backtrace() const { return {}; } #endif + +#ifdef SEASTAR_TASK_CONTEXT + task_context* context_ptr() const noexcept { return _context.get(); } + void set_context(lw_shared_ptr ctx) noexcept { + _context = std::move(ctx); + } +#else + task_context* context_ptr() const noexcept { return nullptr; } + void set_context(lw_shared_ptr) noexcept {} +#endif }; @@ -65,4 +111,39 @@ void schedule(task* t) noexcept; void schedule_checked(task* t) noexcept; void schedule_urgent(task* t) noexcept; +namespace internal { +#ifdef SEASTAR_TASK_CONTEXT +#ifndef SEASTAR_BUILD_SHARED_LIBS +inline task_context*& current_task_context_ref() noexcept { + static thread_local task_context* ptr = nullptr; + return ptr; +} +#endif + +inline lw_shared_ptr inherit_task_context() noexcept { + auto* p = current_task_context_ref(); + if (__builtin_expect(p != nullptr, false)) { + return p->shared_from_this(); + } + return {}; +} +#endif +} // namespace internal + +/// Returns the current task context, or nullptr if none is set. +#ifdef SEASTAR_TASK_CONTEXT +inline task_context* current_task_context() noexcept { + return internal::current_task_context_ref(); +} +/// Set the current task context (TLS). Prefer \ref with_context to +/// scope a context to a function call; this setter is primarily for +/// building custom RAII or awaitable helpers on top of the primitive. +inline void set_current_task_context(task_context* ctx) noexcept { + internal::current_task_context_ref() = ctx; +} +#else +inline task_context* current_task_context() noexcept { return nullptr; } +inline void set_current_task_context(task_context*) noexcept {} +#endif + } diff --git a/include/seastar/core/with_context.hh b/include/seastar/core/with_context.hh new file mode 100644 index 00000000000..c0c1c21ed5e --- /dev/null +++ b/include/seastar/core/with_context.hh @@ -0,0 +1,62 @@ +/* + * This file is open source software, licensed to you under the terms + * of the Apache License, Version 2.0 (the "License"). See the NOTICE file + * distributed with this work for additional information regarding copyright + * ownership. 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. + */ +/* + * Copyright 2026 Redpanda Data, Inc. + */ + +#pragma once + +#include +#include +#include + +namespace seastar { + +/// \addtogroup future-util +/// @{ + +/// \brief Run a callable with \c ctx installed as the current task context. +/// +/// TLS is set to \c ctx for the duration of \c func's synchronous +/// execution, so tasks created inline (coroutine promises, \c .then() +/// continuations, etc.) inherit \c ctx at construction and carry it +/// through their own suspensions via their own \c task._context. The +/// caller's \c task._context is never modified, so this is safe +/// regardless of what task is running when invoked. +/// +/// \param ctx context to install; empty clears the context for the scope. +/// \param func callable returning a future; its async work inherits \c ctx. +/// \param args forwarded to \c func. +template +inline +auto +with_context([[maybe_unused]] const lw_shared_ptr& ctx, + Func&& func, Args&&... args) noexcept { +#ifdef SEASTAR_TASK_CONTEXT + auto* prev = current_task_context(); + set_current_task_context(ctx.get()); + auto restore = defer([prev] () noexcept { + set_current_task_context(prev); + }); +#endif + return futurize_invoke(std::forward(func), std::forward(args)...); +} + +/// @} + +} // namespace seastar diff --git a/src/core/reactor.cc b/src/core/reactor.cc index b70a0a4cc41..ef4276eff54 100644 --- a/src/core/reactor.cc +++ b/src/core/reactor.cc @@ -2744,6 +2744,9 @@ bool reactor::task_queue::run_tasks() { STAP_PROBE(seastar, reactor_run_tasks_single_start); internal::task_histogram_add_task(*tsk); 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; STAP_PROBE(seastar, reactor_run_tasks_single_end); @@ -2771,6 +2774,9 @@ bool reactor::task_queue::run_tasks() { } } +#ifdef SEASTAR_TASK_CONTEXT + internal::current_task_context_ref() = nullptr; +#endif return !_q.empty(); } @@ -5176,6 +5182,14 @@ internal::current_scheduling_group_ptr() noexcept { static thread_local scheduling_group sg; return &sg; } + +#ifdef SEASTAR_TASK_CONTEXT +task_context*& +internal::current_task_context_ref() noexcept { + static thread_local task_context* ptr = nullptr; + return ptr; +} +#endif #endif const sstring& @@ -5469,6 +5483,9 @@ void log_timer_callback_exception(std::exception_ptr ex) noexcept { void set_current_task(task* t) { local_engine->_current_task = t; +#ifdef SEASTAR_TASK_CONTEXT + internal::current_task_context_ref() = t ? t->context_ptr() : nullptr; +#endif } } diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt index 413ca0b0c5d..6020809e4d8 100644 --- a/tests/unit/CMakeLists.txt +++ b/tests/unit/CMakeLists.txt @@ -840,6 +840,10 @@ seastar_add_test (gate SOURCES gate_test.cc) +seastar_add_test (task_context + SOURCES + task_context_test.cc) + seastar_add_test (test_fixture SOURCES test_fixture_test.cc) diff --git a/tests/unit/task_context_test.cc b/tests/unit/task_context_test.cc new file mode 100644 index 00000000000..0fd4cda0765 --- /dev/null +++ b/tests/unit/task_context_test.cc @@ -0,0 +1,495 @@ +/* + * This file is open source software, licensed to you under the terms + * of the Apache License, Version 2.0 (the "License"). See the NOTICE file + * distributed with this work for additional information regarding copyright + * ownership. 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. + */ +/* + * Copyright 2025 Redpanda Data, Inc. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace seastar; + +#ifdef SEASTAR_TASK_CONTEXT + +namespace { + +struct test_context : task_context { + int value = 0; + explicit test_context(int v) : value(v) {} +}; + +test_context* current_test_context() { + return static_cast(current_task_context()); +} + +lw_shared_ptr make_test_ctx(int value) { + return (new test_context(value))->shared_from_this(); +} + +} // anonymous namespace + +// TLS starts cleared on entry to a fresh task. +SEASTAR_TEST_CASE(test_context_default_null) { + BOOST_REQUIRE(!current_task_context()); + return make_ready_future<>(); +} + +// Direct set/get on a task's _context via the public accessors. +SEASTAR_TEST_CASE(test_context_task_set_and_get) { + auto* t = engine().current_task(); + auto ctx = make_test_ctx(42); + t->set_context(ctx); + BOOST_REQUIRE_EQUAL(t->context_ptr(), ctx.get()); + t->set_context({}); + BOOST_REQUIRE(!t->context_ptr()); + return make_ready_future<>(); +} + +// Direct set/get on TLS via the public accessors. +SEASTAR_TEST_CASE(test_context_tls_set_and_get) { + auto ctx = make_test_ctx(7); + set_current_task_context(ctx.get()); + BOOST_REQUIRE_EQUAL(current_task_context(), ctx.get()); + set_current_task_context(nullptr); + BOOST_REQUIRE(!current_task_context()); + return make_ready_future<>(); +} + +// lw_shared_ptr refcount semantics on task._context. +SEASTAR_TEST_CASE(test_context_refcount) { + auto* raw = new test_context(600); + auto ctx = raw->shared_from_this(); + BOOST_REQUIRE_EQUAL(raw->use_count(), 1); + + engine().current_task()->set_context(ctx); + BOOST_REQUIRE_EQUAL(raw->use_count(), 2); + + engine().current_task()->set_context({}); + BOOST_REQUIRE_EQUAL(raw->use_count(), 1); + + return make_ready_future<>(); +} + +// Polymorphic deletion via virtual destructor on task_context. +SEASTAR_TEST_CASE(test_context_virtual_destructor) { + bool destroyed = false; + + struct destructor_tracker : task_context { + bool& flag; + explicit destructor_tracker(bool& f) : flag(f) {} + ~destructor_tracker() override { flag = true; } + }; + + { + auto* raw = new destructor_tracker(destroyed); + auto ctx = raw->shared_from_this(); + BOOST_REQUIRE(!destroyed); + } + BOOST_REQUIRE(destroyed); + return make_ready_future<>(); +} + +// with_context installs ctx as TLS while func runs. Inline children +// inherit ctx via task construction. +SEASTAR_TEST_CASE(test_with_context_installs_tls_during_func) { + auto ctx = make_test_ctx(1); + BOOST_REQUIRE(!current_task_context()); + co_await with_context(ctx, [] () -> future<> { + BOOST_REQUIRE(current_test_context()); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1); + co_return; + }); + BOOST_REQUIRE(!current_task_context()); +} + +// If func is a coroutine, its promise is constructed with TLS=ctx, +// so promise._context=ctx from birth. Suspension + resume reads +// promise._context, so TLS comes back as ctx on resume. +SEASTAR_TEST_CASE(test_with_context_survives_co_await_in_func) { + auto ctx = make_test_ctx(2); + co_await with_context(ctx, [] () -> future<> { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 2); + co_await coroutine::maybe_yield(); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 2); + co_await yield(); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 2); + }); + BOOST_REQUIRE(!current_task_context()); +} + +// .then() continuations created inside func inherit ctx at +// construction via TLS, so they propagate correctly even after +// with_context returns. +SEASTAR_TEST_CASE(test_with_context_then_chain_inherits) { + auto ctx = make_test_ctx(3); + co_await with_context(ctx, [] { + return yield().then([] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 3); + }).then([] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 3); + }); + }); + BOOST_REQUIRE(!current_task_context()); +} + +// with_context must not modify the current task's _context, otherwise +// long-lived driver tasks would leak context across iterations. +SEASTAR_TEST_CASE(test_with_context_does_not_pollute_current_task) { + auto* t = engine().current_task(); + auto original = t->context_ptr(); + auto ctx = make_test_ctx(4); + co_await with_context(ctx, [] () -> future<> { + co_return; + }); + BOOST_REQUIRE_EQUAL(t->context_ptr(), original); + BOOST_REQUIRE(!current_task_context()); +} + +SEASTAR_TEST_CASE(test_with_context_nested) { + auto outer = make_test_ctx(10); + auto inner = make_test_ctx(20); + co_await with_context(outer, [inner] () -> future<> { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 10); + co_await with_context(inner, [] () -> future<> { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 20); + co_await coroutine::maybe_yield(); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 20); + }); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 10); + }); + BOOST_REQUIRE(!current_task_context()); +} + +SEASTAR_TEST_CASE(test_with_context_restores_tls_on_exception) { + auto ctx = make_test_ctx(5); + try { + co_await with_context(ctx, [] () -> future<> { + throw std::runtime_error("oops"); + co_return; + }); + BOOST_FAIL("expected exception"); + } catch (const std::runtime_error&) { + // expected + } + BOOST_REQUIRE(!current_task_context()); +} + +// Passing an empty context clears TLS for the duration of func, +// then restores it on return. +SEASTAR_TEST_CASE(test_with_context_empty_ctx) { + auto outer = make_test_ctx(99); + co_await with_context(outer, [] () -> future<> { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 99); + co_await with_context(lw_shared_ptr{}, [] () -> future<> { + BOOST_REQUIRE(!current_task_context()); + co_return; + }); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 99); + }); + BOOST_REQUIRE(!current_task_context()); +} + +// parallel_for_each sub-task branches see the enclosing context. +// Regression for the set_current_task sync added to ensure inline +// resumption via parallel_for_each doesn't drop TLS. +SEASTAR_TEST_CASE(test_with_context_parallel_for_each) { + auto ctx = make_test_ctx(6); + co_await with_context(ctx, [] () -> future<> { + std::vector items{1, 2, 3}; + co_await parallel_for_each(items, [] (int) -> future<> { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 6); + co_return; + }); + }); + BOOST_REQUIRE(!current_task_context()); +} + +// Context propagates into with_gate lambdas. +SEASTAR_TEST_CASE(test_with_context_gate) { + auto ctx = make_test_ctx(500); + co_await with_context(ctx, [] () -> future<> { + gate g; + auto holder = g.hold(); + auto f = with_gate(g, [] () -> future<> { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 500); + return make_ready_future<>(); + }); + co_await std::move(f); + holder.release(); + co_await g.close(); + }); + BOOST_REQUIRE(!current_task_context()); +} + +// Context propagates into seastar::async. +SEASTAR_TEST_CASE(test_with_context_async) { + auto ctx = make_test_ctx(400); + co_await with_context(ctx, [] () -> future<> { + co_await async([] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 400); + }); + }); + BOOST_REQUIRE(!current_task_context()); +} + +// Context does NOT cross shards: lw_shared_ptr is not thread-safe, +// and work_item uses no_context_tag to avoid inheriting. +SEASTAR_TEST_CASE(test_context_does_not_propagate_across_shards) { + if (smp::count < 2) { + co_return; + } + + auto ctx = make_test_ctx(300); + co_await with_context(ctx, [] () -> future<> { + auto other_shard = (this_shard_id() + 1) % smp::count; + bool remote_has_context = co_await smp::submit_to(other_shard, [] { + auto* t = engine().current_task(); + return t && t->context_ptr() != nullptr; + }); + BOOST_REQUIRE(!remote_has_context); + }); +} + +// A child task's captured lw_shared_ptr keeps the context alive +// after with_context returns. The child task's own _context owns a +// reference from its construction. +SEASTAR_TEST_CASE(test_context_child_outlives_caller) { + promise<> p; + auto f = make_ready_future<>(); + + co_await with_context(make_test_ctx(42), [&] () -> future<> { + f = p.get_future().then([] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 42); + }); + co_return; + }); + // TLS restored after with_context returns. + BOOST_REQUIRE(!current_task_context()); + p.set_value(); + co_await std::move(f); +} + +// Clearing the TLS mid-execution stops propagation to subsequently +// created tasks. +SEASTAR_TEST_CASE(test_context_clear_stops_propagation) { + auto ctx = make_test_ctx(700); + co_await with_context(ctx, [] () -> future<> { + set_current_task_context(nullptr); + co_await yield().then([] { + BOOST_REQUIRE(!current_test_context()); + }); + }); +} + +// Clearing context mid-.then() chain does NOT affect subsequent +// continuations because they inherited at construction time. +// Fixing this would require schedule-time propagation in add_task, +// which causes cross-contamination from unrelated promise resolvers. +SEASTAR_TEST_CASE(test_context_clear_mid_then_chain_does_not_propagate) { + auto ctx = make_test_ctx(800); + return with_context(ctx, [] { + return yield().then([] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 800); + set_current_task_context(nullptr); + }).then([] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 800); + }); + }); +} + +// All when_any branches share the same context object via pointer +// identity. Context stays alive while any branch holds a reference. +SEASTAR_TEST_CASE(test_context_when_any) { + auto* raw = new test_context(1100); + auto ctx = raw->shared_from_this(); + + co_await with_context(ctx, [raw] () -> future<> { + promise p1; + promise p2; + + auto f1 = p1.get_future().then([raw] (int v) { + BOOST_REQUIRE_EQUAL(current_test_context(), raw); + return v; + }); + auto f2 = p2.get_future().then([raw] (int v) { + BOOST_REQUIRE_EQUAL(current_test_context(), raw); + return v; + }); + + auto any = when_any(std::move(f1), std::move(f2)); + p1.set_value(10); + auto result = co_await std::move(any); + (void)result; + p2.set_value(20); + }); +} + +// Context survives an I/O completion (sleep-based here). +SEASTAR_TEST_CASE(test_context_survives_io_completion) { + auto ctx = make_test_ctx(1200); + co_await with_context(ctx, [] () -> future<> { + co_await sleep(std::chrono::milliseconds(1)); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1200); + }); +} + +// Context persists through a not-ready future (set_coroutine path). +SEASTAR_TEST_CASE(test_context_awaiter_not_ready) { + auto ctx = make_test_ctx(1300); + co_await with_context(ctx, [] () -> future<> { + co_await yield(); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1300); + }); +} + +// Context persists through a ready future (preemption/schedule path). +SEASTAR_TEST_CASE(test_context_awaiter_ready_preempt) { + auto ctx = make_test_ctx(1400); + co_await with_context(ctx, [] () -> future<> { + co_await make_ready_future<>(); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1400); + }); +} + +// Context persists through coroutine::as_future with a not-ready future. +SEASTAR_TEST_CASE(test_context_as_future_not_ready) { + auto ctx = make_test_ctx(1500); + co_await with_context(ctx, [] () -> future<> { + auto f = co_await coroutine::as_future(yield().then([] { return 42; })); + BOOST_REQUIRE(!f.failed()); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1500); + }); +} + +// Context persists through coroutine::as_future with a ready future. +SEASTAR_TEST_CASE(test_context_as_future_ready) { + auto ctx = make_test_ctx(1600); + co_await with_context(ctx, [] () -> future<> { + auto f = co_await coroutine::as_future(make_ready_future(42)); + BOOST_REQUIRE(!f.failed()); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1600); + }); +} + +// Context persists through coroutine::maybe_yield. +SEASTAR_TEST_CASE(test_context_maybe_yield) { + auto ctx = make_test_ctx(1700); + co_await with_context(ctx, [] () -> future<> { + co_await coroutine::maybe_yield(); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1700); + }); +} + +// Context persists across scheduling group switches. +SEASTAR_TEST_CASE(test_context_switch_to) { + auto ctx = make_test_ctx(1800); + co_await with_context(ctx, [] () -> future<> { + auto sg = co_await create_scheduling_group("test_ctx_sg", 100); + auto prev_sg = co_await coroutine::switch_to(sg); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1800); + co_await coroutine::switch_to(prev_sg); + co_await destroy_scheduling_group(sg); + }); +} + +// Context persists through coroutine::try_future with a not-ready +// future. +SEASTAR_TEST_CASE(test_context_try_future_not_ready) { + auto ctx = make_test_ctx(1900); + co_await with_context(ctx, [] () -> future<> { + co_await coroutine::try_future(yield().then([] { return 42; })); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1900); + }); +} + +// parallel_for_each resumes coroutines via set_current_task, not +// through the reactor loop. Regression for the TLS sync in +// set_current_task. +SEASTAR_TEST_CASE(test_context_parallel_for_each_set_current_task) { + auto ctx = make_test_ctx(2000); + co_await with_context(ctx, [] () -> future<> { + std::vector items{1, 2, 3}; + co_await parallel_for_each(items, [] (int) -> future<> { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 2000); + co_await yield(); + BOOST_REQUIRE_EQUAL(current_test_context()->value, 2000); + }); + }); +} + +// thread_context::switch_in calls set_current_task(nullptr). After +// a blocking .get() inside async, TLS must reflect the null state. +// Regression for the thread.cc sync. +SEASTAR_TEST_CASE(test_context_async_after_blocking_get) { + auto ctx = make_test_ctx(2100); + co_await with_context(ctx, [] () -> future<> { + co_await async([] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 2100); + yield().get(); + // After .get() resumes, set_current_task(nullptr) was called. + // TLS should be null. + BOOST_REQUIRE(!current_test_context()); + }); + }); +} + +// do_with keeps a with_context-wrapped continuation chain alive. +SEASTAR_TEST_CASE(test_with_context_do_with_chain) { + auto ctx = make_test_ctx(1000); + return with_context(ctx, [] { + return do_with(int{0}, [] (int& counter) { + return yield().then([&counter] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1000); + ++counter; + }).then([&counter] { + BOOST_REQUIRE_EQUAL(current_test_context()->value, 1000); + ++counter; + BOOST_REQUIRE_EQUAL(counter, 2); + }); + }); + }); +} + +#else // !SEASTAR_TASK_CONTEXT + +// When disabled, accessors compile and return null. +SEASTAR_TEST_CASE(test_context_noop_when_disabled) { + auto* t = engine().current_task(); + BOOST_REQUIRE(t); + BOOST_REQUIRE(!t->context_ptr()); + t->set_context({}); + BOOST_REQUIRE(!current_task_context()); + set_current_task_context(nullptr); + return make_ready_future<>(); +} + +#endif // SEASTAR_TASK_CONTEXT