From 906d2b817e552924f333c19a560d2cac1c8c0e0d Mon Sep 17 00:00:00 2001 From: Ferenc Szili Date: Wed, 6 May 2026 11:27:19 +0200 Subject: [PATCH 1/2] service: allow draining with forced capacity-based balancing When force_capacity_based_balancing is enabled, the tablet allocator balances by node and shard capacity rather than by tablet sizes. When the data needed for load balancing is incomplete, the balancer fails and waits until load_stats is available and correct for all the nodes. An exception to this is when a node is being drained and excluded: it is unreachable, and will not return. In this case the balancer has to do its best and ignore the missing data. This patch fixes a bug where forcing capacity based balancing made the balancer not ignore missing data in these cases, and instead abort the balancing. --- service/tablet_allocator.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/service/tablet_allocator.cc b/service/tablet_allocator.cc index 9b80e6a158..c3d0e27ed0 100644 --- a/service/tablet_allocator.cc +++ b/service/tablet_allocator.cc @@ -4242,10 +4242,10 @@ public: } } - // For size based balancing, only excluded nodes are allowed to have incomplete tablet stats + // Only excluded nodes are allowed to have incomplete tablet stats for (auto& [host, node] : nodes) { if (!_load_sketch->has_complete_data(host)) { - if (!_force_capacity_based_balancing && node.drained && node.node->is_excluded()) { + if (node.drained && node.node->is_excluded()) { _load_sketch->ignore_incomplete_data(host); } else { lblogger.info("Cannot balance because node {} (or more) has incomplete tablet stats", host); From f7bc8f5fa71c5a7a980088ce0f1bd0f685842e3f Mon Sep 17 00:00:00 2001 From: Ferenc Szili Date: Wed, 6 May 2026 11:27:19 +0200 Subject: [PATCH 2/2] test: boost: add drain test for forced capacity-based balancing Add a Boost unit test that forces capacity-based balancing through configuration and verifies that a drained and excluded node will be drained of its tablets when tablet size stats are missing. The test covers the regression where the allocator rejected the plan due to incomplete tablet stats, even though forced capacity-based balancing does not depend on tablet sizes. --- test/boost/tablets_test.cc | 52 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/test/boost/tablets_test.cc b/test/boost/tablets_test.cc index 0eaa5ab6c5..0797c37dd1 100644 --- a/test/boost/tablets_test.cc +++ b/test/boost/tablets_test.cc @@ -4136,6 +4136,58 @@ SEASTAR_THREAD_TEST_CASE(test_load_balancing_with_asymmetric_node_capacity) { }).get(); } +SEASTAR_THREAD_TEST_CASE(test_drain_with_forced_capacity_based_balancing_with_incomplete_data) { + auto cfg = tablet_cql_test_config(); + cfg.db_config->force_capacity_based_balancing.set(true); + + do_with_cql_env_thread([] (auto& e) { + logging::logger_registry().set_logger_level("load_balancer", logging::log_level::debug); + topology_builder topo(e); + + auto host1 = topo.add_node(node_state::removing, 8); + e.get_storage_service().local().mark_excluded({host1}).get(); + auto host2 = topo.add_node(node_state::normal, 1); + auto host3 = topo.add_node(node_state::normal, 7); + + const uint64_t capacity_unit = 100UL * 1024UL * 1024UL * 1024UL; + topo.get_shared_load_stats().set_capacity(host2, capacity_unit); + topo.get_shared_load_stats().set_capacity(host3, capacity_unit * 7); + + auto ks_name = add_keyspace(e, {{topo.dc(), 1}}, 16); + auto table1 = add_table(e, ks_name).get(); + + mutate_tablets(e, [&] (tablet_metadata& tmeta) -> future<> { + tablet_map tmap(16); + for (auto tid: tmap.tablet_ids()) { + tmap.set_tablet(tid, tablet_info { + tablet_replica_set { + tablet_replica {host1, 0}, + } + }); + } + tmeta.set_tablet_map(table1, std::move(tmap)); + co_return; + }); + + auto until_nodes_drained = [] (const migration_plan& plan) { + return !plan.has_nodes_to_drain(); + }; + + auto& stm = e.shared_token_metadata().local(); + + rebalance_tablets(e, &topo.get_shared_load_stats(), {}, until_nodes_drained); + + load_sketch load(stm.get()); + load.populate().get(); + + for (auto h: {host2, host3}) { + testlog.info("Checking host {}", h); + BOOST_REQUIRE_EQUAL(load.get_avg_tablet_count(h), 2); // 16 tablets / 8 shards = 2 tablets / shard + BOOST_REQUIRE_EQUAL(load.get_shard_tablet_count_imbalance(h), 0); + } + }, std::move(cfg)).get(); +} + SEASTAR_THREAD_TEST_CASE(test_load_balancer_disabling) { do_with_cql_env_thread([] (auto& e) { topology_builder topo(e);