From 4d8537a6227a2cba69348d0283ca09075e0cd87f Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 4 Nov 2026 06:52:10 +0000 Subject: [PATCH] refactor(meta): separate migration construction and drain child tracking (#10130) refactor(meta): separate migration construction and drain child tracking (#10098) refactor(meta): separate migration construction and own drain child tracking (cherry picked from commit 6a731a444fa9416aff4e30b12f4b34ede6f0d965) Signed-off-by: WenyXu --- src/meta-srv/src/procedure/drain_datanode.rs | 8 +- .../procedure/drain_datanode/check_drained.rs | 2 +- .../src/procedure/drain_datanode/dispatch.rs | 459 ++++++++++++++++-- .../src/procedure/drain_datanode/migration.rs | 390 +-------------- .../src/procedure/region_migration/manager.rs | 286 +++++++---- .../src/procedure/region_migration/utils.rs | 9 +- 6 files changed, 634 insertions(+), 520 deletions(-) diff --git a/src/meta-srv/src/procedure/drain_datanode.rs b/src/meta-srv/src/procedure/drain_datanode.rs index 87da8108428..76a389657b6 100644 --- a/src/meta-srv/src/procedure/drain_datanode.rs +++ b/src/meta-srv/src/procedure/drain_datanode.rs @@ -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, /// 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, - /// The migration sub-procedures of the current batch that are not settled yet. - pub inflight_subprocedures: Vec, /// 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. diff --git a/src/meta-srv/src/procedure/drain_datanode/check_drained.rs b/src/meta-srv/src/procedure/drain_datanode/check_drained.rs index 28db9bcab49..e0e3fd21de1 100644 --- a/src/meta-srv/src/procedure/drain_datanode/check_drained.rs +++ b/src/meta-srv/src/procedure/drain_datanode/check_drained.rs @@ -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 { diff --git a/src/meta-srv/src/procedure/drain_datanode/dispatch.rs b/src/meta-srv/src/procedure/drain_datanode/dispatch.rs index e5f89a6ffb1..5424f3cbfa0 100644 --- a/src/meta-srv/src/procedure/drain_datanode/dispatch.rs +++ b/src/meta-srv/src/procedure/drain_datanode/dispatch.rs @@ -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, +} + +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 ®ion_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::().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::().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::().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::().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::().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::().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()); + } + } } diff --git a/src/meta-srv/src/procedure/drain_datanode/migration.rs b/src/meta-srv/src/procedure/drain_datanode/migration.rs index 594106fc2e3..daac15c6bc4 100644 --- a/src/meta-srv/src/procedure/drain_datanode/migration.rs +++ b/src/meta-srv/src/procedure/drain_datanode/migration.rs @@ -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 ®ion_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()); - } - } } diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index 998e13fad20..75bd28d6b9e 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -179,6 +179,14 @@ pub struct SubmitRegionMigrationTaskResult { pub procedure_id: Option, } +/// 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, +} + 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 { - 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 { + 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::>(); 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, - catalog_and_schema: Vec<(String, String)>, - ) -> Result { + analysis.pending = task.region_ids.clone(); + let pending = task.region_ids.iter().copied().collect::>(); + 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; diff --git a/src/meta-srv/src/procedure/region_migration/utils.rs b/src/meta-srv/src/procedure/region_migration/utils.rs index 91af7ccf179..83728401ee0 100644 --- a/src/meta-srv/src/procedure/region_migration/utils.rs +++ b/src/meta-srv/src/procedure/region_migration/utils.rs @@ -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, /// Regions already migrated to the `to_peer`. pub(crate) migrated: Vec, /// 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![], } ); }