feat engine: move local variables inheritance to TaskContext

Changelog entry:

* Allowed to mark any @ref engine::TaskInheritedVariable as @ref engine::TaskInheritedVariablePriority::kBackground, which means that it spills even into @ref utils::AsyncBackground tasks. Previously, only baggage inherited variable had this quality.
* Added @ref engine::TaskLocalVariable::GetOrEmplace, which allows to perform essentially non-default initialization of variables, and more.

Internal engine refactorings (probably not suitable for changelog):

* Moved handling the inheritance of `engine::thread_local::Storage` from `SpanWrapCall` to `TaskContext` constructor.
  * It allowed to properly order variable inheritance w.r.t. plugin hooks.
  * It opens the door to add optional support for inherited variables for `utils::TaskBuilder::NoTracing`.
commit_hash:dbba4ee958e1d24b261025caa3b815fd38889927

Author: Anton3

.mapping.json

@@ -1192,6 +1192,7 @@
   "core/include/userver/engine/task/cancel.hpp":"taxi/uservices/userver/core/include/userver/engine/task/cancel.hpp",
   "core/include/userver/engine/task/current_task.hpp":"taxi/uservices/userver/core/include/userver/engine/task/current_task.hpp",
   "core/include/userver/engine/task/inherited_variable.hpp":"taxi/uservices/userver/core/include/userver/engine/task/inherited_variable.hpp",
+  "core/include/userver/engine/task/inherited_variable_options.hpp":"taxi/uservices/userver/core/include/userver/engine/task/inherited_variable_options.hpp",
   "core/include/userver/engine/task/local_variable.hpp":"taxi/uservices/userver/core/include/userver/engine/task/local_variable.hpp",
   "core/include/userver/engine/task/shared_task.hpp":"taxi/uservices/userver/core/include/userver/engine/task/shared_task.hpp",
   "core/include/userver/engine/task/shared_task_with_result.hpp":"taxi/uservices/userver/core/include/userver/engine/task/shared_task_with_result.hpp",

core/include/userver/baggage/baggage.hpp

@@ -149,7 +149,7 @@ bool HasInvalidSymbols(const T& obj) {
            }) != obj.end();
 }
 
-inline engine::TaskInheritedVariable<Baggage> kInheritedBaggage;
+inline engine::TaskInheritedVariable<Baggage> kInheritedBaggage{engine::TaskInheritedVariablePriority::kBackground};
 
 }  // namespace baggage
 

core/include/userver/engine/async.hpp

