service: extend tablet_migration_info to handle rebuilds

Make tablet_migration_info::{src,dst} optional, so that it can be
reused by rebuild, for respectively leaving and pending replica.
This commit is contained in:
Aleksandra Martyniuk
2026-03-10 17:05:59 +01:00
parent d41c5a7db4
commit bcdab2e012
3 changed files with 31 additions and 20 deletions

View File

@@ -153,19 +153,27 @@ struct hash<locator::range_based_tablet_id> {
namespace locator {
/// Creates a new replica set with old_replica replaced by new_replica.
/// If there is no old_replica, the set is returned unchanged.
/// Returns a copy of the replica set with the following modifications:
/// - If both old_replica and new_replica are set, old_replica is substituted
/// with new_replica. If old_replica is not found in rs, the set is returned as-is.
/// - If only old_replica is set, it is removed from the result.
/// - If only new_replica is set, it is appended to the result.
inline
tablet_replica_set replace_replica(const tablet_replica_set& rs, tablet_replica old_replica, tablet_replica new_replica) {
tablet_replica_set replace_replica(const tablet_replica_set& rs, std::optional<tablet_replica> old_replica, std::optional<tablet_replica> new_replica) {
tablet_replica_set result;
result.reserve(rs.size());
for (auto&& r : rs) {
if (r == old_replica) {
result.push_back(new_replica);
if (old_replica.has_value() && r == old_replica.value()) {
if (new_replica.has_value()) {
result.push_back(new_replica.value());
}
} else {
result.push_back(r);
}
}
if (!old_replica.has_value() && new_replica.has_value()) {
result.push_back(new_replica.value());
}
return result;
}
@@ -383,8 +391,8 @@ bool is_post_cleanup(tablet_replica replica, const tablet_info& tinfo, const tab
struct tablet_migration_info {
locator::tablet_transition_kind kind;
locator::global_tablet_id tablet;
locator::tablet_replica src;
locator::tablet_replica dst;
std::optional<locator::tablet_replica> src;
std::optional<locator::tablet_replica> dst;
};
class tablet_map;

View File

@@ -1997,10 +1997,10 @@ static
future<> apply_plan(token_metadata& tm, const migration_plan& plan, service::topology& topology, shared_load_stats* load_stats) {
for (auto&& mig : plan.migrations()) {
co_await tm.tablets().mutate_tablet_map_async(mig.tablet.table, [&] (tablet_map& tmap) {
if (load_stats) {
if (load_stats && mig.src && mig.dst) {
global_tablet_id gid {mig.tablet.table, mig.tablet.tablet};
dht::token_range trange {tmap.get_token_range(mig.tablet.tablet)};
auto new_stats = load_stats->stats.migrate_tablet_size(mig.src.host, mig.dst.host, gid, trange);
auto new_stats = load_stats->stats.migrate_tablet_size(mig.src->host, mig.dst->host, gid, trange);
if (new_stats) {
load_stats->stats = std::move(*new_stats);
}
@@ -2661,13 +2661,13 @@ SEASTAR_THREAD_TEST_CASE(test_rack_list_conversion) {
for (auto& mig : plan.migrations()) {
testlog.info("Rack list colocation migration: {}", mig);
BOOST_REQUIRE(mig.kind == locator::tablet_transition_kind::migration);
BOOST_REQUIRE(mig.src.host == host3 || mig.src.host == host4);
if (mig.src.host == host3) {
BOOST_REQUIRE(mig.src && (mig.src->host == host3 || mig.src->host == host4));
if (mig.src->host == host3) {
BOOST_REQUIRE(mig.tablet.tablet == A);
BOOST_REQUIRE(mig.dst.host == host5 || mig.dst.host == host6);
BOOST_REQUIRE(mig.dst && (mig.dst->host == host5 || mig.dst->host == host6));
} else {
BOOST_REQUIRE(mig.tablet.tablet == B);
BOOST_REQUIRE(mig.dst.host == host1 || mig.dst.host == host2);
BOOST_REQUIRE(mig.dst && (mig.dst->host == host1 || mig.dst->host == host2));
}
}
}).get();
@@ -2745,8 +2745,9 @@ SEASTAR_THREAD_TEST_CASE(test_colocation_skipped_on_excluded_nodes) {
rebalance_tablets_as_in_progress(e, topo.get_shared_load_stats(), [&] (const migration_plan& plan) {
// Verify that only rebuilding migrations involve the excluded host.
for (auto&& mig : plan.migrations()) {
BOOST_REQUIRE_NE(mig.dst.host, host1);
if (mig.src.host == host1) {
BOOST_REQUIRE(mig.src && mig.dst);
BOOST_REQUIRE_NE(mig.dst->host, host1);
if (mig.src->host == host1) {
BOOST_REQUIRE(mig.kind == tablet_transition_kind::rebuild_v2);
}
}
@@ -2790,8 +2791,9 @@ SEASTAR_THREAD_TEST_CASE(test_no_intranode_migration_on_draining_node) {
rebalance_tablets_as_in_progress(e, topo.get_shared_load_stats(), [&] (const migration_plan& plan) {
// Verify no intra-node migrations on the draining host.
for (auto&& mig : plan.migrations()) {
if (mig.src.host == host1) {
BOOST_REQUIRE_NE(mig.dst.host, host1);
BOOST_REQUIRE(mig.src && mig.dst);
if (mig.src->host == host1) {
BOOST_REQUIRE_NE(mig.dst->host, host1);
}
}
return false;
@@ -4945,7 +4947,8 @@ SEASTAR_THREAD_TEST_CASE(test_load_balancer_ignores_hosts_with_incomplete_stats)
BOOST_REQUIRE(!plan.empty());
BOOST_REQUIRE(!plan.migrations().empty());
for (auto&& mig : plan.migrations()) {
BOOST_REQUIRE_EQUAL(mig.src.host, host2);
BOOST_REQUIRE(mig.src);
BOOST_REQUIRE_EQUAL(mig.src->host, host2);
}
}
}).get();

View File

@@ -122,10 +122,10 @@ future<> apply_plan(token_metadata& tm, const migration_plan& plan, locator::loa
return make_ready_future();
});
// Move tablet size in load_stats to account for the migration
if (mig.src.host != mig.dst.host) {
if (mig.src && mig.dst && mig.src->host != mig.dst->host) {
auto& tmap = tm.tablets().get_tablet_map(mig.tablet.table);
const dht::token_range trange = tmap.get_token_range(mig.tablet.tablet);
lw_shared_ptr<locator::load_stats> new_stats = load_stats.migrate_tablet_size(mig.src.host, mig.dst.host, mig.tablet, trange);
lw_shared_ptr<locator::load_stats> new_stats = load_stats.migrate_tablet_size(mig.src->host, mig.dst->host, mig.tablet, trange);
if (new_stats) {
load_stats = std::move(*new_stats);