Fix unbounded recursion in T[Cancelable]BoundedConcurrencyRunner

commit_hash:59473e0eed6e5a362e144c261c9db74a7bc0d2cc
ydb-platform · Jan 15, 2025 · 52092f0 · 52092f0
1 parent 8dff5c1
commit 52092f0
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 35 deletions.
diff --git a/yt/yt/core/actions/future-inl.h b/yt/yt/core/actions/future-inl.h
@@ -2461,7 +2461,7 @@ class TCancelableBoundedConcurrencyRunner
         , ConcurrencyLimit_(concurrencyLimit)
         , Futures_(Callbacks_.size(), VoidFuture)
         , Results_(Callbacks_.size())
-        , CurrentIndex_(std::min<int>(ConcurrencyLimit_, ssize(Callbacks_)))
+        , CurrentIndex_(std::min<int>(ConcurrencyLimit_, std::ssize(Callbacks_)))
         , FailOnFirstError_(failOnError)
     { }
 
@@ -2499,56 +2499,70 @@ class TCancelableBoundedConcurrencyRunner
 
     void RunCallback(int index)
     {
-        auto future = Callbacks_[index]();
+        // ORD-2002: Avoid calling OnResult directly to prevent unbounded
+        // recursive chains like RunCallback -> OnResult -> RunCallback...
+        while (true) {
+            auto future = Callbacks_[index]();
+            if (!future.IsSet()) {
+                {
+                    auto guard = Guard(SpinLock_);
+                    if (Error_) {
+                        guard.Release();
+                        future.Cancel(*Error_);
+                        return;
+                    }
 
-        if (future.IsSet()) {
-            OnResult(index, std::move(future.Get()));
-            return;
-        }
+                    Futures_[index] = future.template As<void>();
+                }
 
-        {
-            auto guard = Guard(SpinLock_);
-            if (Error_) {
-                guard.Release();
-                future.Cancel(*Error_);
-                return;
+                future.Subscribe(
+                    BIND_NO_PROPAGATE(&TCancelableBoundedConcurrencyRunner::OnResult, MakeStrong(this), index)
+                        // NB: Sync invoker protects from unbounded recursion.
+                        .Via(GetSyncInvoker()));
+                break;
             }
 
-            Futures_[index] = future.template As<void>();
-        }
+            auto suggestedIndex = HandleResultAndSuggestNextIndex(index, std::move(future.Get()));
+            if (!suggestedIndex) {
+                break;
+            }
 
-        future.Subscribe(
-            BIND_NO_PROPAGATE(&TCancelableBoundedConcurrencyRunner::OnResult, MakeStrong(this), index)
-                // NB: Sync invoker protects from unbounded recursion.
-                .Via(GetSyncInvoker()));
+            index = *suggestedIndex;
+        }
     }
 
-    void OnResult(int index, const NYT::TErrorOr<T>& result)
+    [[nodiscard]]
+    std::optional<int> HandleResultAndSuggestNextIndex(int index, const TErrorOr<T>& result)
     {
         if (FailOnFirstError_ && !result.IsOK()) {
             OnError(result);
-            return;
+            return std::nullopt;
         }
 
         int newIndex;
         int finishedCount;
         {
             auto guard = Guard(SpinLock_);
             if (Error_) {
-                return;
+                return std::nullopt;
             }
 
             newIndex = CurrentIndex_++;
             finishedCount = ++FinishedCount_;
             Results_[index] = result;
         }
 
-        if (finishedCount == ssize(Callbacks_)) {
+        if (finishedCount == std::ssize(Callbacks_)) {
             Promise_.TrySet(Results_);
         }
 
-        if (newIndex < ssize(Callbacks_)) {
-            RunCallback(newIndex);
+        return newIndex < std::ssize(Callbacks_) ? std::optional(newIndex) : std::nullopt;
+    }
+
+    void OnResult(int index, const TErrorOr<T>& result)
+    {
+        if (auto suggestedIndex = HandleResultAndSuggestNextIndex(index, result)) {
+            RunCallback(*suggestedIndex);
         }
     }
 
@@ -2616,26 +2630,44 @@ class TBoundedConcurrencyRunner
 
     void RunCallback(int index)
     {
-        auto future = Callbacks_[index]();
-        if (future.IsSet()) {
-            OnResult(index, future.Get());
-        } else {
-            future.Subscribe(
-                BIND_NO_PROPAGATE(&TBoundedConcurrencyRunner::OnResult, MakeStrong(this), index));
+        // ORD-2002: Avoid calling OnResult directly to prevent unbounded
+        // recursive chains like RunCallback -> OnResult -> RunCallback...
+        while (true) {
+            auto future = Callbacks_[index]();
+            if (!future.IsSet()) {
+                future.Subscribe(
+                    BIND_NO_PROPAGATE(&TBoundedConcurrencyRunner::OnResult, MakeStrong(this), index)
+                        // NB: Sync invoker protects from unbounded recursion.
+                        .Via(GetSyncInvoker()));
+                        break;
+            }
+
+            auto suggestedIndex = HandleResultAndSuggestNextIndex(index, future.Get());
+            if (!suggestedIndex) {
+                break;
+            }
+
+            index = *suggestedIndex;
         }
     }
 
-    void OnResult(int index, const NYT::TErrorOr<T>& result)
+    [[nodiscard]]
+    std::optional<int> HandleResultAndSuggestNextIndex(int index, const TErrorOr<T>& result)
     {
         Results_[index] = result;
 
         int newIndex = CurrentIndex_++;
-        if (newIndex < static_cast<ssize_t>(Callbacks_.size())) {
-            RunCallback(newIndex);
+        if (++FinishedCount_ == std::ssize(Callbacks_)) {
+            Promise_.Set(Results_);
         }
 
-        if (++FinishedCount_ == static_cast<ssize_t>(Callbacks_.size())) {
-            Promise_.Set(Results_);
+        return newIndex < std::ssize(Callbacks_) ? std::optional(newIndex) : std::nullopt;
+    }
+
+    void OnResult(int index, const TErrorOr<T>& result)
+    {
+        if (auto suggestedIndex = HandleResultAndSuggestNextIndex(index, result)) {
+            RunCallback(*suggestedIndex);
         }
     }
 };

diff --git a/yt/yt/core/actions/unittests/actions_ut.cpp b/yt/yt/core/actions/unittests/actions_ut.cpp
@@ -94,6 +94,25 @@ TEST(TCancelableRunWithBoundedConcurrencyTest, Cancelation)
     EXPECT_EQ(canceledCount, 4);
 }
 
+TEST(TCancelableRunWithBoundedConcurrencyTest, RecurseRunner)
+{
+    auto threadPool = CreateThreadPool(4, "ThreadPool");
+
+    std::vector<TCallback<TFuture<void>()>> callbacks;
+    for (int i = 0; i < 50'000; ++i) {
+        callbacks.push_back(BIND([] {
+            return VoidFuture;
+        }));
+    }
+
+    auto future = CancelableRunWithBoundedConcurrency<void>(
+        std::move(callbacks),
+        /*concurrencyLimit*/ 5);
+
+    WaitFor(future)
+        .ThrowOnError();
+}
+
 TEST(TAllSucceededBoundedConcurrencyTest, CancelOthers)
 {
     using TCounter = std::atomic<int>;