@@ -38,7 +38,10 @@ namespace engine {
 template <typename Function, typename... Args>
 [[nodiscard]] auto AsyncNoTracing(TaskProcessor& task_processor, Function&& f, Args&&... args) {
     return impl::MakeTaskWithResult<TaskWithResult>(
-        impl::TaskConfig{.task_processor = &task_processor},
+        impl::TaskConfig{
+            .task_processor = &task_processor,
+            .inherited_variables_priority = TaskInheritedVariablePriority::kNone,
+        },
         std::forward<Function>(f),
         std::forward<Args>(args)...
     );
@@ -47,8 +50,11 @@ template <typename Function, typename... Args>
 /// @overload
 template <typename Function, typename... Args>
 [[nodiscard]] auto AsyncNoTracing(Function&& f, Args&&... args) {
-    return impl::MakeTaskWithResult<
-        TaskWithResult>(impl::TaskConfig{}, std::forward<Function>(f), std::forward<Args>(args)...);
+    return impl::MakeTaskWithResult<TaskWithResult>(
+        impl::TaskConfig{.inherited_variables_priority = TaskInheritedVariablePriority::kNone},
+        std::forward<Function>(f),
+        std::forward<Args>(args)...
+    );
 }
 
 /// @overload
@@ -59,6 +65,7 @@ template <typename Function, typename... Args>
         impl::TaskConfig{
             .task_processor = &task_processor,
             .importance = Task::Importance::kCritical,
+            .inherited_variables_priority = TaskInheritedVariablePriority::kNone,
         },
         std::forward<Function>(f),
         std::forward<Args>(args)...
@@ -70,7 +77,10 @@ template <typename Function, typename... Args>
 template <typename Function, typename... Args>
 [[nodiscard]] auto CriticalAsyncNoTracing(Function&& f, Args&&... args) {
     return impl::MakeTaskWithResult<TaskWithResult>(
-        impl::TaskConfig{.importance = Task::Importance::kCritical},
+        impl::TaskConfig{
+            .importance = Task::Importance::kCritical,
+            .inherited_variables_priority = TaskInheritedVariablePriority::kNone,
+        },
         std::forward<Function>(f),
         std::forward<Args>(args)...
     );

core/include/userver/engine/impl/task_context_factory.hpp

@@ -6,6 +6,9 @@
 #include <utility>
 
 #include <userver/engine/impl/task_context_holder.hpp>
+#include <userver/engine/impl/task_local_storage.hpp>
+#include <userver/engine/task/current_task.hpp>
+#include <userver/engine/task/inherited_variable_options.hpp>
 #include <userver/engine/task/task.hpp>
 #include <userver/utils/fast_scope_guard.hpp>
 #include <userver/utils/impl/wrapped_call.hpp>
@@ -25,6 +28,9 @@ struct TaskConfig final {
     Task::Importance importance{Task::Importance::kNormal};
     Task::WaitMode wait_mode{Task::WaitMode::kSingleAwaiter};
     engine::Deadline deadline{};
+    // The lower bound of priority of variables that the new task inherits from the creating task.
+    // The parent storage itself is captured synchronously on the creating thread.
+    TaskInheritedVariablePriority inherited_variables_priority{TaskInheritedVariablePriority::kNormal};
 };
 
 [[nodiscard]] TaskContext& PlacementNewTaskContext(

core/include/userver/engine/impl/task_local_storage.hpp

@@ -6,6 +6,7 @@
 #include <type_traits>
 #include <typeinfo>
 
+#include <userver/engine/task/inherited_variable_options.hpp>
 #include <userver/utils/fast_pimpl.hpp>
 
 USERVER_NAMESPACE_BEGIN
@@ -14,7 +15,7 @@ namespace engine::impl::task_local {
 
 using Key = std::size_t;
 
-enum class VariableKind { kNormal, kInherited };
+enum class VariableKind : std::uint8_t { kNormal, kInherited };
 
 class DataBase {
 public:
@@ -89,40 +90,29 @@ class Storage final {
     Storage(Storage&&) = delete;
     Storage& operator=(Storage&&) = delete;
 
-    // Requires that the variables have already been destroyed using
-    // DestroyVariables, or that there are none.
+    // Requires that the normal (local) variables have already been destroyed
+    // using DestroyVariables. Any remaining inherited variables (which can be
+    // left over if the task never finished normally) are destroyed here.
     ~Storage();
 
     // Destroys the variables in reverse-initialization order.
-    // Must be called before the destructor. Must be called in the coroutine
-    // owning `this`: the destructors may want to sleep, and the storage must
-    // still be usable (e.g. via GetCurrentStorage) while they run.
+    // Must be called for normal (local) variables before the destructor, in the
+    // coroutine owning `this`: their destructors may want to sleep, and the
+    // storage must still be usable (e.g. via GetCurrentStorage) while they run.
     void DestroyVariables() noexcept;
 
-    // Copies pointers to inherited variables from 'other'
-    // 'this' must not contain any variables
-    void InheritFrom(Storage& other);
+    // Copies pointers to inherited variables from 'other'.
+    // 'this' must not contain any variables.
+    // A variable is copied when its priority is not less than 'priority'.
+    void InheritFrom(Storage& other, TaskInheritedVariablePriority priority);
 
-    // Copies pointers to specific inherited variables from 'other'
-    // 'this' must not contain the variable being copied
-    // does nothing if there is nothing to copy
-    void InheritNodeIfExists(Storage& other, Key key);
-
-    // Moves other's variables into 'this', leaving 'other' with no variables.
-    // 'this' must not contain inherited variables. It may contain normal
-    // variables (e.g. task-locals initialized by engine plugin hooks before
-    // the task payload starts); those are kept. In that case 'other' must not
-    // contain normal variables.
-    void InitializeFrom(Storage&& other) noexcept;
-
-    // Variable accessors must be called with the same T, Kind, key.
-    // Otherwise it is UB.
-    template <typename T, VariableKind Kind>
-    T& GetOrEmplace(Key key) {
+    // Multiple calls with the same key should have the same T and Kind, otherwise it is UB.
+    template <typename T, VariableKind Kind, typename... Args>
+    T& GetOrEmplace(Key key, Args&&... args) {
         DataBase* const old_data = GetGeneric(key);
         if (!old_data) {
             const bool has_existing_variable = false;
-            return DoEmplace<T, Kind>(key, has_existing_variable);
+            return DoEmplace<T, Kind>(key, has_existing_variable, std::forward<Args>(args)...);
         }
         return static_cast<DataImpl<T, Kind>&>(*old_data).Get();
     }
@@ -188,7 +178,7 @@ class Storage final {
 
 class Variable final {
 public:
-    Variable();
+    explicit Variable(TaskInheritedVariablePriority priority = TaskInheritedVariablePriority::kNormal);
 
     Variable(const Variable&) = delete;
     Variable(Variable&&) = delete;
@@ -203,10 +193,6 @@ class Variable final {
 
 Storage& GetCurrentStorage() noexcept;
 
-struct InternalTag final {
-    explicit InternalTag() = default;
-};
-
 }  // namespace engine::impl::task_local
 
 USERVER_NAMESPACE_END

core/include/userver/engine/task/inherited_variable.hpp

@@ -7,6 +7,7 @@
 #include <utility>
 
 #include <userver/engine/impl/task_local_storage.hpp>
+#include <userver/engine/task/inherited_variable_options.hpp>
 
 USERVER_NAMESPACE_BEGIN
 
@@ -19,13 +20,26 @@ namespace engine {
 /// These are like engine::TaskLocalVariable, but the variable instances are
 /// inherited by child tasks created via utils::Async.
 ///
+/// By default a variable is inherited only by regular child tasks, not @ref utils::AsyncBackground.
+/// This behavior can be changed by passing a different priority to the constructor.
+///
 /// The order of destruction of task-inherited variables is unspecified.
 template <typename T>
 class TaskInheritedVariable final {
     static_assert(!std::is_reference_v<T>);
     static_assert(!std::is_const_v<T>);
 
 public:
+    /// @brief Create a task-inherited variable with @ref TaskInheritedVariablePriority::kNormal.
+    TaskInheritedVariable()
+        : TaskInheritedVariable(TaskInheritedVariablePriority::kNormal)
+    {}
+
+    /// @brief Create a task-inherited variable with the specified inheritance priority.
+    explicit TaskInheritedVariable(TaskInheritedVariablePriority priority)
+        : impl_(priority)
+    {}
+
     /// @brief Get the variable instance for the current task.
     /// @returns the variable or `nullptr` if variable was not set.
     const T* GetOptional() const noexcept { return Storage().GetOptional<T, kVariableKind>(impl_.GetKey()); }
@@ -51,14 +65,6 @@ class TaskInheritedVariable final {
     /// @note The variable might not actually be destroyed immediately.
     void Erase() { Storage().Erase<T, kVariableKind>(impl_.GetKey()); }
 
-    /// @cond
-    // For internal use only
-    // inherits data to another storage
-    void InheritTo(impl::task_local::Storage& other, impl::task_local::InternalTag) {
-        other.InheritNodeIfExists(Storage(), impl_.GetKey());
-    }
-    /// @endcond
-
 private:
     static constexpr auto kVariableKind = impl::task_local::VariableKind::kInherited;
 

core/include/userver/engine/task/inherited_variable_options.hpp

@@ -0,0 +1,29 @@
+#pragma once
+
+/// @file userver/engine/task/inherited_variable_options.hpp
+/// @brief @copybrief engine::TaskInheritedVariablePriority
+
+#include <cstdint>
+
+USERVER_NAMESPACE_BEGIN
+
+namespace engine {
+
+/// @brief Controls whether a @ref engine::TaskInheritedVariable instance is inherited from the creating task.
+enum class TaskInheritedVariablePriority : std::uint8_t {
+    /// Default priority for @ref engine::TaskInheritedVariable. Propagates to normal tasks.
+    kNormal = 0,
+    /// The minimum priority for @ref engine::TaskInheritedVariable instances to propagate to background tasks
+    /// (e.g. @ref utils::AsyncBackground).
+    kBackground = 1,
+    /// The minimum priority for @ref engine::TaskInheritedVariable instances to propagate to no-tracing tasks
+    /// (e.g. @ref engine::AsyncNoTracing).
+    /// As of now, this priority cannot be assigned to a variable.
+    kNoTracing = 2,
+    /// Do not inherit any variables.
+    kNone = 2,
+};
+
+}  // namespace engine
+
+USERVER_NAMESPACE_END

core/include/userver/engine/task/local_variable.hpp

@@ -34,6 +34,15 @@ class TaskLocalVariable final {
     /// @overload
     T* operator->();
 
+    /// @brief Get the variable instance for the current task, or emplace a new
+    /// one with the given arguments.
+    /// @note Must be called from a coroutine, otherwise it is UB.
+    template <typename... Args>
+    T& GetOrEmplace(Args&&... args) {
+        return impl::task_local::GetCurrentStorage()
+            .GetOrEmplace<T, kVariableKind>(impl_.GetKey(), std::forward<Args>(args)...);
+    }
+
     /// @brief Get the variable instance for the current task.
     /// @returns the variable or `nullptr` if the variable was not initialized,
     /// was already destroyed, or is being destroyed right now. That is,

core/include/userver/utils/async.hpp

@@ -74,7 +74,9 @@ template <typename Function, typename... Args>
 template <typename Function, typename... Args>
 [[nodiscard]] auto Async(engine::TaskProcessor& task_processor, std::string name, Function&& f, Args&&... args) {
     return engine::impl::MakeTaskWithResult<engine::TaskWithResult>(
-        engine::impl::TaskConfig{.task_processor = &task_processor},
+        engine::impl::TaskConfig{
+            .task_processor = &task_processor,
+        },
         utils::impl::SpanLazyPrvalue(std::move(name)),
         std::forward<Function>(f),
         std::forward<Args>(args)...
@@ -98,11 +100,7 @@ template <typename Function, typename... Args>
 [[nodiscard]] auto AsyncHideSpan(Function&& f, Args&&... args) {
     return engine::impl::MakeTaskWithResult<engine::TaskWithResult>(
         engine::impl::TaskConfig{},
-        utils::impl::SpanLazyPrvalue(
-            std::string{},
-            utils::impl::SpanWrapCall::InheritVariables::kYes,
-            utils::impl::SpanWrapCall::HideSpan::kYes
-        ),
+        utils::impl::SpanLazyPrvalue(std::string{}, utils::impl::SpanWrapCall::HideSpan::kYes),
         std::forward<Function>(f),
         std::forward<Args>(args)...
     );
@@ -125,12 +123,10 @@ template <typename Function, typename... Args>
 template <typename Function, typename... Args>
 [[nodiscard]] auto AsyncHideSpan(engine::TaskProcessor& task_processor, Function&& f, Args&&... args) {
     return engine::impl::MakeTaskWithResult<engine::TaskWithResult>(
-        engine::impl::TaskConfig{.task_processor = &task_processor},
-        utils::impl::SpanLazyPrvalue(
-            std::string{},
-            utils::impl::SpanWrapCall::InheritVariables::kYes,
-            utils::impl::SpanWrapCall::HideSpan::kYes
-        ),
+        engine::impl::TaskConfig{
+            .task_processor = &task_processor,
+        },
+        utils::impl::SpanLazyPrvalue(std::string{}, utils::impl::SpanWrapCall::HideSpan::kYes),
         std::forward<Function>(f),
         std::forward<Args>(args)...
     );
@@ -179,7 +175,9 @@ template <typename Function, typename... Args>
 template <typename Function, typename... Args>
 [[nodiscard]] auto SharedAsync(engine::TaskProcessor& task_processor, std::string name, Function&& f, Args&&... args) {
     return engine::impl::MakeTaskWithResult<engine::SharedTaskWithResult>(
-        engine::impl::TaskConfig{.task_processor = &task_processor},
+        engine::impl::TaskConfig{
+            .task_processor = &task_processor,
+        },
         utils::impl::SpanLazyPrvalue(std::move(name)),
         std::forward<Function>(f),
         std::forward<Args>(args)...
@@ -199,7 +197,9 @@ template <typename Function, typename... Args>
 template <typename Function, typename... Args>
 [[nodiscard]] auto CriticalAsync(std::string name, Function&& f, Args&&... args) {
     return engine::impl::MakeTaskWithResult<engine::TaskWithResult>(
-        engine::impl::TaskConfig{.importance = engine::Task::Importance::kCritical},
+        engine::impl::TaskConfig{
+            .importance = engine::Task::Importance::kCritical,
+        },
         utils::impl::SpanLazyPrvalue(std::move(name)),
         std::forward<Function>(f),
         std::forward<Args>(args)...
@@ -219,7 +219,9 @@ template <typename Function, typename... Args>
 template <typename Function, typename... Args>
 [[nodiscard]] auto SharedCriticalAsync(std::string name, Function&& f, Args&&... args) {
     return engine::impl::MakeTaskWithResult<engine::SharedTaskWithResult>(
-        engine::impl::TaskConfig{.importance = engine::Task::Importance::kCritical},
+        engine::impl::TaskConfig{
+            .importance = engine::Task::Importance::kCritical,
+        },
         utils::impl::SpanLazyPrvalue(std::move(name)),
         std::forward<Function>(f),
         std::forward<Args>(args)...
@@ -260,7 +262,9 @@ template <typename Function, typename... Args>
 template <typename Function, typename... Args>
 [[nodiscard]] auto Async(std::string name, engine::Deadline deadline, Function&& f, Args&&... args) {
     return engine::impl::MakeTaskWithResult<engine::TaskWithResult>(
-        engine::impl::TaskConfig{.deadline = deadline},
+        engine::impl::TaskConfig{
+            .deadline = deadline,
+        },
         utils::impl::SpanLazyPrvalue(std::move(name)),
         std::forward<Function>(f),
         std::forward<Args>(args)...
@@ -305,12 +309,11 @@ template <typename Function, typename... Args>
     Args&&... args
 ) {
     return engine::impl::MakeTaskWithResult<engine::TaskWithResult>(
-        engine::impl::TaskConfig{.task_processor = &task_processor},
-        utils::impl::SpanLazyPrvalue(
-            std::move(name),
-            utils::impl::SpanWrapCall::InheritVariables::kNo,
-            utils::impl::SpanWrapCall::HideSpan::kNo
-        ),
+        engine::impl::TaskConfig{
+            .task_processor = &task_processor,
+            .inherited_variables_priority = engine::TaskInheritedVariablePriority::kBackground,
+        },
+        utils::impl::SpanLazyPrvalue(std::move(name)),
         std::forward<Function>(f),
         std::forward<Args>(args)...
     );
@@ -340,12 +343,9 @@ template <typename Function, typename... Args>
         engine::impl::TaskConfig{
             .task_processor = &task_processor,
             .importance = engine::Task::Importance::kCritical,
+            .inherited_variables_priority = engine::TaskInheritedVariablePriority::kBackground,
         },
-        utils::impl::SpanLazyPrvalue(
-            std::move(name),
-            utils::impl::SpanWrapCall::InheritVariables::kNo,
-            utils::impl::SpanWrapCall::HideSpan::kNo
-        ),
+        utils::impl::SpanLazyPrvalue(std::move(name)),
         std::forward<Function>(f),
         std::forward<Args>(args)...
     );