Merge pull request #29959 from redpanda-data/stephan/greedy-improvements-fix

cluster: Improve greedy leader balancer

Author: StephanDollberg

src/v/cluster/scheduling/leader_balancer.cc

@@ -646,16 +646,26 @@ ss::future<ss::stop_iteration> leader_balancer::balance() {
           std::move(muted_index),
           std::move(preference_index));
         break;
-    case model::leader_balancer_mode::greedy:
+    case model::leader_balancer_mode::greedy: {
         vlog(clusterlog.debug, "using greedy strategy");
+        // Collect non-user topic IDs so the greedy strategy can exclude
+        // them from cross-topic global counts.
+        absl::flat_hash_set<leader_balancer_types::topic_id_t> internal_topics;
+        for (const auto& t : _topics.topics_map()) {
+            if (!model::is_user_topic(t.first)) {
+                internal_topics.emplace(t.second.get_revision());
+            }
+        }
         strategy = std::make_unique<
           leader_balancer_types::greedy_topic_aware_strategy>(
           _members.node_count(),
           std::move(index),
           std::move(group_id_to_topic),
+          std::move(internal_topics),
           std::move(muted_index),
           std::move(preference_index));
         break;
+    }
     case model::leader_balancer_mode::calibrated:
         vlog(clusterlog.debug, "using calibrated_hill_climbing strategy");
         strategy = std::make_unique<

src/v/cluster/scheduling/leader_balancer_greedy.cc

@@ -24,13 +24,15 @@ greedy_topic_aware_strategy::greedy_topic_aware_strategy(
   size_t node_count,
   index_type index,
   group_id_to_topic_id group_to_topic,
+  absl::flat_hash_set<topic_id_t> internal_topics,
   muted_index muted_index_value,
   std::optional<preference_index> preference_idx)
   : _muted_index(std::move(muted_index_value))
   , _group_to_topic(std::move(group_to_topic))
   , _shard_index(std::move(index))
   , _topic_distribution_constraint(_group_to_topic, _shard_index, _muted_index)
   , _shard_load_constraint(_shard_index, _muted_index)
