fix engine: support accessing local variables in destructor of another local variable

* Fixed UB in case a @ref engine::TaskLocalVariable is accessed in the destructor of another @ref engine::TaskLocalVariable. This use case is now fully supported, just like C++ supports this for `thread_local` variables.
commit_hash:92859d762135387d41b26ea16d6fcee7e39fe28c

Author: Anton3

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

@@ -6,7 +6,6 @@
 #include <boost/intrusive/list_hook.hpp>
 
 #include <userver/engine/impl/awaiter_fwd.hpp>
-#include <userver/engine/impl/epoch.hpp>
 #include <userver/utils/assert.hpp>
 
 USERVER_NAMESPACE_BEGIN

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

@@ -88,8 +88,17 @@ 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.
     ~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.
+    void DestroyVariables() noexcept;
+
     // Copies pointers to inherited variables from 'other'
     // 'this' must not contain any variables
     void InheritFrom(Storage& other);
@@ -99,8 +108,11 @@ class Storage final {
     // does nothing if there is nothing to copy
     void InheritNodeIfExists(Storage& other, Key key);
 
-    // Moves other's variables into 'this', leaving 'other' in an empty state
-    // 'this' must not contain any variables
+    // 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.

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

@@ -35,7 +35,11 @@ class TaskLocalVariable final {
     T* operator->();
 
     /// @brief Get the variable instance for the current task.
-    /// @returns the variable or `nullptr` if variable was not initialized.
+    /// @returns the variable or `nullptr` if the variable was not initialized,
+    /// was already destroyed, or is being destroyed right now. That is,
+    /// a non-null result is guaranteed to point to a variable whose
+    /// destruction has not started (the variable is unset before its
+    /// destructor is invoked, as in POSIX `pthread_getspecific`).
     T* GetOptional() noexcept {
         return impl::task_local::GetCurrentStorage().GetOptional<T, kVariableKind>(impl_.GetKey());
     }

core/src/engine/impl/task_local_storage.cpp

@@ -1,6 +1,7 @@
 #include <userver/engine/impl/task_local_storage.hpp>
 
 #include <ranges>
+#include <utility>
 
 #include <fmt/format.h>
 #include <boost/intrusive/list.hpp>
@@ -82,20 +83,42 @@ struct Storage::Impl final {
 Storage::Storage() { utils::impl::AssertStaticRegistrationFinished(); }
 
 Storage::~Storage() {
+    UASSERT_MSG(
+        impl_->normal_data_storage.empty() && impl_->inherited_data_storage.empty(),
+        "Storage::DestroyVariables must be called before the Storage destructor. The destructors of task-local "
+        "variables may sleep and access the storage, so they must not run while the Storage (typically stored in "
+        "a std::optional) is being destroyed"
+    );
+}
+
+void Storage::DestroyVariables() noexcept {
     const auto disposer = [](DataPtr* node_ptr) noexcept {
         UASSERT(node_ptr->ptr);
-        node_ptr->ptr->DeleteSelf();
+        // Unset the variable before running its destructor (POSIX
+        // pthread_getspecific-style). This way GetOptional, when called from a
+        // destructor of another task-local variable (they may sleep, letting
+        // arbitrary engine code run), returns nullptr instead of a pointer to
+        // a destroyed or currently-being-destroyed object.
+        auto* const data = std::exchange(node_ptr->ptr, nullptr);
+        data->DeleteSelf();
     };
 
     // By default, boost::intrusive containers don't own their elements (nodes),
     // so we need to destroy them explicitly. The variables are destroyed
     // front-to-back, in reverse-initialization order.
-    while (!impl_->normal_data_storage.empty()) {
-        impl_->normal_data_storage.pop_front_and_dispose(disposer);
-    }
-
-    while (!impl_->inherited_data_storage.empty()) {
-        impl_->inherited_data_storage.pop_front_and_dispose(disposer);
+    //
+    // A destructor may initialize other task-local variables. Same as for
+    // `thread_local` ([basic.stc.thread]/2): a variable initialized after
+    // a destructor has started executing is destroyed after that destructor
+    // completes. Freshly initialized variables are pushed to the list front,
+    // so the inner loops pick them up; the outer loop handles a destructor of
+    // a variable of one kind initializing a variable of the other kind.
+    while (!impl_->normal_data_storage.empty() || !impl_->inherited_data_storage.empty()) {
+        if (!impl_->normal_data_storage.empty()) {
+            impl_->normal_data_storage.pop_front_and_dispose(disposer);
+        } else {
+            impl_->inherited_data_storage.pop_front_and_dispose(disposer);
+        }
     }
 }
 

core/src/engine/plugin.hpp

@@ -37,6 +37,9 @@ class PluginBase {
     /// Callback that is called in coroutine when the task finishes executing its
     /// payload (once per task, after completion or cancellation). The terminal
     /// state is available via task.GetPendingFinalState().
+    /// Task-local variables are destroyed after this callback is called.
+    /// Note: the task may sleep during the destruction of task-local variables
+    /// after this callback is called.
     virtual void HookTaskStop(impl::TaskContext& /*task*/) noexcept {}
 };
 

core/src/engine/task/async_benchmark.cpp

@@ -111,6 +111,7 @@ void WrapCallAndPerform(benchmark::State& state) {
                 // Perform requires that task-local storage is empty, then fills it
                 engine::impl::task_local::Storage discarded_storage;
                 discarded_storage.InitializeFrom(std::move(engine::impl::task_local::GetCurrentStorage()));
+                discarded_storage.DestroyVariables();
             }
             wrapped_call.Perform();
         }

core/src/engine/task/local_variable_test.cpp

@@ -1,14 +1,19 @@
 #include <userver/utest/utest.hpp>
 
+#include <functional>
 #include <optional>
 #include <utility>
 
 #include <userver/engine/async.hpp>
 #include <userver/engine/single_use_event.hpp>
 #include <userver/engine/sleep.hpp>
+#include <userver/engine/task/current_task.hpp>
+#include <userver/engine/task/inherited_variable.hpp>
 #include <userver/engine/task/local_variable.hpp>
 #include <userver/utils/async.hpp>
 
+#include <engine/task/task_processor.hpp>
+
 USERVER_NAMESPACE_BEGIN
 
 namespace {
@@ -198,4 +203,203 @@ UTEST(TaskLocalVariable, WaitInDestructorCancelled) {
     UEXPECT_THROW(task.Get(), engine::TaskCancelledException);
 }
 
+namespace {
+
+struct DtorCallback final {
+    std::function<void()> on_destroy;
+
+    DtorCallback() = default;
+    DtorCallback(const DtorCallback&) = delete;
+    DtorCallback& operator=(const DtorCallback&) = delete;
+
+    ~DtorCallback() {
+        if (on_destroy) {
+            on_destroy();
+        }
+    }
+};
+
+engine::TaskLocalVariable<DtorCallback> kDtorCallback;
+engine::TaskLocalVariable<int> kObservedInt;
+
+}  // namespace
+
+// Situation 1a: GetOptional on a live (not yet destroyed) variable.
+UTEST(TaskLocalVariable, GetOptionalAlive) {
+    EXPECT_EQ(kObservedInt.GetOptional(), nullptr);
+
+    *kObservedInt = 42;
+    ASSERT_NE(kObservedInt.GetOptional(), nullptr);
+    EXPECT_EQ(*kObservedInt.GetOptional(), 42);
+}
+
+// Situation 1b: a variable that is initialized earlier (thus destroyed later)
+// is still observable from the destructor of another task-local variable.
+UTEST(TaskLocalVariable, GetOptionalAliveFromOtherVariableDestructor) {
+    bool checked = false;
+
+    engine::AsyncNoTracing([&] {
+        // kObservedInt is initialized first => destroyed after kDtorCallback.
+        *kObservedInt = 42;
+        kDtorCallback->on_destroy = [&] {
+            auto* const observed = kObservedInt.GetOptional();
+            ASSERT_NE(observed, nullptr);
+            EXPECT_EQ(*observed, 42);
+            checked = true;
+        };
+    }).Get();
+
+    EXPECT_TRUE(checked);
+}
+
+// Situation 2: GetOptional on an already-destroyed variable must return
+// nullptr, not a dangling pointer.
+UTEST(TaskLocalVariable, GetOptionalAfterVariableDestroyed) {
+    bool checked = false;
+
+    engine::AsyncNoTracing([&] {
+        kDtorCallback->on_destroy = [&] {
+            EXPECT_EQ(kObservedInt.GetOptional(), nullptr);
+            checked = true;
+        };
+        // kObservedInt is initialized after kDtorCallback => destroyed before
+        // it, so ~DtorCallback observes an already-destroyed variable.
+        *kObservedInt = 42;
+    }).Get();
+
+    EXPECT_TRUE(checked);
+}
+
+// Situation 3: GetOptional on a variable whose destructor is currently
+// running must return nullptr (POSIX pthread_getspecific-style: the variable
+// is unset before its destructor is invoked). The destructor itself can use
+// `this` if needed; other code must not observe a half-destroyed object.
+UTEST(TaskLocalVariable, GetOptionalDuringOwnDestruction) {
+    bool checked = false;
+
+    engine::AsyncNoTracing([&] {
+        kDtorCallback->on_destroy = [&] {
+            EXPECT_EQ(kDtorCallback.GetOptional(), nullptr);
+            checked = true;
+        };
+    }).Get();
+
+    EXPECT_TRUE(checked);
+}
+
+namespace {
+
+struct CallbackHolder final {
+    std::function<void()> on_destroy;
+
+    explicit CallbackHolder(std::function<void()> cb)
+        : on_destroy(std::move(cb))
+    {}
+
+    CallbackHolder(const CallbackHolder&) = delete;
+    CallbackHolder& operator=(const CallbackHolder&) = delete;
+
+    ~CallbackHolder() {
+        if (on_destroy) {
+            on_destroy();
+        }
+    }
+};
+
+engine::TaskInheritedVariable<CallbackHolder> kInheritedCallback;
+
+}  // namespace
+
+// A task-local variable initialized from a destructor of another task-local
+// variable is destroyed after that destructor completes. Same as for
+// `thread_local`, [basic.stc.thread]/2: "If a variable with thread storage
+// duration is initialized after a thread-local destructor has started
+// executing, its destructor is scheduled to run after that destructor
+// completes".
+UTEST(TaskLocalVariable, DtorInitializesLocalVariable) {
+    std::string order;
+
+    engine::AsyncNoTracing([&] {
+        kDtorCallback->on_destroy = [&] {
+            kGuardX->emplace(order, "x");
+            order += "a";
+        };
+    }).Get();
+
+    // "a" (the initializing destructor completes) strictly before "x"
+    EXPECT_EQ(order, "ax");
+}
+
+// A task-inherited variable initialized from a destructor of a task-local
+// variable is destroyed after that destructor completes.
+UTEST(TaskLocalVariable, LocalDtorInitializesInheritedVariable) {
+    std::string order;
+
+    engine::AsyncNoTracing([&] {
+        kDtorCallback->on_destroy = [&] {
+            kInheritedCallback.Emplace([&order] { order += "i"; });
+            order += "a";
+        };
+    }).Get();
+
+    EXPECT_EQ(order, "ai");
+}
+
+// A task-local variable initialized from a destructor of a task-inherited
+// variable is destroyed after that destructor completes.
+UTEST(TaskLocalVariable, InheritedDtorInitializesLocalVariable) {
+    std::string order;
+
+    engine::AsyncNoTracing([&] {
+        kInheritedCallback.Emplace([&] {
+            kGuardX->emplace(order, "x");
+            order += "a";
+        });
+    }).Get();
+
+    EXPECT_EQ(order, "ax");
+}
+
+// A task-inherited variable initialized from a destructor of another
+// task-inherited variable is destroyed after that destructor completes.
+UTEST(TaskLocalVariable, InheritedDtorInitializesInheritedVariable) {
+    std::string order;
+
+    engine::AsyncNoTracing([&] {
+        kInheritedCallback.Emplace([&] {
+            kInheritedCallback.Emplace([&order] { order += "i"; });
+            order += "a";
+        });
+    }).Get();
+
+    EXPECT_EQ(order, "ai");
+}
+
+// A child task inherits a variable on construction and keeps it alive even if:
+// - the parent erases the variable before the child starts,
+// - the child is cancelled before start (finishes as State::kCancelled).
+// The variable is destroyed only after the child task finishes.
+UTEST(TaskInheritedVariable, ErasedInParentWhileChildCancelledBeforeStart) {
+    // With more than 1 worker thread the child task would start running
+    // concurrently, racing with RequestCancel and Erase below.
+    ASSERT_EQ(engine::current_task::GetTaskProcessor().GetWorkerCount(), 1);
+
+    bool destroyed = false;
+    kInheritedCallback.Emplace([&destroyed] { destroyed = true; });
+
+    // The child inherits the variable on construction...
+    auto task = utils::Async("child", [] { FAIL() << "the task must not start"; });
+    // ...and is cancelled before it gets a chance to start.
+    task.RequestCancel();
+
+    // Hide the variable from the parent. The child still holds a reference.
+    kInheritedCallback.Erase();
+    EXPECT_FALSE(destroyed);
+
+    task.Wait();
+    EXPECT_EQ(task.GetState(), engine::TaskBase::State::kCancelled);
+    // The child dropped the last reference when it finished.
+    EXPECT_TRUE(destroyed);
+}
+
 USERVER_NAMESPACE_END

core/src/engine/task/task_context.cpp

@@ -518,7 +518,15 @@ class TaskContext::LocalStorageGuard {
         context_.local_storage_.emplace();
     }
 
-    ~LocalStorageGuard() { context_.local_storage_.reset(); }
+    ~LocalStorageGuard() {
+        // Destroy the variables while the storage is still alive and reachable
+        // through GetCurrentStorage: the destructors may sleep, letting
+        // arbitrary engine code (e.g. plugin hooks) access task-locals.
+        // Calling this from ~Storage during optional::reset would make that
+        // access UB.
+        context_.local_storage_->DestroyVariables();
+        context_.local_storage_.reset();
+    }
 
 private:
     TaskContext& context_;

core/src/utils/impl/span_wrap_call.cpp

@@ -48,7 +48,11 @@ void SpanWrapCall::DoBeforeInvoke() {
     pimpl_->span.Get().AttachToCoroStack();
 }
 
-SpanWrapCall::~SpanWrapCall() = default;
+SpanWrapCall::~SpanWrapCall() {
+    // Non-empty if the task was never started (e.g. cancelled before Perform):
+    // drop the references to the inherited variables.
+    pimpl_->storage.DestroyVariables();
+}
 
 }  // namespace utils::impl