YT-24270: Added overdraft for TAsyncSemaphore

* Changelog entry
  Type: feature
  Component: core

Add slots overdraft into TAsyncSemaphore
commit_hash:3a25c8c48546671fece656d75b2b5e4d37f829cd

Author: cherepashka

yt/yt/core/concurrency/async_semaphore.cpp

@@ -11,6 +11,14 @@ TAsyncSemaphore::TAsyncSemaphore(i64 totalSlots)
     YT_VERIFY(TotalSlots_ >= 0);
 }
 
+TAsyncSemaphore::TAsyncSemaphore(i64 totalSlots, bool enableOverdraft)
+    : TotalSlots_(totalSlots)
+    , FreeSlots_(totalSlots)
+    , EnableOverdraft_(enableOverdraft)
+{
+    YT_VERIFY(TotalSlots_ >= 0);
+}
+
 void TAsyncSemaphore::SetTotal(i64 totalSlots)
 {
     YT_VERIFY(totalSlots >= 0);
@@ -33,6 +41,10 @@ void TAsyncSemaphore::Release(i64 slots)
         auto guard = WriterGuard(SpinLock_);
 
         FreeSlots_ += slots;
+        if (EnableOverdraft_) {
+            FreeSlots_ = std::min(FreeSlots_, TotalSlots_);
+        }
+
         YT_VERIFY(FreeSlots_ <= TotalSlots_);
 
         if (Releasing_) {
@@ -48,9 +60,17 @@ void TAsyncSemaphore::Release(i64 slots)
 
         {
             auto guard = WriterGuard(SpinLock_);
+            auto frontWaiterOverflowsSlots = [&] {
+                return EnableOverdraft_ && !Waiters_.empty() && Waiters_.front().Slots > TotalSlots_ && FreeSlots_ == TotalSlots_;
+            };
 
-            while (!Waiters_.empty() && FreeSlots_ >= Waiters_.front().Slots) {
+            while (!Waiters_.empty() && FreeSlots_ >= Waiters_.front().Slots || frontWaiterOverflowsSlots()) {
                 auto& waiter = Waiters_.front();
+                if (frontWaiterOverflowsSlots()) {
+                    // For "fat" request we need to acquire all total slots in semaphore.
+                    YT_ASSERT(FreeSlots_ == TotalSlots_);
+                    waiter.Slots = FreeSlots_;
+                }
                 FreeSlots_ -= waiter.Slots;
                 waitersToRelease.push_back(std::move(waiter));
                 Waiters_.pop();
@@ -82,6 +102,10 @@ bool TAsyncSemaphore::Acquire(i64 slots)
     YT_VERIFY(slots >= 0);
 
     auto guard = WriterGuard(SpinLock_);
+    if (EnableOverdraft_) {
+        slots = std::min(slots, TotalSlots_);
+    }
+
     FreeSlots_ -= slots;
 
     return FreeSlots_ >= 0;
@@ -92,6 +116,10 @@ bool TAsyncSemaphore::TryAcquire(i64 slots)
     YT_VERIFY(slots >= 0);
 
     auto guard = WriterGuard(SpinLock_);
+    if (EnableOverdraft_) {
+        slots = std::min(slots, TotalSlots_);
+    }
+
     if (FreeSlots_ < slots) {
         return false;
     }
@@ -104,6 +132,10 @@ TFuture<TAsyncSemaphoreGuard> TAsyncSemaphore::AsyncAcquire(i64 slots)
     YT_VERIFY(slots >= 0);
 
     auto guard = WriterGuard(SpinLock_);
+    if (EnableOverdraft_) {
+        slots = std::min(slots, TotalSlots_);
+    }
+
     if (FreeSlots_ >= slots) {
         FreeSlots_ -= slots;
         return MakeFuture(TAsyncSemaphoreGuard(this, slots));

yt/yt/core/concurrency/async_semaphore.h

@@ -54,6 +54,7 @@ class TAsyncSemaphore
 {
 public:
     explicit TAsyncSemaphore(i64 totalSlots);
+    TAsyncSemaphore(i64 totalSlots, bool enableOverdraft);
 
     //! Updates the total number of slots.
     void SetTotal(i64 totalSlots);
@@ -95,6 +96,7 @@ class TAsyncSemaphore
     i64 TotalSlots_;
     i64 FreeSlots_;
 
+    bool EnableOverdraft_ = false;
     bool Releasing_ = false;
 
     TPromise<void> ReadyEvent_;

yt/yt/core/concurrency/unittests/async_semaphore_ut.cpp

@@ -4,6 +4,7 @@
 
 #include <yt/yt/core/concurrency/async_semaphore.h>
 #include <yt/yt/core/concurrency/scheduler_api.h>
+#include <yt/yt/core/concurrency/thread_pool.h>
 
 namespace NYT {
 namespace {
@@ -26,6 +27,78 @@ TEST(TAsyncSemaphoreTest, CancelReadyEvent)
     EXPECT_TRUE(WaitFor(readyTwo).IsOK());
 }
 
+TEST(TAsyncSemaphoreTest, OverdraftSlots)
+{
+    constexpr static int ThreadCount = 4;
+    constexpr static int RequestCount = 10;
+    constexpr static int SemaphoreTotalSlots = 1;
+    constexpr static int RequestWeight = 100;
+
+    auto threadPool = CreateThreadPool(ThreadCount, "SemaphoreAcqusition");
+    auto semaphore = New<TAsyncSemaphore>(SemaphoreTotalSlots, /*enableOverdraft*/ true);
+
+    {
+        std::vector<TFuture<void>> futures;
+        auto barrierPromise = NewPromise<void>();
+        for (int i = 0; i < RequestCount; ++i) {
+            futures.push_back(BIND([
+                &semaphore,
+                barrierFuture = barrierPromise.ToFuture()
+            ] {
+                WaitForFast(barrierFuture)
+                    .ThrowOnError();
+
+                WaitFor(semaphore->AsyncAcquire(RequestWeight).AsVoid())
+                    .ThrowOnError();
+            })
+                .AsyncVia(threadPool->GetInvoker())
+                .Run());
+        }
+
+        barrierPromise.Set();
+
+        WaitFor(AllSucceeded(std::move(futures)))
+            .ThrowOnError();
+    }
+
+    {
+        semaphore->SetTotal(RequestWeight);
+        std::vector<TFuture<void>> futures;
+        auto barrierPromise = NewPromise<void>();
+        for (int i = 0; i < RequestCount; ++i) {
+            futures.push_back(BIND([
+                &semaphore,
+                barrierFuture = barrierPromise.ToFuture()
+            ] {
+                WaitForFast(barrierFuture)
+                    .ThrowOnError();
+
+                WaitFor(semaphore->AsyncAcquire(RequestWeight).AsVoid())
+                    .ThrowOnError();
+            })
+                .AsyncVia(threadPool->GetInvoker())
+                .Run());
+        }
+
+        futures.push_back(BIND([
+            &semaphore,
+            barrierFuture = barrierPromise.ToFuture()
+        ] {
+            WaitForFast(barrierFuture)
+                .ThrowOnError();
+
+            semaphore->SetTotal(SemaphoreTotalSlots);
+        })
+            .AsyncVia(threadPool->GetInvoker())
+            .Run());
+
+        barrierPromise.Set();
+
+        WaitFor(AllSucceeded(std::move(futures)))
+            .ThrowOnError();
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace