task: add task_context propagation support (v2)

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)

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'),

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;

include/seastar/core/task.hh

@@ -22,19 +22,44 @@
 #pragma once
 
 #include <seastar/core/scheduling.hh>
+#include <seastar/core/shared_ptr.hh>
 #include <seastar/util/backtrace.hh>
 
 #include <utility>
 
 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<task_context> {
+    virtual ~task_context() = default;
+};
+
+namespace internal {
+#ifdef SEASTAR_TASK_CONTEXT
+task_context*& current_task_context_ref() noexcept;
+lw_shared_ptr<task_context> inherit_task_context() noexcept;
+#endif
+} // namespace internal
+
 class task {
 protected:
     scheduling_group _sg;
 private:
 #ifdef SEASTAR_TASK_BACKTRACE
     shared_backtrace _bt;
 #endif
+#ifdef SEASTAR_TASK_CONTEXT
+    lw_shared_ptr<task_context> _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,11 +94,56 @@ 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<task_context> ctx) noexcept {
+        _context = std::move(ctx);
+    }
+#else
+    task_context* context_ptr() const noexcept { return nullptr; }
+    void set_context(lw_shared_ptr<task_context>) noexcept {}
+#endif
 };
 
 
 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<task_context> 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
+
 }

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 <seastar/core/future.hh>
+#include <seastar/core/task.hh>
+#include <seastar/util/defer.hh>
+
+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 <typename Func, typename... Args>
+inline
+auto
+with_context([[maybe_unused]] const lw_shared_ptr<task_context>& 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>(func), std::forward<Args>(args)...);
+}
+
+/// @}
+
+} // namespace seastar

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
 }
 
 }

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)