intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 83 insertions, 0 deletions

diff --git a/library/cpp/messagebus/misc/tokenquota.h b/library/cpp/messagebus/misc/tokenquota.h
new file mode 100644
index 0000000000..190547fa54
--- /dev/null
+++ b/library/cpp/messagebus/misc/tokenquota.h
@@ -0,0 +1,83 @@
+#pragma once
+
+#include <util/system/atomic.h>
+
+namespace NBus {
+    /* Consumer and feeder quota model impl.
+
+        Consumer thread only calls:
+            Acquire(),  fetches tokens for usage from bucket;
+            Consume(),  eats given amount of tokens, must not
+                        be greater than Value() items;
+
+        Other threads (feeders) calls:
+            Return(),   put used tokens back to bucket;
+     */
+
+    class TTokenQuota {
+    public:
+        TTokenQuota(bool enabled, size_t tokens, size_t wake)
+            : Enabled(tokens > 0 ? enabled : false)
+            , Acquired(0)
+            , WakeLev(wake < 1 ? Max<size_t>(1, tokens / 2) : 0)
+            , Tokens_(tokens)
+        {
+            Y_UNUSED(padd_);
+        }
+
+        bool Acquire(TAtomic level = 1, bool force = false) {
+            level = Max(TAtomicBase(level), TAtomicBase(1));
+
+            if (Enabled && (Acquired < level || force)) {
+                Acquired += AtomicSwap(&Tokens_, 0);
+            }
+
+            return !Enabled || Acquired >= level;
+        }
+
+        void Consume(size_t items) {
+            if (Enabled) {
+                Y_ASSERT(Acquired >= TAtomicBase(items));
+
+                Acquired -= items;
+            }
+        }
+
+        bool Return(size_t items_) noexcept {
+            if (!Enabled || items_ == 0)
+                return false;
+
+            const TAtomic items = items_;
+            const TAtomic value = AtomicAdd(Tokens_, items);
+
+            return (value - items < WakeLev && value >= WakeLev);
+        }
+
+        bool IsEnabled() const noexcept {
+            return Enabled;
+        }
+
+        bool IsAboveWake() const noexcept {
+            return !Enabled || (WakeLev <= AtomicGet(Tokens_));
+        }
+
+        size_t Tokens() const noexcept {
+            return Acquired + AtomicGet(Tokens_);
+        }
+
+        size_t Check(const TAtomic level) const noexcept {
+            return !Enabled || level <= Acquired;
+        }
+
+    private:
+        bool Enabled;
+        TAtomicBase Acquired;
+        const TAtomicBase WakeLev;
+        TAtomic Tokens_;
+
+        /* This padd requires for align Tokens_ member on its own
+            CPU cacheline. */
+
+        ui64 padd_;
+    };
+}