Ydb stable 23-1-1923.1.19

x-stable-origin-commit: c5d5a396e89d0a72e0267a55e93d8404d4fb54fe

1 files changed, 67 insertions, 0 deletions

diff --git a/library/cpp/unified_agent_client/throttling.cpp b/library/cpp/unified_agent_client/throttling.cpp
new file mode 100644
index 0000000000..271f7b0e7e
--- /dev/null
+++ b/library/cpp/unified_agent_client/throttling.cpp
@@ -0,0 +1,67 @@
+#include "throttling.h"
+
+#include <util/datetime/cputimer.h>
+
+namespace NUnifiedAgent {
+    TThrottler::TThrottler(double rate, TDuration updatePeriod)
+        : CyclesPerMillisecond(GetCyclesPerMillisecond())
+        , UpdatePeriod(updatePeriod.MilliSeconds() * CyclesPerMillisecond)
+        , PeriodTokens(updatePeriod.SecondsFloat() * rate)
+        , AvailableTokens(0)
+        , ExpirationTime(0)
+    {
+    }
+
+    TThrottler::TThrottler(double rate, double burst)
+        : TThrottler(rate, TDuration::Seconds(burst / rate))
+    {
+    }
+
+    void TThrottler::Consume(double& tokens, TFMaybe<TDuration>& nextCheckDelay) {
+        const auto updateTime = UpdateTokens();
+
+        if (tokens <= AvailableTokens) {
+            AvailableTokens -= tokens;
+            tokens = 0.0;
+            nextCheckDelay = Nothing();
+        } else {
+            tokens -= AvailableTokens;
+            AvailableTokens = 0.0;
+            nextCheckDelay = TDuration::MicroSeconds((ExpirationTime - updateTime) * 1000 / CyclesPerMillisecond + 1);
+        }
+    }
+
+    bool TThrottler::TryConsume(double tokens) {
+        UpdateTokens();
+
+        if (tokens >  AvailableTokens) {
+            return false;
+        }
+        AvailableTokens -= tokens;
+        return true;
+    }
+
+    void TThrottler::ConsumeAndWait(double tokens) {
+        TFMaybe<TDuration> nextCheckDelay;
+        while (true) {
+            Consume(tokens, nextCheckDelay);
+            if (!nextCheckDelay.Defined()) {
+                return;
+            }
+            Sleep(*nextCheckDelay);
+        }
+    }
+
+    ui64 TThrottler::UpdateTokens() {
+        const auto updateTime = GetCycleCount();
+        if (updateTime >= ExpirationTime) {
+            if (ExpirationTime == 0) {
+                ExpirationTime = updateTime + UpdatePeriod;
+            } else {
+                ExpirationTime += ((updateTime - ExpirationTime) / UpdatePeriod + 1) * UpdatePeriod;
+            }
+            AvailableTokens = PeriodTokens;
+        }
+        return updateTime;
+    }
+}