c/shard_balancer: enable persistence when the feature barrier allows

redpanda-data · May 17, 2024 · faa76fd · faa76fd
1 parent caa89ca
commit faa76fd
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/src/v/cluster/controller.cc b/src/v/cluster/controller.cc
@@ -442,6 +442,7 @@ ss::future<> controller::start(
       .then([this] {
           return _shard_balancer.start_single(
             std::ref(_shard_placement),
+            std::ref(_feature_table),
             std::ref(_tp_state),
             std::ref(_backend));
       })

diff --git a/src/v/cluster/shard_balancer.cc b/src/v/cluster/shard_balancer.cc
@@ -20,9 +20,11 @@ namespace cluster {
 
 shard_balancer::shard_balancer(
   ss::sharded<shard_placement_table>& spt,
+  ss::sharded<features::feature_table>& features,
   ss::sharded<topic_table>& topics,
   ss::sharded<controller_backend>& cb)
   : _shard_placement(spt.local())
+  , _features(features.local())
   , _topics(topics)
   , _controller_backend(cb)
   , _self(*config::node().node_id()) {}
@@ -94,6 +96,11 @@ ss::future<> shard_balancer::start() {
         co_await do_assign_ntps();
     } else {
         co_await _shard_placement.initialize_from_topic_table(_topics, _self);
+
+        if (_features.is_active(
+              features::feature::shard_placement_persistence)) {
+            co_await _shard_placement.enable_persistence();
+        }
     }
 
     vassert(
@@ -146,6 +153,12 @@ ss::future<> shard_balancer::assign_fiber() {
             co_return;
         }
 
+        if (
+          _features.is_active(features::feature::shard_placement_persistence)
+          && !_shard_placement.is_persistence_enabled()) {
+            co_await _shard_placement.enable_persistence();
+        }
+
         co_await do_assign_ntps();
     }
 }

diff --git a/src/v/cluster/shard_balancer.h b/src/v/cluster/shard_balancer.h
@@ -32,6 +32,7 @@ class shard_balancer {
 
     shard_balancer(
       ss::sharded<shard_placement_table>&,
+      ss::sharded<features::feature_table>&,
       ss::sharded<topic_table>&,
       ss::sharded<controller_backend>&);
 
@@ -47,6 +48,7 @@ class shard_balancer {
 
 private:
     shard_placement_table& _shard_placement;
+    features::feature_table& _features;
     ss::sharded<topic_table>& _topics;
     ss::sharded<controller_backend>& _controller_backend;
     model::node_id _self;