utest: partition balancer planner cancelation

redpanda-data · Jul 18, 2022 · 7f29bca · 7f29bca
1 parent 8bb7725
commit 7f29bca
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 0 deletions.
diff --git a/src/v/cluster/tests/partition_balancer_planner_fixture.h b/src/v/cluster/tests/partition_balancer_planner_fixture.h
@@ -133,6 +133,22 @@ struct partition_balancer_planner_fixture {
         }
     }
 
+    cluster::move_partition_replicas_cmd make_move_partition_replicas_cmd(
+      model::ntp ntp, std::vector<model::broker_shard> replica_set) {
+        return cluster::move_partition_replicas_cmd(
+          std::move(ntp), std::move(replica_set));
+    }
+
+    void move_partition_replicas(cluster::ntp_reassignments& reassignment) {
+        auto cmd = make_move_partition_replicas_cmd(
+          reassignment.ntp,
+          reassignment.allocation_units.get_assignments().front().replicas);
+        auto res = workers.dispatcher
+                     .apply_update(serialize_cmd(std::move(cmd)).get())
+                     .get();
+        BOOST_REQUIRE_EQUAL(res, cluster::errc::success);
+    }
+
     std::vector<raft::follower_metrics>
     create_follower_metrics(const std::set<size_t>& unavailable_nodes = {}) {
         std::vector<raft::follower_metrics> metrics;

diff --git a/src/v/cluster/tests/partition_balancer_planner_test.cc b/src/v/cluster/tests/partition_balancer_planner_test.cc
@@ -616,3 +616,58 @@ FIXTURE_TEST(test_lot_of_partitions, partition_balancer_planner_fixture) {
     BOOST_REQUIRE_EQUAL(node_3_counter, node_4_counter);
     BOOST_REQUIRE_EQUAL(node_4_counter, movement_batch_partitions_amount / 2);
 }
+
+/*
+ * 5 nodes; 1 topic; 1 node down
+ * Node 3 is down after planning
+ * Actual
+ *   node_0: partitions: 1; down: True; disk: unfilled;
+ *   node_1: partitions: 1; down: False; disk: unfilled;
+ *   node_2: partitions: 1; down: False; disk: unfilled;
+ *   node_3: partitions: 0; down: False; disk: unfilled;
+ * Expected
+ *   node_0: partitions: 1;
+ *   node_1: partitions: 1;
+ *   node_2: partitions: 1;
+ *   node_3: partitions: 0;
+ */
+FIXTURE_TEST(test_node_cancelation, partition_balancer_planner_fixture) {
+    vlog(logger.debug, "test_node_cancelation");
+    allocator_register_nodes(3);
+    create_topic("topic-1", 1, 3);
+    allocator_register_nodes(1);
+
+    auto hr = create_health_report();
+
+    std::set<size_t> unavailable_nodes = {0};
+    auto fm = create_follower_metrics(unavailable_nodes);
+
+    auto planner_result = planner.plan_reassignments(hr, fm);
+
+    BOOST_REQUIRE_EQUAL(planner_result.reassignments.size(), 1);
+
+    auto ntp = planner_result.reassignments.front().ntp;
+
+    std::unordered_set<model::node_id> expected_nodes(
+      {model::node_id(1), model::node_id(2), model::node_id(3)});
+
+    auto new_replicas = planner_result.reassignments.front()
+                          .allocation_units.get_assignments()
+                          .front()
+                          .replicas;
+    check_expected_assignments(new_replicas, expected_nodes);
+
+    for (auto& reassignment : planner_result.reassignments) {
+        move_partition_replicas(reassignment);
+    }
+
+    hr = create_health_report();
+
+    unavailable_nodes = {0, 3};
+    fm = create_follower_metrics(unavailable_nodes);
+
+    planner_result = planner.plan_reassignments(hr, fm);
+    BOOST_REQUIRE(planner_result.reassignments.size() == 0);
+    BOOST_REQUIRE(planner_result.cancellations.size() == 1);
+    BOOST_REQUIRE(planner_result.cancellations.front() == ntp);
+}