partition_balancer: propagate an abort source into planner

.. to make the planner tick abortable.

src/v/cluster/partition_balancer_backend.cc

@@ -34,6 +34,12 @@ namespace cluster {
 static constexpr std::chrono::seconds controller_stm_sync_timeout = 10s;
 static constexpr std::chrono::seconds add_move_cmd_timeout = 10s;
 
+class balancer_tick_aborted_exception final : public std::runtime_error {
+public:
+    explicit balancer_tick_aborted_exception(const std::string& msg)
+      : std::runtime_error(msg) {}
+};
+
 partition_balancer_backend::partition_balancer_backend(
   consensus_ptr raft0,
   ss::sharded<controller_stm>& controller_stm,
@@ -138,12 +144,21 @@ void partition_balancer_backend::on_topic_table_update() {
 }
 
 void partition_balancer_backend::tick() {
-    ssx::background = ssx::spawn_with_gate_then(_gate, [this] {
-                          return do_tick().finally(
-                            [this] { maybe_rearm_timer(); });
-                      }).handle_exception([](const std::exception_ptr& e) {
-        vlog(clusterlog.warn, "tick error: {}", e);
-    });
+    ssx::background
+      = ssx::spawn_with_gate_then(
+          _gate,
+          [this] {
+              return do_tick().finally([this] {
+                  _tick_in_progress = {};
+                  maybe_rearm_timer();
+              });
+          })
+          .handle_exception_type([](balancer_tick_aborted_exception& e) {
+              vlog(clusterlog.info, "tick aborted, reason: {}", e.what());
+          })
+          .handle_exception([](const std::exception_ptr& e) {
+              vlog(clusterlog.warn, "tick error: {}", e);
+          });
 }
 
 ss::future<> partition_balancer_backend::stop() {
@@ -152,6 +167,10 @@ ss::future<> partition_balancer_backend::stop() {
     _state.members().unregister_members_updated_notification(_member_updates);
     _timer.cancel();
     _lock.broken();
+    if (_tick_in_progress) {
+        _tick_in_progress->request_abort_ex(
+          balancer_tick_aborted_exception{"shutting down"});
+    }
     return _gate.close();
 }
 
@@ -176,6 +195,8 @@ ss::future<> partition_balancer_backend::do_tick() {
         co_return;
     }
 
+    _tick_in_progress = ss::abort_source{};
+
     auto health_report = co_await _health_monitor.get_cluster_health(
       cluster_report_filter{},
       force_refresh::no,
@@ -213,7 +234,7 @@ ss::future<> partition_balancer_backend::do_tick() {
             .node_responsiveness_timeout = node_responsiveness_timeout},
           _state,
           _partition_allocator)
-          .plan_actions(health_report.value());
+          .plan_actions(health_report.value(), _tick_in_progress.value());
 
     _last_leader_term = _raft0->term();
     _last_tick_time = clock_t::now();
@@ -252,6 +273,7 @@ ss::future<> partition_balancer_backend::do_tick() {
 
     co_await ss::max_concurrent_for_each(
       plan_data.cancellations, 32, [this, current_term](model::ntp& ntp) {
+          _tick_in_progress->check();
           auto f = _topics_frontend.cancel_moving_partition_replicas(
             ntp,
             model::timeout_clock::now() + add_move_cmd_timeout,
@@ -272,6 +294,7 @@ ss::future<> partition_balancer_backend::do_tick() {
       plan_data.reassignments,
       32,
       [this, current_term](ntp_reassignment& reassignment) {
+          _tick_in_progress->check();
           auto f = _topics_frontend.move_partition_replicas(
             reassignment.ntp,
             reassignment.allocated.replicas(),

src/v/cluster/partition_balancer_backend.h

@@ -110,6 +110,7 @@ class partition_balancer_backend {
     ss::timer<clock_t> _timer;
     notification_id_type _topic_table_updates;
     notification_id_type _member_updates;
+    std::optional<ss::abort_source> _tick_in_progress;
 };
 
 } // namespace cluster

src/v/cluster/partition_balancer_planner.cc

@@ -107,13 +107,17 @@ class partition_balancer_planner::request_context {
         }
     }
 
-    ss::future<> maybe_yield() { co_await ss::coroutine::maybe_yield(); }
+    ss::future<> maybe_yield() {
+        co_await ss::coroutine::maybe_yield();
+        _as.check();
+    }
 
 private:
     friend class partition_balancer_planner;
 
-    request_context(partition_balancer_planner& parent)
-      : _parent(parent) {}
+    request_context(partition_balancer_planner& parent, ss::abort_source& as)
+      : _parent(parent)
+      , _as(as) {}
 
     bool all_reports_received() const;
 
@@ -128,7 +132,6 @@ class partition_balancer_planner::request_context {
     // returns true if the failure can be logged
     bool increment_failure_count();
 
-private:
     partition_balancer_planner& _parent;
     absl::node_hash_map<model::ntp, size_t> _ntp2size;
     absl::node_hash_map<model::ntp, absl::flat_hash_map<model::node_id, size_t>>
@@ -137,6 +140,7 @@ class partition_balancer_planner::request_context {
     uint64_t _planned_moves_size_bytes = 0;
     size_t _failed_actions_count = 0;
     absl::node_hash_set<model::ntp> _cancellations;
+    ss::abort_source& _as;
 };
 
 void partition_balancer_planner::init_per_node_state(
@@ -1198,8 +1202,8 @@ void partition_balancer_planner::request_context::collect_actions(
 
 ss::future<partition_balancer_planner::plan_data>
 partition_balancer_planner::plan_actions(
-  const cluster_health_report& health_report) {
-    request_context ctx(*this);
+  const cluster_health_report& health_report, ss::abort_source& as) {
+    request_context ctx(*this, as);
     plan_data result;
 
     init_per_node_state(health_report, ctx, result);

src/v/cluster/partition_balancer_planner.h

@@ -74,7 +74,8 @@ class partition_balancer_planner {
         status status = status::empty;
     };
 
-    ss::future<plan_data> plan_actions(const cluster_health_report&);
+    ss::future<plan_data>
+    plan_actions(const cluster_health_report&, ss::abort_source&);
 
 private:
     class request_context;

src/v/cluster/tests/partition_balancer_bench.cc

@@ -32,8 +32,9 @@ PERF_TEST_C(partition_balancer_planner_fixture, unavailable_nodes) {
 
     auto planner = make_planner();
 
+    abort_source as;
     perf_tests::start_measuring_time();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
     perf_tests::stop_measuring_time();
 
     const auto& reassignments = plan_data.reassignments;

src/v/cluster/tests/partition_balancer_planner_fixture.h

@@ -356,4 +356,5 @@ struct partition_balancer_planner_fixture {
 
     controller_workers workers;
     int last_node_idx{};
+    ss::abort_source as;
 };

src/v/cluster/tests/partition_balancer_planner_test.cc

@@ -72,7 +72,7 @@ FIXTURE_TEST(test_stable, partition_balancer_planner_fixture) {
     populate_node_status_table().get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
     check_violations(plan_data, {}, {});
     BOOST_REQUIRE_EQUAL(plan_data.reassignments.size(), 0);
 }
@@ -102,7 +102,7 @@ FIXTURE_TEST(test_node_down, partition_balancer_planner_fixture) {
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
 
@@ -140,7 +140,7 @@ FIXTURE_TEST(test_no_quorum_for_partition, partition_balancer_planner_fixture) {
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
     BOOST_REQUIRE_EQUAL(plan_data.reassignments.size(), 0);
 }
 
@@ -174,7 +174,7 @@ FIXTURE_TEST(
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
 
@@ -213,7 +213,7 @@ FIXTURE_TEST(
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
 
@@ -252,7 +252,7 @@ FIXTURE_TEST(
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, full_nodes);
 
@@ -284,7 +284,7 @@ FIXTURE_TEST(test_move_from_full_node, partition_balancer_planner_fixture) {
     populate_node_status_table().get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, {}, full_nodes);
 
@@ -325,7 +325,7 @@ FIXTURE_TEST(
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
 
@@ -373,7 +373,7 @@ FIXTURE_TEST(
     populate_node_status_table().get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
     check_violations(plan_data, {}, full_nodes);
 
     const auto& reassignments = plan_data.reassignments;
@@ -420,7 +420,7 @@ FIXTURE_TEST(
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
     check_violations(plan_data, unavailable_nodes, full_nodes);
 
     const auto& reassignments = plan_data.reassignments;
@@ -471,7 +471,7 @@ FIXTURE_TEST(test_move_part_of_replicas, partition_balancer_planner_fixture) {
     hr.node_reports[2].local_state.data_disk.free -= 2_MiB;
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, {}, full_nodes);
 
@@ -529,7 +529,7 @@ FIXTURE_TEST(
     }
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, {}, full_nodes);
 
@@ -579,7 +579,7 @@ FIXTURE_TEST(test_lot_of_partitions, partition_balancer_planner_fixture) {
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
     check_violations(plan_data, unavailable_nodes, {});
 
     const auto& reassignments = plan_data.reassignments;
@@ -638,7 +638,7 @@ FIXTURE_TEST(test_node_cancelation, partition_balancer_planner_fixture) {
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto planner_result = planner.plan_actions(hr).get();
+    auto planner_result = planner.plan_actions(hr, as).get();
 
     BOOST_REQUIRE_EQUAL(planner_result.reassignments.size(), 1);
 
@@ -660,7 +660,7 @@ FIXTURE_TEST(test_node_cancelation, partition_balancer_planner_fixture) {
     unavailable_nodes = {0, 3};
     populate_node_status_table(unavailable_nodes).get();
 
-    planner_result = planner.plan_actions(hr).get();
+    planner_result = planner.plan_actions(hr, as).get();
     BOOST_REQUIRE(planner_result.reassignments.size() == 0);
     BOOST_REQUIRE(planner_result.cancellations.size() == 1);
     BOOST_REQUIRE(planner_result.cancellations.front() == ntp);
@@ -699,7 +699,7 @@ FIXTURE_TEST(test_rack_awareness, partition_balancer_planner_fixture) {
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
 
@@ -734,7 +734,7 @@ FIXTURE_TEST(
     set_maintenance_mode(model::node_id{3});
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
     BOOST_REQUIRE_EQUAL(plan_data.reassignments.size(), 0);
@@ -764,7 +764,7 @@ FIXTURE_TEST(
     set_decommissioning(model::node_id{3});
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
     BOOST_REQUIRE_EQUAL(plan_data.reassignments.size(), 0);
@@ -794,7 +794,7 @@ FIXTURE_TEST(
     set_decommissioning(model::node_id{0});
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, unavailable_nodes, {});
     BOOST_REQUIRE_EQUAL(plan_data.reassignments.size(), 1);
@@ -866,7 +866,7 @@ FIXTURE_TEST(test_rack_awareness_repair, partition_balancer_planner_fixture) {
     populate_node_status_table().get();
 
     auto planner = make_planner();
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, {}, {});
     BOOST_REQUIRE_EQUAL(plan_data.reassignments.size(), 2);
@@ -897,7 +897,7 @@ FIXTURE_TEST(balancing_modes, partition_balancer_planner_fixture) {
     populate_node_status_table(unavailable_nodes).get();
 
     auto planner = make_planner(model::partition_autobalancing_mode::node_add);
-    auto plan_data = planner.plan_actions(hr).get();
+    auto plan_data = planner.plan_actions(hr, as).get();
 
     check_violations(plan_data, {}, {});
     BOOST_REQUIRE_EQUAL(plan_data.reassignments.size(), 0);