diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index ca0f3b74057..278afb7f1da 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -681,6 +681,8 @@ impl StartCommand { )); let ddl_context = DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: layered_cache_registry.clone(), memory_region_keeper: Arc::new(MemoryRegionKeeper::default()), diff --git a/src/common/meta/AGENTS.md b/src/common/meta/AGENTS.md index 6c9710b9e45..69d76d3c5ca 100644 --- a/src/common/meta/AGENTS.md +++ b/src/common/meta/AGENTS.md @@ -11,6 +11,7 @@ and durable DDL procedures. Metasrv process and service wiring live in | Area | Path | Entry point | | --- | --- | --- | +| Affinity | `src/common/meta/src/affinity.rs` | Requirements, reference resolution, node matching, and allocation-context provider | | Metadata model | `src/common/meta/src/key.rs`, `src/common/meta/src/key/` | Typed keys, values, and `TableMetadataManager` | | KV storage | `src/common/meta/src/kv_backend.rs`, `src/common/meta/src/kv_backend/` | `KvBackend`, transactions, memory/etcd/RDS backends | | DDL procedures | `src/common/meta/src/ddl.rs`, `src/common/meta/src/ddl/`, `src/common/meta/src/ddl_manager.rs` | Durable DDL state machines and task dispatch | @@ -31,6 +32,10 @@ and durable DDL procedures. Metasrv process and service wiring live in `src/table/src/requests.rs` — classification, validation, logical-table eligibility, and the metadata-only alter flow come with it. Do not add per-option special cases to the alter conversion or DDL routing. +- Create-table preparation obtains allocation inputs from the optional query + provider in `DdlContext`. Install providers before registering procedure + loaders so restored procedures use the same runtime dependencies. Affinity + requirements and fallback policy are runtime inputs, not persisted fields. - `KvBackend` behavior changes should extend the shared tests in `src/common/meta/src/kv_backend/test.rs`. diff --git a/src/common/meta/src/affinity.rs b/src/common/meta/src/affinity.rs new file mode 100644 index 00000000000..2bcb0e934f0 --- /dev/null +++ b/src/common/meta/src/affinity.rs @@ -0,0 +1,254 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Resolve reference attributes into requirements, then evaluate candidates. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::DatanodeId; +use crate::cluster::NodeInfo; +use crate::error::Result; +use crate::peer::PeerAllocContext; +use crate::rpc::ddl::QueryContext; + +/// Named attributes used to express affinity requirements. +pub type AffinityAttributes = HashMap; + +/// Identifies the node whose attributes define affinity requirements. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum AffinityReference { + /// Uses the datanode's last-known attributes, even if its lease has expired. + Datanode(DatanodeId), + /// Uses the frontend identified by its advertised address. + Frontend(String), +} + +/// Behavior when affinity cannot be satisfied. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum AffinityFallback { + /// Returns an error instead of selecting without affinity. + #[default] + Error, + /// Allows selection among eligible nodes without affinity. + Unconstrained, +} + +/// Builds allocation inputs from the query that initiated an operation. +#[async_trait::async_trait] +pub trait PeerAllocContextProvider: Send + Sync { + /// Resolves query-specific allocation inputs using metadata, without selecting peers. + async fn build(&self, query_context: &QueryContext) -> Result; +} + +/// Shared allocation-context provider. +pub type PeerAllocContextProviderRef = Arc; + +/// Nonempty attributes a candidate must match. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AffinityRequirements { + attributes: AffinityAttributes, +} + +impl AffinityRequirements { + /// Returns no requirements when the supplied attributes are empty. + pub fn new(attributes: AffinityAttributes) -> Option { + (!attributes.is_empty()).then_some(Self { attributes }) + } + + /// Returns the attributes that candidates are evaluated against. + pub fn attributes(&self) -> &AffinityAttributes { + &self.attributes + } +} + +/// Resolves references using current metadata, without selecting candidates. +#[async_trait::async_trait] +pub trait AffinityResolver: Send + Sync { + /// Returns one result per reference, in order. + /// None means requirements are unavailable; Err means resolution failed. + async fn resolve( + &self, + references: &[AffinityReference], + ) -> Result>>; +} + +/// Shared requirement resolver. +pub type AffinityResolverRef = Arc; + +/// Evaluates candidate nodes without I/O or scheduling decisions. +pub trait Affinity: Send + Sync { + /// Evaluates a node snapshot; implementations choose which fields to inspect. + fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool; +} + +/// Accepts every candidate without imposing affinity constraints. +pub struct NoopAffinity; + +impl Affinity for NoopAffinity { + fn matches(&self, _: &AffinityRequirements, _: &NodeInfo) -> bool { + true + } +} + +/// Shared affinity evaluator. +pub type AffinityRef = Arc; + +#[cfg(test)] +mod tests { + use std::sync::RwLock; + + use super::*; + use crate::cluster::{DatanodeStatus, NodeStatus}; + use crate::peer::Peer; + + struct NodeAffinityResolver { + attributes: RwLock>, + } + + #[async_trait::async_trait] + impl AffinityResolver for NodeAffinityResolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> Result>> { + let snapshot = self.attributes.read().unwrap(); + Ok(references + .iter() + .map(|reference| { + snapshot + .get(reference) + .cloned() + .and_then(AffinityRequirements::new) + }) + .collect()) + } + } + + struct ExactAffinity; + + impl Affinity for ExactAffinity { + fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool { + requirements + .attributes() + .iter() + .all(|(key, expected)| candidate.env_vars.get(key) == Some(expected)) + } + } + + fn zone(name: &str) -> AffinityAttributes { + HashMap::from([("AZ".to_string(), name.to_string())]) + } + + fn node(id: DatanodeId, env_vars: AffinityAttributes) -> NodeInfo { + NodeInfo { + peer: Peer::empty(id), + last_activity_ts: 0, + status: NodeStatus::Datanode(DatanodeStatus { + rcus: 0, + wcus: 0, + leader_regions: 0, + follower_regions: 0, + workloads: Default::default(), + }), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars, + } + } + + #[tokio::test] + async fn test_resolve_reference_and_match_candidates() { + let resolver = NodeAffinityResolver { + attributes: RwLock::new(HashMap::from([(AffinityReference::Datanode(1), zone("a"))])), + }; + let references = [AffinityReference::Datanode(1)]; + let resolved = resolver.resolve(&references).await.unwrap(); + assert_eq!(resolved, vec![AffinityRequirements::new(zone("a"))]); + let requirements = resolved[0].as_ref().unwrap(); + let candidates = [ + node(2, zone("a")), + node(3, zone("b")), + node(4, HashMap::new()), + ]; + assert_eq!( + candidates + .iter() + .map(|node| (node.peer.id, ExactAffinity.matches(requirements, node))) + .collect::>(), + vec![(2, true), (3, false), (4, false),] + ); + + assert!( + candidates + .iter() + .all(|node| NoopAffinity.matches(requirements, node)) + ); + + // A new allocation attempt resolves current attributes rather than a saved requirement. + resolver + .attributes + .write() + .unwrap() + .insert(AffinityReference::Datanode(1), zone("b")); + let resolved = resolver.resolve(&references).await.unwrap(); + let requirements = resolved[0].as_ref().unwrap(); + let matched: Vec<_> = candidates + .iter() + .filter(|node| ExactAffinity.matches(requirements, node)) + .map(|node| node.peer.id) + .collect(); + assert_eq!(matched, vec![3]); + } + + #[tokio::test] + async fn test_resolution_preserves_order_and_unavailable_requirements() { + let resolver = NodeAffinityResolver { + attributes: RwLock::new(HashMap::from([ + (AffinityReference::Datanode(1), zone("a")), + (AffinityReference::Datanode(2), HashMap::new()), + ( + AffinityReference::Frontend("frontend:4001".into()), + zone("b"), + ), + ])), + }; + let references = [ + AffinityReference::Frontend("missing:4001".into()), + AffinityReference::Datanode(1), + AffinityReference::Frontend("frontend:4001".into()), + AffinityReference::Datanode(2), + AffinityReference::Datanode(3), + AffinityReference::Frontend("frontend:4001".into()), + ]; + assert_eq!( + resolver.resolve(&references).await.unwrap(), + vec![ + None, + AffinityRequirements::new(zone("a")), + AffinityRequirements::new(zone("b")), + None, + None, + AffinityRequirements::new(zone("b")), + ] + ); + assert!(AffinityRequirements::new(HashMap::new()).is_none()); + } +} diff --git a/src/common/meta/src/cluster.rs b/src/common/meta/src/cluster.rs index 149e5a059d5..4e7ffb7f920 100644 --- a/src/common/meta/src/cluster.rs +++ b/src/common/meta/src/cluster.rs @@ -106,7 +106,7 @@ fn calculate_node_id(addr: &str) -> u64 { } /// The information of a node in the cluster. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct NodeInfo { /// The peer information. [node_id, address] pub peer: Peer, @@ -148,7 +148,7 @@ pub enum Role { Metasrv, } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum NodeStatus { Datanode(DatanodeStatus), Frontend(FrontendStatus), @@ -171,7 +171,7 @@ impl NodeStatus { } /// The status of a datanode. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct DatanodeStatus { /// The read capacity units during this period. pub rcus: i64, @@ -186,7 +186,7 @@ pub struct DatanodeStatus { } /// The status of a frontend. -#[derive(Debug, Default, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct FrontendStatus { /// The workloads of the frontend. #[serde(default)] @@ -194,7 +194,7 @@ pub struct FrontendStatus { } /// The status of a flownode. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct FlownodeStatus { /// The workloads of the flownode. #[serde(default)] @@ -202,7 +202,7 @@ pub struct FlownodeStatus { } /// The status of a metasrv. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct MetasrvStatus { pub is_leader: bool, } diff --git a/src/common/meta/src/ddl.rs b/src/common/meta/src/ddl.rs index f70ba2b3861..3eea16fcfa5 100644 --- a/src/common/meta/src/ddl.rs +++ b/src/common/meta/src/ddl.rs @@ -18,6 +18,7 @@ use std::time::Duration; use store_api::storage::{RegionId, TableId}; use crate::DatanodeId; +use crate::affinity::{AffinityResolverRef, PeerAllocContextProviderRef}; use crate::cache_invalidator::CacheInvalidatorRef; use crate::ddl::flow_meta::FlowMetadataAllocatorRef; use crate::ddl::table_meta::TableMetadataAllocatorRef; @@ -105,6 +106,10 @@ impl RegionFailureDetectorController for NoopRegionFailureDetectorControl { /// The context of ddl. #[derive(Clone)] pub struct DdlContext { + /// Optional resolver for automatic allocation affinity. + pub affinity_resolver: Option, + /// Builds allocation inputs from the initiating query. + pub peer_alloc_context_provider: Option, /// Sends querying and requests to nodes. pub node_manager: NodeManagerRef, /// Cache invalidation. diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index ce3e1aa060a..122b0f24030 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -165,7 +165,10 @@ impl CreateTableProcedure { return Ok(Status::done_with_output(table_id)); } - self.data.state = CreateTableState::DatanodeCreateRegions; + let alloc_context = match &self.context.peer_alloc_context_provider { + Some(provider) => provider.build(&self.data.query_context).await?, + None => PeerAllocContext::default(), + }; let TableMetadata { table_id, table_route, @@ -173,14 +176,10 @@ impl CreateTableProcedure { } = self .context .table_metadata_allocator - .create_with_context( - &self.data.task, - &PeerAllocContext { - extensions: self.data.query_context.extensions.clone(), - }, - ) + .create_with_context(&self.data.task, &alloc_context) .await?; self.set_allocated_metadata(table_id, table_route, region_wal_options); + self.data.state = CreateTableState::DatanodeCreateRegions; Ok(Status::executing(true)) } diff --git a/src/common/meta/src/ddl/tests/create_table.rs b/src/common/meta/src/ddl/tests/create_table.rs index c97e61926c0..d39346e02da 100644 --- a/src/common/meta/src/ddl/tests/create_table.rs +++ b/src/common/meta/src/ddl/tests/create_table.rs @@ -398,3 +398,143 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { .contains(datanode_id, region_id) ); } + +#[tokio::test] +async fn test_on_prepare_uses_peer_alloc_context_provider() { + use std::sync::Mutex; + + use crate::affinity::{AffinityRequirements, PeerAllocContextProvider}; + use crate::ddl::table_meta::TableMetadataAllocator; + use crate::error; + use crate::peer::{PeerAllocContext, PeerAllocator}; + use crate::rpc::ddl::QueryContext; + + #[derive(Clone, Copy, PartialEq, Eq)] + enum ProviderMode { + Absent, + Available, + Error, + } + + struct Provider { + mode: ProviderMode, + calls: Mutex>, + } + + #[async_trait::async_trait] + impl PeerAllocContextProvider for Provider { + async fn build(&self, query_context: &QueryContext) -> Result { + self.calls.lock().unwrap().push(query_context.clone()); + if self.mode == ProviderMode::Error { + return error::UnexpectedSnafu { + err_msg: "allocation context lookup failed", + } + .fail(); + } + Ok(PeerAllocContext { + affinity: AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])), + }) + } + } + + #[derive(Default)] + struct Allocator(Mutex>); + + #[async_trait::async_trait] + impl PeerAllocator for Allocator { + async fn alloc(&self, num: usize, ctx: &PeerAllocContext) -> Result> { + self.0.lock().unwrap().push(ctx.clone()); + Ok(vec![Peer::empty(2); num]) + } + } + + for (mode, table_exists) in [ + (ProviderMode::Available, false), + (ProviderMode::Absent, false), + (ProviderMode::Error, false), + (ProviderMode::Error, true), + ] { + let provider = Arc::new(Provider { + mode, + calls: Mutex::new(Vec::new()), + }); + let allocator = Arc::new(Allocator::default()); + let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + if mode != ProviderMode::Absent { + context.peer_alloc_context_provider = Some(provider.clone()); + } + context.table_metadata_allocator = Arc::new(TableMetadataAllocator::with_peer_allocator( + context.table_metadata_allocator.table_id_allocator(), + context.table_metadata_allocator.wal_options_allocator(), + allocator.clone(), + )); + let mut task = test_create_table_task("provider_table"); + task.create_table.create_if_not_exists = true; + if table_exists { + context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + } + let query_context = QueryContext { + current_catalog: "test_catalog".into(), + current_schema: "test_schema".into(), + timezone: "UTC".into(), + extensions: HashMap::from([("custom-key".into(), "custom-value".into())]), + channel: 1, + snapshot_seqs: HashMap::from([(1, 10)]), + sst_min_sequences: HashMap::from([(1, 2)]), + }; + let mut procedure = + CreateTableProcedure::new_with_query_context(task, query_context.clone(), context) + .unwrap(); + let result = procedure.on_prepare().await; + if table_exists { + assert_matches!(result.unwrap(), Status::Done { .. }); + assert!(provider.calls.lock().unwrap().is_empty()); + assert!(allocator.0.lock().unwrap().is_empty()); + continue; + } + let expected_calls = if mode == ProviderMode::Absent { + vec![] + } else { + vec![query_context] + }; + assert_eq!(*provider.calls.lock().unwrap(), expected_calls); + if mode == ProviderMode::Error { + let error = result.unwrap_err(); + assert_matches!(error, Error::Unexpected { .. }); + assert!( + error + .to_string() + .contains("allocation context lookup failed") + ); + assert_matches!(procedure.data.state, CreateTableState::Prepare); + assert!(allocator.0.lock().unwrap().is_empty()); + } else { + result.unwrap(); + assert_matches!( + procedure.data.state, + CreateTableState::DatanodeCreateRegions + ); + let calls = allocator.0.lock().unwrap(); + assert_eq!(calls.len(), 1); + let expected = if mode == ProviderMode::Available { + PeerAllocContext { + affinity: AffinityRequirements::new(HashMap::from([( + "pool".into(), + "a".into(), + )])), + } + } else { + PeerAllocContext::default() + }; + assert_eq!(calls[0].affinity, expected.affinity); + } + } +} diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index eddb4148767..0a1abde30f7 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -33,6 +33,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::{RegionId, TableId}; use table::table_name::TableName; +use crate::affinity::PeerAllocContextProviderRef; use crate::ddl::alter_database::AlterDatabaseProcedure; use crate::ddl::alter_logical_tables::AlterLogicalTablesProcedure; use crate::ddl::alter_table::{AlterTableProcedure, RegionRouteChanged, only_sets_skip_wal}; @@ -300,6 +301,15 @@ impl DdlManager { self } + /// Sets the provider used to build allocation inputs from queries. + pub fn with_peer_alloc_context_provider( + mut self, + provider: PeerAllocContextProviderRef, + ) -> Self { + self.ddl_context.peer_alloc_context_provider = Some(provider); + self + } + pub fn with_create_database_metadata_committer( mut self, committer: CreateDatabaseMetadataCommitterRef, @@ -1875,6 +1885,8 @@ mod tests { let ddl_manager = DdlManager::new( DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), table_metadata_manager, @@ -1941,6 +1953,8 @@ mod tests { DdlManager::new( DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), table_metadata_manager, diff --git a/src/common/meta/src/lib.rs b/src/common/meta/src/lib.rs index ef5109dc03e..22fea83563e 100644 --- a/src/common/meta/src/lib.rs +++ b/src/common/meta/src/lib.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod affinity; pub mod cache; pub mod cache_invalidator; pub mod cluster; diff --git a/src/common/meta/src/peer.rs b/src/common/meta/src/peer.rs index 18ffd92f8bb..0e1f53ac462 100644 --- a/src/common/meta/src/peer.rs +++ b/src/common/meta/src/peer.rs @@ -12,12 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; use std::sync::Arc; pub use api::v1::meta::Peer; use api::v1::meta::heartbeat_request::NodeWorkloads; +use crate::affinity::AffinityRequirements; use crate::cluster::NodeInfo; use crate::error::Error; use crate::{DatanodeId, FlownodeId}; @@ -80,7 +80,8 @@ pub type PeerDiscoveryRef = Arc; #[derive(Debug, Clone, Default)] pub struct PeerAllocContext { - pub extensions: HashMap, + /// Resolved requirements passed to node selection. + pub affinity: Option, } /// [`PeerAllocator`] allocates [`Peer`]s for creating region or flow. diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index 618603c13e5..e574161d367 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -198,6 +198,8 @@ pub fn new_ddl_context_with_kv_backend( .build(), ))); DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager, cache_invalidator: Arc::new(DummyCacheInvalidator), memory_region_keeper: Arc::new(MemoryRegionKeeper::new()), diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index 5f9f573978c..863b2a2aab4 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -64,6 +64,12 @@ API under `src/meta-srv/src/service/admin/`. must be idempotent for crash recovery. - **Cache invalidation**: new metadata types need matching invalidation in `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. - **Selector stats fields**: keep `selector/` implementations in sync with the stats they consume. diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index 86aee9f8b82..f9b7a8dd5d7 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -26,6 +26,7 @@ use common_base::readable_size::ReadableSize; use common_config::{Configurable, DEFAULT_DATA_HOME}; use common_event_recorder::EventRecorderOptions; use common_greptimedb_telemetry::GreptimeDBTelemetryTask; +use common_meta::affinity::{AffinityFallback, AffinityRef, AffinityResolverRef}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::ddl::allocator::resource_id::ResourceIdAllocatorRef; use common_meta::ddl_manager::DdlManagerRef; @@ -486,6 +487,10 @@ impl Display for SelectTarget { #[derive(Clone)] pub struct SelectorContext { pub peer_discovery: PeerDiscoveryRef, + /// Evaluator shared by datanode and flownode selection. + pub affinity: AffinityRef, + /// Whether selection may retry without affinity when candidates are insufficient. + pub affinity_fallback: AffinityFallback, } pub type SelectorRef = Arc>>; @@ -507,6 +512,14 @@ pub struct SelectorFactoryContext { /// Builds the final datanode selector metasrv should use. pub trait SelectorFactory: Send + Sync { fn build(&self, ctx: SelectorFactoryContext) -> SelectorRef; + + /// Builds the resolver shared by allocation and failover callers. + fn build_affinity_resolver( + &self, + _meta_peer_client: MetaPeerClientRef, + ) -> Option { + None + } } /// Shared selector factory plugin registered through [`common_base::Plugins`]. diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 57e7deb3db6..96cd2a8adb5 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -22,6 +22,7 @@ use client::inserter::InsertOptions; use common_base::Plugins; use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID}; use common_event_recorder::{DEFAULT_COMPACTION_TIME_WINDOW, EventRecorderImpl, EventRecorderRef}; +use common_meta::affinity::{AffinityFallback, AffinityRef, NoopAffinity}; use common_meta::ddl::flow_meta::FlowMetadataAllocator; use common_meta::ddl::table_meta::{TableMetadataAllocator, TableMetadataAllocatorRef}; use common_meta::ddl::{ @@ -66,7 +67,8 @@ use crate::handler::region_lease_handler::{CustomizedRegionLeaseRenewerRef, Regi use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatMailbox, Pushers}; use crate::metasrv::{ ElectionRef, FLOW_ID_SEQ, METASRV_DATA_DIR, Metasrv, MetasrvInfo, MetasrvOptions, - RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorRef, TABLE_ID_SEQ, + RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorFactoryRef, SelectorRef, + TABLE_ID_SEQ, }; use crate::peer::MetasrvPeerAllocator; use crate::procedure::region_migration::DefaultContextFactory; @@ -274,6 +276,14 @@ impl MetasrvBuilder { )); let selector_ctx = SelectorContext { + affinity_fallback: plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .unwrap_or_default(), + affinity: plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .unwrap_or_else(|| Arc::new(NoopAffinity)), peer_discovery: meta_peer_client.clone(), }; @@ -385,6 +395,10 @@ impl MetasrvBuilder { ), )); region_migration_manager.try_start()?; + let affinity_resolver = plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .and_then(|factory| factory.build_affinity_resolver(meta_peer_client.clone())); let region_supervisor_selector = plugins .as_ref() .and_then(|plugins| plugins.get::()); @@ -422,6 +436,8 @@ impl MetasrvBuilder { let topic_stats_registry = Arc::new(TopicStatsRegistry::default()); let ddl_context = DdlContext { + peer_alloc_context_provider: None, + affinity_resolver, node_manager: node_manager.clone(), cache_invalidator: cache_invalidator.clone(), memory_region_keeper: memory_region_keeper.clone(), @@ -795,10 +811,99 @@ mod tests { options.gc.enable = false; options.gc.max_concurrent_tables = 0; - MetasrvBuilder::new() + let metasrv = MetasrvBuilder::new() .options(options) .build() .await .unwrap(); + let context = metasrv.ddl_manager().create_context(); + assert!(context.affinity_resolver.is_none()); + assert!(context.peer_alloc_context_provider.is_none()); + } + + #[tokio::test] + async fn test_builder_injects_affinity_components() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::{ + Affinity, AffinityRef, AffinityReference, AffinityRequirements, AffinityResolver, + AffinityResolverRef, + }; + use common_meta::cluster::NodeInfo; + use common_meta::error::Result as MetaResult; + + use crate::metasrv::{SelectorFactory, SelectorFactoryContext}; + + struct TestAffinity; + + impl Affinity for TestAffinity { + fn matches(&self, _: &AffinityRequirements, _: &NodeInfo) -> bool { + true + } + } + + struct TestResolver; + + #[async_trait::async_trait] + impl AffinityResolver for TestResolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + Ok(vec![None; references.len()]) + } + } + + struct RecordingFactory { + calls: Arc, + resolver: AffinityResolverRef, + } + + impl SelectorFactory for RecordingFactory { + fn build_affinity_resolver(&self, _: MetaPeerClientRef) -> Option { + self.calls.fetch_add(1, Ordering::Relaxed); + Some(self.resolver.clone()) + } + + fn build(&self, ctx: SelectorFactoryContext) -> SelectorRef { + ctx.base_selector + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let resolver: AffinityResolverRef = Arc::new(TestResolver); + let affinity: AffinityRef = Arc::new(TestAffinity); + let plugins = Plugins::new(); + plugins.insert(affinity.clone()); + plugins.insert(AffinityFallback::Unconstrained); + plugins.insert(Arc::new(RecordingFactory { + calls: calls.clone(), + resolver: resolver.clone(), + }) as SelectorFactoryRef); + let metasrv = MetasrvBuilder::new() + .options(MetasrvOptions { + enable_telemetry: false, + ..Default::default() + }) + .plugins(plugins) + .build() + .await + .unwrap(); + + assert_eq!( + metasrv.selector_ctx().affinity_fallback, + AffinityFallback::Unconstrained + ); + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert!(Arc::ptr_eq(&metasrv.selector_ctx().affinity, &affinity)); + assert!(Arc::ptr_eq( + metasrv + .ddl_manager() + .create_context() + .affinity_resolver + .as_ref() + .unwrap(), + &resolver + )); } } diff --git a/src/meta-srv/src/peer.rs b/src/meta-srv/src/peer.rs index 9572006182b..8843d58ad87 100644 --- a/src/meta-srv/src/peer.rs +++ b/src/meta-srv/src/peer.rs @@ -21,8 +21,8 @@ use common_meta::peer::{Peer, PeerAllocContext, PeerAllocator}; use snafu::{ResultExt, ensure}; use crate::discovery::utils::accept_ingest_workload; -use crate::error::{Result, TooManyPartitionsSnafu}; -use crate::metasrv::{SelectorContext, SelectorRef}; +use crate::error::{self, Result, TooManyPartitionsSnafu}; +use crate::metasrv::{SelectTarget, SelectorContext, SelectorRef}; use crate::selector::SelectorOptions; pub struct MetasrvPeerAllocator { @@ -67,18 +67,23 @@ impl MetasrvPeerAllocator { ); } - self.selector - .select( - &self.ctx, - SelectorOptions { - min_required_items, - allow_duplication: true, - exclude_peer_ids: HashSet::new(), - workload_filter: Some(accept_ingest_workload), - extensions: alloc_context.extensions.clone(), - }, - ) - .await + let options = SelectorOptions { + affinity: alloc_context.affinity.as_ref(), + min_required_items, + allow_duplication: true, + exclude_peer_ids: HashSet::new(), + workload_filter: Some(accept_ingest_workload), + }; + let peers = self.selector.select(&self.ctx, options).await?; + ensure!( + peers.len() >= min_required_items, + error::NoEnoughAvailableNodeSnafu { + required: min_required_items, + available: peers.len(), + select_target: SelectTarget::Datanode, + } + ); + Ok(peers) } } @@ -91,3 +96,187 @@ impl PeerAllocator for MetasrvPeerAllocator { .context(ExternalSnafu) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::ptr; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::{ + Affinity, AffinityFallback, AffinityRef, AffinityRequirements, NoopAffinity, + }; + use common_meta::cluster::NodeInfo; + + use super::*; + use crate::selector::Selector; + use crate::selector::round_robin::RoundRobinSelector; + use crate::test_util::{create_meta_peer_client, put_datanodes}; + + struct OnlySecondNode; + + impl Affinity for OnlySecondNode { + fn matches(&self, _: &AffinityRequirements, node: &NodeInfo) -> bool { + node.peer.id == 2 + } + } + + #[derive(Clone, Copy)] + enum Outcome { + Base, + BackendError, + Short, + } + + struct CountingSelector { + base: RoundRobinSelector, + calls: AtomicUsize, + requirements_address: AtomicUsize, + outcome: Outcome, + } + + #[async_trait] + impl Selector for CountingSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &SelectorContext, + opts: SelectorOptions<'_>, + ) -> Result { + self.calls.fetch_add(1, Ordering::Relaxed); + self.requirements_address.store( + opts.affinity + .map_or(0, |requirements| ptr::from_ref(requirements) as usize), + Ordering::Relaxed, + ); + assert!(opts.allow_duplication); + assert!(opts.exclude_peer_ids.is_empty()); + assert!(opts.workload_filter.is_some()); + assert_eq!(opts.min_required_items, 2); + match self.outcome { + Outcome::Base => self.base.select(ctx, opts).await, + Outcome::BackendError => error::UnexpectedSnafu { + violated: "selection backend failed", + } + .fail(), + Outcome::Short => Ok(vec![Peer::empty(1)]), + } + } + } + + #[tokio::test] + async fn test_allocator_delegates_fallback_and_validates_results() { + let discovery = create_meta_peer_client(); + put_datanodes(&discovery, vec![Peer::empty(1)]).await; + for (has_requirements, fallback, noop, outcome, expected_calls, succeeds) in [ + ( + true, + AffinityFallback::Error, + false, + Outcome::Base, + 1, + false, + ), + ( + true, + AffinityFallback::Unconstrained, + false, + Outcome::Base, + 1, + true, + ), + (true, AffinityFallback::Error, true, Outcome::Base, 1, true), + ( + false, + AffinityFallback::Unconstrained, + false, + Outcome::Base, + 1, + true, + ), + ( + true, + AffinityFallback::Unconstrained, + false, + Outcome::BackendError, + 1, + false, + ), + ( + true, + AffinityFallback::Error, + false, + Outcome::Short, + 1, + false, + ), + ( + true, + AffinityFallback::Unconstrained, + false, + Outcome::Short, + 1, + false, + ), + ( + false, + AffinityFallback::Unconstrained, + false, + Outcome::Short, + 1, + false, + ), + ] { + let selector = Arc::new(CountingSelector { + base: RoundRobinSelector::default(), + calls: AtomicUsize::new(0), + requirements_address: AtomicUsize::new(0), + outcome, + }); + let affinity: AffinityRef = if noop { + Arc::new(NoopAffinity) + } else { + Arc::new(OnlySecondNode) + }; + let allocator = MetasrvPeerAllocator::new( + SelectorContext { + affinity_fallback: fallback, + peer_discovery: discovery.clone(), + affinity, + }, + selector.clone(), + ); + let alloc_context = PeerAllocContext { + affinity: has_requirements.then(|| { + AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])).unwrap() + }), + }; + let result = allocator.alloc(2, &alloc_context).await; + assert_eq!( + selector.requirements_address.load(Ordering::Relaxed), + alloc_context + .affinity + .as_ref() + .map_or(0, |requirements| { ptr::from_ref(requirements) as usize }), + ); + assert_eq!(selector.calls.load(Ordering::Relaxed), expected_calls); + if succeeds { + assert_eq!(result.unwrap(), vec![Peer::empty(1); 2]); + } else { + match outcome { + Outcome::BackendError => assert!(matches!( + result.unwrap_err(), + error::Error::Unexpected { .. } + )), + _ => assert!(matches!( + result.unwrap_err(), + error::Error::NoEnoughAvailableNode { .. } + )), + } + } + } + } +} diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index 49964de24b7..1f975b2f9c0 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -535,6 +535,8 @@ pub mod test_data { Arc::new(SequenceBuilder::new("test", kv_backend).build()), )); DdlContext { + peer_alloc_context_provider: None, + affinity_resolver: None, node_manager, cache_invalidator: Arc::new(MetasrvCacheInvalidator::new( mailbox, diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 99c68ac23c8..93624a61892 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -667,7 +667,7 @@ impl RegionSupervisor { allow_duplication: true, exclude_peer_ids, workload_filter: Some(accept_ingest_workload), - extensions: Default::default(), + affinity: None, }; let peers = selector.select(&self.selector_context, opt).await?; ensure!( diff --git a/src/meta-srv/src/selector.rs b/src/meta-srv/src/selector.rs index a02d294d822..e5c643b7b7c 100644 --- a/src/meta-srv/src/selector.rs +++ b/src/meta-srv/src/selector.rs @@ -20,9 +20,10 @@ pub mod round_robin; pub(crate) mod test_utils; pub mod weight_compute; pub mod weighted_choose; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use api::v1::meta::heartbeat_request::NodeWorkloads; +use common_meta::affinity::AffinityRequirements; use serde::{Deserialize, Serialize}; use store_api::storage::RegionId; use strum::AsRefStr; @@ -35,7 +36,7 @@ pub trait Selector: Send + Sync { type Context; type Output; - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result; + async fn select(&self, ctx: &Self::Context, opts: SelectorOptions<'_>) -> Result; } /// A selector that aware of region statistics @@ -56,8 +57,10 @@ pub trait RegionStatAwareSelector: Send + Sync { ) -> Result; } -#[derive(Debug)] -pub struct SelectorOptions { +#[derive(Debug, Clone)] +pub struct SelectorOptions<'a> { + /// Resolved requirements used to filter candidate nodes. + pub affinity: Option<&'a AffinityRequirements>, /// Minimum number of selected results. pub min_required_items: usize, /// Whether duplicates are allowed in the selected result, default false. @@ -66,18 +69,16 @@ pub struct SelectorOptions { pub exclude_peer_ids: HashSet, /// The filter to select the peers based on their workloads. pub workload_filter: Option bool>, - /// Additional placement context, such as the frontend identity that originated the request. - pub extensions: HashMap, } -impl Default for SelectorOptions { +impl Default for SelectorOptions<'_> { fn default() -> Self { Self { + affinity: None, min_required_items: 1, allow_duplication: false, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), } } } diff --git a/src/meta-srv/src/selector/common.rs b/src/meta-srv/src/selector/common.rs index 2a31c1ce86d..00e72b23096 100644 --- a/src/meta-srv/src/selector/common.rs +++ b/src/meta-srv/src/selector/common.rs @@ -32,7 +32,7 @@ pub fn filter_out_excluded_peers( } /// According to the `opts`, choose peers from the `weight_array` through `weighted_choose`. -pub fn choose_items(opts: &SelectorOptions, weighted_choose: &mut W) -> Result> +pub fn choose_items(opts: &SelectorOptions<'_>, weighted_choose: &mut W) -> Result> where W: WeightedChoose, { @@ -86,13 +86,36 @@ where #[cfg(test)] mod tests { + use std::any::Any; use std::collections::{HashMap, HashSet}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; - use common_meta::peer::Peer; + use api::v1::meta::DatanodeWorkloads; + use api::v1::meta::heartbeat_request::NodeWorkloads; + use common_meta::affinity::{Affinity, AffinityFallback, AffinityRequirements, NoopAffinity}; + use common_meta::cluster::{DatanodeStatus, NodeInfo, NodeStatus}; + use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; + use common_meta::error::{Error as MetaError, Result as MetaResult}; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend, TxnService}; + use common_meta::peer::{Peer, PeerDiscovery}; + use common_meta::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, + BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutRequest, + PutResponse, RangeRequest, RangeResponse, + }; - use crate::selector::SelectorOptions; + use crate::cluster::MetaPeerClientBuilder; + use crate::discovery::utils::accept_ingest_workload; + use crate::metasrv::{SelectTarget, SelectorContext, SelectorRef}; use crate::selector::common::{choose_items, filter_out_excluded_peers}; + use crate::selector::lease_based::LeaseBasedSelector; + use crate::selector::load_based::LoadBasedSelector; + use crate::selector::round_robin::RoundRobinSelector; + use crate::selector::weight_compute::RegionNumsBasedWeightCompute; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; + use crate::selector::{Selector, SelectorOptions}; #[test] fn test_choose_peers() { @@ -136,11 +159,11 @@ mod tests { for i in 1..=5 { let opts = SelectorOptions { + affinity: None, min_required_items: i, allow_duplication: false, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), }; let selected_peers: HashSet<_> = @@ -153,11 +176,11 @@ mod tests { } let opts = SelectorOptions { + affinity: None, min_required_items: 6, allow_duplication: false, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), }; let selected_result = @@ -166,11 +189,11 @@ mod tests { for i in 1..=50 { let opts = SelectorOptions { + affinity: None, min_required_items: i, allow_duplication: true, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), }; let selected_peers = @@ -205,4 +228,508 @@ mod tests { assert_eq!(weight_array.len(), 1); assert_eq!(weight_array[0].item.id, 2); } + + struct ExactAffinity; + + impl Affinity for ExactAffinity { + fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool { + requirements + .attributes() + .iter() + .all(|(key, value)| candidate.env_vars.get(key) == Some(value)) + } + } + + fn requirements() -> AffinityRequirements { + AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])).unwrap() + } + + fn node(id: u64, pool: Option<&str>, accepts_ingest: bool) -> NodeInfo { + NodeInfo { + peer: Peer::empty(id), + last_activity_ts: 0, + status: NodeStatus::Datanode(DatanodeStatus { + rcus: 0, + wcus: 0, + leader_regions: 0, + follower_regions: 0, + workloads: DatanodeWorkloads { + types: if accepts_ingest { vec![0] } else { vec![] }, + }, + }), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars: pool + .map(|pool| HashMap::from([("pool".into(), pool.into())])) + .unwrap_or_default(), + } + } + + struct Discovery { + snapshots: Vec>, + fail: bool, + reads: AtomicUsize, + } + + #[async_trait::async_trait] + impl PeerDiscovery for Discovery { + async fn active_frontends(&self) -> MetaResult> { + Ok(vec![]) + } + + async fn active_datanodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> MetaResult> { + let index = self.reads.fetch_add(1, Ordering::Relaxed); + if self.fail { + return common_meta::error::UnexpectedSnafu { + err_msg: "discovery unavailable", + } + .fail(); + } + let index = index.min(self.snapshots.len() - 1); + Ok(self.snapshots[index] + .iter() + .filter(|node| { + let NodeStatus::Datanode(status) = &node.status else { + unreachable!() + }; + filter.is_none_or(|filter| { + filter(&NodeWorkloads::Datanode(status.workloads.clone())) + }) + }) + .cloned() + .collect()) + } + + async fn active_flownodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> MetaResult> { + self.active_datanodes(filter).await + } + } + + #[tokio::test] + async fn test_affinity_constrains_base_selectors() { + let affinity_requirements = requirements(); + for lease_based in [false, true] { + for (use_exact_affinity, has_requirements) in + [(true, false), (true, true), (false, true)] + { + let enabled = use_exact_affinity && has_requirements; + let selector: SelectorRef = if lease_based { + Arc::new(LeaseBasedSelector) + } else { + Arc::new(RoundRobinSelector::new(SelectTarget::Datanode)) + }; + let ctx = SelectorContext { + affinity_fallback: AffinityFallback::Error, + affinity: if use_exact_affinity { + Arc::new(ExactAffinity) + } else { + Arc::new(NoopAffinity) + }, + peer_discovery: Arc::new(Discovery { + snapshots: vec![vec![ + node(1, Some("b"), true), + node(2, Some("a"), true), + node(3, Some("a"), true), + node(4, None, true), + node(5, Some("a"), false), + ]], + reads: AtomicUsize::new(0), + fail: false, + }), + }; + // Selecting a complete round makes the lease-based result deterministic as a set. + let peers = selector + .select( + &ctx, + SelectorOptions { + affinity: has_requirements.then_some(&affinity_requirements), + min_required_items: if enabled { 1 } else { 3 }, + exclude_peer_ids: HashSet::from([3]), + workload_filter: Some(accept_ingest_workload), + ..Default::default() + }, + ) + .await + .unwrap(); + let ids = peers + .into_iter() + .map(|peer| peer.id) + .collect::>(); + assert_eq!( + ids, + if enabled { + HashSet::from([2]) + } else { + HashSet::from([1, 2, 4]) + } + ); + } + } + } + + #[tokio::test] + async fn test_round_robin_filters_each_discovery_read_for_both_targets() { + let affinity_requirements = requirements(); + for target in [SelectTarget::Datanode, SelectTarget::Flownode] { + let discovery = Arc::new(Discovery { + snapshots: vec![ + vec![node(1, Some("b"), true), node(2, Some("a"), true)], + vec![ + node(1, Some("b"), true), + node(2, Some("a"), true), + node(3, None, true), + ], + ], + reads: AtomicUsize::new(0), + fail: false, + }); + let ctx = SelectorContext { + affinity_fallback: AffinityFallback::Error, + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + }; + let selector = RoundRobinSelector::new(target); + let opts = SelectorOptions { + affinity: Some(&affinity_requirements), + ..Default::default() + }; + for _ in 0..2 { + assert_eq!( + selector.select(&ctx, opts.clone()).await.unwrap(), + vec![Peer::empty(2)] + ); + } + assert_eq!(discovery.reads.load(Ordering::Relaxed), 2); + } + } + + #[tokio::test] + async fn test_round_robin_excludes_peers_for_both_targets() { + let affinity_requirements = requirements(); + for target in [SelectTarget::Datanode, SelectTarget::Flownode] { + let ctx = SelectorContext { + affinity_fallback: AffinityFallback::Error, + peer_discovery: Arc::new(Discovery { + snapshots: vec![vec![node(1, Some("a"), true), node(2, Some("a"), true)]], + reads: AtomicUsize::new(0), + fail: false, + }), + affinity: Arc::new(ExactAffinity), + }; + for affinity in [None, Some(&affinity_requirements)] { + // Keep the topology fixed and vary only the exclusion set. + for (excluded, expected) in [(vec![], vec![1, 2]), (vec![1], vec![2, 2])] { + let selector = RoundRobinSelector::new(target); + let peers = selector + .select( + &ctx, + SelectorOptions { + affinity, + min_required_items: 2, + allow_duplication: true, + exclude_peer_ids: excluded.into_iter().collect(), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!( + peers.into_iter().map(|peer| peer.id).collect::>(), + expected + ); + } + let selector = RoundRobinSelector::new(target); + let error = selector + .select( + &ctx, + SelectorOptions { + affinity, + exclude_peer_ids: HashSet::from([1, 2]), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!(matches!( + error, + crate::error::Error::NoEnoughAvailableNode { available: 0, .. } + )); + } + } + } + + #[tokio::test] + async fn test_base_selector_affinity_fallback() { + let affinity_requirements = requirements(); + for lease_based in [false, true] { + for fallback in [AffinityFallback::Error, AffinityFallback::Unconstrained] { + for has_requirements in [false, true] { + for required in [2, 3] { + let discovery = Arc::new(Discovery { + snapshots: vec![vec![ + node(1, Some("a"), true), + node(2, Some("b"), true), + node(3, Some("a"), true), + node(4, Some("a"), false), + ]], + reads: AtomicUsize::new(0), + fail: false, + }); + let ctx = SelectorContext { + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + affinity_fallback: fallback, + }; + let selector: SelectorRef = if lease_based { + Arc::new(LeaseBasedSelector) + } else { + Arc::new(RoundRobinSelector::new(SelectTarget::Datanode)) + }; + let result = selector + .select( + &ctx, + SelectorOptions { + affinity: has_requirements.then_some(&affinity_requirements), + min_required_items: required, + allow_duplication: false, + exclude_peer_ids: HashSet::from([3]), + workload_filter: Some(accept_ingest_workload), + }, + ) + .await; + let retries = + has_requirements && fallback == AffinityFallback::Unconstrained; + assert_eq!( + discovery.reads.load(Ordering::Relaxed), + 1 + usize::from(retries) + ); + let constrained = has_requirements && !retries; + if required == 2 && !constrained { + assert_eq!( + result + .unwrap() + .into_iter() + .map(|peer| peer.id) + .collect::>(), + HashSet::from([1, 2]) + ); + } else { + let available = if constrained { 1 } else { 2 }; + assert!( + matches!(result.unwrap_err(), crate::error::Error::NoEnoughAvailableNode { + required: actual_required, available: actual_available, .. + } if actual_required == required && actual_available == available) + ); + } + } + } + } + } + } + + #[tokio::test] + async fn test_base_selector_discovery_error_does_not_fallback() { + let affinity_requirements = requirements(); + for lease_based in [false, true] { + let discovery = Arc::new(Discovery { + snapshots: vec![], + reads: AtomicUsize::new(0), + fail: true, + }); + let ctx = SelectorContext { + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + affinity_fallback: AffinityFallback::Unconstrained, + }; + let selector: SelectorRef = if lease_based { + Arc::new(LeaseBasedSelector) + } else { + Arc::new(RoundRobinSelector::new(SelectTarget::Datanode)) + }; + let error = selector + .select( + &ctx, + SelectorOptions { + affinity: Some(&affinity_requirements), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!(matches!( + error, + crate::error::Error::ListActiveDatanodes { .. } + )); + assert_eq!(discovery.reads.load(Ordering::Relaxed), 1); + } + } + + struct CountingStatBackend { + inner: MemoryKvBackend, + reads: AtomicUsize, + fail: bool, + } + + #[async_trait::async_trait] + impl TxnService for CountingStatBackend { + type Error = MetaError; + } + + #[async_trait::async_trait] + impl KvBackend for CountingStatBackend { + fn name(&self) -> &str { + "counting-stat-backend" + } + + fn as_any(&self) -> &dyn Any { + self + } + + async fn batch_get(&self, req: BatchGetRequest) -> MetaResult { + self.reads.fetch_add(1, Ordering::Relaxed); + if self.fail { + return common_meta::error::UnexpectedSnafu { + err_msg: "stat backend unavailable", + } + .fail(); + } + self.inner.batch_get(req).await + } + + async fn range(&self, req: RangeRequest) -> MetaResult { + self.inner.range(req).await + } + + async fn put(&self, req: PutRequest) -> MetaResult { + self.inner.put(req).await + } + + async fn batch_put(&self, req: BatchPutRequest) -> MetaResult { + self.inner.batch_put(req).await + } + + async fn delete_range(&self, req: DeleteRangeRequest) -> MetaResult { + self.inner.delete_range(req).await + } + + async fn batch_delete(&self, req: BatchDeleteRequest) -> MetaResult { + self.inner.batch_delete(req).await + } + } + + impl ResettableKvBackend for CountingStatBackend { + fn reset(&self) { + self.inner.reset(); + } + + fn as_kv_backend_ref(self: Arc) -> KvBackendRef { + self + } + } + + #[tokio::test] + async fn test_load_based_affinity_fallback_queries() { + let affinity_requirements = requirements(); + for (has_requirements, fallback, required, fail_stats, expected_reads, available) in [ + (false, AffinityFallback::Unconstrained, 2, false, 1, 2), + (true, AffinityFallback::Error, 2, false, 1, 1), + (true, AffinityFallback::Unconstrained, 2, false, 2, 2), + (true, AffinityFallback::Unconstrained, 3, false, 2, 2), + (true, AffinityFallback::Unconstrained, 2, true, 1, 0), + ] { + let backend = Arc::new(CountingStatBackend { + inner: MemoryKvBackend::new(), + reads: AtomicUsize::new(0), + fail: fail_stats, + }); + for id in 1..=4 { + let value = DatanodeStatValue { + stats: vec![Stat { + id, + addr: format!("node-{id}"), + region_num: id, + ..Default::default() + }], + }; + backend + .put( + PutRequest::new() + .with_key(Vec::::from(DatanodeStatKey { node_id: id })) + .with_value(Vec::::try_from(value).unwrap()), + ) + .await + .unwrap(); + } + let client = Arc::new( + MetaPeerClientBuilder::default() + .election(None) + .in_memory(backend.clone()) + .build() + .unwrap(), + ); + let selector = LoadBasedSelector::new(RegionNumsBasedWeightCompute, client); + let discovery = Arc::new(Discovery { + snapshots: vec![vec![ + node(1, Some("a"), true), + node(2, Some("b"), true), + node(3, Some("a"), true), + node(4, Some("a"), false), + ]], + reads: AtomicUsize::new(0), + fail: false, + }); + let ctx = SelectorContext { + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + affinity_fallback: fallback, + }; + let result = selector + .select( + &ctx, + SelectorOptions { + affinity: has_requirements.then_some(&affinity_requirements), + min_required_items: required, + allow_duplication: false, + exclude_peer_ids: HashSet::from([3]), + workload_filter: Some(accept_ingest_workload), + }, + ) + .await; + assert_eq!(discovery.reads.load(Ordering::Relaxed), expected_reads); + assert_eq!(backend.reads.load(Ordering::Relaxed), expected_reads); + if fail_stats { + assert!(matches!( + result.unwrap_err(), + crate::error::Error::KvBackend { .. } + )); + } else if available < required { + assert!( + matches!(result.unwrap_err(), crate::error::Error::NoEnoughAvailableNode { + required: actual_required, available: actual_available, .. + } if actual_required == required && actual_available == available) + ); + } else { + assert_eq!( + result + .unwrap() + .into_iter() + .map(|peer| peer.id) + .collect::>(), + HashSet::from([1, 2]) + ); + } + } + } } diff --git a/src/meta-srv/src/selector/lease_based.rs b/src/meta-srv/src/selector/lease_based.rs index 4a44bfca1de..b9ddd852dc4 100644 --- a/src/meta-srv/src/selector/lease_based.rs +++ b/src/meta-srv/src/selector/lease_based.rs @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; use snafu::ResultExt; -use crate::error::{ListActiveDatanodesSnafu, Result}; +use crate::error::{self, ListActiveDatanodesSnafu, Result}; use crate::metasrv::SelectorContext; use crate::selector::common::{choose_items, filter_out_excluded_peers}; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; @@ -25,12 +26,12 @@ use crate::selector::{Selector, SelectorOptions}; #[derive(Default)] pub struct LeaseBasedSelector; -#[async_trait::async_trait] -impl Selector for LeaseBasedSelector { - type Context = SelectorContext; - type Output = Vec; - - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result { +impl LeaseBasedSelector { + async fn select_once( + &self, + ctx: &SelectorContext, + opts: &SelectorOptions<'_>, + ) -> Result> { // 1. get alive datanodes. let alive_datanodes = ctx .peer_discovery @@ -38,6 +39,10 @@ impl Selector for LeaseBasedSelector { .await .context(ListActiveDatanodesSnafu)? .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) .map(|node| node.peer) .collect::>(); @@ -53,8 +58,35 @@ impl Selector for LeaseBasedSelector { // 3. choose peers by weight_array. filter_out_excluded_peers(&mut weight_array, &opts.exclude_peer_ids); let mut weighted_choose = RandomWeightedChoose::new(weight_array); - let selected = choose_items(&opts, &mut weighted_choose)?; + let selected = choose_items(opts, &mut weighted_choose)?; Ok(selected) } } + +#[async_trait::async_trait] +impl Selector for LeaseBasedSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &Self::Context, + mut opts: SelectorOptions<'_>, + ) -> Result { + match self.select_once(ctx, &opts).await { + Err(error::Error::NoEnoughAvailableNode { + required, + available, + .. + }) if opts.affinity.is_some() + && ctx.affinity_fallback == AffinityFallback::Unconstrained + && available < required => + { + opts.affinity = None; + self.select_once(ctx, &opts).await + } + result => result, + } + } +} diff --git a/src/meta-srv/src/selector/load_based.rs b/src/meta-srv/src/selector/load_based.rs index 9f36ba051f9..84196f3973e 100644 --- a/src/meta-srv/src/selector/load_based.rs +++ b/src/meta-srv/src/selector/load_based.rs @@ -14,13 +14,14 @@ use std::collections::HashMap; +use common_meta::affinity::AffinityFallback; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue}; use common_meta::peer::Peer; use common_telemetry::debug; use snafu::ResultExt; use crate::cluster::MetaPeerClientRef; -use crate::error::{ListActiveDatanodesSnafu, Result}; +use crate::error::{self, ListActiveDatanodesSnafu, Result}; use crate::metasrv::SelectorContext; use crate::selector::common::{choose_items, filter_out_excluded_peers}; use crate::selector::weight_compute::WeightCompute; @@ -41,15 +42,15 @@ impl LoadBasedSelector { } } -#[async_trait::async_trait] -impl Selector for LoadBasedSelector +impl LoadBasedSelector where C: WeightCompute>, { - type Context = SelectorContext; - type Output = Vec; - - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result { + async fn select_once( + &self, + ctx: &SelectorContext, + opts: &SelectorOptions<'_>, + ) -> Result> { // 1. get alive datanodes. let alive_datanodes = ctx .peer_discovery @@ -57,6 +58,10 @@ where .await .context(ListActiveDatanodesSnafu)? .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) .map(|node| node.peer) .collect::>(); @@ -77,7 +82,7 @@ where filter_out_excluded_peers(&mut weight_array, &opts.exclude_peer_ids); // 5. choose peers by weight_array. let mut weighted_choose = RandomWeightedChoose::new(weight_array); - let selected = choose_items(&opts, &mut weighted_choose)?; + let selected = choose_items(opts, &mut weighted_choose)?; debug!( "LoadBasedSelector select peers: {:?}, opts: {:?}.", @@ -88,6 +93,36 @@ where } } +#[async_trait::async_trait] +impl Selector for LoadBasedSelector +where + C: WeightCompute>, +{ + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &Self::Context, + mut opts: SelectorOptions<'_>, + ) -> Result { + match self.select_once(ctx, &opts).await { + Err(error::Error::NoEnoughAvailableNode { + required, + available, + .. + }) if opts.affinity.is_some() + && ctx.affinity_fallback == AffinityFallback::Unconstrained + && available < required => + { + opts.affinity = None; + self.select_once(ctx, &opts).await + } + result => result, + } + } +} + fn filter_out_expired_datanode( mut stat_kvs: HashMap, datanodes: &[Peer], diff --git a/src/meta-srv/src/selector/round_robin.rs b/src/meta-srv/src/selector/round_robin.rs index f9cf71e949b..05921921145 100644 --- a/src/meta-srv/src/selector/round_robin.rs +++ b/src/meta-srv/src/selector/round_robin.rs @@ -14,11 +14,12 @@ use std::sync::atomic::AtomicUsize; +use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; use snafu::{ResultExt, ensure}; use crate::error::{ - ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result, + self, ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result, }; use crate::metasrv::{SelectTarget, SelectorContext}; use crate::selector::{Selector, SelectorOptions}; @@ -52,35 +53,42 @@ impl RoundRobinSelector { } } - async fn get_peers(&self, opts: &SelectorOptions, ctx: &SelectorContext) -> Result> { + async fn get_peers( + &self, + opts: &SelectorOptions<'_>, + ctx: &SelectorContext, + ) -> Result> { let mut peers = match self.select_target { - SelectTarget::Datanode => { - // 1. get alive datanodes. - let alive_datanodes = ctx - .peer_discovery - .active_datanodes(opts.workload_filter) - .await - .context(ListActiveDatanodesSnafu)?; - - // 2. filter out excluded datanodes. - alive_datanodes - .into_iter() - .map(|node| node.peer) - .filter(|p| !opts.exclude_peer_ids.contains(&p.id)) - .collect::>() - } + SelectTarget::Datanode => ctx + .peer_discovery + .active_datanodes(opts.workload_filter) + .await + .context(ListActiveDatanodesSnafu)? + .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) + .map(|node| node.peer) + .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id)) + .collect::>(), SelectTarget::Flownode => ctx .peer_discovery .active_flownodes(opts.workload_filter) .await .context(ListActiveFlownodesSnafu)? .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) .map(|node| node.peer) + .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id)) .collect::>(), }; ensure!( - !peers.is_empty(), + !peers.is_empty() && (opts.allow_duplication || peers.len() >= opts.min_required_items), NoEnoughAvailableNodeSnafu { required: opts.min_required_items, available: peers.len(), @@ -93,15 +101,13 @@ impl RoundRobinSelector { Ok(peers) } -} -#[async_trait::async_trait] -impl Selector for RoundRobinSelector { - type Context = SelectorContext; - type Output = Vec; - - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result> { - let peers = self.get_peers(&opts, ctx).await?; + async fn select_once( + &self, + ctx: &SelectorContext, + opts: &SelectorOptions<'_>, + ) -> Result> { + let peers = self.get_peers(opts, ctx).await?; // choose peers let mut selected = Vec::with_capacity(opts.min_required_items); for _ in 0..opts.min_required_items { @@ -116,9 +122,39 @@ impl Selector for RoundRobinSelector { } } +#[async_trait::async_trait] +impl Selector for RoundRobinSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &Self::Context, + mut opts: SelectorOptions<'_>, + ) -> Result { + match self.select_once(ctx, &opts).await { + Err(error::Error::NoEnoughAvailableNode { + required, + available, + .. + }) if opts.affinity.is_some() + && ctx.affinity_fallback == AffinityFallback::Unconstrained + && available < required => + { + opts.affinity = None; + self.select_once(ctx, &opts).await + } + result => result, + } + } +} + #[cfg(test)] mod test { use std::collections::HashSet; + use std::sync::Arc; + + use common_meta::affinity::NoopAffinity; use super::*; use crate::test_util::{create_meta_peer_client, put_datanodes}; @@ -128,6 +164,8 @@ mod test { let selector = RoundRobinSelector::default(); let meta_peer_client = create_meta_peer_client(); let ctx = SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), peer_discovery: meta_peer_client.clone(), }; // add three nodes @@ -150,11 +188,11 @@ mod test { .select( &ctx, SelectorOptions { + affinity: None, min_required_items: 4, allow_duplication: true, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: Default::default(), }, ) .await @@ -169,11 +207,11 @@ mod test { .select( &ctx, SelectorOptions { + affinity: None, min_required_items: 2, allow_duplication: true, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: Default::default(), }, ) .await @@ -187,6 +225,8 @@ mod test { let selector = RoundRobinSelector::new(SelectTarget::Datanode); let meta_peer_client = create_meta_peer_client(); let ctx = SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), peer_discovery: meta_peer_client.clone(), }; // add three nodes @@ -212,11 +252,11 @@ mod test { .select( &ctx, SelectorOptions { + affinity: None, min_required_items: 1, allow_duplication: true, exclude_peer_ids: HashSet::from([2, 5]), workload_filter: None, - extensions: Default::default(), }, ) .await diff --git a/src/meta-srv/src/selector/test_utils.rs b/src/meta-srv/src/selector/test_utils.rs index 6a992e967fc..9b79a435e90 100644 --- a/src/meta-srv/src/selector/test_utils.rs +++ b/src/meta-srv/src/selector/test_utils.rs @@ -14,6 +14,7 @@ use std::sync::Arc; +use common_meta::affinity::NoopAffinity; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::peer::Peer; use rand::prelude::SliceRandom; @@ -34,6 +35,8 @@ pub fn new_test_selector_context() -> SelectorContext { .unwrap(); SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), peer_discovery: meta_peer_client, } } @@ -54,7 +57,11 @@ impl Selector for RandomNodeSelector { type Context = SelectorContext; type Output = Vec; - async fn select(&self, _ctx: &Self::Context, _opts: SelectorOptions) -> Result { + async fn select( + &self, + _ctx: &Self::Context, + _opts: SelectorOptions<'_>, + ) -> Result { let mut rng = rand::rng(); let mut nodes = self.nodes.clone(); nodes.shuffle(&mut rng); diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 074be282543..f8e24cb001d 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -281,6 +281,8 @@ impl GreptimeDbStandaloneBuilder { let repartition_procedure_factory = Arc::new(StandaloneRepartitionProcedureFactory); let ddl_manager = Arc::new(DdlManager::new( DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: cache_registry.clone(), memory_region_keeper: Arc::new(MemoryRegionKeeper::default()), diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 0df42ee1e88..d50905637db 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -1224,7 +1224,7 @@ impl Selector for ConstNodeSelector { async fn select( &self, _ctx: &Self::Context, - _opts: SelectorOptions, + _opts: SelectorOptions<'_>, ) -> MetaResult { Ok(self.peers.clone()) }