rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 14:56:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(meta): separate migration construction and drain child tracking (#10130)

Browse Source 
refactor(meta): separate migration construction and drain child tracking (#10098)

refactor(meta): separate migration construction and own drain child tracking


(cherry picked from commit 6a731a444f)

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-11-04 06:52:10 +00:00
committed by GitHub
parent 25003cea16
commit 4d8537a622
6 changed files with 634 additions and 520 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/meta-srv/src/procedure/drain_datanode.rs
+3 -5
View File
@@ -96,10 +96,10 @@ pub struct PersistentContext {
}
/// A dispatched migration sub-procedure that is not settled yet.
#[derive(Debug, Clone, Copy)]
#[derive(Debug, Clone)]
pub struct SubprocedureMeta {
/// The region migrated by the sub-procedure.
pub region_id: RegionId,
/// The regions migrated by the sub-procedure.
pub region_ids: Vec<RegionId>,
/// The procedure id of the sub-procedure.
pub procedure_id: ProcedureId,
}
@@ -118,8 +118,6 @@ pub struct VolatileContext {
/// its migration succeeded, when it was given up by [`DrainDatanodeStrategy::TryAll`],
/// when the drained datanode stopped leading it, or when another migration owns it.
pub plan: Option<DrainDatanodePlan>,
/// The migration sub-procedures of the current batch that are not settled yet.
pub inflight_subprocedures: Vec<SubprocedureMeta>,
/// The number of regions confirmed to have left the source while collecting batch results.
pub migrated: usize,
/// The regions given up by [`DrainDatanodeStrategy::TryAll`] in this run.
src/meta-srv/src/procedure/drain_datanode/check_drained.rs
+1 -1
View File
@@ -69,7 +69,7 @@ impl State for CheckDrained {
node_id, remaining_leader_regions, rescan_round
);
Ok((Box::new(Dispatch), Status::executing(false)))
Ok((Box::new(Dispatch::default()), Status::executing(false)))
}
fn as_any(&self) -> &dyn Any {
src/meta-srv/src/procedure/drain_datanode/dispatch.rs
+425 -34
View File
@@ -14,9 +14,10 @@
use std::any::Any;
use common_procedure::{Context as ProcedureContext, Status, SubprocedureRecovery};
use common_procedure::{Context as ProcedureContext, Status, SubprocedureRecovery, watcher};
use common_telemetry::{error, info, warn};
use serde::{Deserialize, Serialize};
use snafu::ResultExt;
use store_api::storage::RegionId;
use crate::error::{self, Result};
@@ -28,13 +29,94 @@ use crate::procedure::drain_datanode::{Context, DrainDatanodeStrategy, State, Su
/// Migrates the leader regions off the drained datanode in batches.
///
/// The state first settles the previous batch, then dispatches the next one. Sub-procedures
/// are tracked in the volatile context, and the suspension is not persisted, so a metasrv
/// are tracked in this state, and the suspension is not persisted, so a metasrv
/// restart rebuilds the plan by rescanning the datanode. Recovered migrations run
/// independently and are excluded by the migration tracker; regions that already left
/// the datanode are skipped by the route check in
/// [`Context::build_migration_procedure`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct Dispatch;
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub(crate) struct Dispatch {
/// Dispatched children whose results have not been collected yet.
#[serde(skip)]
inflight_subprocedures: Vec<SubprocedureMeta>,
}
impl Dispatch {
/// Waits for every migration in the batch before applying the drain strategy.
///
/// Keep the batch and its results unchanged until all fallible route checks succeed,
/// so a metadata retry cannot lose failures or apply progress twice.
async fn wait_for_inflight_subprocedures(
&mut self,
ctx: &mut Context,
procedure_ctx: &ProcedureContext,
strategy: DrainDatanodeStrategy,
) -> Result<()> {
let node_id = ctx.persistent_ctx.node_id;
let mut migrated = Vec::new();
let mut failed = Vec::new();
for meta in &self.inflight_subprocedures {
let receiver = procedure_ctx
.provider
.procedure_state_receiver(meta.procedure_id)
.await
.inspect_err(|e| {
error!(e; "Failed to get state receiver for migration sub-procedure {}", meta.procedure_id);
})
.context(error::QueryProcedureSnafu)?;
let Some(mut receiver) = receiver else {
warn!(
"Migration sub-procedure {} of regions {:?} is unknown, keep the region pending",
meta.procedure_id, meta.region_ids
);
continue;
};
match watcher::wait(&mut receiver).await {
Ok(_) => migrated.extend(meta.region_ids.iter().copied()),
Err(e) => {
error!(e; "Migration sub-procedure {} of regions {:?} failed", meta.procedure_id, meta.region_ids);
failed.extend(meta.region_ids.iter().copied());
}
}
}
let failed_count = failed.len();
let mut failed_regions = Vec::new();
for region_id in failed {
// A concurrent failover may have achieved the goal despite the child failure.
if utils::leads_region(&ctx.table_metadata_manager, node_id, region_id).await? {
failed_regions.push(region_id);
} else {
migrated.push(region_id);
}
}
METRIC_META_REGION_MIGRATION_FAIL.inc_by(failed_count as u64);
for region_id in migrated {
ctx.forget_leader_region(region_id);
ctx.volatile_ctx.migrated += 1;
info!("Region {} left datanode {}", region_id, node_id);
}
self.inflight_subprocedures.clear();
match strategy {
DrainDatanodeStrategy::FailFast if !failed_regions.is_empty() => {
error::DrainDatanodeAbortSnafu {
node_id,
failed_regions,
}
.fail()
}
DrainDatanodeStrategy::TryAll => {
for &region_id in &failed_regions {
ctx.forget_leader_region(region_id);
}
ctx.volatile_ctx.failed_regions.extend(failed_regions);
Ok(())
}
DrainDatanodeStrategy::FailFast => Ok(()),
}
}
}
#[async_trait::async_trait]
#[typetag::serde]
@@ -47,10 +129,10 @@ impl State for Dispatch {
let node_id = ctx.persistent_ctx.node_id;
// Collect the previous batch's results before dispatching more migrations.
if !ctx.volatile_ctx.inflight_subprocedures.is_empty() {
ctx.wait_for_inflight_subprocedures(procedure_ctx, ctx.persistent_ctx.strategy)
if !self.inflight_subprocedures.is_empty() {
self.wait_for_inflight_subprocedures(ctx, procedure_ctx, ctx.persistent_ctx.strategy)
.await?;
return Ok((Box::new(Dispatch), Status::executing(false)));
return Ok((Box::new(Dispatch::default()), Status::executing(false)));
}
utils::ensure_drain_deadline(ctx)?;
@@ -81,7 +163,7 @@ impl State for Dispatch {
region_id, node_id, child.id
);
inflight.push(SubprocedureMeta {
region_id,
region_ids: vec![region_id],
procedure_id: child.id,
});
targets.push(target_node_id);
@@ -152,12 +234,12 @@ impl State for Dispatch {
}
}
}
ctx.volatile_ctx.inflight_subprocedures = inflight;
self.inflight_subprocedures = inflight;
if children.is_empty() {
// The whole batch was handled without a sub-procedure, so keep draining the
// remaining leaders instead of spending a rescan round on them.
return Ok((Box::new(Dispatch), Status::executing(false)));
return Ok((Box::new(Dispatch::default()), Status::executing(false)));
}
for target_node_id in targets {
@@ -172,7 +254,9 @@ impl State for Dispatch {
// Recovered children continue independently. The parent rebuilds its plan and
// consults the migration tracker instead of restoring the batch's child links.
Ok((
Box::new(Dispatch),
Box::new(Self {
inflight_subprocedures: std::mem::take(&mut self.inflight_subprocedures),
}),
Status::Suspended {
subprocedures: children,
recovery: SubprocedureRecovery::Detach,
@@ -194,6 +278,8 @@ mod tests {
use crate::procedure::test_util::{procedure_context_with_receivers, procedure_state_receiver};
use api::v1::meta::heartbeat_request::NodeWorkloads;
use common_error::mock::MockError;
use common_error::status_code::StatusCode;
use common_meta::cluster::NodeInfo;
use common_meta::error::Result as MetaResult;
use common_meta::key::table_route::TableRouteKey;
@@ -201,8 +287,9 @@ mod tests {
use common_meta::kv_backend::test_util::MockKvBackendBuilder;
use common_meta::peer::{Peer, PeerDiscovery, PeerDiscoveryRef};
use common_meta::rpc::store::BatchGetResponse;
use common_procedure::{ProcedureId, ProcedureState, Status};
use common_procedure::{Error as ProcedureError, ProcedureId, ProcedureState, Status};
use store_api::storage::RegionId;
use tokio::sync::watch;
use crate::error;
use crate::key::NodeSchedulingRole;
@@ -221,6 +308,7 @@ mod tests {
#[tokio::test]
async fn test_drain_datanode_dispatch_transitions_to_check_drained() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
// Only a follower region: there is no leader region to migrate.
env.create_table(vec![follower_region(RegionId::new(1024, 1), 2, 1)])
@@ -232,7 +320,7 @@ mod tests {
));
let procedure_ctx = TestingEnv::procedure_context();
seed_plan(&mut ctx, vec![]);
let (next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
let (next, status) = dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
assert!(next.as_any().downcast_ref::<CheckDrained>().is_some());
assert!(matches!(status, Status::Executing { persist: false, .. }));
@@ -270,10 +358,11 @@ mod tests {
assert!(!persist);
assert_eq!(common_procedure::SubprocedureRecovery::Detach, recovery);
assert_eq!(1, subprocedures.len());
assert_eq!(1, ctx.volatile_ctx.inflight_subprocedures.len());
let dispatched = next.as_any().downcast_ref::<Dispatch>().unwrap();
assert_eq!(1, dispatched.inflight_subprocedures.len());
assert_eq!(
region_id,
ctx.volatile_ctx.inflight_subprocedures[0].region_id
dispatched.inflight_subprocedures[0].region_ids[0]
);
// The region is only removed after its migration is settled.
assert_eq!(vec![region_id], leader_regions(&ctx));
@@ -282,6 +371,7 @@ mod tests {
#[tokio::test]
async fn test_drain_datanode_dispatch_skips_region_moved_away() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
// The plan lists the region, but the table route says another datanode leads it.
@@ -294,12 +384,12 @@ mod tests {
seed_plan(&mut ctx, vec![region_id]);
let procedure_ctx = TestingEnv::procedure_context();
let (next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
let (next, status) = dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
assert!(matches!(status, Status::Executing { .. }));
assert!(next.as_any().downcast_ref::<Dispatch>().is_some());
assert!(leader_regions(&ctx).is_empty());
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
assert!(dispatch.inflight_subprocedures.is_empty());
}
/// Changes the route after the drain's first read, during target discovery.
@@ -366,6 +456,7 @@ mod tests {
}
async fn assert_drop_after_route_read_is_skipped(strategy: DrainDatanodeStrategy) {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
@@ -417,7 +508,7 @@ mod tests {
ctx.region_migration_manager = manager.clone();
seed_plan(&mut ctx, vec![region_id]);
let (next, status) = Dispatch
let (next, status) = dispatch
.next(&mut ctx, &TestingEnv::procedure_context())
.await
.unwrap();
@@ -426,7 +517,7 @@ mod tests {
assert!(matches!(status, Status::Executing { persist: false, .. }));
assert!(leader_regions(&ctx).is_empty());
assert!(ctx.volatile_ctx.failed_regions.is_empty());
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
assert!(dispatch.inflight_subprocedures.is_empty());
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(!manager.tracker().contains(region_id));
assert!(
@@ -441,6 +532,7 @@ mod tests {
#[tokio::test]
async fn test_drain_datanode_dispatch_skips_region_moved_or_dropped_during_discovery() {
let mut dispatch = Dispatch::default();
for strategy in [
DrainDatanodeStrategy::FailFast,
DrainDatanodeStrategy::TryAll,
@@ -460,7 +552,7 @@ mod tests {
});
seed_plan(&mut ctx, vec![region_id]);
let (next, status) = Dispatch
let (next, status) = dispatch
.next(&mut ctx, &TestingEnv::procedure_context())
.await
.unwrap();
@@ -469,7 +561,7 @@ mod tests {
assert!(matches!(status, Status::Executing { persist: false, .. }));
assert!(leader_regions(&ctx).is_empty());
assert!(ctx.volatile_ctx.failed_regions.is_empty());
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
assert!(dispatch.inflight_subprocedures.is_empty());
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(!env.region_migration_manager.tracker().contains(region_id));
}
@@ -478,6 +570,7 @@ mod tests {
#[tokio::test]
async fn test_drain_datanode_skips_region_whose_migration_is_running() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
@@ -506,7 +599,7 @@ mod tests {
.await
.expect("the manual migration should be buildable");
let (next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
let (next, status) = dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
// The region is ignored for this round instead of failing the drain: it leaves the plan
// and is neither migrated nor given up. The rescan re-adds it when the datanode still
@@ -515,7 +608,7 @@ mod tests {
assert!(next.as_any().downcast_ref::<Dispatch>().is_some());
assert!(leader_regions(&ctx).is_empty());
assert!(ctx.volatile_ctx.failed_regions.is_empty());
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
assert!(dispatch.inflight_subprocedures.is_empty());
}
#[tokio::test]
@@ -549,6 +642,7 @@ mod tests {
#[tokio::test]
async fn test_drain_datanode_waits_for_migrated_batch() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
@@ -558,22 +652,22 @@ mod tests {
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
region_id,
dispatch.inflight_subprocedures = vec![SubprocedureMeta {
region_ids: vec![region_id],
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
dispatch.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::Done { output: None }),
)]));
let (mut next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
let (mut next, status) = dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
assert!(matches!(status, Status::Executing { .. }));
assert!(next.as_any().downcast_ref::<Dispatch>().is_some());
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(leader_regions(&ctx).is_empty());
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
assert!(dispatch.inflight_subprocedures.is_empty());
// The next step checks the remaining regions.
let (next, _) = next.next(&mut ctx, &procedure_ctx).await.unwrap();
assert!(next.as_any().downcast_ref::<CheckDrained>().is_some());
@@ -581,6 +675,7 @@ mod tests {
#[tokio::test]
async fn test_drain_datanode_exceeded_deadline() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
@@ -591,7 +686,7 @@ mod tests {
seed_plan(&mut ctx, vec![region_id]);
let procedure_ctx = TestingEnv::procedure_context();
let error = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap_err();
let error = dispatch.next(&mut ctx, &procedure_ctx).await.unwrap_err();
assert!(matches!(error, error::Error::ExceededDeadline { .. }));
assert!(!error.is_retryable());
@@ -600,6 +695,7 @@ mod tests {
#[tokio::test]
async fn test_drain_datanode_deadline_does_not_interrupt_batch() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
@@ -608,22 +704,317 @@ mod tests {
persistent_ctx.deadline_ms = 0;
let mut ctx = env.context(persistent_ctx);
seed_plan(&mut ctx, vec![region_id]);
ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
region_id,
dispatch.inflight_subprocedures = vec![SubprocedureMeta {
region_ids: vec![region_id],
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
dispatch.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::Done { output: None }),
)]));
// The in-flight batch is settled before the deadline is checked.
let (_, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
let (_, status) = dispatch.next(&mut ctx, &procedure_ctx).await.unwrap();
assert!(matches!(status, Status::Executing { .. }));
assert_eq!(1, ctx.volatile_ctx.migrated);
// The next step observes the exceeded deadline.
let error = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap_err();
let error = dispatch.next(&mut ctx, &procedure_ctx).await.unwrap_err();
assert!(matches!(error, error::Error::ExceededDeadline { .. }));
}
#[tokio::test]
async fn test_drain_datanode_wait_failed_child_fail_fast() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
dispatch.inflight_subprocedures = vec![SubprocedureMeta {
region_ids: vec![region_id],
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
dispatch.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
MockError::new(StatusCode::Internal),
)))),
)]));
let error = dispatch
.wait_for_inflight_subprocedures(
&mut ctx,
&procedure_ctx,
DrainDatanodeStrategy::FailFast,
)
.await
.unwrap_err();
// The whole batch is settled before the procedure aborts.
assert!(matches!(error, error::Error::DrainDatanodeAbort { .. }));
assert_eq!(0, ctx.volatile_ctx.migrated);
assert_eq!(vec![region_id], leader_regions(&ctx));
assert!(dispatch.inflight_subprocedures.is_empty());
}
#[tokio::test]
async fn test_drain_datanode_wait_failed_child_try_all() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
dispatch.inflight_subprocedures = vec![SubprocedureMeta {
region_ids: vec![region_id],
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
dispatch.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
MockError::new(StatusCode::Internal),
)))),
)]));
dispatch
.wait_for_inflight_subprocedures(
&mut ctx,
&procedure_ctx,
DrainDatanodeStrategy::TryAll,
)
.await
.unwrap();
assert_eq!(vec![region_id], ctx.volatile_ctx.failed_regions);
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(leader_regions(&ctx).is_empty());
}
#[tokio::test]
async fn test_drain_datanode_wait_failed_child_counts_migrated_when_moved() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
// The route reports another datanode as the leader, which is what a failover that
// raced with the migration leaves behind.
env.create_table(vec![leader_region(region_id, 2)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
dispatch.inflight_subprocedures = vec![SubprocedureMeta {
region_ids: vec![region_id],
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
dispatch.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
MockError::new(StatusCode::Internal),
)))),
)]));
dispatch
.wait_for_inflight_subprocedures(
&mut ctx,
&procedure_ctx,
DrainDatanodeStrategy::FailFast,
)
.await
.unwrap();
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.failed_regions.is_empty());
assert!(leader_regions(&ctx).is_empty());
}
#[tokio::test]
async fn test_drain_datanode_wait_unknown_child_keeps_pending() {
let mut dispatch = Dispatch::default();
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
dispatch.inflight_subprocedures = vec![SubprocedureMeta {
region_ids: vec![region_id],
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::new());
dispatch
.wait_for_inflight_subprocedures(
&mut ctx,
&procedure_ctx,
DrainDatanodeStrategy::FailFast,
)
.await
.unwrap();
// The region stays pending and is dispatched again instead of being reported as
// failed.
assert_eq!(vec![region_id], leader_regions(&ctx));
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.failed_regions.is_empty());
}
#[tokio::test]
async fn test_drain_datanode_waits_for_entire_batch() {
let mut dispatch = Dispatch::default();
for strategy in [
DrainDatanodeStrategy::FailFast,
DrainDatanodeStrategy::TryAll,
] {
let env = TestEnv::new().await;
let failed_region = RegionId::new(1024, 1);
let succeeded_region = RegionId::new(1024, 2);
env.create_table(vec![
leader_region(failed_region, 1),
leader_region(succeeded_region, 1),
])
.await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![failed_region, succeeded_region]);
let failed_id = ProcedureId::random();
let pending_id = ProcedureId::random();
dispatch.inflight_subprocedures = vec![
SubprocedureMeta {
region_ids: vec![failed_region],
procedure_id: failed_id,
},
SubprocedureMeta {
region_ids: vec![succeeded_region],
procedure_id: pending_id,
},
];
let (sender, receiver) = watch::channel(ProcedureState::Running);
let procedure_ctx = procedure_context_with_receivers(HashMap::from([
(
failed_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(
ProcedureError::external(MockError::new(StatusCode::Internal)),
))),
),
(pending_id, receiver),
]));
let mut waiting = Box::pin(dispatch.wait_for_inflight_subprocedures(
&mut ctx,
&procedure_ctx,
strategy,
));
assert!(futures::poll!(&mut waiting).is_pending());
sender.send(ProcedureState::Done { output: None }).unwrap();
let result = waiting.await;
match strategy {
DrainDatanodeStrategy::FailFast => assert!(
matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region])
),
DrainDatanodeStrategy::TryAll => {
result.unwrap();
assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions);
}
}
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(dispatch.inflight_subprocedures.is_empty());
}
}
#[tokio::test]
async fn test_drain_datanode_wait_preserves_batch_on_metadata_retry() {
let mut dispatch = Dispatch::default();
use common_meta::key::TableMetadataManager;
use common_meta::kv_backend::test_util::MockKvBackendBuilder;
for strategy in [
DrainDatanodeStrategy::FailFast,
DrainDatanodeStrategy::TryAll,
] {
let env = TestEnv::new().await;
let failed_region = RegionId::new(1024, 1);
let succeeded_region = RegionId::new(1024, 2);
env.create_table(vec![
leader_region(failed_region, 1),
leader_region(succeeded_region, 1),
])
.await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![failed_region, succeeded_region]);
let failed_id = ProcedureId::random();
let succeeded_id = ProcedureId::random();
dispatch.inflight_subprocedures = vec![
SubprocedureMeta {
region_ids: vec![succeeded_region],
procedure_id: succeeded_id,
},
SubprocedureMeta {
region_ids: vec![failed_region],
procedure_id: failed_id,
},
];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([
(
succeeded_id,
procedure_state_receiver(ProcedureState::Done { output: None }),
),
(
failed_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(
ProcedureError::external(MockError::new(StatusCode::Internal)),
))),
),
]));
ctx.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new(
MockKvBackendBuilder::default()
.range_fn(Arc::new(|_| {
common_meta::error::ElectionNoLeaderSnafu.fail()
}))
.build()
.unwrap(),
)));
let error = dispatch
.wait_for_inflight_subprocedures(&mut ctx, &procedure_ctx, strategy)
.await
.unwrap_err();
assert!(error.is_retryable());
assert_eq!(2, dispatch.inflight_subprocedures.len());
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.failed_regions.is_empty());
assert_eq!(vec![failed_region, succeeded_region], leader_regions(&ctx));
ctx.table_metadata_manager = env.table_metadata_manager().clone();
let result = dispatch
.wait_for_inflight_subprocedures(&mut ctx, &procedure_ctx, strategy)
.await;
match strategy {
DrainDatanodeStrategy::FailFast => assert!(
matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region])
),
DrainDatanodeStrategy::TryAll => {
result.unwrap();
assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions);
}
}
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(dispatch.inflight_subprocedures.is_empty());
}
}
}
src/meta-srv/src/procedure/drain_datanode/migration.rs
+4 -386
View File
@@ -16,96 +16,18 @@ use common_meta::DatanodeId;
use common_meta::key::table_route::TableRouteValue;
use common_meta::peer::Peer;
use common_procedure::{
Context as ProcedureContext, ProcedureContext as SubmissionContext, ProcedureWithId, watcher,
Context as ProcedureContext, ProcedureContext as SubmissionContext, ProcedureWithId,
};
use common_telemetry::{error, info, warn};
use snafu::ResultExt;
use store_api::storage::RegionId;
use crate::error::{self, Result};
use crate::metrics::METRIC_META_REGION_MIGRATION_FAIL;
use crate::procedure::drain_datanode::{Context, DrainDatanodeStrategy, utils};
use crate::procedure::drain_datanode::{Context, utils};
use crate::procedure::region_migration::DEFAULT_REGION_MIGRATION_TIMEOUT;
use crate::procedure::region_migration::manager::{
RegionMigrationProcedureTask, RegionMigrationTriggerReason,
};
impl Context {
/// Waits for every migration in the batch before applying the drain strategy.
///
/// Keep the batch and its results unchanged until all fallible route checks succeed,
/// so a metadata retry cannot lose failures or apply progress twice.
pub async fn wait_for_inflight_subprocedures(
&mut self,
procedure_ctx: &ProcedureContext,
strategy: DrainDatanodeStrategy,
) -> Result<()> {
let node_id = self.persistent_ctx.node_id;
let mut migrated = Vec::new();
let mut failed = Vec::new();
for meta in &self.volatile_ctx.inflight_subprocedures {
let receiver = procedure_ctx
.provider
.procedure_state_receiver(meta.procedure_id)
.await
.inspect_err(|e| {
error!(e; "Failed to get state receiver for migration sub-procedure {}", meta.procedure_id);
})
.context(error::QueryProcedureSnafu)?;
let Some(mut receiver) = receiver else {
warn!(
"Migration sub-procedure {} of region {} is unknown, keep the region pending",
meta.procedure_id, meta.region_id
);
continue;
};
match watcher::wait(&mut receiver).await {
Ok(_) => migrated.push(meta.region_id),
Err(e) => {
error!(e; "Migration sub-procedure {} of region {} failed", meta.procedure_id, meta.region_id);
failed.push(meta.region_id);
}
}
}
let failed_count = failed.len();
let mut failed_regions = Vec::new();
for region_id in failed {
// A concurrent failover may have achieved the goal despite the child failure.
if utils::leads_region(&self.table_metadata_manager, node_id, region_id).await? {
failed_regions.push(region_id);
} else {
migrated.push(region_id);
}
}
METRIC_META_REGION_MIGRATION_FAIL.inc_by(failed_count as u64);
for region_id in migrated {
self.forget_leader_region(region_id);
self.volatile_ctx.migrated += 1;
info!("Region {} left datanode {}", region_id, node_id);
}
self.volatile_ctx.inflight_subprocedures.clear();
match strategy {
DrainDatanodeStrategy::FailFast if !failed_regions.is_empty() => {
error::DrainDatanodeAbortSnafu {
node_id,
failed_regions,
}
.fail()
}
DrainDatanodeStrategy::TryAll => {
for &region_id in &failed_regions {
self.forget_leader_region(region_id);
}
self.volatile_ctx.failed_regions.extend(failed_regions);
Ok(())
}
DrainDatanodeStrategy::FailFast => Ok(()),
}
}
}
impl Context {
/// Builds the migration sub-procedure of `region_id`.
///
@@ -169,18 +91,11 @@ impl Context {
#[cfg(test)]
mod tests {
use crate::error;
use crate::procedure::drain_datanode::DrainDatanodeStrategy;
use crate::procedure::drain_datanode::test_util::{
TestEnv, leader_region, leader_regions, new_persistent_context, seed_plan,
TestEnv, leader_region, new_persistent_context,
};
use crate::procedure::drain_datanode::{DrainDatanodeStrategy, SubprocedureMeta};
use crate::procedure::test_util::{procedure_context_with_receivers, procedure_state_receiver};
use common_error::mock::MockError;
use common_error::status_code::StatusCode;
use common_procedure::{Error as ProcedureError, ProcedureId, ProcedureState};
use std::collections::HashMap;
use std::sync::Arc;
use store_api::storage::RegionId;
use tokio::sync::watch;
#[tokio::test]
async fn test_build_migration_procedure_checks_route_and_guard() {
for leader_id in [1, 2] {
@@ -215,301 +130,4 @@ mod tests {
assert!(!env.region_migration_manager.tracker().contains(region_id));
}
}
#[tokio::test]
async fn test_drain_datanode_waits_for_migrated_batch() {
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
region_id,
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::Done { output: None }),
)]));
ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
.await
.unwrap();
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(leader_regions(&ctx).is_empty());
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
}
#[tokio::test]
async fn test_drain_datanode_wait_failed_child_fail_fast() {
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
region_id,
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
MockError::new(StatusCode::Internal),
)))),
)]));
let error = ctx
.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
.await
.unwrap_err();
// The whole batch is settled before the procedure aborts.
assert!(matches!(error, error::Error::DrainDatanodeAbort { .. }));
assert_eq!(0, ctx.volatile_ctx.migrated);
assert_eq!(vec![region_id], leader_regions(&ctx));
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
}
#[tokio::test]
async fn test_drain_datanode_wait_failed_child_try_all() {
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
region_id,
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
MockError::new(StatusCode::Internal),
)))),
)]));
ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::TryAll)
.await
.unwrap();
assert_eq!(vec![region_id], ctx.volatile_ctx.failed_regions);
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(leader_regions(&ctx).is_empty());
}
#[tokio::test]
async fn test_drain_datanode_wait_failed_child_counts_migrated_when_moved() {
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
// The route reports another datanode as the leader, which is what a failover that
// raced with the migration leaves behind.
env.create_table(vec![leader_region(region_id, 2)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
region_id,
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
MockError::new(StatusCode::Internal),
)))),
)]));
ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
.await
.unwrap();
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.failed_regions.is_empty());
assert!(leader_regions(&ctx).is_empty());
}
#[tokio::test]
async fn test_drain_datanode_wait_unknown_child_keeps_pending() {
let env = TestEnv::new().await;
let region_id = RegionId::new(1024, 1);
env.create_table(vec![leader_region(region_id, 1)]).await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![region_id]);
ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
region_id,
procedure_id: ProcedureId::random(),
}];
let procedure_ctx = procedure_context_with_receivers(HashMap::new());
ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
.await
.unwrap();
// The region stays pending and is dispatched again instead of being reported as
// failed.
assert_eq!(vec![region_id], leader_regions(&ctx));
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.failed_regions.is_empty());
}
#[tokio::test]
async fn test_drain_datanode_waits_for_entire_batch() {
for strategy in [
DrainDatanodeStrategy::FailFast,
DrainDatanodeStrategy::TryAll,
] {
let env = TestEnv::new().await;
let failed_region = RegionId::new(1024, 1);
let succeeded_region = RegionId::new(1024, 2);
env.create_table(vec![
leader_region(failed_region, 1),
leader_region(succeeded_region, 1),
])
.await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![failed_region, succeeded_region]);
let failed_id = ProcedureId::random();
let pending_id = ProcedureId::random();
ctx.volatile_ctx.inflight_subprocedures = vec![
SubprocedureMeta {
region_id: failed_region,
procedure_id: failed_id,
},
SubprocedureMeta {
region_id: succeeded_region,
procedure_id: pending_id,
},
];
let (sender, receiver) = watch::channel(ProcedureState::Running);
let procedure_ctx = procedure_context_with_receivers(HashMap::from([
(
failed_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(
ProcedureError::external(MockError::new(StatusCode::Internal)),
))),
),
(pending_id, receiver),
]));
let mut waiting =
Box::pin(ctx.wait_for_inflight_subprocedures(&procedure_ctx, strategy));
assert!(futures::poll!(&mut waiting).is_pending());
sender.send(ProcedureState::Done { output: None }).unwrap();
let result = waiting.await;
match strategy {
DrainDatanodeStrategy::FailFast => assert!(
matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region])
),
DrainDatanodeStrategy::TryAll => {
result.unwrap();
assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions);
}
}
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
}
}
#[tokio::test]
async fn test_drain_datanode_wait_preserves_batch_on_metadata_retry() {
use common_meta::key::TableMetadataManager;
use common_meta::kv_backend::test_util::MockKvBackendBuilder;
for strategy in [
DrainDatanodeStrategy::FailFast,
DrainDatanodeStrategy::TryAll,
] {
let env = TestEnv::new().await;
let failed_region = RegionId::new(1024, 1);
let succeeded_region = RegionId::new(1024, 2);
env.create_table(vec![
leader_region(failed_region, 1),
leader_region(succeeded_region, 1),
])
.await;
let mut ctx = env.context(new_persistent_context(
Some(2),
DrainDatanodeStrategy::FailFast,
));
seed_plan(&mut ctx, vec![failed_region, succeeded_region]);
let failed_id = ProcedureId::random();
let succeeded_id = ProcedureId::random();
ctx.volatile_ctx.inflight_subprocedures = vec![
SubprocedureMeta {
region_id: succeeded_region,
procedure_id: succeeded_id,
},
SubprocedureMeta {
region_id: failed_region,
procedure_id: failed_id,
},
];
let procedure_ctx = procedure_context_with_receivers(HashMap::from([
(
succeeded_id,
procedure_state_receiver(ProcedureState::Done { output: None }),
),
(
failed_id,
procedure_state_receiver(ProcedureState::failed(Arc::new(
ProcedureError::external(MockError::new(StatusCode::Internal)),
))),
),
]));
ctx.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new(
MockKvBackendBuilder::default()
.range_fn(Arc::new(|_| {
common_meta::error::ElectionNoLeaderSnafu.fail()
}))
.build()
.unwrap(),
)));
let error = ctx
.wait_for_inflight_subprocedures(&procedure_ctx, strategy)
.await
.unwrap_err();
assert!(error.is_retryable());
assert_eq!(2, ctx.volatile_ctx.inflight_subprocedures.len());
assert_eq!(0, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.failed_regions.is_empty());
assert_eq!(vec![failed_region, succeeded_region], leader_regions(&ctx));
ctx.table_metadata_manager = env.table_metadata_manager().clone();
let result = ctx
.wait_for_inflight_subprocedures(&procedure_ctx, strategy)
.await;
match strategy {
DrainDatanodeStrategy::FailFast => assert!(
matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region])
),
DrainDatanodeStrategy::TryAll => {
result.unwrap();
assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions);
}
}
assert_eq!(1, ctx.volatile_ctx.migrated);
assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
}
}
}
src/meta-srv/src/procedure/region_migration/manager.rs
+194 -92
View File
@@ -179,6 +179,14 @@ pub struct SubmitRegionMigrationTaskResult {
pub procedure_id: Option<ProcedureId>,
}
/// A constructed migration batch that has not been submitted.
pub(crate) struct BuiltMigrationBatch {
/// Region classifications from tracker, route, and table-info checks.
pub(crate) analysis: RegionMigrationAnalysis,
/// The procedure owning guards for the pending regions, if any remain.
pub(crate) procedure: Option<ProcedureWithId>,
}
impl RegionMigrationManager {
/// Returns new [`RegionMigrationManager`]
pub(crate) fn new(
@@ -393,33 +401,90 @@ impl RegionMigrationManager {
pub async fn submit_region_migration_task(
&self,
mut task: RegionMigrationTaskBatch,
task: RegionMigrationTaskBatch,
) -> Result<SubmitRegionMigrationTaskResult> {
let (migrating_region_ids, procedure_guards) = self.extract_running_regions(&mut task);
let procedure_context = ProcedureContext {
actor: None,
event_context: Some(PersistentEventContext::new(
task.trigger_reason.to_trigger_reason(),
)),
};
let from_peer_id = task.from_peer.id;
let to_peer_id = task.to_peer.id;
let BuiltMigrationBatch {
analysis,
procedure,
} = self.build_batch_procedure(procedure_context, task).await?;
let procedure_id = procedure.as_ref().map(|procedure| procedure.id);
if let Some(procedure) = procedure {
let procedure_id = procedure.id;
info!(
"Starting region migration procedure {procedure_id}, regions: {:?}, from_peer: {}, to_peer: {}",
analysis.pending, from_peer_id, to_peer_id
);
let mut watcher = self
.procedure_manager
.submit(procedure)
.await
.context(error::SubmitProcedureSnafu)?;
let num_regions = analysis.pending.len();
common_runtime::spawn_global(async move {
METRIC_META_REGION_MIGRATION_DATANODES
.with_label_values(&["src", &from_peer_id.to_string()])
.inc_by(num_regions as u64);
METRIC_META_REGION_MIGRATION_DATANODES
.with_label_values(&["desc", &to_peer_id.to_string()])
.inc_by(num_regions as u64);
if let Err(e) = watcher::wait(&mut watcher).await {
error!(e; "Failed to wait region migration procedure {procedure_id}");
METRIC_META_REGION_MIGRATION_FAIL.inc();
return;
}
info!("Region migration procedure {procedure_id} finished successfully");
});
}
let RegionMigrationAnalysis {
migrating,
migrated,
leader_changed,
peer_conflict,
mut table_not_found,
table_not_found,
region_not_found,
pending,
} = analyze_region_migration_task(&task, &self.context_factory.table_metadata_manager)
.await?;
if pending.is_empty() {
return Ok(SubmitRegionMigrationTaskResult {
migrated,
leader_changed,
peer_conflict,
table_not_found,
region_not_found,
migrating: migrating_region_ids,
submitted: vec![],
procedure_id: None,
} = analysis;
Ok(SubmitRegionMigrationTaskResult {
migrated,
leader_changed,
peer_conflict,
table_not_found,
region_not_found,
migrating,
submitted: pending,
procedure_id,
})
}
/// Builds a batch without submitting it. The returned procedure owns the guards
/// of its pending regions; skipped regions are reported separately.
pub(crate) async fn build_batch_procedure(
&self,
procedure_context: ProcedureContext,
mut task: RegionMigrationTaskBatch,
) -> Result<BuiltMigrationBatch> {
let (migrating, procedure_guards) = self.extract_running_regions(&mut task);
let mut analysis =
analyze_region_migration_task(&task, &self.context_factory.table_metadata_manager)
.await?;
analysis.migrating = migrating;
if analysis.pending.is_empty() {
return Ok(BuiltMigrationBatch {
analysis,
procedure: None,
});
}
// Updates the region ids to the pending region ids.
task.region_ids = pending;
task.region_ids = std::mem::take(&mut analysis.pending);
let table_regions = task.table_regions();
let table_ids = table_regions.keys().cloned().collect::<Vec<_>>();
let table_info_values = self
@@ -442,99 +507,41 @@ impl RegionMigrationManager {
}
None => {
task.region_ids.retain(|id| id.table_id() != table_id);
table_not_found.extend(regions);
analysis.table_not_found.extend(regions);
}
}
}
if task.region_ids.is_empty() {
return Ok(SubmitRegionMigrationTaskResult {
migrated,
leader_changed,
peer_conflict,
table_not_found,
region_not_found,
migrating: migrating_region_ids,
submitted: vec![],
procedure_id: None,
return Ok(BuiltMigrationBatch {
analysis,
procedure: None,
});
}
let submitting_region_ids = task.region_ids.clone();
// TODO(weny): Pass the actor when region migration task submission propagates it.
let procedure_context = ProcedureContext {
actor: None,
event_context: Some(PersistentEventContext::new(
task.trigger_reason.to_trigger_reason(),
)),
};
let procedure_id = self
.submit_procedure_inner(
procedure_context,
task,
procedure_guards,
catalog_and_schema,
)
.await?;
Ok(SubmitRegionMigrationTaskResult {
migrated,
leader_changed,
peer_conflict,
table_not_found,
region_not_found,
migrating: migrating_region_ids,
submitted: submitting_region_ids,
procedure_id: Some(procedure_id),
})
}
async fn submit_procedure_inner(
&self,
procedure_context: ProcedureContext,
task: RegionMigrationTaskBatch,
procedure_guards: Vec<RegionMigrationProcedureGuard>,
catalog_and_schema: Vec<(String, String)>,
) -> Result<ProcedureId> {
analysis.pending = task.region_ids.clone();
let pending = task.region_ids.iter().copied().collect::<HashSet<_>>();
let procedure_guards = procedure_guards
.into_iter()
.filter(|guard| pending.contains(&guard.region_id))
.collect();
let procedure = RegionMigrationProcedure::new(
PersistentContext::new(
catalog_and_schema,
task.from_peer.clone(),
task.to_peer.clone(),
task.region_ids.clone(),
task.from_peer,
task.to_peer,
task.region_ids,
task.timeout,
task.trigger_reason,
),
self.context_factory.clone(),
procedure_guards,
);
let procedure_with_id =
let procedure =
ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context);
let procedure_id = procedure_with_id.id;
info!("Starting region migration procedure {procedure_id} for {task}");
let mut watcher = self
.procedure_manager
.submit(procedure_with_id)
.await
.context(error::SubmitProcedureSnafu)?;
let num_region = task.region_ids.len();
common_runtime::spawn_global(async move {
METRIC_META_REGION_MIGRATION_DATANODES
.with_label_values(&["src", &task.from_peer.id.to_string()])
.inc_by(num_region as u64);
METRIC_META_REGION_MIGRATION_DATANODES
.with_label_values(&["desc", &task.to_peer.id.to_string()])
.inc_by(num_region as u64);
if let Err(e) = watcher::wait(&mut watcher).await {
error!(e; "Failed to wait region migration procedure {procedure_id} for {task}");
METRIC_META_REGION_MIGRATION_FAIL.inc();
return;
}
info!("Region migration procedure {procedure_id} for {task} is finished successfully!");
});
Ok(procedure_id)
Ok(BuiltMigrationBatch {
analysis,
procedure: Some(procedure),
})
}
/// Builds a new region migration procedure for `task` without submitting it.
@@ -790,6 +797,101 @@ mod test {
procedure_manager.stop().await.unwrap();
}
#[tokio::test]
async fn test_build_batch_procedure_retains_only_pending_guards() {
let env = TestingEnv::new();
let manager =
RegionMigrationManager::new(env.procedure_manager().clone(), env.context_factory());
let pending = RegionId::new(1024, 1);
let migrated = RegionId::new(1024, 2);
let migrating = RegionId::new(2048, 1);
let running_guard = manager
.insert_running_procedure(&RegionMigrationProcedureTask::new(
migrating,
Peer::empty(1),
Peer::empty(2),
Duration::from_secs(10),
RegionMigrationTriggerReason::Manual,
))
.unwrap();
env.create_physical_table_metadata(
new_test_table_info(1024),
vec![
RegionRoute {
region: Region::new_test(pending),
leader_peer: Some(Peer::empty(1)),
..Default::default()
},
RegionRoute {
region: Region::new_test(migrated),
leader_peer: Some(Peer::empty(2)),
..Default::default()
},
],
)
.await;
let batch = manager
.build_batch_procedure(
ProcedureContext::default(),
RegionMigrationTaskBatch {
region_ids: vec![pending, migrated, migrating],
from_peer: Peer::empty(1),
to_peer: Peer::empty(2),
timeout: Duration::from_secs(10),
trigger_reason: RegionMigrationTriggerReason::Manual,
},
)
.await
.unwrap();
assert_eq!(
RegionMigrationAnalysis {
pending: vec![pending],
migrated: vec![migrated],
migrating: vec![migrating],
..Default::default()
},
batch.analysis
);
let child = batch.procedure.unwrap();
assert!(manager.tracker.contains(pending));
assert!(!manager.tracker.contains(migrated));
assert!(
env.procedure_manager()
.list_procedures()
.await
.unwrap()
.is_empty()
);
drop(child);
assert!(!manager.tracker.contains(pending));
assert!(manager.tracker.contains(migrating));
let batch = manager
.build_batch_procedure(
ProcedureContext::default(),
RegionMigrationTaskBatch {
region_ids: vec![migrating],
from_peer: Peer::empty(1),
to_peer: Peer::empty(2),
timeout: Duration::from_secs(10),
trigger_reason: RegionMigrationTriggerReason::Manual,
},
)
.await
.unwrap();
assert_eq!(
RegionMigrationAnalysis {
migrating: vec![migrating],
..Default::default()
},
batch.analysis
);
assert!(batch.procedure.is_none());
assert!(manager.tracker.contains(migrating));
drop(running_guard);
assert!(!manager.tracker.contains(migrating));
}
#[tokio::test]
async fn test_submit_procedure_registration() {
check_submission_registration(false).await;
src/meta-srv/src/procedure/region_migration/utils.rs
+7 -2
View File
@@ -94,9 +94,12 @@ impl RegionMigrationTaskBatch {
}
}
/// Represents the result of analyzing a migration task.
/// Classifies regions during migration batch construction.
/// Route analysis fills metadata categories; the builder adds tracker-owned regions.
#[derive(Debug, Clone, Default, PartialEq)]
pub(crate) struct RegionMigrationAnalysis {
/// Regions already owned by another migration, populated by the batch builder.
pub(crate) migrating: Vec<RegionId>,
/// Regions already migrated to the `to_peer`.
pub(crate) migrated: Vec<RegionId>,
/// Regions where the leader peer has changed.
@@ -174,7 +177,8 @@ fn update_result_with_region_route(
/// Analyzes the migration task and categorizes regions by their current state.
///
/// Returns a [`RegionMigrationAnalysis`] describing the migration status.
/// Returns a [`RegionMigrationAnalysis`] based on table routes.
/// The `migrating` category remains empty because this function does not access the tracker.
pub async fn analyze_region_migration_task(
task: &RegionMigrationTaskBatch,
table_metadata_manager: &TableMetadataManagerRef,
@@ -504,6 +508,7 @@ mod tests {
peer_conflict: vec![RegionId::new(1024, 3)],
region_not_found: vec![RegionId::new(1024, 5)],
table_not_found: vec![RegionId::new(1025, 1)],
migrating: vec![],
}
);
}