From 4859e57e0bec44a871168be0f364fee262ddd827 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 15 Oct 2026 14:29:10 +0000 Subject: [PATCH] feat(meta): apply affinity to failover and repartition (#9604) * feat(meta): apply affinity to failover and repartition Signed-off-by: WenyXu * refactor(meta): share repartition affinity until group allocation Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/cmd/src/standalone.rs | 1 + src/common/meta/src/ddl.rs | 4 +- src/common/meta/src/ddl_manager.rs | 2 + src/common/meta/src/test_util.rs | 1 + src/meta-srv/AGENTS.md | 10 +- src/meta-srv/src/metasrv/builder.rs | 11 + src/meta-srv/src/procedure/repartition.rs | 5 + .../procedure/repartition/allocate_region.rs | 576 +++++++++++++++++- src/meta-srv/src/procedure/utils.rs | 1 + src/meta-srv/src/region/supervisor.rs | 318 +++++++++- src/meta-srv/src/selector.rs | 1 + tests-integration/src/standalone.rs | 1 + 12 files changed, 911 insertions(+), 20 deletions(-) diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 278afb7f1da..3dca89f6c25 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -682,6 +682,7 @@ impl StartCommand { let ddl_context = DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: layered_cache_registry.clone(), diff --git a/src/common/meta/src/ddl.rs b/src/common/meta/src/ddl.rs index 3eea16fcfa5..82b43cd832c 100644 --- a/src/common/meta/src/ddl.rs +++ b/src/common/meta/src/ddl.rs @@ -18,7 +18,7 @@ use std::time::Duration; use store_api::storage::{RegionId, TableId}; use crate::DatanodeId; -use crate::affinity::{AffinityResolverRef, PeerAllocContextProviderRef}; +use crate::affinity::{AffinityFallback, AffinityResolverRef, PeerAllocContextProviderRef}; use crate::cache_invalidator::CacheInvalidatorRef; use crate::ddl::flow_meta::FlowMetadataAllocatorRef; use crate::ddl::table_meta::TableMetadataAllocatorRef; @@ -108,6 +108,8 @@ impl RegionFailureDetectorController for NoopRegionFailureDetectorControl { pub struct DdlContext { /// Optional resolver for automatic allocation affinity. pub affinity_resolver: Option, + /// Fallback policy for allocation without query-specific inputs. + pub affinity_fallback: AffinityFallback, /// Builds allocation inputs from the initiating query. pub peer_alloc_context_provider: Option, /// Sends querying and requests to nodes. diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 0a1abde30f7..7b490895a4b 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -1886,6 +1886,7 @@ mod tests { let ddl_manager = DdlManager::new( DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), @@ -1954,6 +1955,7 @@ mod tests { DdlManager::new( DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index e574161d367..c1c3af8f5e9 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -199,6 +199,7 @@ pub fn new_ddl_context_with_kv_backend( ))); DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager, cache_invalidator: Arc::new(DummyCacheInvalidator), diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index 863b2a2aab4..f1153b4cc35 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -66,10 +66,12 @@ API under `src/meta-srv/src/service/admin/`. `cache_invalidator.rs` and the heartbeat publish handler. - **Affinity**: `common-meta/affinity.rs` defines requirements, matching, batch resolution, and the query allocation-context provider. The selector factory - optionally supplies a resolver for allocation inputs. DDL configurators can - install a query provider before procedure loaders are registered. Each selector - allows one unconstrained retry on candidate shortage only when explicitly - configured. Procedure state is unchanged. + optionally supplies a resolver shared by DDL allocation and the region + supervisor; without one, allocation skips source lookup. DDL configurators can + install a query provider before procedure loaders are registered. Repartition + and failover resolve source requirements before selection. Each selector allows + one unconstrained retry on candidate shortage only when explicitly configured; + metadata errors propagate. Procedure state is unchanged. - **Selector stats fields**: keep `selector/` implementations in sync with the stats they consume. diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 96cd2a8adb5..93ce228a156 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -423,6 +423,11 @@ impl MetasrvBuilder { leader_cached_kv_backend.clone(), ) .with_state(state.clone()); + let region_supervisor = if let Some(resolver) = &affinity_resolver { + region_supervisor.with_affinity_resolver(resolver.clone()) + } else { + region_supervisor + }; Some(RegionFailureHandler::new( region_supervisor, @@ -437,6 +442,7 @@ impl MetasrvBuilder { let ddl_context = DdlContext { peer_alloc_context_provider: None, + affinity_fallback: selector_ctx.affinity_fallback, affinity_resolver, node_manager: node_manager.clone(), cache_invalidator: cache_invalidator.clone(), @@ -817,6 +823,7 @@ mod tests { .await .unwrap(); let context = metasrv.ddl_manager().create_context(); + assert_eq!(context.affinity_fallback, AffinityFallback::Error); assert!(context.affinity_resolver.is_none()); assert!(context.peer_alloc_context_provider.is_none()); } @@ -894,6 +901,10 @@ mod tests { metasrv.selector_ctx().affinity_fallback, AffinityFallback::Unconstrained ); + assert_eq!( + metasrv.ddl_manager().create_context().affinity_fallback, + AffinityFallback::Unconstrained + ); assert_eq!(calls.load(Ordering::Relaxed), 1); assert!(Arc::ptr_eq(&metasrv.selector_ctx().affinity, &affinity)); assert!(Arc::ptr_eq( diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index 008c0ce05bb..b137eac2166 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -30,6 +30,7 @@ use std::fmt::{Debug, Display}; use std::time::{Duration, Instant}; use common_error::ext::BoxedError; +use common_meta::affinity::{AffinityFallback, AffinityResolverRef}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::ddl::DdlContext; use common_meta::ddl::allocator::region_routes::RegionRoutesAllocatorRef; @@ -148,6 +149,8 @@ impl PersistentContext { #[derive(Clone)] pub struct Context { + pub affinity_resolver: Option, + pub affinity_fallback: AffinityFallback, pub persistent_ctx: PersistentContext, pub volatile_ctx: VolatileContext, pub table_metadata_manager: TableMetadataManagerRef, @@ -248,6 +251,8 @@ impl Context { persistent_ctx: PersistentContext, ) -> Self { Self { + affinity_resolver: ddl_ctx.affinity_resolver.clone(), + affinity_fallback: ddl_ctx.affinity_fallback, persistent_ctx, table_metadata_manager: ddl_ctx.table_metadata_manager.clone(), memory_region_keeper: ddl_ctx.memory_region_keeper.clone(), diff --git a/src/meta-srv/src/procedure/repartition/allocate_region.rs b/src/meta-srv/src/procedure/repartition/allocate_region.rs index 68000682142..edbae33c17d 100644 --- a/src/meta-srv/src/procedure/repartition/allocate_region.rs +++ b/src/meta-srv/src/procedure/repartition/allocate_region.rs @@ -15,6 +15,9 @@ use std::any::Any; use std::collections::{HashMap, HashSet}; +use common_meta::affinity::{ + AffinityFallback, AffinityReference, AffinityRequirements, AffinityResolver, +}; use common_meta::ddl::create_table::executor::CreateTableExecutor; use common_meta::ddl::create_table::template::{ CreateRequestBuilder, build_template_from_raw_table_info_for_physical_table, @@ -155,18 +158,37 @@ impl ExecutePlan { let region_number_and_partition_exprs = AllocateRegion::prepare_region_allocation_data(&allocate_regions)?; let table_info_value = ctx.get_table_info_value().await?; - let new_allocated_region_routes = ctx - .region_routes_allocator - .allocate( - table_id, - ®ion_number_and_partition_exprs + let affinity = if let Some(resolver) = &ctx.affinity_resolver { + Some(resolve_allocation_affinity(ctx, resolver.as_ref()).await?) + } else { + None + }; + let groups = if let Some(affinity) = &affinity { + group_allocations(table_id, ®ion_number_and_partition_exprs, affinity) + } else { + vec![AllocationGroup { + affinity: None, + regions: region_number_and_partition_exprs .iter() - .map(|(n, p)| (*n, p.as_str())) - .collect::>(), - &PeerAllocContext::default(), - ) - .await - .context(error::AllocateRegionRoutesSnafu { table_id })?; + .map(|(number, partition)| (*number, partition.as_str())) + .collect(), + }] + }; + let mut new_allocated_region_routes = Vec::with_capacity(allocate_regions.len()); + for AllocationGroup { affinity, regions } in groups { + let routes = ctx + .region_routes_allocator + .allocate( + table_id, + ®ions, + &PeerAllocContext { + affinity: affinity.cloned(), + }, + ) + .await + .context(error::AllocateRegionRoutesSnafu { table_id })?; + new_allocated_region_routes.extend(routes); + } let skip_wal = table_info_value.table_info.meta.options.skip_wal; let allocate_noop = if skip_wal { let table_route_value = ctx.get_table_route_value().await?; @@ -439,8 +461,146 @@ impl AllocateRegion { } } +/// Selects the lowest-ID contributing source for a target spanning multiple sources. +/// Legacy plans without a transition map fall back to the lowest-ID source in the group. +fn collect_allocation_sources( + plans: &[RepartitionPlanEntry], + routes: &[RegionRoute], +) -> Vec<(RegionId, Option)> { + let leaders: HashMap<_, _> = routes + .iter() + .filter_map(|route| { + route + .leader_peer + .as_ref() + .map(|peer| (route.region.id, peer.id)) + }) + .collect(); + let mut sources = Vec::new(); + for plan in plans { + let allocated: HashSet<_> = plan.allocated_region_ids.iter().copied().collect(); + let mut target_sources = vec![None; plan.target_regions.len()]; + if plan.transition_map.is_empty() { + let source = plan + .source_regions + .iter() + .map(|source| source.region_id()) + .min(); + target_sources.fill(source); + } else { + for (source, targets) in plan.source_regions.iter().zip(&plan.transition_map) { + for target_index in targets { + if let Some(target_source) = target_sources.get_mut(*target_index) { + let source_id = source.region_id(); + *target_source = + Some(target_source.map_or(source_id, |current| current.min(source_id))); + } + } + } + } + for (target, source) in plan.target_regions.iter().zip(target_sources) { + if allocated.contains(&target.region_id) { + sources.push(( + target.region_id, + source.and_then(|source| leaders.get(&source)).copied(), + )); + } + } + } + sources +} + +/// Owns requirements once per source; targets refer to them by index. +#[derive(Debug)] +struct AllocationAffinity { + requirements: Vec>, + target_sources: HashMap, +} + +async fn resolve_allocation_affinity( + ctx: &Context, + resolver: &dyn AffinityResolver, +) -> Result { + let table_id = ctx.persistent_ctx.table_id; + let routes = ctx.get_table_route_value().await?; + let sources = + collect_allocation_sources(&ctx.persistent_ctx.plans, routes.region_routes().unwrap()); + let mut source_indices = HashMap::new(); + let mut references = Vec::new(); + for (_, source_id) in &sources { + if let Some(source_id) = source_id { + source_indices.entry(*source_id).or_insert_with(|| { + let index = references.len(); + references.push(AffinityReference::Datanode(*source_id)); + index + }); + } + } + let resolved = resolver + .resolve(&references) + .await + .context(error::AllocateRegionRoutesSnafu { table_id })?; + snafu::ensure!( + resolved.len() == source_indices.len(), + error::UnexpectedSnafu { + violated: "affinity resolution must preserve the number and order of references", + } + ); + snafu::ensure!( + ctx.affinity_fallback == AffinityFallback::Unconstrained + || sources.iter().all(|(_, source)| { + source.is_some_and(|id| resolved[source_indices[&id]].is_some()) + }), + error::UnexpectedSnafu { + violated: "source affinity requirements are unavailable", + } + ); + Ok(AllocationAffinity { + requirements: resolved, + target_sources: sources + .into_iter() + .filter_map(|(region_id, source_id)| { + source_id.map(|source_id| (region_id, source_indices[&source_id])) + }) + .collect(), + }) +} + +#[derive(Debug, PartialEq, Eq)] +struct AllocationGroup<'a> { + affinity: Option<&'a AffinityRequirements>, + regions: Vec<(RegionNumber, &'a str)>, +} + +fn group_allocations<'a>( + table_id: TableId, + regions: &'a [(RegionNumber, String)], + affinity: &'a AllocationAffinity, +) -> Vec> { + let mut groups: Vec> = Vec::new(); + for (number, partition) in regions { + let requirements = affinity + .target_sources + .get(&RegionId::new(table_id, *number)) + .and_then(|index| affinity.requirements[*index].as_ref()); + if let Some(group) = groups + .iter_mut() + .find(|group| group.affinity == requirements) + { + group.regions.push((*number, partition.as_str())); + } else { + groups.push(AllocationGroup { + affinity: requirements, + regions: vec![(*number, partition.as_str())], + }); + } + } + groups +} + #[cfg(test)] mod tests { + use std::ptr; use std::sync::Arc; use api::v1::region::region_request::Body; @@ -1035,4 +1195,398 @@ mod tests { AllocateRegion::Build(_) => panic!("expected execute plan"), } } + #[test] + fn test_collect_allocation_sources_use_target_overlap() { + let table_id = 1024; + let mut plan = RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![ + create_region_descriptor(table_id, 2, "x", 100, 200), + create_region_descriptor(table_id, 1, "x", 0, 100), + ], + target_regions: vec![ + create_target_region_descriptor(table_id, 3, "x", 0, 50), + create_target_region_descriptor(table_id, 4, "x", 50, 150), + create_target_region_descriptor(table_id, 5, "x", 150, 200), + ], + allocated_region_ids: (3..=5).map(|n| RegionId::new(table_id, n)).collect(), + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![1, 2], vec![0, 1]], + original_target_routes: vec![], + }; + let mut routes = create_current_region_routes(table_id, &[1, 2]); + routes[0].leader_peer = Some(Peer::empty(10)); + routes[1].leader_peer = Some(Peer::empty(20)); + assert_eq!( + collect_allocation_sources(&[plan.clone()], &routes), + vec![ + (RegionId::new(table_id, 3), Some(10)), + (RegionId::new(table_id, 4), Some(10)), + (RegionId::new(table_id, 5), Some(20)), + ] + ); + // Source enumeration does not change the tie-break for overlapping targets. + plan.source_regions.reverse(); + plan.transition_map.reverse(); + assert_eq!( + collect_allocation_sources(std::slice::from_ref(&plan), &routes)[1], + (RegionId::new(table_id, 4), Some(10)) + ); + plan.transition_map.clear(); + assert_eq!( + collect_allocation_sources(std::slice::from_ref(&plan), &routes), + (3..=5) + .map(|number| (RegionId::new(table_id, number), Some(10))) + .collect::>() + ); + plan.source_regions.clear(); + assert_eq!( + collect_allocation_sources(&[plan], &routes), + (3..=5) + .map(|number| (RegionId::new(table_id, number), None)) + .collect::>() + ); + } + + #[test] + fn test_group_allocations_preserves_region_partition_mapping() { + use common_meta::affinity::AffinityRequirements; + let table_id = 1024; + let requirements = |zone: &str| { + AffinityRequirements::new(HashMap::from([("ZONE".into(), zone.into())])).unwrap() + }; + let region_requirements = AllocationAffinity { + requirements: vec![ + Some(requirements("a")), + Some(requirements("b")), + Some(requirements("a")), + ], + target_sources: HashMap::from([ + (RegionId::new(table_id, 1), 0), + (RegionId::new(table_id, 2), 1), + (RegionId::new(table_id, 3), 2), + ]), + }; + let regions = vec![ + (1, "p1".into()), + (2, "p2".into()), + (3, "p3".into()), + (4, "p4".into()), + ]; + let groups = group_allocations(table_id, ®ions, ®ion_requirements); + assert!(ptr::eq( + groups[0].affinity.unwrap(), + region_requirements.requirements[0].as_ref().unwrap(), + )); + assert!(!ptr::eq( + region_requirements.requirements[0].as_ref().unwrap(), + region_requirements.requirements[2].as_ref().unwrap(), + )); + assert_eq!( + groups, + vec![ + AllocationGroup { + affinity: region_requirements.requirements[0].as_ref(), + regions: vec![(1, "p1"), (3, "p3")], + }, + AllocationGroup { + affinity: region_requirements.requirements[1].as_ref(), + regions: vec![(2, "p2")], + }, + AllocationGroup { + affinity: None, + regions: vec![(4, "p4")], + }, + ] + ); + } + + #[tokio::test] + async fn test_execute_refreshes_affinity_on_retry() { + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::AffinityResolver; + use common_meta::error::{self as meta_error, Result as MetaResult}; + use common_meta::key::table_route::TableRouteKey; + use common_meta::peer::PeerAllocator; + + #[derive(Default)] + struct ChangingResolver { + calls: AtomicUsize, + } + + #[async_trait::async_trait] + impl AffinityResolver for ChangingResolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert_eq!( + references, + &[ + AffinityReference::Datanode(1), + AffinityReference::Datanode(2) + ] + ); + let generation = self.calls.fetch_add(1, Ordering::Relaxed); + Ok(references + .iter() + .map(|_| { + AffinityRequirements::new(HashMap::from([( + "ZONE".into(), + format!("zone-{generation}"), + )])) + }) + .collect()) + } + } + + #[derive(Default)] + struct RecordingAllocator(Mutex>>); + + #[async_trait::async_trait] + impl PeerAllocator for RecordingAllocator { + async fn alloc(&self, _num: usize, ctx: &PeerAllocContext) -> MetaResult> { + self.0.lock().unwrap().push(ctx.affinity.clone()); + meta_error::UnexpectedSnafu { + err_msg: "stop before creating regions", + } + .fail() + } + } + + let env = TestingEnv::new(); + let table_id = 1024; + let mut routes = create_current_region_routes(table_id, &[1, 2]); + routes[1].leader_peer = Some(Peer::empty(2)); + env.create_physical_table_metadata_for_repartition( + table_id, + routes, + test_region_wal_options(&[1, 2]), + ) + .await; + let (sender, _receiver) = mpsc::channel(1); + let mut ctx = new_parent_context( + &env, + Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(sender))), + table_id, + ); + let resolver = Arc::new(ChangingResolver::default()); + let allocator = Arc::new(RecordingAllocator::default()); + ctx.affinity_resolver = Some(resolver.clone()); + ctx.region_routes_allocator = allocator.clone(); + ctx.persistent_ctx.plans = vec![RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![ + create_region_descriptor(table_id, 1, "x", 0, 100), + create_region_descriptor(table_id, 2, "x", 100, 200), + ], + target_regions: vec![ + create_target_region_descriptor(table_id, 3, "x", 0, 50), + create_target_region_descriptor(table_id, 4, "x", 50, 100), + create_target_region_descriptor(table_id, 5, "x", 100, 200), + ], + allocated_region_ids: (3..=5).map(|n| RegionId::new(table_id, n)).collect(), + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![0, 1], vec![2]], + original_target_routes: vec![], + }]; + let procedure_ctx = ProcedureContext { + procedure_id: ProcedureId::random(), + provider: Arc::new(MockContextProvider::default()), + event_context: None, + }; + for _ in 0..2 { + assert!(ExecutePlan.next(&mut ctx, &procedure_ctx).await.is_err()); + } + assert_eq!(resolver.calls.load(Ordering::Relaxed), 2); + assert_eq!( + *allocator.0.lock().unwrap(), + (0..2) + .map(|generation| AffinityRequirements::new(HashMap::from([( + "ZONE".into(), + format!("zone-{generation}") + ),]))) + .collect::>() + ); + // Without a resolver, allocation does not need source route metadata. + // Removing it makes any accidental lookup fail before the recording allocator. + ctx.affinity_resolver = None; + env.kv_backend + .delete(TableRouteKey::new(table_id).to_string().as_bytes(), false) + .await + .unwrap(); + assert!(ExecutePlan.next(&mut ctx, &procedure_ctx).await.is_err()); + assert_eq!(resolver.calls.load(Ordering::Relaxed), 2); + let allocations = allocator.0.lock().unwrap(); + assert_eq!(allocations.len(), 3); + assert_eq!(allocations.last(), Some(&None)); + } + + #[tokio::test] + async fn test_resolved_affinity_preserves_source_associations() { + use common_meta::affinity::AffinityResolver; + use common_meta::error::Result as MetaResult; + + struct Resolver(Vec>); + + #[async_trait::async_trait] + impl AffinityResolver for Resolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert_eq!( + references, + &[ + AffinityReference::Datanode(2), + AffinityReference::Datanode(1) + ] + ); + Ok(self.0.clone()) + } + } + + let env = TestingEnv::new(); + let table_id = 1024; + let mut routes = create_current_region_routes(table_id, &[1, 2]); + routes[1].leader_peer = Some(Peer::empty(2)); + env.create_physical_table_metadata_for_repartition( + table_id, + routes, + test_region_wal_options(&[1, 2]), + ) + .await; + let (sender, _receiver) = mpsc::channel(1); + let mut ctx = new_parent_context( + &env, + Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(sender))), + table_id, + ); + ctx.persistent_ctx.plans = vec![RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![ + create_region_descriptor(table_id, 1, "x", 0, 100), + create_region_descriptor(table_id, 2, "x", 100, 200), + ], + target_regions: vec![ + create_target_region_descriptor(table_id, 3, "x", 100, 200), + create_target_region_descriptor(table_id, 4, "x", 0, 50), + create_target_region_descriptor(table_id, 5, "x", 50, 100), + ], + allocated_region_ids: (3..=5).map(|n| RegionId::new(table_id, n)).collect(), + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![1, 2], vec![0]], + original_target_routes: vec![], + }]; + let a = AffinityRequirements::new(HashMap::from([("ZONE".into(), "a".into())])).unwrap(); + let b = AffinityRequirements::new(HashMap::from([("ZONE".into(), "b".into())])).unwrap(); + assert!(matches!( + resolve_allocation_affinity(&ctx, &Resolver(vec![None, Some(a.clone())])).await, + Err(error::Error::Unexpected { .. }) + )); + ctx.affinity_fallback = AffinityFallback::Unconstrained; + for (resolved, expected) in [ + ( + vec![Some(b.clone()), Some(a.clone())], + HashMap::from([ + (RegionId::new(table_id, 3), b.clone()), + (RegionId::new(table_id, 4), a.clone()), + (RegionId::new(table_id, 5), a.clone()), + ]), + ), + ( + vec![None, Some(a.clone())], + HashMap::from([ + (RegionId::new(table_id, 4), a.clone()), + (RegionId::new(table_id, 5), a), + ]), + ), + ( + vec![Some(b.clone()), None], + HashMap::from([(RegionId::new(table_id, 3), b)]), + ), + ] { + let affinity = resolve_allocation_affinity(&ctx, &Resolver(resolved)) + .await + .unwrap(); + assert_eq!(affinity.requirements.len(), 2); + assert_eq!( + affinity.target_sources[&RegionId::new(table_id, 4)], + affinity.target_sources[&RegionId::new(table_id, 5)] + ); + let actual = affinity + .target_sources + .iter() + .filter_map(|(target, index)| { + affinity.requirements[*index] + .as_ref() + .map(|requirements| (*target, requirements)) + }) + .collect::>(); + if let Some(first) = actual.get(&RegionId::new(table_id, 4)) { + let second = actual[&RegionId::new(table_id, 5)]; + assert!(ptr::eq(*first, second)); + } + assert_eq!( + actual, + expected + .iter() + .map(|(target, requirements)| (*target, requirements)) + .collect() + ); + } + } + + #[tokio::test] + async fn test_missing_source_leader_requires_explicit_fallback() { + use common_meta::error::Result as MetaResult; + + struct Resolver; + + #[async_trait::async_trait] + impl AffinityResolver for Resolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert!(references.is_empty()); + Ok(vec![]) + } + } + + let env = TestingEnv::new(); + let table_id = 1024; + let mut routes = create_current_region_routes(table_id, &[1]); + routes[0].leader_peer = None; + env.create_physical_table_metadata_for_repartition( + table_id, + routes, + test_region_wal_options(&[1]), + ) + .await; + let mut ctx = new_parent_context(&env, Arc::new(MockDatanodeManager::new(())), table_id); + ctx.persistent_ctx.plans = vec![RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![create_region_descriptor(table_id, 1, "x", 0, 100)], + target_regions: vec![create_target_region_descriptor(table_id, 2, "x", 0, 100)], + allocated_region_ids: vec![RegionId::new(table_id, 2)], + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![0]], + original_target_routes: vec![], + }]; + assert!(matches!( + resolve_allocation_affinity(&ctx, &Resolver).await, + Err(error::Error::Unexpected { .. }) + )); + ctx.affinity_fallback = AffinityFallback::Unconstrained; + assert!( + resolve_allocation_affinity(&ctx, &Resolver) + .await + .unwrap() + .target_sources + .is_empty() + ); + } } diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index 1f975b2f9c0..ff88d1b38ba 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -536,6 +536,7 @@ pub mod test_data { )); DdlContext { peer_alloc_context_provider: None, + affinity_fallback: Default::default(), affinity_resolver: None, node_manager, cache_invalidator: Arc::new(MetasrvCacheInvalidator::new( diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 93624a61892..1c431ef9960 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -18,7 +18,9 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use async_trait::async_trait; +use common_error::ext::BoxedError; use common_meta::DatanodeId; +use common_meta::affinity::{AffinityFallback, AffinityReference, AffinityResolverRef}; use common_meta::datanode::Stat; use common_meta::ddl::{DetectingRegion, RegionFailureDetectorController}; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; @@ -290,6 +292,8 @@ pub struct RegionSupervisor { selector_context: SelectorContext, /// Candidate node selector. selector: RegionSupervisorSelector, + /// Resolves the failed node attributes before selection. + affinity_resolver: Option, /// Region migration manager. region_migration_manager: RegionMigrationManagerRef, /// The maintenance mode manager. @@ -408,6 +412,7 @@ impl RegionSupervisor { receiver: event_receiver, selector_context, selector, + affinity_resolver: None, region_migration_manager, runtime_switch_manager, peer_resolver, @@ -416,6 +421,12 @@ impl RegionSupervisor { } } + /// Sets the resolver used to prepare failover affinity requirements. + pub(crate) fn with_affinity_resolver(mut self, resolver: AffinityResolverRef) -> Self { + self.affinity_resolver = Some(resolver); + self + } + /// Sets the meta state. pub(crate) fn with_state(mut self, state: StateRef) -> Self { self.state = Some(state); @@ -659,17 +670,41 @@ impl RegionSupervisor { regions: &[RegionId], failure_datanodes: &[DatanodeId], ) -> Result> { - let exclude_peer_ids = HashSet::from_iter(failure_datanodes.iter().cloned()); + let affinity = if let Some(resolver) = &self.affinity_resolver { + let mut resolved = resolver + .resolve(&[AffinityReference::Datanode(from_peer_id)]) + .await + .map_err(BoxedError::new) + .context(error::OtherSnafu)?; + ensure!( + resolved.len() == 1, + error::UnexpectedSnafu { + violated: "affinity resolution must preserve the number and order of references", + } + ); + let affinity = resolved.pop().flatten(); + ensure!( + affinity.is_some() + || self.selector_context.affinity_fallback == AffinityFallback::Unconstrained, + error::UnexpectedSnafu { + violated: "source affinity requirements are unavailable" + } + ); + affinity + } else { + None + }; + let exclude_peer_ids = HashSet::from_iter(failure_datanodes.iter().copied()); match &self.selector { RegionSupervisorSelector::NaiveSelector(selector) => { - let opt = SelectorOptions { + let opts = SelectorOptions { + affinity: affinity.as_ref(), min_required_items: regions.len(), allow_duplication: true, exclude_peer_ids, workload_filter: Some(accept_ingest_workload), - affinity: None, }; - let peers = selector.select(&self.selector_context, opt).await?; + let peers = selector.select(&self.selector_context, opts).await?; ensure!( peers.len() == regions.len(), error::NoEnoughAvailableNodeSnafu { @@ -693,6 +728,7 @@ impl RegionSupervisor { from_peer_id, regions, exclude_peer_ids, + affinity.as_ref(), ) .await?; ensure!( @@ -962,6 +998,280 @@ pub(crate) mod tests { ) } + #[tokio::test] + async fn test_failover_passes_affinity_and_propagates_errors() { + use std::collections::HashSet; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::{ + AffinityFallback, AffinityReference, AffinityRequirements, AffinityResolver, + }; + use common_meta::error::{self as meta_error, Result as MetaResult}; + + use crate::error::{self, Result}; + use crate::metasrv::{SelectTarget, SelectorContext}; + use crate::selector::{RegionStatAwareSelector, Selector, SelectorOptions}; + + #[derive(Clone, Copy, PartialEq)] + enum Resolution { + Disabled, + Available, + Missing, + Error, + } + + struct Resolver { + result: Resolution, + calls: AtomicUsize, + } + + #[async_trait::async_trait] + impl AffinityResolver for Resolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert_eq!(references, &[AffinityReference::Datanode(7)]); + self.calls.fetch_add(1, Ordering::Relaxed); + match self.result { + Resolution::Available => { + Ok(vec![AffinityRequirements::new(HashMap::from([( + "pool".into(), + "a".into(), + )]))]) + } + Resolution::Disabled => unreachable!(), + Resolution::Missing => Ok(vec![None]), + Resolution::Error => meta_error::UnexpectedSnafu { + err_msg: "source lookup failed", + } + .fail(), + } + } + } + + #[derive(Clone, Copy)] + enum Outcome { + Success, + Short, + Exhausted, + Error, + } + struct FailoverSelector { + outcome: Outcome, + expected_fallback: AffinityFallback, + calls: Mutex>>, + } + impl FailoverSelector { + fn select_peers(&self, affinity: Option<&AffinityRequirements>) -> Result> { + self.calls.lock().unwrap().push(affinity.cloned()); + if affinity.is_some() { + match self.outcome { + Outcome::Success => {} + Outcome::Short => return Ok(vec![Peer::empty(9)]), + Outcome::Exhausted => { + return error::NoEnoughAvailableNodeSnafu { + required: 2usize, + available: 0usize, + select_target: SelectTarget::Datanode, + } + .fail(); + } + Outcome::Error => { + return error::UnexpectedSnafu { + violated: "selection backend failed", + } + .fail(); + } + } + } + Ok(vec![Peer::empty(9); 2]) + } + } + #[async_trait::async_trait] + impl Selector for FailoverSelector { + type Context = SelectorContext; + type Output = Vec; + async fn select( + &self, + ctx: &SelectorContext, + opts: SelectorOptions<'_>, + ) -> Result { + assert_eq!(ctx.affinity_fallback, self.expected_fallback); + assert_eq!(opts.exclude_peer_ids, HashSet::from([8])); + assert!(opts.workload_filter.is_some()); + assert!(opts.allow_duplication); + assert_eq!(opts.min_required_items, 2); + self.select_peers(opts.affinity) + } + } + #[async_trait::async_trait] + impl RegionStatAwareSelector for FailoverSelector { + type Context = SelectorContext; + type Output = Vec<(RegionId, Peer)>; + async fn select( + &self, + ctx: &SelectorContext, + source: u64, + regions: &[RegionId], + excluded: HashSet, + affinity: Option<&AffinityRequirements>, + ) -> Result { + assert_eq!(ctx.affinity_fallback, self.expected_fallback); + assert_eq!(source, 7); + assert_eq!(excluded, HashSet::from([8])); + Ok(regions + .iter() + .copied() + .zip(self.select_peers(affinity)?) + .collect()) + } + } + for stat_aware in [false, true] { + for (resolution, fallback, outcome, expected_calls, succeeds) in [ + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Success, + 1, + true, + ), + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Short, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Exhausted, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Error, + 1, + false, + ), + ( + Resolution::Missing, + AffinityFallback::Error, + Outcome::Success, + 0, + false, + ), + ( + Resolution::Error, + AffinityFallback::Error, + Outcome::Success, + 0, + false, + ), + ( + Resolution::Available, + AffinityFallback::Unconstrained, + Outcome::Short, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Unconstrained, + Outcome::Exhausted, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Unconstrained, + Outcome::Error, + 1, + false, + ), + ( + Resolution::Missing, + AffinityFallback::Unconstrained, + Outcome::Success, + 1, + true, + ), + ( + Resolution::Error, + AffinityFallback::Unconstrained, + Outcome::Success, + 0, + false, + ), + ( + Resolution::Disabled, + AffinityFallback::Unconstrained, + Outcome::Success, + 1, + true, + ), + ] { + let resolver = Arc::new(Resolver { + result: resolution, + calls: AtomicUsize::new(0), + }); + let selector = Arc::new(FailoverSelector { + outcome, + expected_fallback: fallback, + calls: Mutex::new(vec![]), + }); + let (mut supervisor, _sender) = new_test_supervisor(); + supervisor.selector_context.affinity_fallback = fallback; + if resolution != Resolution::Disabled { + supervisor = supervisor.with_affinity_resolver(resolver.clone()); + } + supervisor.selector = if stat_aware { + RegionSupervisorSelector::RegionStatAwareSelector(selector.clone()) + } else { + RegionSupervisorSelector::NaiveSelector(selector.clone()) + }; + let regions = [RegionId::new(1, 1), RegionId::new(1, 2)]; + // Preserve the caller's exclusion set without adding the source. + let result = supervisor.select_peers(7, ®ions, &[8]).await; + if succeeds { + assert_eq!( + result.unwrap(), + vec![(regions[0], Peer::empty(9)), (regions[1], Peer::empty(9))] + ); + } else { + let error = result.unwrap_err(); + if resolution == Resolution::Error { + assert!(matches!(error, error::Error::Other { .. })); + } else if resolution == Resolution::Missing { + assert!(matches!(error, error::Error::Unexpected { .. })); + } else { + match outcome { + Outcome::Short | Outcome::Exhausted => { + assert!(matches!(error, error::Error::NoEnoughAvailableNode { .. })) + } + Outcome::Error => { + assert!(matches!(error, error::Error::Unexpected { .. })) + } + Outcome::Success => unreachable!(), + } + } + } + assert_eq!( + resolver.calls.load(Ordering::Relaxed), + usize::from(resolution != Resolution::Disabled) + ); + let calls = selector.calls.lock().unwrap(); + assert_eq!(calls.len(), expected_calls); + if let Some(affinity) = calls.first() { + assert_eq!(affinity.is_some(), resolution == Resolution::Available); + } + } + } + } + fn filler_heartbeat() -> DatanodeHeartbeat { DatanodeHeartbeat { datanode_id: 1, diff --git a/src/meta-srv/src/selector.rs b/src/meta-srv/src/selector.rs index e5c643b7b7c..c4d29b8ecea 100644 --- a/src/meta-srv/src/selector.rs +++ b/src/meta-srv/src/selector.rs @@ -54,6 +54,7 @@ pub trait RegionStatAwareSelector: Send + Sync { from_peer_id: u64, region_ids: &[RegionId], exclude_peer_ids: HashSet, + affinity: Option<&AffinityRequirements>, ) -> Result; } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index e0113f90425..5bd0d970de5 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -291,6 +291,7 @@ impl GreptimeDbStandaloneBuilder { let ddl_manager = Arc::new(DdlManager::new( DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: cache_registry.clone(),