feat engine: support WaitAny in SharedTask

Anton3 · Anton3 · commit 10dbb36acd7d · 2026-07-03T22:28:16.000+03:00
* Made @ref engine::SharedTask an @ref engine::Awaitable, making it usable in @ref engine::MakeWaitAny and friends. commit_hash:4ba17b631cda4f2abf17ae56228a8537f1b99664
diff --git a/core/include/userver/engine/task/shared_task.hpp b/core/include/userver/engine/task/shared_task.hpp
@@ -51,6 +51,9 @@ class [[nodiscard]] SharedTask : public TaskBase {
     /// other task into this, leaving the other in an invalid state.
     SharedTask& operator=(SharedTask&& other) noexcept;
 
+    /// Satisfies @ref engine::Awaitable, for use with @ref engine::WaitAnyContext and friends.
+    AwaitableToken GetAwaitableToken() noexcept USERVER_IMPL_LIFETIME_BOUND;
+
     /// @cond
     static constexpr WaitMode kWaitMode = WaitMode::kMultipleAwaiters;
 
diff --git a/core/src/engine/task/shared_task.cpp b/core/src/engine/task/shared_task.cpp
@@ -47,6 +47,14 @@ SharedTask& SharedTask::operator=(SharedTask&& other) noexcept {
     return *this;
 }
 
+AwaitableToken SharedTask::GetAwaitableToken() noexcept USERVER_IMPL_LIFETIME_BOUND {
+    if (!IsValid()) {
+        return {};
+    }
+
+    return AwaitableToken{utils::impl::InternalTag{}, &GetContext()};
+}
+
 void SharedTask::DecrementSharedUsages() noexcept {
     if (!IsValid()) {
         return;
diff --git a/core/src/engine/task/shared_task_with_result_test.cpp b/core/src/engine/task/shared_task_with_result_test.cpp
@@ -1,10 +1,14 @@
 #include <userver/utest/utest.hpp>
 
+#include <array>
 #include <type_traits>
 
+#include <userver/engine/future.hpp>
 #include <userver/engine/single_consumer_event.hpp>
 #include <userver/engine/sleep.hpp>
 #include <userver/engine/task/shared_task_with_result.hpp>
+#include <userver/engine/wait_all_checked.hpp>
+#include <userver/engine/wait_any.hpp>
 #include <userver/utils/async.hpp>
 
 USERVER_NAMESPACE_BEGIN
@@ -392,4 +396,123 @@ static_assert(std::is_move_assignable_v<engine::SharedTaskWithResult<void>>);
 static_assert(std::is_copy_constructible_v<engine::SharedTaskWithResult<void>>);
 static_assert(std::is_copy_assignable_v<engine::SharedTaskWithResult<void>>);
 
+namespace {
+
+constexpr int kFirstSharedTaskValue = 2;
+constexpr int kSecondSharedTaskValue = 3;
+
+using IntSharedTask = engine::SharedTaskWithResult<int>;
+
+std::array<int, 2> WaitAnyAndGetResults(IntSharedTask task1, IntSharedTask task2) {
+    auto wait_any = engine::MakeWaitAny(task1, task2);
+    std::array<int, 2> values{};
+    std::array<bool, 2> completed{};
+
+    for (std::size_t i = 0; i < completed.size(); ++i) {
+        const auto index = wait_any.Wait();
+        EXPECT_TRUE(index.has_value());
+        if (index != 0 && index != 1) {
+            ADD_FAILURE() << "Unexpected task index";
+            return values;
+        }
+        EXPECT_FALSE(completed[*index]);
+        completed[*index] = true;
+        values[*index] = *index == 0 ? task1.Get() : task2.Get();
+    }
+
+    return values;
+}
+
+}  // namespace
+
+UTEST(SharedTask, WaitAnyWithMultipleAwaiters) {
+    engine::Promise<int> promise1;
+    engine::Promise<int> promise2;
+
+    auto future1 = promise1.get_future();
+    auto future2 = promise2.get_future();
+
+    auto task1 = utils::SharedAsync("shared_task_1", [future = std::move(future1)]() mutable { return future.get(); });
+    auto task2 = utils::SharedAsync("shared_task_2", [future = std::move(future2)]() mutable { return future.get(); });
+
+    auto sum_waiter = utils::Async(
+        "sum_waiter",
+        [](IntSharedTask task1, IntSharedTask task2) {
+            const auto values = WaitAnyAndGetResults(std::move(task1), std::move(task2));
+            return values[0] + values[1];
+        },
+        task1,
+        task2
+    );
+
+    auto product_waiter = utils::Async(
+        "product_waiter",
+        [](IntSharedTask task1, IntSharedTask task2) {
+            const auto values = WaitAnyAndGetResults(std::move(task1), std::move(task2));
+            return values[0] * values[1];
+        },
+        task1,
+        task2
+    );
+
+    engine::Yield();
+    EXPECT_FALSE(task1.IsFinished());
+    EXPECT_FALSE(task2.IsFinished());
+
+    promise1.set_value(kFirstSharedTaskValue);
+    promise2.set_value(kSecondSharedTaskValue);
+
+    const auto sum = sum_waiter.Get();
+    const auto product = product_waiter.Get();
+
+    EXPECT_EQ(sum, kFirstSharedTaskValue + kSecondSharedTaskValue);
+    EXPECT_EQ(product, kFirstSharedTaskValue * kSecondSharedTaskValue);
+    EXPECT_NE(sum, product);
+}
+
+UTEST(SharedTask, WaitAllCheckedWithMultipleAwaiters) {
+    engine::Promise<int> promise1;
+    engine::Promise<int> promise2;
+
+    auto future1 = promise1.get_future();
+    auto future2 = promise2.get_future();
+
+    auto task1 = utils::SharedAsync("shared_task_1", [future = std::move(future1)]() mutable { return future.get(); });
+    auto task2 = utils::SharedAsync("shared_task_2", [future = std::move(future2)]() mutable { return future.get(); });
+
+    auto sum_waiter = utils::Async(
+        "sum_waiter",
+        [](IntSharedTask task1, IntSharedTask task2) {
+            engine::WaitAllChecked(task1, task2);
+            return task1.Get() + task2.Get();
+        },
+        task1,
+        task2
+    );
+
+    auto product_waiter = utils::Async(
+        "product_waiter",
+        [](IntSharedTask task1, IntSharedTask task2) {
+            engine::WaitAllChecked(task1, task2);
+            return task1.Get() * task2.Get();
+        },
+        task1,
+        task2
+    );
+
+    engine::Yield();
+    EXPECT_FALSE(task1.IsFinished());
+    EXPECT_FALSE(task2.IsFinished());
+
+    promise1.set_value(kFirstSharedTaskValue);
+    promise2.set_value(kSecondSharedTaskValue);
+
+    const auto sum = sum_waiter.Get();
+    const auto product = product_waiter.Get();
+
+    EXPECT_EQ(sum, kFirstSharedTaskValue + kSecondSharedTaskValue);
+    EXPECT_EQ(product, kFirstSharedTaskValue * kSecondSharedTaskValue);
+    EXPECT_NE(sum, product);
+}
+
 USERVER_NAMESPACE_END
diff --git a/core/src/engine/wait_all_checked.cpp b/core/src/engine/wait_all_checked.cpp
@@ -9,6 +9,8 @@
 #include <userver/engine/exception.hpp>
 #include <userver/engine/impl/awaiter.hpp>
 #include <userver/engine/single_consumer_event.hpp>
+#include <userver/utils/fast_scope_guard.hpp>
+#include <userver/utils/fixed_array.hpp>
 #include <userver/utils/impl/intrusive_ref_counter_one.hpp>
 #include <userver/utils/impl/userver_experiments.hpp>
 
@@ -45,8 +47,11 @@ class WaitAllCheckedContext final {
     FutureStatus WaitUntil(Deadline deadline) { return impl_->WaitUntil(deadline); }
 
 private:
-    // NOLINTNEXTLINE(fuchsia-multiple-inheritance)
-    class Impl final : public impl::PolymorphicAwaiter, public utils::impl::IntrusiveRefCounterOne<Impl> {
+    // The refcounter of Impl counts:
+    // 1. the reference from WaitAllCheckedContext (impl_);
+    // 2. one reference per subscribed Node. This way, Impl stays alive until all in-flight
+    //    notifications complete, even if WaitAllCheckedContext is destroyed concurrently with them.
+    class Impl final : public utils::impl::IntrusiveRefCounterOne<Impl> {
     public:
         explicit Impl(std::vector<AwaitableToken>&& targets);
 
@@ -61,15 +66,51 @@ class WaitAllCheckedContext final {
         void Unsubscribe() noexcept;
 
     private:
+        // Unlike Impl itself, a Node is an Awaiter, so it has its own wait list hook, allowing
+        // multiple Nodes of a single WaitAllCheckedContext to be linked into the WaitList
+        // of a single shared target (as well as into wait lists of distinct targets).
+        //
+        // Nodes are owned by Impl (nodes_). The AwaiterPtr ownership passed to wait lists
+        // is logical: instead of destroying the Node, its dispose hooks release
+        // the subscription's reference to Impl.
+        struct Node final : public impl::PolymorphicAwaiter {
+            Node(Impl& owner, AwaitableToken target, std::size_t index)
+                : owner(owner),
+                  target(target),
+                  index(index)
+            {}
+
+            void NotifyAndDispose(std::uintptr_t context) noexcept override {
+                UASSERT(context == kUnusedContext);
+
+                owner.HandleNotification(index);
+                // Release the subscription's reference to Impl, see above. After this point,
+                // Impl (including `*this`) may be destroyed (if WaitAllCheckedContext
+                // is destroyed concurrently).
+                intrusive_ptr_release(&owner);
+            }
+
+            void DisposeWithoutNotification() noexcept override {
+                // The Node is removed from a wait list without notification. Release
+                // the subscription's reference to Impl.
+                intrusive_ptr_release(&owner);
+            }
+
+            Impl& owner;
+            AwaitableToken target;
+            std::size_t index;
+            bool subscribed{false};
+        };
+
+        // Nodes do not need to pass any information through the context parameter of Awaiter:
+        // a Node itself identifies the subscription.
+        static constexpr std::uintptr_t kUnusedContext = 0;
         static constexpr std::size_t kNoExceptionSource = std::numeric_limits<std::size_t>::max();
 
-        AwaiterPtr AsAwaiterPtr() noexcept;
+        void HandleNotification(std::size_t index) noexcept;
 
-        void NotifyAndDispose(std::uintptr_t context) noexcept override;
-
-        void DisposeWithoutNotification() noexcept override { intrusive_ptr_release(this); }
-
-        std::vector<AwaitableToken> targets_;
+        // One node per target; nodes of empty targets stay unused.
+        utils::FixedArray<Node> nodes_;
         std::size_t next_subscription_index_{0};
         std::atomic<std::size_t> pending_notifications_{0};
         std::atomic<std::size_t> exception_source_index_{kNoExceptionSource};
@@ -80,24 +121,42 @@ class WaitAllCheckedContext final {
 };
 
 WaitAllCheckedContext::Impl::Impl(std::vector<AwaitableToken>&& targets)
-    : targets_(std::move(targets))
+    : nodes_(utils::GenerateFixedArray(
+          targets.size(),
+          [&](std::size_t index) { return Node(*this, targets[index], index); }
+      ))
 {}
 
 FutureStatus WaitAllCheckedContext::Impl::WaitUntil(Deadline deadline) {
-    while (next_subscription_index_ < targets_.size()) {
-        const auto target = targets_[next_subscription_index_];
+    while (next_subscription_index_ < nodes_.size()) {
+        const auto index = next_subscription_index_++;
+        auto& node = nodes_[index];
+        const auto target = node.target;
         if (target.IsEmpty()) {
-            ++next_subscription_index_;
             continue;
         }
         auto& awaitable = target.GetAwaitable(utils::impl::InternalTag{});
 
         pending_notifications_.fetch_add(1, std::memory_order_relaxed);
-        auto awaiter_ptr = AsAwaiterPtr();
-        awaitable.TryAppendAwaiter(awaiter_ptr, next_subscription_index_++);
-        if (awaiter_ptr != nullptr) {  // target is already ready.
-            pending_notifications_.fetch_sub(1, std::memory_order_relaxed);
-            RethrowError(target);
+        // Acquire the subscription's reference to Impl in advance.
+        intrusive_ptr_add_ref(this);
+
+        AwaiterPtr awaiter_ptr{&node};
+
+        {
+            const utils::FastScopeGuard finalize_subscription([&]() noexcept {
+                if (awaiter_ptr != nullptr) {
+                    // The target is already ready: deliver the missed notification in-place.
+                    // If the target has completed with an exception, it will be rethrown
+                    // after the subscription loop.
+                    UASSERT(awaiter_ptr.get() == &node);
+                    impl::NotifyAndDispose(std::move(awaiter_ptr), kUnusedContext);
+                } else {
+                    node.subscribed = true;
+                }
+            });
+
+            awaitable.TryAppendAwaiter(awaiter_ptr, kUnusedContext);
         }
     }
 
@@ -110,7 +169,7 @@ FutureStatus WaitAllCheckedContext::Impl::WaitUntil(Deadline deadline) {
     }
     auto exception_source_index = exception_source_index_.exchange(kNoExceptionSource, std::memory_order_relaxed);
     if (exception_source_index != kNoExceptionSource) {
-        RethrowError(targets_[exception_source_index]);
+        RethrowError(nodes_[exception_source_index].target);
         utils::AbortWithStacktrace("Should never be there");
     }
     UASSERT(pending_notifications_.load(std::memory_order_relaxed) == 0);
@@ -119,23 +178,27 @@ FutureStatus WaitAllCheckedContext::Impl::WaitUntil(Deadline deadline) {
 
 void WaitAllCheckedContext::Impl::Unsubscribe() noexcept {
     for (std::size_t i = 0; i < next_subscription_index_; ++i) {
-        if (targets_[i].IsEmpty()) {
+        auto& node = nodes_[i];
+        if (!node.subscribed) {
             continue;
         }
-        auto& awaitable = targets_[i].GetAwaitable(utils::impl::InternalTag{});
-        awaitable.RemoveAwaiter(*this, i);
+        auto& awaitable = node.target.GetAwaitable(utils::impl::InternalTag{});
+        auto removed = awaitable.RemoveAwaiter(node, kUnusedContext);
+        if (removed != nullptr) {
+            UASSERT(removed.get() == &node);
+        }
+        // Dropping `removed` calls DisposeWithoutNotification, releasing the subscription's
+        // reference to Impl. This never destroys Impl here, because the caller
+        // (WaitAllCheckedContext) still holds its own reference. If the node was already
+        // notified, `removed` is null, and the in-flight NotifyAndDispose releases
+        // the subscription's reference.
+        removed.reset();
+        node.subscribed = false;
     }
 }
 
-AwaiterPtr WaitAllCheckedContext::Impl::AsAwaiterPtr() noexcept {
-    intrusive_ptr_add_ref(this);
-    return AwaiterPtr{this};
-}
-
-void WaitAllCheckedContext::Impl::NotifyAndDispose(std::uintptr_t context) noexcept {
-    UASSERT(context <= std::numeric_limits<std::size_t>::max());
-    const auto index = static_cast<std::size_t>(context);
-    const auto target = targets_[index];
+void WaitAllCheckedContext::Impl::HandleNotification(std::size_t index) noexcept {
+    const auto target = nodes_[index].target;
     auto& awaitable = target.GetAwaitable(utils::impl::InternalTag{});
     UASSERT(awaitable.IsReady());
     UASSERT(pending_notifications_.load(std::memory_order_relaxed) > 0);
@@ -153,8 +216,6 @@ void WaitAllCheckedContext::Impl::NotifyAndDispose(std::uintptr_t context) noexc
     if (ready) {
         result_ready_.Send();
     }
-
-    intrusive_ptr_release(this);
 }
 
 FutureStatus DoWaitAllChecked(std::vector<AwaitableToken>&& targets, Deadline deadline) {
diff --git a/core/src/engine/wait_any.cpp b/core/src/engine/wait_any.cpp
