Allow to change (reduce) max queue size for common elastic queue

Add same method as in TFastElasticQueue, and test.
commit_hash:0a2b618325e57c32fd269254a7dbe912849c3f10

3 files changed, 28 insertions, 1 deletions

diff --git a/library/cpp/threading/equeue/equeue.cpp b/library/cpp/threading/equeue/equeue.cpp
index 54b088635d5..bc0b6fc8144 100644
--- a/library/cpp/threading/equeue/equeue.cpp
+++ b/library/cpp/threading/equeue/equeue.cpp
@@ -10,7 +10,7 @@ size_t TElasticQueue::ObjectCount() const {
 }
 
 bool TElasticQueue::TryIncCounter() {
-    if (++GuardCount_ > MaxQueueSize_) {
+    if (++GuardCount_ > CurrentMaxQueueSize_) {
         --GuardCount_;
         return false;
     }
@@ -68,6 +68,7 @@ bool TElasticQueue::Add(IObjectInQueue* obj) {
 
 void TElasticQueue::Start(size_t threadCount, size_t maxQueueSize) {
     MaxQueueSize_ = maxQueueSize;
+    CurrentMaxQueueSize_ = maxQueueSize;
     SlaveQueue_->Start(threadCount, maxQueueSize);
 }
 
diff --git a/library/cpp/threading/equeue/equeue.h b/library/cpp/threading/equeue/equeue.h
index 4d03ad051db..a76271c27ea 100644
--- a/library/cpp/threading/equeue/equeue.h
+++ b/library/cpp/threading/equeue/equeue.h
@@ -17,13 +17,20 @@ public:
     void Stop() noexcept override;
 
     size_t ObjectCount() const;
+
+    void SetCurrentMaxQueueSize(size_t v) {
+        Y_ENSURE(v <= MaxQueueSize_);
+        CurrentMaxQueueSize_ = v;
+    }
 private:
     class TDecrementingWrapper;
 
     bool TryIncCounter();
 private:
     THolder<IThreadPool> SlaveQueue_;
+
     size_t MaxQueueSize_ = 0;
+    std::atomic<size_t> CurrentMaxQueueSize_ = 0;
     std::atomic<size_t> ObjectCount_ = 0;
     std::atomic<size_t> GuardCount_ = 0;
 };
diff --git a/library/cpp/threading/equeue/equeue_ut.cpp b/library/cpp/threading/equeue/equeue_ut.cpp
index 0b66128c73d..615d3d877d5 100644
--- a/library/cpp/threading/equeue/equeue_ut.cpp
+++ b/library/cpp/threading/equeue/equeue_ut.cpp
@@ -88,6 +88,25 @@ Y_UNIT_TEST_SUITE(TElasticQueueTest) {
         UNIT_ASSERT_VALUES_EQUAL(0u, TEnv<T>::Queue->ObjectCount());
         UNIT_ASSERT_VALUES_EQUAL(0u, TEnv<T>::Queue->Size());
         UNIT_ASSERT_VALUES_EQUAL(Counters.Processed.load(), enqueued);
+
+        Counters.Reset();
+
+        enqueued = 0;
+        {
+            TLocalSetup setup;
+            TEnv<T>::Queue->SetCurrentMaxQueueSize(MaxQueueSize / 2);
+
+            while (TEnv<T>::Queue->Add(&job) && enqueued < MaxQueueSize + 100) {
+                ++enqueued;
+            }
+
+            UNIT_ASSERT_VALUES_EQUAL(enqueued, MaxQueueSize / 2);
+            UNIT_ASSERT_VALUES_EQUAL(enqueued, TEnv<T>::Queue->ObjectCount());
+        }
+
+        UNIT_ASSERT_VALUES_EQUAL(0u, TEnv<T>::Queue->ObjectCount());
+        UNIT_ASSERT_VALUES_EQUAL(0u, TEnv<T>::Queue->Size());
+        UNIT_ASSERT_VALUES_EQUAL(Counters.Processed.load(), enqueued);
     }
 
     Y_UNIT_TEST(FillTest) {