Use BucketSpaceStateMap to track cluster state and distribution in the top-level distributor.

This replaces the previous hack (needed in legacy mode) that used
DistributorBucketSpaceRepo to achieve the same.

3 files changed, 17 insertions, 17 deletions

diff --git a/storage/src/tests/distributor/top_level_bucket_db_updater_test.cpp b/storage/src/tests/distributor/top_level_bucket_db_updater_test.cpp
index d57ce228908..fe8a607c9ae 100644
--- a/storage/src/tests/distributor/top_level_bucket_db_updater_test.cpp
+++ b/storage/src/tests/distributor/top_level_bucket_db_updater_test.cpp
@@ -364,7 +364,7 @@ public:
             OutdatedNodesMap outdated_nodes_map;
             state = PendingClusterState::createForClusterStateChange(
                     clock, cluster_info, sender,
-                    owner.top_level_bucket_space_repo(),
+                    owner.bucket_space_states(),
                     cmd, outdated_nodes_map, api::Timestamp(1));
         }
 
@@ -374,7 +374,7 @@ public:
         {
             auto cluster_info = owner.create_cluster_info(old_cluster_state);
             state = PendingClusterState::createForDistributionChange(
-                    clock, cluster_info, sender, owner.top_level_bucket_space_repo(), api::Timestamp(1));
+                    clock, cluster_info, sender, owner.bucket_space_states(), api::Timestamp(1));
         }
     };
 
@@ -1389,7 +1389,7 @@ TopLevelBucketDBUpdaterTest::get_sent_nodes_distribution_changed(const std::stri
     auto cluster_info = create_cluster_info(old_cluster_state);
     std::unique_ptr<PendingClusterState> state(
             PendingClusterState::createForDistributionChange(
-                    clock, cluster_info, sender, top_level_bucket_space_repo(), api::Timestamp(1)));
+                    clock, cluster_info, sender, bucket_space_states(), api::Timestamp(1)));
 
     sort_sent_messages_by_index(sender);
 
@@ -1514,7 +1514,7 @@ TEST_F(TopLevelBucketDBUpdaterTest, pending_cluster_state_receive) {
     OutdatedNodesMap outdated_nodes_map;
     std::unique_ptr<PendingClusterState> state(
             PendingClusterState::createForClusterStateChange(
-                    clock, cluster_info, sender, top_level_bucket_space_repo(),
+                    clock, cluster_info, sender, bucket_space_states(),
                     cmd, outdated_nodes_map, api::Timestamp(1)));
 
     ASSERT_EQ(message_count(3), sender.commands().size());
@@ -1670,7 +1670,7 @@ TopLevelBucketDBUpdaterTest::merge_bucket_lists(
         auto cluster_info = create_cluster_info("cluster:d");
 
         auto state = PendingClusterState::createForClusterStateChange(
-                clock, cluster_info, sender, top_level_bucket_space_repo(),
+                clock, cluster_info, sender, bucket_space_states(),
                 cmd, outdated_nodes_map, before_time);
 
         parse_input_data(existing_data, before_time, *state, include_bucket_info);
@@ -1690,7 +1690,7 @@ TopLevelBucketDBUpdaterTest::merge_bucket_lists(
         auto cluster_info = create_cluster_info(old_state.toString());
 
         auto state = PendingClusterState::createForClusterStateChange(
-                        clock, cluster_info, sender, top_level_bucket_space_repo(),
+                        clock, cluster_info, sender, bucket_space_states(),
                         cmd, outdated_nodes_map, after_time);
 
         parse_input_data(new_data, after_time, *state, include_bucket_info);
diff --git a/storage/src/tests/distributor/top_level_distributor_test_util.cpp b/storage/src/tests/distributor/top_level_distributor_test_util.cpp
index e4002b3d1cb..6a0aa015ba4 100644
--- a/storage/src/tests/distributor/top_level_distributor_test_util.cpp
+++ b/storage/src/tests/distributor/top_level_distributor_test_util.cpp
@@ -265,16 +265,16 @@ TopLevelDistributorTestUtil::get_bucket(const document::BucketId& bId) const
     return stripe_bucket_database(stripe_index_of_bucket(bId)).get(bId);
 }
 
-DistributorBucketSpaceRepo&
-TopLevelDistributorTestUtil::top_level_bucket_space_repo() noexcept
+BucketSpaceStateMap&
+TopLevelDistributorTestUtil::bucket_space_states() noexcept
 {
-    return _distributor->_component.bucket_space_repo();
+    return _distributor->_component.bucket_space_states();
 }
 
-const DistributorBucketSpaceRepo&
-TopLevelDistributorTestUtil::top_level_bucket_space_repo() const noexcept
+const BucketSpaceStateMap&
+TopLevelDistributorTestUtil::bucket_space_states() const noexcept
 {
-    return _distributor->_component.bucket_space_repo();
+    return _distributor->_component.bucket_space_states();
 }
 
 std::unique_ptr<StripeAccessGuard>
diff --git a/storage/src/tests/distributor/top_level_distributor_test_util.h b/storage/src/tests/distributor/top_level_distributor_test_util.h
index e8794eb4199..6efc36a8215 100644
--- a/storage/src/tests/distributor/top_level_distributor_test_util.h
+++ b/storage/src/tests/distributor/top_level_distributor_test_util.h
@@ -17,7 +17,7 @@ namespace framework { struct TickingThreadPool; }
 
 namespace distributor {
 
-class TopLevelDistributor;
+class BucketSpaceStateMap;
 class DistributorBucketSpace;
 class DistributorBucketSpaceRepo;
 class DistributorMetricSet;
@@ -26,10 +26,11 @@ class DistributorStripe;
 class DistributorStripeComponent;
 class DistributorStripeOperationContext;
 class DistributorStripePool;
-class StripeAccessGuard;
 class IdealStateMetricSet;
 class Operation;
+class StripeAccessGuard;
 class TopLevelBucketDBUpdater;
+class TopLevelDistributor;
 
 class TopLevelDistributorTestUtil : private DoneInitializeHandler
 {
@@ -60,9 +61,8 @@ public:
     // As the above, but always inserts into default bucket space
     void add_nodes_to_stripe_bucket_db(const document::BucketId& id, const std::string& nodeStr);
 
-    // TODO STRIPE replace with BucketSpaceStateMap once legacy is gone
-    DistributorBucketSpaceRepo& top_level_bucket_space_repo() noexcept;
-    const DistributorBucketSpaceRepo& top_level_bucket_space_repo() const noexcept;
+    BucketSpaceStateMap& bucket_space_states() noexcept;
+    const BucketSpaceStateMap& bucket_space_states() const noexcept;
 
     std::unique_ptr<StripeAccessGuard> acquire_stripe_guard();