Use MonitoredRefCount/RetainGuard instead of atomic/sleep

2 files changed, 12 insertions, 13 deletions

diff --git a/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.cpp b/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.cpp
index 0a24efd0e66..b9cbdab1c0a 100644
--- a/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.cpp
+++ b/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.cpp
@@ -40,7 +40,7 @@ WarmupIndexCollection::WarmupIndexCollection(const WarmupConfig & warmupConfig,
     _warmupDone(warmupDone),
     _warmupEndTime(vespalib::steady_clock::now() + warmupConfig.getDuration()),
     _handledTerms(std::make_unique<FieldTermMap>()),
-    _pendingTasks(0)
+    _pendingTasks()
 {
     if (_next->valid()) {
         setCurrentIndex(_next->getCurrentIndex());
@@ -81,7 +81,7 @@ WarmupIndexCollection::~WarmupIndexCollection()
     if (_warmupEndTime != vespalib::steady_time()) {
         LOG(info, "Warmup aborted due to new state change or application shutdown");
     }
-    assert(_pendingTasks == 0);
+    assert(_pendingTasks.has_zero_ref_count());
 }
 
 const ISourceSelector &
@@ -220,23 +220,19 @@ WarmupIndexCollection::getSearchableSP(uint32_t i) const
 
 void
 WarmupIndexCollection::drainPending() {
-    while (_pendingTasks > 0) {
-        std::this_thread::sleep_for(1ms);
-    }
+    _pendingTasks.waitForZeroRefCount();
 }
 
 WarmupIndexCollection::WarmupTask::WarmupTask(std::unique_ptr<MatchData> md, std::shared_ptr<WarmupIndexCollection> warmup)
     : _warmup(std::move(warmup)),
+      _retainGuard(_warmup->_pendingTasks),
       _matchData(std::move(md)),
       _bluePrint(),
       _requestContext()
 {
-    _warmup->_pendingTasks++;
 }
 
-WarmupIndexCollection::WarmupTask::~WarmupTask() {
-    _warmup->_pendingTasks--;
-}
+WarmupIndexCollection::WarmupTask::~WarmupTask() = default;
 
 void
 WarmupIndexCollection::WarmupTask::run()
diff --git a/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.h b/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.h
index 5292fb36298..b0b2952bee8 100644
--- a/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.h
+++ b/searchcorespi/src/vespa/searchcorespi/index/warmupindexcollection.h
@@ -5,6 +5,8 @@
 #include "isearchableindexcollection.h"
 #include "warmupconfig.h"
 #include <vespa/vespalib/util/threadexecutor.h>
+#include <vespa/vespalib/util/monitored_refcount.h>
+#include <vespa/vespalib/util/retain_guard.h>
 #include <vespa/searchlib/queryeval/fake_requestcontext.h>
 
 namespace searchcorespi {
@@ -85,9 +87,10 @@ private:
     private:
         void run() override;
         std::shared_ptr<WarmupIndexCollection>  _warmup;
-        std::unique_ptr<MatchData>   _matchData;
-        Blueprint::UP                _bluePrint;
-        FakeRequestContext           _requestContext;
+        vespalib::RetainGuard                   _retainGuard;
+        std::unique_ptr<MatchData>              _matchData;
+        Blueprint::UP                           _bluePrint;
+        FakeRequestContext                      _requestContext;
     };
 
     void fireWarmup(Task::UP task);
@@ -102,7 +105,7 @@ private:
     vespalib::steady_time              _warmupEndTime;
     std::mutex                         _lock;
     std::unique_ptr<FieldTermMap>      _handledTerms;
-    std::atomic<uint64_t>              _pendingTasks;
+    vespalib::MonitoredRefCount        _pendingTasks;
 };
 
 }  // namespace searchcorespi