Add distribution to content bucket space.

5 files changed, 36 insertions, 2 deletions

diff --git a/storage/src/vespa/storage/common/content_bucket_space.cpp b/storage/src/vespa/storage/common/content_bucket_space.cpp
index b78be81c9de..0f84a2d0e38 100644
--- a/storage/src/vespa/storage/common/content_bucket_space.cpp
+++ b/storage/src/vespa/storage/common/content_bucket_space.cpp
@@ -5,8 +5,24 @@
 namespace storage {
 
 ContentBucketSpace::ContentBucketSpace()
-    : _bucketDatabase()
+    : _bucketDatabase(),
+      _lock(),
+      _distribution()
 {
 }
 
+void
+ContentBucketSpace::setDistribution(std::shared_ptr<const lib::Distribution> distribution)
+{
+    std::lock_guard<std::mutex> guard(_lock);
+    _distribution = std::move(distribution);
+}
+
+std::shared_ptr<const lib::Distribution>
+ContentBucketSpace::getDistribution() const
+{
+    std::lock_guard<std::mutex> guard(_lock);
+    return _distribution;
+}
+
 }
diff --git a/storage/src/vespa/storage/common/content_bucket_space.h b/storage/src/vespa/storage/common/content_bucket_space.h
index 2efb2eca06d..65524121ca6 100644
--- a/storage/src/vespa/storage/common/content_bucket_space.h
+++ b/storage/src/vespa/storage/common/content_bucket_space.h
@@ -2,20 +2,27 @@
 #pragma once
 
 #include <vespa/storage/bucketdb/storbucketdb.h>
+#include <mutex>
 
 namespace storage {
 
+namespace lib { class Distribution; }
+
 /**
  * Class representing a bucket space (with associated bucket database) on a content node.
  */
 class ContentBucketSpace {
 private:
     StorBucketDatabase _bucketDatabase;
+    mutable std::mutex _lock;
+    std::shared_ptr<const lib::Distribution> _distribution;
 
 public:
     using UP = std::unique_ptr<ContentBucketSpace>;
     ContentBucketSpace();
     StorBucketDatabase &bucketDatabase() { return _bucketDatabase; }
+    void setDistribution(std::shared_ptr<const lib::Distribution> distribution);
+    std::shared_ptr<const lib::Distribution> getDistribution() const;
 };
 
 }
diff --git a/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.cpp b/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.cpp
index da734c07c2d..6fd75e06d9d 100644
--- a/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.cpp
+++ b/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.cpp
@@ -36,4 +36,14 @@ ServiceLayerComponentRegisterImpl::setDiskCount(uint16_t count)
     }
 }
 
+void
+ServiceLayerComponentRegisterImpl::setDistribution(lib::Distribution::SP distribution)
+{
+    // For now, copy distribution to all content bucket spaces
+    for (const auto &elem : _bucketSpaceRepo) {
+        elem.second->setDistribution(distribution);
+    }
+    StorageComponentRegisterImpl::setDistribution(distribution);
+}
+
 } // storage
diff --git a/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.h b/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.h
index 5b3e54e3831..df4047c92c3 100644
--- a/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.h
+++ b/storage/src/vespa/storage/frameworkimpl/component/servicelayercomponentregisterimpl.h
@@ -37,6 +37,7 @@ public:
 
     void registerServiceLayerComponent(ServiceLayerManagedComponent&) override;
     void setDiskCount(uint16_t count);
+    virtual void setDistribution(lib::Distribution::SP distribution) override;
 };
 
 } // storage
diff --git a/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp b/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp
index dfd04b271b2..cf41a297541 100644
--- a/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp
+++ b/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp
@@ -978,7 +978,7 @@ FileStorManager::updateState()
     }
     for (const auto &elem : _component.getBucketSpaceRepo()) {
         BucketSpace bucketSpace(elem.first);
-        spi::ClusterState spiState(*state, _component.getIndex(), *_component.getDistribution());
+        spi::ClusterState spiState(*state, _component.getIndex(), *elem.second->getDistribution());
         _provider->setClusterState(bucketSpace, spiState);
     }
     _nodeUpInLastNodeStateSeenByProvider = nodeUp;