Use ContentBucketSpace to get Distribution.

4 files changed, 10 insertions, 9 deletions

diff --git a/storage/src/vespa/storage/bucketdb/bucketmanager.cpp b/storage/src/vespa/storage/bucketdb/bucketmanager.cpp
index 5e7cf4af046..00fa5c95c9b 100644
--- a/storage/src/vespa/storage/bucketdb/bucketmanager.cpp
+++ b/storage/src/vespa/storage/bucketdb/bucketmanager.cpp
@@ -104,14 +104,14 @@ namespace {
         DistributorStateCache _state;
         std::unordered_map<uint16_t, ResultArray>& _result;
         const document::BucketIdFactory& _factory;
-        std::shared_ptr<lib::Distribution> _storageDistribution;
+        std::shared_ptr<const lib::Distribution> _storageDistribution;
 
     public:
         DistributorInfoGatherer(
                 const lib::ClusterState& systemState,
                 std::unordered_map<uint16_t, ResultArray>& result,
                 const document::BucketIdFactory& factory,
-                std::shared_ptr<lib::Distribution> distribution)
+                std::shared_ptr<const lib::Distribution> distribution)
             : _state(*distribution, systemState),
               _result(result),
               _factory(factory),
@@ -513,7 +513,7 @@ BucketManager::processRequestBucketInfoCommands(document::BucketSpace bucketSpac
     typedef std::shared_ptr<api::RequestBucketInfoCommand> RBISP;
     std::map<uint16_t, RBISP> requests;
 
-    lib::Distribution::SP distribution(_component.getDistribution());
+    auto distribution(_component.getBucketSpaceRepo().get(bucketSpace).getDistribution());
     lib::ClusterState::CSP clusterState(
             _component.getStateUpdater().getSystemState());
     assert(clusterState.get());
diff --git a/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.cpp b/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.cpp
index eaebba9ccc5..c4f566f0286 100644
--- a/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.cpp
+++ b/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.cpp
@@ -68,7 +68,6 @@ StorageBucketDBInitializer::System::System(
       _partitions(partitions),
       _bucketSpaceRepo(_component.getBucketSpaceRepo()),
       _nodeIndex(_component.getIndex()),
-      _distribution(*_component.getDistribution()),
       _nodeState()
 {
     // Is this correct? We should get the node state from the node state updater
@@ -376,6 +375,7 @@ StorageBucketDBInitializer::registerBucket(const document::Bucket &bucket,
                                            api::BucketInfo bucketInfo)
 {
     document::BucketId bucketId(bucket.getBucketId());
+    const auto &contentBucketSpace(_system._bucketSpaceRepo.get(bucket.getBucketSpace()));
     StorBucketDatabase::WrappedEntry entry(_system.getBucketDatabase(bucket.getBucketSpace()).get(
                 bucketId, "StorageBucketDBInitializer::registerBucket",
                 StorBucketDatabase::CREATE_IF_NONEXISTING));
@@ -403,7 +403,8 @@ StorageBucketDBInitializer::registerBucket(const document::Bucket &bucket,
             return;
         }
         uint32_t keepOnDisk, joinFromDisk;
-        if (_system._distribution.getPreferredAvailableDisk(
+        auto distribution(contentBucketSpace.getDistribution());
+        if (distribution->getPreferredAvailableDisk(
                 _system._nodeState, _system._nodeIndex,
                 bucketId.stripUnused()) == partition)
         {
@@ -429,7 +430,8 @@ StorageBucketDBInitializer::registerBucket(const document::Bucket &bucket,
             bucketId.toString().c_str(), int(partition));
         entry->disk = partition;
         entry.write();
-        uint16_t disk(_system._distribution.getIdealDisk(
+        auto distribution(contentBucketSpace.getDistribution());
+        uint16_t disk(distribution->getIdealDisk(
                 _system._nodeState, _system._nodeIndex, bucketId.stripUnused(),
                 lib::Distribution::IDEAL_DISK_EVEN_IF_DOWN));
         if (disk != partition) {
diff --git a/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.h b/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.h
index 15fbd5172cb..eeb86eaff55 100644
--- a/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.h
+++ b/storage/src/vespa/storage/bucketdb/storagebucketdbinitializer.h
@@ -80,7 +80,6 @@ class StorageBucketDBInitializer : public StorageLink,
         const spi::PartitionStateList& _partitions;
         const ContentBucketSpaceRepo& _bucketSpaceRepo;
         uint32_t _nodeIndex;
-        lib::Distribution& _distribution;
         lib::NodeState _nodeState; // Disk info for ideal state calculations
         framework::Thread::UP _thread;
 
diff --git a/storage/src/vespa/storage/common/servicelayercomponent.cpp b/storage/src/vespa/storage/common/servicelayercomponent.cpp
index 68c41536f97..11311a4d189 100644
--- a/storage/src/vespa/storage/common/servicelayercomponent.cpp
+++ b/storage/src/vespa/storage/common/servicelayercomponent.cpp
@@ -28,7 +28,7 @@ ServiceLayerComponent::getBucketDatabase(BucketSpace bucketSpace) const
 uint16_t
 ServiceLayerComponent::getIdealPartition(const document::Bucket& bucket) const
 {
-    return getDistribution()->getIdealDisk(
+    return _bucketSpaceRepo->get(bucket.getBucketSpace()).getDistribution()->getIdealDisk(
             *getStateUpdater().getReportedNodeState(), getIndex(), bucket.getBucketId(),
             lib::Distribution::IDEAL_DISK_EVEN_IF_DOWN);
 }
@@ -37,7 +37,7 @@ uint16_t
 ServiceLayerComponent::getPreferredAvailablePartition(
         const document::Bucket& bucket) const
 {
-    return getDistribution()->getPreferredAvailableDisk(
+    return _bucketSpaceRepo->get(bucket.getBucketSpace()).getDistribution()->getPreferredAvailableDisk(
             *getStateUpdater().getReportedNodeState(), getIndex(), bucket.getBucketId());
 }