YT-23585: Fix race between Cancel and GetCanceledError in Canceler

commit_hash:2c1b31d95037975cf4703cb90f81232f880a37e6
ytsaurus · Nov 21, 2024 · 9994cd4 · 9994cd4
1 parent 668dc9b
commit 9994cd4
Showing 1 changed file with 14 additions and 11 deletions.
diff --git a/yt/yt/core/concurrency/fiber_scheduler_thread.cpp b/yt/yt/core/concurrency/fiber_scheduler_thread.cpp
@@ -621,7 +621,17 @@ class TCanceler
     void SetFuture(TFuture<void> awaitable)
     {
         auto guard = Guard(Lock_);
-        Future_ = std::move(awaitable);
+        if (!IsCanceled()) {
+            Future_ = std::move(awaitable);
+            return;
+        }
+
+        guard.Release();
+
+        ErrorSet_.Wait();
+
+        YT_ASSERT(!CancelationError_.IsOK());
+        awaitable.Cancel(CancelationError_);
     }
 
     void ResetFuture()
@@ -644,6 +654,8 @@ class TCanceler
             future = std::move(Future_);
         }
 
+        ErrorSet_.NotifyAll();
+
         if (future) {
             YT_LOG_DEBUG("Sending cancelation to fiber, propagating to the awaited future (TargetFiberId: %x)",
                 FiberId_);
@@ -654,12 +666,6 @@ class TCanceler
         }
     }
 
-    TError GetCancelationError() const
-    {
-        auto guard = Guard(Lock_);
-        return CancelationError_;
-    }
-
     void Run(const TError& error)
     {
         Cancel(error);
@@ -680,6 +686,7 @@ class TCanceler
     const TFiberId FiberId_;
 
     std::atomic<bool> Canceled_ = false;
+    NThreading::TEvent ErrorSet_;
     NThreading::TSpinLock Lock_;
     TError CancelationError_;
     TFuture<void> Future_;
@@ -1173,10 +1180,6 @@ void WaitUntilSet(TFuture<void> future, IInvokerPtr invoker)
     GetCurrentFiberCanceler();
 
     const auto& canceler = NDetail::GetFiberSwitchHandler()->Canceler();
-    if (canceler->IsCanceled()) {
-        future.Cancel(canceler->GetCancelationError());
-    }
-
     canceler->SetFuture(future);
     auto finally = Finally([&] {
         canceler->ResetFuture();