Pass distributor bucket space repo to pending cluster state.

7 files changed, 45 insertions, 14 deletions

diff --git a/storage/src/tests/distributor/bucketdbupdatertest.cpp b/storage/src/tests/distributor/bucketdbupdatertest.cpp
index d1a54c04359..fce2ff9b985 100644
--- a/storage/src/tests/distributor/bucketdbupdatertest.cpp
+++ b/storage/src/tests/distributor/bucketdbupdatertest.cpp
@@ -538,7 +538,7 @@ public:
 
             std::unordered_set<uint16_t> outdatedNodes;
             state = PendingClusterState::createForClusterStateChange(
-                    clock, clusterInfo, sender, cmd, outdatedNodes,
+                    clock, clusterInfo, sender, owner.getBucketSpaceRepo(), cmd, outdatedNodes,
                     api::Timestamp(1));
         }
 
@@ -551,7 +551,7 @@ public:
 
             std::unordered_set<uint16_t> outdatedNodes;
             state = PendingClusterState::createForDistributionChange(
-                    clock, clusterInfo, sender, api::Timestamp(1));
+                    clock, clusterInfo, sender, owner.getBucketSpaceRepo(), api::Timestamp(1));
         }
     };
 
@@ -1475,7 +1475,7 @@ BucketDBUpdaterTest::getSentNodesDistributionChanged(
     ClusterInformation::CSP clusterInfo(createClusterInfo(oldClusterState));
     std::unique_ptr<PendingClusterState> state(
             PendingClusterState::createForDistributionChange(
-                    clock, clusterInfo, sender, api::Timestamp(1)));
+                    clock, clusterInfo, sender, getBucketSpaceRepo(), api::Timestamp(1)));
 
     sortSentMessagesByIndex(sender);
 
@@ -1640,7 +1640,7 @@ BucketDBUpdaterTest::testPendingClusterStateReceive()
     std::unordered_set<uint16_t> outdatedNodes;
     std::unique_ptr<PendingClusterState> state(
             PendingClusterState::createForClusterStateChange(
-                    clock, clusterInfo, sender, cmd, outdatedNodes,
+                    clock, clusterInfo, sender, getBucketSpaceRepo(), cmd, outdatedNodes,
                     api::Timestamp(1)));
 
     CPPUNIT_ASSERT_EQUAL(3, (int)sender.commands.size());