+  , _internal_topics(std::move(internal_topics))
   , _node_count(node_count) {
     if (preference_idx) {
         _pinning_constraint.emplace(
@@ -85,38 +87,36 @@ std::vector<shard_load> greedy_topic_aware_strategy::stats() const {
 }
 
 // Lower-is-better on every element:
-// 1) excess above per-topic floor (0 while at or below
-//    fair share; positive once exceeded — this gates the
-//    remainder allocation so global counts can steer it)
+// 1) at_quota: 0 while broker is below its per-topic fair share,
+//    1 once reached — guarantees each broker fills to floor_quota
+//    before any gets more, so shard preferences in fields 2-3
+//    cannot cause broker imbalance.
 // 2) fewest topic leaders on this broker-shard
-// 3) fewest global leaders on this broker
-// 4) fewest topic leaders on this broker
-// 5) fewest global leaders on this broker-shard
-// 6) deterministic tie-break by broker-shard ordering
+// 3) fewest topic leaders on this broker (tie-break to spread
+//    across nodes when shards tie)
+// 4) fewest global leaders on this broker-shard
+// 5) deterministic tie-break by broker-shard ordering
 struct leader_preference {
-    size_t excess;
+    size_t at_quota;
     size_t shard_count;
-    size_t global_broker_count;
     size_t broker_count;
-    size_t global_shard_counts;
+    size_t global_shard_count;
     model::broker_shard replica;
 
     auto operator<=>(const leader_preference&) const = default;
 };
 
 void greedy_topic_aware_strategy::build_target_assignment() {
     // `partitions_by_topic` is the input to the greedy planner grouped by
-    // topic. Each entry stores the raft groups for one topic along with the
-    // broker-shards that can legally host leadership for those groups.
+    // topic. Each entry stores the raft groups for one topic along with
+    // the broker-shards that can legally host leadership for those groups.
     std::map<topic_id_t, chunked_vector<partition_info>> partitions_by_topic;
 
-    // `global_shard_counts` and `global_broker_counts` track how many target
-    // leaders have been assigned to each broker-shard and broker respectively
-    // across previously processed topics. They serve as cross-topic
-    // tie-breakers: when two replicas have equal per-topic counts, the one on
-    // the globally least-loaded broker (then shard) wins.
+    // `global_shard_counts` tracks how many target leaders have been
+    // assigned to each broker-shard across all topics processed so far.
+    // It serves as a cross-topic tie-breaker so that topics processed
+    // later inherit shard awareness from earlier ones.
     chunked_hash_map<model::broker_shard, size_t> global_shard_counts;
-    chunked_hash_map<model::node_id, size_t> global_broker_counts;
 
     for (const auto& [leader, groups] : _shard_index.shards()) {
         for (const auto& [group, replicas] : groups) {
@@ -126,18 +126,18 @@ void greedy_topic_aware_strategy::build_target_assignment() {
                   clusterlog.warn, "missing topic mapping for group {}", group);
                 continue;
             }
-
             partitions_by_topic[topic_iterator->second].emplace_back(
               group, topic_iterator->second, leader, replicas);
         }
     }
 
-    // Build targets topic-by-topic. For each topic the greedy walk assigns
-    // each partition to the replica that minimises a scoring tuple. A
-    // per-topic floor (partitions / brokers) gates the first element so
-    // that a broker that has already received its fair share is penalised,
-    // letting globally under-loaded brokers absorb the remainder.
+    // Build targets topic-by-topic. For each topic the greedy walk
+    // assigns each partition to the replica that minimises the
+    // leader_preference scoring tuple defined above. The at_quota
+    // gate guarantees broker balance while shard_count optimises
+    // shard placement within that constraint.
     for (auto& [topic, partitions] : partitions_by_topic) {
+        bool is_internal = _internal_topics.contains(topic);
         std::ranges::sort(
           partitions, std::ranges::less{}, &partition_info::group);
 
@@ -155,48 +155,35 @@ void greedy_topic_aware_strategy::build_target_assignment() {
             std::optional<leader_preference> best_assignment;
             for (const model::broker_shard& replica : partition.replicas) {
                 size_t broker_count = assigned_broker_counts[replica.node_id];
-                size_t excess = broker_count > floor_quota
-                                  ? broker_count - floor_quota
-                                  : 0;
-                auto candidate_assignment = leader_preference{
-                  .excess = excess,
+                auto candidate = leader_preference{
+                  .at_quota = broker_count >= floor_quota ? 1u : 0u,
                   .shard_count = assigned_shard_counts[replica],
-                  .global_broker_count = global_broker_counts[replica.node_id],
                   .broker_count = broker_count,
-                  .global_shard_counts = global_shard_counts[replica],
+                  .global_shard_count = global_shard_counts[replica],
                   .replica = replica};
-                if (
-                  !best_assignment || candidate_assignment < *best_assignment) {
-                    best_assignment = candidate_assignment;
+                if (!best_assignment || candidate < *best_assignment) {
+                    best_assignment = candidate;
                 }
             }
 
             if (!best_assignment.has_value()) {
                 // group has no replicas
                 continue;
             }
-            model::broker_shard selected_leader_shard
-              = best_assignment->replica;
-            if (partition.leader != selected_leader_shard) {
+            model::broker_shard selected = best_assignment->replica;
+            if (partition.leader != selected) {
                 _pending_moves.emplace_back(
-                  partition.group, partition.leader, selected_leader_shard);
+                  partition.group, partition.leader, selected);
+            }
+            assigned_broker_counts[selected.node_id] += 1;
+            assigned_shard_counts[selected] += 1;
+
+            // Internal topics (e.g. id_allocator, tx_manager) are
+            // excluded from global counts so their fixed placement
+            // does not distort the cross-topic balance of user topics.
+            if (!is_internal) {
+                global_shard_counts[selected] += 1;
             }
-            assigned_broker_counts[selected_leader_shard.node_id] += 1;
-            assigned_shard_counts[selected_leader_shard] += 1;
-        }
-
-        // Commit this topic's assignments to the global counters after
-        // the entire topic is processed, not per-partition.  If globals
-        // were updated per-partition, the first assignment to an
-        // under-loaded broker (e.g. global {3,3,2} → {3,3,3}) would
-        // erase the imbalance before the remaining partitions are
-        // scored, preventing global_broker_counts from steering subsequent
-        // remainder assignments to that broker.
-        for (const auto& [node, count] : assigned_broker_counts) {
-            global_broker_counts[node] += count;
-        }
-        for (const auto& [shard, count] : assigned_shard_counts) {
-            global_shard_counts[shard] += count;
         }
     }
 }

src/v/cluster/scheduling/leader_balancer_greedy.h

@@ -35,6 +35,7 @@ class greedy_topic_aware_strategy final : public leader_balancer_strategy {
       size_t node_count,
       index_type index,
       group_id_to_topic_id group_to_topic,
+      absl::flat_hash_set<topic_id_t> internal_topics,
       muted_index muted_index_value,
       std::optional<preference_index> preference_idx);
 
@@ -66,6 +67,8 @@ class greedy_topic_aware_strategy final : public leader_balancer_strategy {
 
     std::optional<pinning_constraint> _pinning_constraint;
 
+    absl::flat_hash_set<topic_id_t> _internal_topics;
+
     chunked_vector<reassignment> _pending_moves;
     size_t _next_pending{0};
     size_t _node_count;

src/v/cluster/tests/leader_balancer_greedy_test.cc

@@ -23,6 +23,7 @@
 #include <cstdint>
 #include <iterator>
 #include <memory>
+#include <random>
 #include <ranges>
 #include <tuple>
 #include <vector>
@@ -98,6 +99,7 @@ make_strategy_from_partitions(
       nodes.size(),
       std::move(index),
       std::move(group_to_topic),
+      absl::flat_hash_set<cluster::leader_balancer_types::topic_id_t>{},
       cluster::leader_balancer_types::muted_index{muted_nodes, {}},
       std::move(preference_idx));
 
@@ -125,18 +127,38 @@ make_greedy_strategy(
         }
     }
 
+    // Each node independently assigns partitions to shards in a
+    // balanced-but-shuffled order: a permutation of [0..shards-1]
+    // that repeats, so every shard gets the same number of replicas.
+    // Different nodes use different permutations (seeded by node id
+    // and topic) so replicas of the same partition land on different
+    // shards across nodes, matching real node-local shard placement.
+    auto make_shard_sequence = [&](int node, int topic) {
+        std::vector<uint32_t> perm(static_cast<size_t>(shards_per_broker));
+        std::iota(perm.begin(), perm.end(), uint32_t{0});
+        std::mt19937 rng(static_cast<uint32_t>(node * 137 + topic * 31));
+        std::ranges::shuffle(perm, rng);
+        return perm;
+    };
+
     for (int topic : std::views::iota(0, topic_count)) {
+        // Build per-node shard permutations and counters.
+        std::vector<std::vector<uint32_t>> node_perms;
+        std::vector<size_t> node_counters(static_cast<size_t>(broker_count), 0);
+        for (int n = 0; n < broker_count; ++n) {
+            node_perms.push_back(make_shard_sequence(n, topic));
+        }
+
         for (int partition : std::views::iota(0, partitions_per_topic)) {
             std::vector<model::broker_shard> replicas;
-            uint32_t topic_shard_offset = static_cast<uint32_t>(
-              std::max(1, shards_per_broker / 2));
             int broker_offset = partition + topic;
             for (int replica : std::views::iota(0, replication_factor)) {
-                replicas.push_back(bs(
-                  (broker_offset + replica) % broker_count,
-                  static_cast<uint32_t>(
-                    (partition + topic * topic_shard_offset)
-                    % shards_per_broker)));
+                int node = (broker_offset + replica) % broker_count;
+                auto& ctr = node_counters[static_cast<size_t>(node)];
+                auto& perm = node_perms[static_cast<size_t>(node)];
+                auto shard = perm[ctr % perm.size()];
+                ++ctr;
+                replicas.push_back(bs(node, shard));
             }
             partitions.push_back(
               partition_spec{
@@ -419,28 +441,27 @@ TEST(GreedyLeaderBalancerTest, TwoTopicsFourBrokersThreeShards) {
       balanced,
       0,
       std::array<std::array<int, 3>, 4>{{
-        {{2, 2, 1}},
+        {{2, 1, 1}},
         {{1, 2, 2}},
+        {{2, 2, 1}},
         {{1, 1, 2}},
-        {{2, 1, 1}},
       }});
     expect_shard_counts_per_broker(
       balanced,
       1,
       std::array<std::array<int, 3>, 4>{{
+        {{1, 1, 2}},
+        {{1, 2, 1}},
+        {{2, 2, 1}},
         {{2, 2, 1}},
-        {{1, 1, 1}},
-        {{1, 2, 2}},
-        {{2, 1, 2}},
       }});
-    // Not perfect
     expect_combined_shard_counts(
       balanced,
       std::array<std::array<int, 3>, 4>{{
+        {{3, 2, 3}},
+        {{2, 4, 3}},
         {{4, 4, 2}},
-        {{2, 3, 3}},
-        {{2, 3, 4}},
-        {{4, 2, 3}},
+        {{3, 3, 3}},
       }});
 }
 
@@ -468,10 +489,10 @@ TEST(GreedyLeaderBalancerTest, RemainderPartitions) {
     expect_shard_counts_per_broker(
       balanced_leaders,
       1,
-      std::array<std::array<int, 2>, 3>{{{{1, 1}}, {{2, 1}}, {{1, 2}}}});
+      std::array<std::array<int, 2>, 3>{{{{1, 2}}, {{1, 1}}, {{1, 2}}}});
     expect_combined_shard_counts(
       balanced_leaders,
-      std::array<std::array<int, 2>, 3>{{{{3, 2}}, {{3, 3}}, {{2, 3}}}});
+      std::array<std::array<int, 2>, 3>{{{{3, 3}}, {{2, 3}}, {{2, 3}}}});
 }
 
 TEST(GreedyLeaderBalancerTest, RemainderPartitionsThreeTopics) {
@@ -486,11 +507,11 @@ TEST(GreedyLeaderBalancerTest, RemainderPartitionsThreeTopics) {
     expect_shard_counts_per_broker(
       balanced_leaders,
       1,
-      std::array<std::array<int, 2>, 3>{{{{1, 1}}, {{2, 1}}, {{1, 2}}}});
+      std::array<std::array<int, 2>, 3>{{{{1, 2}}, {{1, 1}}, {{1, 2}}}});
     expect_shard_counts_per_broker(
       balanced_leaders,
       2,
-      std::array<std::array<int, 2>, 3>{{{{1, 2}}, {{1, 1}}, {{2, 1}}}});
+      std::array<std::array<int, 2>, 3>{{{{1, 1}}, {{2, 1}}, {{2, 1}}}});
     expect_combined_shard_counts(
       balanced_leaders,
       std::array<std::array<int, 2>, 3>{{{{4, 4}}, {{4, 4}}, {{4, 4}}}});
@@ -506,23 +527,23 @@ TEST(GreedyLeaderBalancerTest, BalancesShardsWithinBroker) {
       0,
       std::array<std::array<int, 4>, 3>{{
         {{2, 1, 1, 2}},
-        {{2, 2, 1, 1}},
-        {{1, 2, 2, 1}},
+        {{1, 1, 2, 2}},
+        {{1, 2, 1, 2}},
       }});
     expect_shard_counts_per_broker(
       balanced_leaders,
       1,
       std::array<std::array<int, 4>, 3>{{
         {{1, 2, 2, 1}},
-        {{1, 1, 2, 2}},
         {{2, 1, 1, 2}},
+        {{2, 2, 1, 1}},
       }});
     expect_combined_shard_counts(
       balanced_leaders,
       std::array<std::array<int, 4>, 3>{{
         {{3, 3, 3, 3}},
-        {{3, 3, 3, 3}},
-        {{3, 3, 3, 3}},
+        {{3, 2, 3, 4}},
+        {{3, 4, 2, 3}},
       }});
 }
 
@@ -536,17 +557,16 @@ TEST(GreedyLeaderBalancerTest, TwoTopicsSixBrokersTwoShardsRfThree) {
       balanced_leaders,
       0,
       std::array<std::array<int, 2>, 6>{
-        {{{2, 1}}, {{1, 2}}, {{2, 1}}, {{1, 2}}, {{2, 1}}, {{1, 2}}}});
+        {{{2, 1}}, {{1, 2}}, {{2, 1}}, {{2, 1}}, {{1, 2}}, {{1, 2}}}});
     expect_shard_counts_per_broker(
       balanced_leaders,
       1,
       std::array<std::array<int, 2>, 6>{
-        {{{2, 2}}, {{2, 1}}, {{1, 2}}, {{1, 1}}, {{1, 2}}, {{2, 1}}}});
+        {{{2, 2}}, {{2, 1}}, {{1, 2}}, {{1, 1}}, {{1, 2}}, {{1, 2}}}});
     expect_combined_shard_counts(
-      // Not perfect
       balanced_leaders,
       std::array<std::array<int, 2>, 6>{
-        {{{4, 3}}, {{3, 3}}, {{3, 3}}, {{2, 3}}, {{3, 3}}, {{3, 3}}}});
+        {{{4, 3}}, {{3, 3}}, {{3, 3}}, {{3, 2}}, {{2, 4}}, {{2, 4}}}});
 }
 
 TEST(GreedyLeaderBalancerTest, SkippedMoveDoesNotDesyncIndex) {