diff options
Diffstat (limited to 'library/cpp/threading/equeue')
| -rw-r--r-- | library/cpp/threading/equeue/equeue.cpp | 3 | ||||
| -rw-r--r-- | library/cpp/threading/equeue/equeue.h | 7 | ||||
| -rw-r--r-- | library/cpp/threading/equeue/equeue_ut.cpp | 19 |
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) { |