@@ -1803,7 +1803,7 @@ BucketDBUpdaterTest::mergeBucketLists(
         ClusterInformation::CSP clusterInfo(createClusterInfo("cluster:d"));
         std::unique_ptr<PendingClusterState> state(
                 PendingClusterState::createForClusterStateChange(
-                        clock, clusterInfo, sender, cmd, outdatedNodes,
+                        clock, clusterInfo, sender, getBucketSpaceRepo(), cmd, outdatedNodes,
                         beforeTime));
 
         parseInputData(existingData, beforeTime, *state, includeBucketInfo);
@@ -1822,7 +1822,7 @@ BucketDBUpdaterTest::mergeBucketLists(
         ClusterInformation::CSP clusterInfo(createClusterInfo(oldState.toString()));
         std::unique_ptr<PendingClusterState> state(
                 PendingClusterState::createForClusterStateChange(
-                        clock, clusterInfo, sender, cmd, outdatedNodes,
+                        clock, clusterInfo, sender, getBucketSpaceRepo(), cmd, outdatedNodes,
                         afterTime));
 
         parseInputData(newData, afterTime, *state, includeBucketInfo);
diff --git a/storage/src/tests/distributor/distributortestutil.cpp b/storage/src/tests/distributor/distributortestutil.cpp
index 5deb31f8579..a0aa8e00070 100644
--- a/storage/src/tests/distributor/distributortestutil.cpp
+++ b/storage/src/tests/distributor/distributortestutil.cpp
@@ -358,6 +358,16 @@ DistributorTestUtil::getBucketDatabase() const {
     return _distributor->getDefaultBucketSpace().getBucketDatabase();
 }
 
+DistributorBucketSpaceRepo &
+DistributorTestUtil::getBucketSpaceRepo() {
+    return _distributor->getBucketSpaceRepo();
+}
+
+const DistributorBucketSpaceRepo &
+DistributorTestUtil::getBucketSpaceRepo() const {
+    return _distributor->getBucketSpaceRepo();
+}
+
 const lib::Distribution&
 DistributorTestUtil::getDistribution() const {
     return _distributor->getDefaultBucketSpace().getDistribution();
diff --git a/storage/src/tests/distributor/distributortestutil.h b/storage/src/tests/distributor/distributortestutil.h
index 4f09c11ac03..19da0483165 100644
--- a/storage/src/tests/distributor/distributortestutil.h
+++ b/storage/src/tests/distributor/distributortestutil.h
@@ -20,6 +20,7 @@ namespace distributor {
 class BucketDBUpdater;
 class Distributor;
 class DistributorBucketSpace;
+class DistributorBucketSpaceRepo;
 class IdealStateManager;
 class ExternalOperationHandler;
 class Operation;
@@ -125,6 +126,8 @@ public:
     DistributorBucketSpace &getDistributorBucketSpace();
     BucketDatabase& getBucketDatabase();
     const BucketDatabase& getBucketDatabase() const;
+    DistributorBucketSpaceRepo &getBucketSpaceRepo();
+    const DistributorBucketSpaceRepo &getBucketSpaceRepo() const;
 
     const lib::Distribution& getDistribution() const;
     // "End to end" distribution change trigger, which will invoke the bucket
diff --git a/storage/src/vespa/storage/distributor/bucketdbupdater.cpp b/storage/src/vespa/storage/distributor/bucketdbupdater.cpp
index 569136b8b10..82d49d68f12 100644
--- a/storage/src/vespa/storage/distributor/bucketdbupdater.cpp
+++ b/storage/src/vespa/storage/distributor/bucketdbupdater.cpp
@@ -169,6 +169,7 @@ BucketDBUpdater::storageDistributionChanged(
             _bucketSpaceComponent.getClock(),
             std::move(clusterInfo),
             _sender,
+            _bucketSpaceComponent.getBucketSpaceRepo(),
             _bucketSpaceComponent.getUniqueTimestamp());
     _outdatedNodes = _pendingClusterState->getOutdatedNodeSet();
 }
@@ -216,6 +217,7 @@ BucketDBUpdater::onSetSystemState(
             _bucketSpaceComponent.getClock(),
             std::move(clusterInfo),
             _sender,
+            _bucketSpaceComponent.getBucketSpaceRepo(),
             cmd,
             _outdatedNodes,
             _bucketSpaceComponent.getUniqueTimestamp());
diff --git a/storage/src/vespa/storage/distributor/distributor.h b/storage/src/vespa/storage/distributor/distributor.h
index f59b47574ba..3cb3408a951 100644
--- a/storage/src/vespa/storage/distributor/distributor.h
+++ b/storage/src/vespa/storage/distributor/distributor.h
@@ -161,6 +161,8 @@ public:
 
     DistributorBucketSpace& getDefaultBucketSpace() noexcept;
     const DistributorBucketSpace& getDefaultBucketSpace() const noexcept;
+    DistributorBucketSpaceRepo &getBucketSpaceRepo() noexcept { return *_bucketSpaceRepo; }
+    const DistributorBucketSpaceRepo &getBucketSpaceRepo() const noexcept { return *_bucketSpaceRepo; }
 
 private:
     friend class Distributor_Test;
diff --git a/storage/src/vespa/storage/distributor/pendingclusterstate.cpp b/storage/src/vespa/storage/distributor/pendingclusterstate.cpp
index 41c8974470f..bcdc160ee00 100644
--- a/storage/src/vespa/storage/distributor/pendingclusterstate.cpp
+++ b/storage/src/vespa/storage/distributor/pendingclusterstate.cpp
@@ -23,6 +23,7 @@ PendingClusterState::PendingClusterState(
         const framework::Clock& clock,
         const ClusterInformation::CSP& clusterInfo,
         DistributorMessageSender& sender,
+        DistributorBucketSpaceRepo &bucketSpaceRepo,
         const std::shared_ptr<api::SetSystemStateCommand>& newStateCmd,
         const std::unordered_set<uint16_t>& outdatedNodes,
         api::Timestamp creationTimestamp)
@@ -35,6 +36,7 @@ PendingClusterState::PendingClusterState(
       _clusterInfo(clusterInfo),
       _creationTimestamp(creationTimestamp),
       _sender(sender),
+      _bucketSpaceRepo(bucketSpaceRepo),
       _bucketOwnershipTransfer(distributorChanged(_prevClusterState, _newClusterState)),
       _pendingTransition()
 {
@@ -45,16 +47,14 @@ PendingClusterState::PendingClusterState(
         updateSetOfNodesThatAreOutdated();
         addAdditionalNodesToOutdatedSet(outdatedNodes);
     }
-    _pendingTransition = std::make_unique<PendingBucketSpaceDbTransition>(*this, _clusterInfo, _newClusterState, _creationTimestamp);
-    if (shouldRequestBucketInfo()) {
-        requestNodes();
-    }
+    constructorHelper();
 }
 
 PendingClusterState::PendingClusterState(
         const framework::Clock& clock,
         const ClusterInformation::CSP& clusterInfo,
         DistributorMessageSender& sender,
+        DistributorBucketSpaceRepo &bucketSpaceRepo,
         api::Timestamp creationTimestamp)
     : _requestedNodes(clusterInfo->getStorageNodeCount()),
       _outdatedNodes(clusterInfo->getStorageNodeCount()),
@@ -64,19 +64,26 @@ PendingClusterState::PendingClusterState(
       _clusterInfo(clusterInfo),
       _creationTimestamp(creationTimestamp),
       _sender(sender),
+      _bucketSpaceRepo(bucketSpaceRepo),
       _bucketOwnershipTransfer(true),
       _pendingTransition()
 {
     logConstructionInformation();
     markAllAvailableNodesAsRequiringRequest();
+    constructorHelper();
+}
+
+PendingClusterState::~PendingClusterState() {}
+
+void
+PendingClusterState::constructorHelper()
+{
     _pendingTransition = std::make_unique<PendingBucketSpaceDbTransition>(*this, _clusterInfo, _newClusterState, _creationTimestamp);
     if (shouldRequestBucketInfo()) {
         requestNodes();
     }
 }
 
-PendingClusterState::~PendingClusterState() {}
-
 void
 PendingClusterState::logConstructionInformation() const
 {
diff --git a/storage/src/vespa/storage/distributor/pendingclusterstate.h b/storage/src/vespa/storage/distributor/pendingclusterstate.h
index b742bd3bf46..a36743c126e 100644
--- a/storage/src/vespa/storage/distributor/pendingclusterstate.h
+++ b/storage/src/vespa/storage/distributor/pendingclusterstate.h
@@ -18,6 +18,7 @@ namespace storage::distributor {
 
 class DistributorMessageSender;
 class PendingBucketSpaceDbTransition;
+class DistributorBucketSpaceRepo;
 
 /**
  * Class used by BucketDBUpdater to track request bucket info
@@ -42,12 +43,13 @@ public:
             const framework::Clock& clock,
             const ClusterInformation::CSP& clusterInfo,
             DistributorMessageSender& sender,
+            DistributorBucketSpaceRepo &bucketSpaceRepo,
             const std::shared_ptr<api::SetSystemStateCommand>& newStateCmd,
             const std::unordered_set<uint16_t>& outdatedNodes,
             api::Timestamp creationTimestamp)
     {
         return std::unique_ptr<PendingClusterState>(
-                new PendingClusterState(clock, clusterInfo, sender, newStateCmd,
+                new PendingClusterState(clock, clusterInfo, sender, bucketSpaceRepo, newStateCmd,
                                         outdatedNodes,
                                         creationTimestamp));
     }
@@ -60,10 +62,11 @@ public:
             const framework::Clock& clock,
             const ClusterInformation::CSP& clusterInfo,
             DistributorMessageSender& sender,
+            DistributorBucketSpaceRepo &bucketSpaceRepo,
             api::Timestamp creationTimestamp)
     {
         return std::unique_ptr<PendingClusterState>(
-                new PendingClusterState(clock, clusterInfo, sender, creationTimestamp));
+                new PendingClusterState(clock, clusterInfo, sender, bucketSpaceRepo, creationTimestamp));
     }
 
     PendingClusterState(const PendingClusterState &) = delete;
@@ -150,6 +153,7 @@ private:
             const framework::Clock&,
             const ClusterInformation::CSP& clusterInfo,
             DistributorMessageSender& sender,
+            DistributorBucketSpaceRepo &bucketSpaceRepo,
             const std::shared_ptr<api::SetSystemStateCommand>& newStateCmd,
             const std::unordered_set<uint16_t>& outdatedNodes,
             api::Timestamp creationTimestamp);
@@ -162,8 +166,10 @@ private:
             const framework::Clock&,
             const ClusterInformation::CSP& clusterInfo,
             DistributorMessageSender& sender,
+            DistributorBucketSpaceRepo &bucketSpaceRepo,
             api::Timestamp creationTimestamp);
 
+    void constructorHelper();
     void logConstructionInformation() const;
     void requestNode(uint16_t node);
     bool distributorChanged(const lib::ClusterState& oldState, const lib::ClusterState& newState);
@@ -209,6 +215,7 @@ private:
     api::Timestamp _creationTimestamp;
 
     DistributorMessageSender& _sender;
+    DistributorBucketSpaceRepo &_bucketSpaceRepo;
 
     bool _distributionChange;
     bool _bucketOwnershipTransfer;