rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-21 19:32:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(meta): apply affinity to failover and repartition (#9604)

Browse Source 
* feat(meta): apply affinity to failover and repartition

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor(meta): share repartition affinity until group allocation

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-10-15 14:29:10 +00:00
committed by GitHub
parent deae75ffee
commit 4859e57e0b
12 changed files with 911 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cmd/src/standalone.rs
+1
View File
@@ -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(),
src/common/meta/src/ddl.rs
+3 -1
View File
@@ -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<AffinityResolverRef>,
/// 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<PeerAllocContextProviderRef>,
/// Sends querying and requests to nodes.
src/common/meta/src/ddl_manager.rs
+2
View File
@@ -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),
src/common/meta/src/test_util.rs
+1
View File
@@ -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),
src/meta-srv/AGENTS.md
+6 -4
View File
@@ -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.
src/meta-srv/src/metasrv/builder.rs
+11
View File
@@ -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(
src/meta-srv/src/procedure/repartition.rs
+5
View File
@@ -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<AffinityResolverRef>,
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(),
src/meta-srv/src/procedure/repartition/allocate_region.rs
+565 -11
View File
@@ -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,
&region_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, &region_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::<Vec<_>>(),
&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,
&regions,
&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<u64>)> {
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<Option<AffinityRequirements>>,
target_sources: HashMap<RegionId, usize>,
}
async fn resolve_allocation_affinity(
ctx: &Context,
resolver: &dyn AffinityResolver,
) -> Result<AllocationAffinity> {
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<AllocationGroup<'a>> {
let mut groups: Vec<AllocationGroup<'a>> = 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::<Vec<_>>()
);
plan.source_regions.clear();
assert_eq!(
collect_allocation_sources(&[plan], &routes),
(3..=5)
.map(|number| (RegionId::new(table_id, number), None))
.collect::<Vec<_>>()
);
}
#[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, &regions, &region_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<Vec<Option<AffinityRequirements>>> {
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<Vec<Option<AffinityRequirements>>>);
#[async_trait::async_trait]
impl PeerAllocator for RecordingAllocator {
async fn alloc(&self, _num: usize, ctx: &PeerAllocContext) -> MetaResult<Vec<Peer>> {
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::<Vec<_>>()
);
// 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<Option<AffinityRequirements>>);
#[async_trait::async_trait]
impl AffinityResolver for Resolver {
async fn resolve(
&self,
references: &[AffinityReference],
) -> MetaResult<Vec<Option<AffinityRequirements>>> {
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::<HashMap<_, _>>();
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<Vec<Option<AffinityRequirements>>> {
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()
);
}
}
src/meta-srv/src/procedure/utils.rs
+1
View File
@@ -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(
src/meta-srv/src/region/supervisor.rs
+314 -4
View File
@@ -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<AffinityResolverRef>,
/// 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<Vec<(RegionId, Peer)>> {
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<Vec<Option<AffinityRequirements>>> {
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<Vec<Option<AffinityRequirements>>>,
}
impl FailoverSelector {
fn select_peers(&self, affinity: Option<&AffinityRequirements>) -> Result<Vec<Peer>> {
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<Peer>;
async fn select(
&self,
ctx: &SelectorContext,
opts: SelectorOptions<'_>,
) -> Result<Self::Output> {
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<u64>,
affinity: Option<&AffinityRequirements>,
) -> Result<Self::Output> {
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, &regions, &[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,
src/meta-srv/src/selector.rs
+1
View File
@@ -54,6 +54,7 @@ pub trait RegionStatAwareSelector: Send + Sync {
from_peer_id: u64,
region_ids: &[RegionId],
exclude_peer_ids: HashSet<u64>,
affinity: Option<&AffinityRequirements>,
) -> Result<Self::Output>;
}
tests-integration/src/standalone.rs
+1
View File
@@ -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(),