Merge pull request #17834 from vespa-engine/geirst/impl-bucket-space-state-map

Implement class that provides mapping from bucket space to state for …

4 files changed, 143 insertions, 2 deletions

diff --git a/storage/src/vespa/storage/distributor/CMakeLists.txt b/storage/src/vespa/storage/distributor/CMakeLists.txt
index 82f5d462dd1..7b048e9f109 100644
--- a/storage/src/vespa/storage/distributor/CMakeLists.txt
+++ b/storage/src/vespa/storage/distributor/CMakeLists.txt
@@ -6,6 +6,7 @@ vespa_add_library(storage_distributor
     bucket_db_prune_elision.cpp
     bucket_space_distribution_configs.cpp
     bucket_space_distribution_context.cpp
+    bucket_space_state_map.cpp
     bucketdbupdater.cpp
     bucketgctimecalculator.cpp
     bucketlistmerger.cpp
diff --git a/storage/src/vespa/storage/distributor/bucket_space_state_map.cpp b/storage/src/vespa/storage/distributor/bucket_space_state_map.cpp
new file mode 100644
index 00000000000..63c408f7e1e
--- /dev/null
+++ b/storage/src/vespa/storage/distributor/bucket_space_state_map.cpp
@@ -0,0 +1,67 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "bucket_space_state_map.h"
+#include <vespa/document/bucket/fixed_bucket_spaces.h>
+
+using document::BucketSpace;
+
+namespace storage::distributor {
+
+BucketSpaceState::BucketSpaceState()
+    : _cluster_state(),
+      _distribution()
+{
+}
+
+void
+BucketSpaceState::set_cluster_state(std::shared_ptr<const lib::ClusterState> cluster_state)
+{
+    _cluster_state = std::move(cluster_state);
+}
+
+void
+BucketSpaceState::set_distribution(std::shared_ptr<const lib::Distribution> distribution)
+{
+    _distribution = distribution;
+}
+
+BucketSpaceStateMap::BucketSpaceStateMap()
+    : _map()
+{
+    _map.emplace(document::FixedBucketSpaces::default_space(), std::make_unique<BucketSpaceState>());
+    _map.emplace(document::FixedBucketSpaces::global_space(), std::make_unique<BucketSpaceState>());
+}
+
+void
+BucketSpaceStateMap::set_cluster_state(std::shared_ptr<const lib::ClusterState> cluster_state)
+{
+    for (auto& space : _map) {
+        space.second->set_cluster_state(cluster_state);
+    }
+}
+
+void
+BucketSpaceStateMap::set_distribution(std::shared_ptr<const lib::Distribution> distribution)
+{
+    for (auto& space : _map) {
+        space.second->set_distribution(distribution);
+    }
+}
+
+const lib::ClusterState&
+BucketSpaceStateMap::get_cluster_state(document::BucketSpace space) const
+{
+    auto itr = _map.find(space);
+    assert(itr != _map.end());
+    return itr->second->get_cluster_state();
+}
+
+const lib::Distribution&
+BucketSpaceStateMap::get_distribution(document::BucketSpace space) const
+{
+    auto itr = _map.find(space);
+    assert(itr != _map.end());
+    return itr->second->get_distribution();
+}
+
+}
diff --git a/storage/src/vespa/storage/distributor/bucket_space_state_map.h b/storage/src/vespa/storage/distributor/bucket_space_state_map.h
new file mode 100644
index 00000000000..57eac9eac0d
--- /dev/null
+++ b/storage/src/vespa/storage/distributor/bucket_space_state_map.h
@@ -0,0 +1,74 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+#pragma once
+
+#include <vespa/document/bucket/bucketspace.h>
+#include <vespa/persistence/spi/clusterstate.h>
+#include <vespa/vdslib/distribution/distribution.h>
+#include <cassert>
+#include <unordered_map>
+
+namespace storage::lib {
+    class ClusterState;
+    class Distribution;
+}
+
+namespace storage::distributor {
+
+/**
+ * Represents cluster state and distribution for a given bucket space.
+ * TODO STRIPE: Make DistributorBucketSpace inherit this class.
+ */
+class BucketSpaceState {
+private:
+    std::shared_ptr<const lib::ClusterState> _cluster_state;
+    std::shared_ptr<const lib::Distribution> _distribution;
+
+public:
+    explicit BucketSpaceState();
+
+    BucketSpaceState(const BucketSpaceState&) = delete;
+    BucketSpaceState& operator=(const BucketSpaceState&) = delete;
+    BucketSpaceState(BucketSpaceState&&) = delete;
+    BucketSpaceState& operator=(BucketSpaceState&&) = delete;
+
+    void set_cluster_state(std::shared_ptr<const lib::ClusterState> cluster_state);
+    void set_distribution(std::shared_ptr<const lib::Distribution> distribution);
+
+    const lib::ClusterState& get_cluster_state() const noexcept {
+        assert(_cluster_state);
+        return *_cluster_state;
+    }
+    const lib::Distribution& get_distribution() const noexcept {
+        assert(_distribution);
+        return *_distribution;
+    }
+};
+
+/**
+ * Provides mapping from bucket space to state for that space.
+ */
+class BucketSpaceStateMap {
+private:
+    using StateMap = std::unordered_map<document::BucketSpace, std::unique_ptr<BucketSpaceState>, document::BucketSpace::hash>;
+
+    StateMap _map;
+
+public:
+    explicit BucketSpaceStateMap();
+
+    BucketSpaceStateMap(const BucketSpaceStateMap&&) = delete;
+    BucketSpaceStateMap& operator=(const BucketSpaceStateMap&) = delete;
+    BucketSpaceStateMap(BucketSpaceStateMap&&) = delete;
+    BucketSpaceStateMap& operator=(BucketSpaceStateMap&&) = delete;
+
+    StateMap::const_iterator begin() const { return _map.begin(); }
+    StateMap::const_iterator end() const { return _map.end(); }
+
+    void set_cluster_state(std::shared_ptr<const lib::ClusterState> cluster_state);
+    void set_distribution(std::shared_ptr<const lib::Distribution> distribution);
+
+    const lib::ClusterState& get_cluster_state(document::BucketSpace space) const;
+    const lib::Distribution& get_distribution(document::BucketSpace space) const;
+};
+
+}
diff --git a/storage/src/vespa/storage/distributor/distributor_component.h b/storage/src/vespa/storage/distributor/distributor_component.h
index e01efeddf3f..68db5a3c483 100644
--- a/storage/src/vespa/storage/distributor/distributor_component.h
+++ b/storage/src/vespa/storage/distributor/distributor_component.h
@@ -22,8 +22,7 @@ class DistributorComponent : public storage::DistributorComponent,
                              public DistributorOperationContext {
 private:
     DistributorInterface& _distributor;
-    // TODO STRIPE: When legacy mode is removed, replace this with mapping from BucketSpace to struct with
-    //              lib::ClusterState and lib::Distribution (need by BucketDBUpdater).
+    // TODO STRIPE: When legacy mode is removed, replace this with BucketSpaceStateMap.
     std::unique_ptr<DistributorBucketSpaceRepo> _bucket_space_repo;
     std::unique_ptr<DistributorBucketSpaceRepo> _read_only_bucket_space_repo;