Use a RetainGuard to ensure DocumentDB is not closed until everything has been drained out.

4 files changed, 9 insertions, 3 deletions

diff --git a/searchcore/src/tests/proton/documentdb/documentbucketmover/documentbucketmover_v2_test.cpp b/searchcore/src/tests/proton/documentdb/documentbucketmover/documentbucketmover_v2_test.cpp
index 450fa7e8318..e844047f686 100644
--- a/searchcore/src/tests/proton/documentdb/documentbucketmover/documentbucketmover_v2_test.cpp
+++ b/searchcore/src/tests/proton/documentdb/documentbucketmover/documentbucketmover_v2_test.cpp
@@ -42,6 +42,7 @@ struct ControllerFixtureBase : public ::testing::Test
     DummyBucketExecutor         _bucketExecutor;
     MyMoveHandler               _moveHandler;
     DocumentDBTaggedMetrics     _metrics;
+    MonitoredRefCount           _refCount;
     std::shared_ptr<BucketMoveJobV2>  _bmj;
     MyCountJobRunner            _runner;
     ControllerFixtureBase(const BlockableMaintenanceJobConfig &blockableConfig, bool storeMoveDoneContexts);
@@ -123,7 +124,7 @@ ControllerFixtureBase::ControllerFixtureBase(const BlockableMaintenanceJobConfig
       _bucketExecutor(4),
       _moveHandler(*_bucketDB, storeMoveDoneContexts),
       _metrics("test", 1),
-      _bmj(BucketMoveJobV2::create(_calc, _moveHandler, _modifiedHandler, _master, _bucketExecutor, _ready._subDb,
+      _bmj(BucketMoveJobV2::create(_calc, RetainGuard(_refCount), _moveHandler, _modifiedHandler, _master, _bucketExecutor, _ready._subDb,
                                    _notReady._subDb, _bucketCreateNotifier,_clusterStateHandler, _bucketHandler,
                                    _diskMemUsageNotifier, blockableConfig, "test", makeBucketSpace())),
       _runner(*_bmj)
diff --git a/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.cpp b/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.cpp
index f592aeab9d2..8386385e7b2 100644
--- a/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.cpp
+++ b/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.cpp
@@ -56,7 +56,7 @@ JobTestBase::init(uint32_t allowedLidBloat,
         _singleExecutor = std::make_unique<vespalib::ThreadStackExecutor>(1, 0x10000);
         _master = std::make_unique<proton::ExecutorThreadService> (*_singleExecutor);
         _bucketExecutor = std::make_unique<storage::spi::dummy::DummyBucketExecutor>(4);
-        _job = lidspace::CompactionJob::create(compactCfg, _handler, _storer, *_master, *_bucketExecutor,
+        _job = lidspace::CompactionJob::create(compactCfg, RetainGuard(_refCount), _handler, _storer, *_master, *_bucketExecutor,
                                                _diskMemUsageNotifier, blockableCfg, _clusterStateHandler, nodeRetired,
                                                document::BucketSpace::placeHolder());
     } else {
diff --git a/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.h b/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.h
index 0a4e5c56acb..e74167db189 100644
--- a/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.h
+++ b/searchcore/src/tests/proton/documentdb/lid_space_compaction/lid_space_jobtest.h
@@ -2,6 +2,7 @@
 
 #include "lid_space_common.h"
 #include <vespa/searchcore/proton/server/blockable_maintenance_job.h>
+#include <vespa/searchcore/proton/common/monitored_refcount.h>
 #include <vespa/persistence/spi/bucketexecutor.h>
 #include <vespa/searchcorespi/index/i_thread_service.h>
 #include <vespa/vespalib/gtest/gtest.h>
@@ -16,6 +17,7 @@ struct JobTestBase : public ::testing::TestWithParam<bool> {
     MyFrozenBucketHandler _frozenHandler;
     test::DiskMemUsageNotifier _diskMemUsageNotifier;
     test::ClusterStateHandler _clusterStateHandler;
+    MonitoredRefCount _refCount;
     std::shared_ptr<BlockableMaintenanceJob> _job;
     JobTestBase();
     ~JobTestBase() override;
diff --git a/searchcore/src/tests/proton/documentdb/maintenancecontroller/maintenancecontroller_test.cpp b/searchcore/src/tests/proton/documentdb/maintenancecontroller/maintenancecontroller_test.cpp
index 7b225571227..bebdeb2a701 100644
--- a/searchcore/src/tests/proton/documentdb/maintenancecontroller/maintenancecontroller_test.cpp
+++ b/searchcore/src/tests/proton/documentdb/maintenancecontroller/maintenancecontroller_test.cpp
@@ -6,6 +6,7 @@
 #include <vespa/searchcore/proton/attribute/i_attribute_manager.h>
 #include <vespa/searchcore/proton/bucketdb/bucket_create_notifier.h>
 #include <vespa/searchcore/proton/common/doctypename.h>
+#include <vespa/searchcore/proton/common/monitored_refcount.h>
 #include <vespa/searchcore/proton/common/transient_resource_usage_provider.h>
 #include <vespa/searchcore/proton/documentmetastore/operation_listener.h>
 #include <vespa/searchcore/proton/documentmetastore/documentmetastore.h>
@@ -376,6 +377,7 @@ public:
     AttributeUsageFilter               _attributeUsageFilter;
     test::DiskMemUsageNotifier         _diskMemUsageNotifier;
     BucketCreateNotifier               _bucketCreateNotifier;
+    MonitoredRefCount                  _refCount;
     MaintenanceController              _mc;
 
     MaintenanceControllerFixture();
@@ -794,7 +796,8 @@ MaintenanceControllerFixture::MaintenanceControllerFixture()
       _notReadyAttributeManager(std::make_shared<MyAttributeManager>()),
       _attributeUsageFilter(),
       _bucketCreateNotifier(),
-      _mc(_threadService, _genericExecutor, _docTypeName)
+      _refCount(),
+      _mc(_threadService, _genericExecutor, _refCount, _docTypeName)
 {
     std::vector<MyDocumentSubDB *> subDBs;
     subDBs.push_back(&_ready);