diff --git a/src/common/meta/src/affinity.rs b/src/common/meta/src/affinity.rs index 2bcb0e934f0..5c8bb73f661 100644 --- a/src/common/meta/src/affinity.rs +++ b/src/common/meta/src/affinity.rs @@ -171,6 +171,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars, + schedulable: None, } } diff --git a/src/common/meta/src/cluster.rs b/src/common/meta/src/cluster.rs index 4e7ffb7f920..2e0631622b8 100644 --- a/src/common/meta/src/cluster.rs +++ b/src/common/meta/src/cluster.rs @@ -138,6 +138,22 @@ pub struct NodeInfo { /// Environment variables reported by the node. #[serde(default)] pub env_vars: HashMap, + /// Whether metasrv may place new regions on this node. + /// + /// `None` means the flag was not resolved (deserialized heartbeat data, tests, + /// standalone); callers treat `None` as schedulable. Filled by the metasrv + /// discovery decorator, never by the heartbeat writer. + #[serde(skip)] + pub schedulable: Option, +} + +impl NodeInfo { + /// Returns `false` only if the node is explicitly marked unschedulable. + /// + /// `None` (unresolved) counts as schedulable. + pub fn is_schedulable(&self) -> bool { + self.schedulable != Some(false) + } } #[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize, PartialOrd, Ord)] @@ -324,6 +340,35 @@ mod tests { use crate::cluster::{DatanodeStatus, FlownodeStatus, NodeInfo, NodeInfoKey, NodeStatus}; use crate::peer::Peer; + fn new_test_node_info() -> NodeInfo { + NodeInfo { + peer: Peer { + id: 1, + addr: "127.0.0.1".to_string(), + }, + last_activity_ts: 123, + status: NodeStatus::Datanode(DatanodeStatus { + rcus: 1, + wcus: 2, + leader_regions: 3, + follower_regions: 4, + workloads: DatanodeWorkloads { + types: vec![DatanodeWorkloadType::Hybrid.to_i32()], + }, + }), + version: "".to_string(), + git_commit: "".to_string(), + start_time_ms: 1, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: "test_hostname".to_string(), + env_vars: Default::default(), + schedulable: None, + } + } + #[test] fn test_node_info_key_round_trip() { let key = NodeInfoKey { @@ -364,10 +409,16 @@ mod tests { memory_usage_bytes: 0, hostname: "test_hostname".to_string(), env_vars: Default::default(), + schedulable: None, }; let node_info_bytes: Vec = node_info.try_into().unwrap(); + assert!( + !String::from_utf8_lossy(&node_info_bytes).contains("schedulable"), + "the runtime scheduling flag must not be persisted" + ); let new_node_info: NodeInfo = node_info_bytes.try_into().unwrap(); + assert!(new_node_info.schedulable.is_none()); assert_matches!( new_node_info, @@ -387,6 +438,31 @@ mod tests { ); } + #[test] + fn test_is_schedulable() { + assert!( + NodeInfo { + schedulable: None, + ..new_test_node_info() + } + .is_schedulable() + ); + assert!( + NodeInfo { + schedulable: Some(true), + ..new_test_node_info() + } + .is_schedulable() + ); + assert!( + !NodeInfo { + schedulable: Some(false), + ..new_test_node_info() + } + .is_schedulable() + ); + } + #[test] fn test_node_info_key_prefix() { let prefix = NodeInfoKey::key_prefix(); @@ -461,6 +537,7 @@ mod tests { memory_usage_bytes: 0, hostname: "test_hostname".to_string(), env_vars: Default::default(), + schedulable: None, }; let node_info_bytes: Vec = node_info.try_into().unwrap(); @@ -523,6 +600,7 @@ mod tests { memory_usage_bytes: 0, hostname: "test".to_string(), env_vars, + schedulable: None, }; let node_info_bytes: Vec = node_info.try_into().unwrap(); diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index 44b6b21dd23..bc147c07e5e 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -509,6 +509,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: node_info.memory_usage_bytes, hostname: node_info.hostname, env_vars: Default::default(), + schedulable: None, } } else { // TODO(zyy17): It's for backward compatibility. Remove this when the deprecated fields are removed from the proto. @@ -525,6 +526,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: 0, hostname: "".to_string(), env_vars: Default::default(), + schedulable: None, } } }) @@ -543,6 +545,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: node_info.memory_usage_bytes, hostname: node_info.hostname, env_vars: Default::default(), + schedulable: None, } } else { // TODO(zyy17): It's for backward compatibility. Remove this when the deprecated fields are removed from the proto. @@ -559,6 +562,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: 0, hostname: "".to_string(), env_vars: Default::default(), + schedulable: None, } } })) diff --git a/src/meta-client/src/client/util.rs b/src/meta-client/src/client/util.rs index 62566f43783..649601741ed 100644 --- a/src/meta-client/src/client/util.rs +++ b/src/meta-client/src/client/util.rs @@ -169,6 +169,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, } } diff --git a/src/meta-srv/src/discovery.rs b/src/meta-srv/src/discovery.rs index b0d14e4a948..0e0e476bb8f 100644 --- a/src/meta-srv/src/discovery.rs +++ b/src/meta-srv/src/discovery.rs @@ -14,6 +14,7 @@ pub mod lease; pub mod node_info; +pub mod node_scheduling; pub mod utils; use api::v1::meta::heartbeat_request::NodeWorkloads; @@ -21,7 +22,7 @@ use common_error::ext::BoxedError; use common_meta::cluster::NodeInfo; use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::error::Result; -use common_meta::peer::{Peer, PeerDiscovery, PeerResolver}; +use common_meta::peer::{Peer, PeerResolver}; use common_meta::{DatanodeId, FlownodeId}; use common_time::util::DefaultSystemTimer; use snafu::ResultExt; @@ -29,9 +30,13 @@ use snafu::ResultExt; use crate::cluster::MetaPeerClient; use crate::discovery::lease::{LeaseValueAccessor, LeaseValueType}; -#[async_trait::async_trait] -impl PeerDiscovery for MetaPeerClient { - async fn active_frontends(&self) -> Result> { +// NOTE: these are inherent methods on purpose. `MetaPeerClient` does not implement +// `PeerDiscovery`, so the only way to obtain a `PeerDiscoveryRef` inside metasrv is +// `discovery::SchedulingAwarePeerDiscovery`, which stamps node scheduling states. Wrapping the raw +// client directly would silently drop the drain fence. +impl MetaPeerClient { + /// Returns all frontends that reported a heartbeat recently. + pub async fn active_frontends(&self) -> Result> { utils::alive_frontend_infos( &DefaultSystemTimer, self, @@ -42,7 +47,8 @@ impl PeerDiscovery for MetaPeerClient { .context(common_meta::error::ExternalSnafu) } - async fn active_datanodes( + /// Returns all datanodes that reported a heartbeat recently, without scheduling states. + pub async fn active_datanodes( &self, filter: Option fn(&'a NodeWorkloads) -> bool>, ) -> Result> { @@ -57,7 +63,8 @@ impl PeerDiscovery for MetaPeerClient { .context(common_meta::error::ExternalSnafu) } - async fn active_flownodes( + /// Returns all flownodes that reported a heartbeat recently. + pub async fn active_flownodes( &self, filter: Option fn(&'a NodeWorkloads) -> bool>, ) -> Result> { diff --git a/src/meta-srv/src/discovery/lease.rs b/src/meta-srv/src/discovery/lease.rs index 585676d41b5..b40bee2dcf2 100644 --- a/src/meta-srv/src/discovery/lease.rs +++ b/src/meta-srv/src/discovery/lease.rs @@ -103,7 +103,7 @@ mod tests { }; use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::kv_backend::ResettableKvBackendRef; - use common_meta::peer::{Peer, PeerDiscovery}; + use common_meta::peer::Peer; use common_meta::rpc::store::PutRequest; use common_time::util::current_time_millis; use common_workload::DatanodeWorkloadType; @@ -157,6 +157,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars, + schedulable: None, }, ) .await; @@ -199,6 +200,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, }, ) .await; @@ -236,6 +238,7 @@ mod tests { memory_usage_bytes: 0, hostname: "test_hostname".to_string(), env_vars: Default::default(), + schedulable: None, }; let key_prefix = NodeInfoKey::key_prefix_with_role(Role::Frontend); diff --git a/src/meta-srv/src/discovery/node_scheduling.rs b/src/meta-srv/src/discovery/node_scheduling.rs new file mode 100644 index 00000000000..bfcaa0e5965 --- /dev/null +++ b/src/meta-srv/src/discovery/node_scheduling.rs @@ -0,0 +1,151 @@ +// 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. + +use api::v1::meta::heartbeat_request::NodeWorkloads; +use common_error::ext::BoxedError; +use common_meta::cluster::NodeInfo; +use common_meta::error::Result; +use common_meta::peer::PeerDiscovery; +use snafu::ResultExt; + +use crate::cluster::MetaPeerClientRef; +use crate::key::NodeSchedulingRole; +use crate::node_scheduling::NodeSchedulingManagerRef; + +/// A [`PeerDiscovery`] decorator that stamps the scheduling state of datanodes +/// onto [`NodeInfo::schedulable`]. +/// +/// It never hides nodes: callers that place regions must filter with +/// `selector::common::active_schedulable_datanodes`. Keeping it additive is what allows it to be +/// the only peer discovery in the placement path (built once by `MetasrvBuilder`), because +/// failover still needs the full active set to compute its exclusions. +/// +/// This is the only `PeerDiscovery` implementation over metasrv's local state: the raw +/// [`MetaPeerClient`](crate::cluster::MetaPeerClient) only exposes inherent discovery methods, so +/// a `PeerDiscoveryRef` cannot be built without the scheduling fence. Do not build a second +/// instance; `MetasrvBuilder` owns the single production one and exposes it as +/// `Metasrv::selector_ctx().peer_discovery`. +pub struct SchedulingAwarePeerDiscovery { + manager: NodeSchedulingManagerRef, + inner: MetaPeerClientRef, +} + +impl SchedulingAwarePeerDiscovery { + pub fn new(manager: NodeSchedulingManagerRef, inner: MetaPeerClientRef) -> Self { + Self { manager, inner } + } +} + +#[async_trait::async_trait] +impl PeerDiscovery for SchedulingAwarePeerDiscovery { + async fn active_frontends(&self) -> Result> { + self.inner.active_frontends().await + } + + async fn active_datanodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> Result> { + let mut datanodes = self.inner.active_datanodes(filter).await?; + if datanodes.is_empty() { + return Ok(datanodes); + } + + let unschedulable = self + .manager + .list_unschedulable_node_ids(NodeSchedulingRole::Datanode) + .await + .map_err(BoxedError::new) + .context(common_meta::error::ExternalSnafu)?; + for datanode in &mut datanodes { + datanode.schedulable = Some(!unschedulable.contains(&datanode.peer.id)); + } + Ok(datanodes) + } + + async fn active_flownodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> Result> { + self.inner.active_flownodes(filter).await + } +} + +#[cfg(test)] +mod tests { + use common_meta::peer::Peer; + + use super::*; + use crate::cluster::MetaPeerClientRef; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; + + async fn new_discovery() -> ( + MetaPeerClientRef, + NodeSchedulingManagerRef, + SchedulingAwarePeerDiscovery, + ) { + let client = create_meta_peer_client(); + put_datanodes( + &client, + vec![Peer::empty(1), Peer::empty(2), Peer::empty(3)], + ) + .await; + let (manager, discovery) = new_node_scheduling_discovery(&client); + (client, manager, discovery) + } + + #[tokio::test] + async fn test_active_datanodes_marks_unschedulable() { + let (_, manager, discovery) = new_discovery().await; + + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + let datanodes = discovery.active_datanodes(None).await.unwrap(); + assert_eq!(3, datanodes.len()); + + let schedulable = |id: u64| { + datanodes + .iter() + .find(|datanode| datanode.peer.id == id) + .unwrap() + .schedulable + }; + assert_eq!(Some(true), schedulable(1)); + assert_eq!(Some(false), schedulable(2)); + assert_eq!(Some(true), schedulable(3)); + } + + #[tokio::test] + async fn test_schedulable_marks_are_not_persisted() { + let (client, manager, discovery) = new_discovery().await; + + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + discovery.active_datanodes(None).await.unwrap(); + + let datanodes = client.active_datanodes(None).await.unwrap(); + assert_eq!(3, datanodes.len()); + assert!( + datanodes + .iter() + .all(|datanode| datanode.schedulable.is_none()) + ); + } +} diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 70997ebe7d1..011576f9b27 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -335,6 +335,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid node scheduling key: {}", key))] + InvalidNodeSchedulingKey { + #[snafu(implicit)] + location: Location, + key: String, + }, + #[snafu(display("Failed to parse lease key from utf8"))] LeaseKeyFromUtf8 { #[snafu(source)] @@ -1217,6 +1224,7 @@ impl ErrorExt for Error { | Error::InvalidLeaseKey { .. } | Error::InvalidStatKey { .. } | Error::InvalidInactiveRegionKey { .. } + | Error::InvalidNodeSchedulingKey { .. } | Error::ParseNum { .. } | Error::ParseBool { .. } | Error::ParseAddr { .. } diff --git a/src/meta-srv/src/handler/collect_cluster_info_handler.rs b/src/meta-srv/src/handler/collect_cluster_info_handler.rs index 12fee395136..f884171a931 100644 --- a/src/meta-srv/src/handler/collect_cluster_info_handler.rs +++ b/src/meta-srv/src/handler/collect_cluster_info_handler.rs @@ -196,6 +196,7 @@ impl NodeInfoUpdate { memory_usage_bytes: info.memory_usage_bytes, hostname: info.hostname, env_vars: self.env_vars, + schedulable: None, }; let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; ctx.in_memory diff --git a/src/meta-srv/src/key.rs b/src/meta-srv/src/key.rs index abdc92ff6d4..eea9573580a 100644 --- a/src/meta-srv/src/key.rs +++ b/src/meta-srv/src/key.rs @@ -14,12 +14,14 @@ mod datanode; mod flownode; +mod node_scheduling; use std::str::FromStr; use api::v1::meta::heartbeat_request::NodeWorkloads; pub use datanode::*; pub use flownode::*; +pub use node_scheduling::*; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; diff --git a/src/meta-srv/src/key/node_scheduling.rs b/src/meta-srv/src/key/node_scheduling.rs new file mode 100644 index 00000000000..1ec398871c7 --- /dev/null +++ b/src/meta-srv/src/key/node_scheduling.rs @@ -0,0 +1,153 @@ +// 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. + +use std::str::FromStr; + +use snafu::OptionExt; + +use crate::error; +use crate::error::Result; + +/// Key prefix of datanode scheduling states. +const DATANODE_SCHEDULING_PREFIX: &str = "__node_scheduling/datanode/"; +/// Key prefix of flownode scheduling states. +const FLOWNODE_SCHEDULING_PREFIX: &str = "__node_scheduling/flownode/"; + +/// The node roles whose scheduling state is tracked. +/// +/// A subset of [`common_meta::cluster::Role`]: frontends and metasrv nodes never receive +/// regions, so they are out of scope. +#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq)] +pub enum NodeSchedulingRole { + Datanode, + Flownode, +} + +impl NodeSchedulingRole { + /// Returns the key prefix of this role, e.g. `__node_scheduling/datanode/`. + fn key_prefix(&self) -> &'static str { + match self { + NodeSchedulingRole::Datanode => DATANODE_SCHEDULING_PREFIX, + NodeSchedulingRole::Flownode => FLOWNODE_SCHEDULING_PREFIX, + } + } +} + +/// Key of a node's scheduling state: `__node_scheduling/{role}/{node_id}`, +/// for example `__node_scheduling/datanode/7`. +#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq)] +pub struct NodeSchedulingKey { + pub role: NodeSchedulingRole, + pub node_id: u64, +} + +impl NodeSchedulingKey { + pub fn new(role: NodeSchedulingRole, node_id: u64) -> NodeSchedulingKey { + NodeSchedulingKey { role, node_id } + } + + /// Returns the prefix covering every node of `role`. + pub fn prefix_key(role: NodeSchedulingRole) -> Vec { + role.key_prefix().as_bytes().to_vec() + } +} + +impl From for Vec { + fn from(key: NodeSchedulingKey) -> Vec { + format!("{}{}", key.role.key_prefix(), key.node_id).into_bytes() + } +} + +impl FromStr for NodeSchedulingKey { + type Err = error::Error; + + fn from_str(key: &str) -> Result { + for role in [NodeSchedulingRole::Datanode, NodeSchedulingRole::Flownode] { + let Some(node_id) = key.strip_prefix(role.key_prefix()) else { + continue; + }; + + let node_id = node_id + .parse::() + .ok() + .context(error::InvalidNodeSchedulingKeySnafu { key })?; + return Ok(NodeSchedulingKey { role, node_id }); + } + + error::InvalidNodeSchedulingKeySnafu { key }.fail() + } +} + +impl TryFrom> for NodeSchedulingKey { + type Error = error::Error; + + fn try_from(bytes: Vec) -> Result { + let key = String::from_utf8(bytes).map_err(|err| { + error::InvalidNodeSchedulingKeySnafu { + key: String::from_utf8_lossy(err.as_bytes()).into_owned(), + } + .build() + })?; + key.parse() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_key_round_trip() { + for (role, expected) in [ + (NodeSchedulingRole::Datanode, "__node_scheduling/datanode/7"), + (NodeSchedulingRole::Flownode, "__node_scheduling/flownode/7"), + ] { + let key = NodeSchedulingKey::new(role, 7); + let bytes: Vec = key.into(); + assert_eq!(expected.as_bytes(), bytes.as_slice()); + + let parsed = NodeSchedulingKey::try_from(bytes).unwrap(); + assert_eq!(key, parsed); + assert_eq!(role, parsed.role); + } + } + + #[test] + fn test_key_prefix() { + assert_eq!( + b"__node_scheduling/datanode/", + NodeSchedulingKey::prefix_key(NodeSchedulingRole::Datanode).as_slice() + ); + assert_eq!( + b"__node_scheduling/flownode/", + NodeSchedulingKey::prefix_key(NodeSchedulingRole::Flownode).as_slice() + ); + } + + #[test] + fn test_invalid_key() { + for key in [ + "__node_scheduling/frontend/7", + "__node_scheduling/datanode/", + "__node_scheduling/datanode/x", + "__node_scheduling", + ] { + let err = key.parse::().unwrap_err(); + assert!( + matches!(err, error::Error::InvalidNodeSchedulingKey { .. }), + "unexpected error for {key}: {err:?}" + ); + } + } +} diff --git a/src/meta-srv/src/lib.rs b/src/meta-srv/src/lib.rs index 948cab240d4..9ae4e6db59b 100644 --- a/src/meta-srv/src/lib.rs +++ b/src/meta-srv/src/lib.rs @@ -26,6 +26,7 @@ pub mod metasrv; pub mod metrics; #[cfg(feature = "mock")] pub mod mocks; +pub mod node_scheduling; pub mod peer; pub mod procedure; pub mod pubsub; diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index f9b7a8dd5d7..fd63deaa75a 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -74,6 +74,7 @@ use crate::error::{ use crate::failure_detector::PhiAccrualFailureDetectorOptions; use crate::gc::{GcSchedulerOptions, GcTickerRef}; use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatHandlerGroupRef}; +use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::ProcedureManagerListenerAdapter; use crate::procedure::region_migration::manager::RegionMigrationManagerRef; use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManagerRef; @@ -611,6 +612,7 @@ pub struct Metasrv { resource_stat: ResourceStatRef, gc_ticker: Option, database_operator: DatabaseOperatorRef, + node_scheduling_manager: NodeSchedulingManagerRef, plugins: Plugins, } @@ -852,6 +854,10 @@ impl Metasrv { &self.kv_backend } + /// Returns the undecorated peer client. + /// + /// Its `active_datanodes` does not know about node scheduling states. Use + /// [`Metasrv::selector_ctx`] for any consumer that places regions or picks datanodes. pub fn meta_peer_client(&self) -> &MetaPeerClientRef { &self.meta_peer_client } @@ -860,6 +866,7 @@ impl Metasrv { &self.selector } + /// Returns the selector context, whose `peer_discovery` stamps node scheduling states. pub fn selector_ctx(&self) -> &SelectorContext { &self.selector_ctx } @@ -924,6 +931,10 @@ impl Metasrv { &self.database_operator } + pub fn node_scheduling_manager(&self) -> &NodeSchedulingManagerRef { + &self.node_scheduling_manager + } + pub fn plugins(&self) -> &Plugins { &self.plugins } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 93ce228a156..55fdf457044 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -56,6 +56,7 @@ use store_api::storage::MAX_REGION_SEQ; use crate::bootstrap::build_default_meta_peer_client; use crate::cache_invalidator::MetasrvCacheInvalidator; use crate::cluster::MetaPeerClientRef; +use crate::discovery::node_scheduling::SchedulingAwarePeerDiscovery; use crate::error::{self, BuildWalProviderSnafu, OtherSnafu, Result}; use crate::event::EventHandlerImpl; use crate::gc::{DefaultGcSchedulerCtx, GcScheduler}; @@ -70,6 +71,7 @@ use crate::metasrv::{ RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorFactoryRef, SelectorRef, TABLE_ID_SEQ, }; +use crate::node_scheduling::NodeSchedulingManager; use crate::peer::MetasrvPeerAllocator; use crate::procedure::region_migration::DefaultContextFactory; use crate::procedure::region_migration::manager::RegionMigrationManager; @@ -238,7 +240,31 @@ impl MetasrvBuilder { let meta_peer_client = meta_peer_client .unwrap_or_else(|| build_default_meta_peer_client(&election, &in_memory)); - let database_operator = Arc::new(DatabaseOperator::new(meta_peer_client.clone())); + + let node_scheduling_manager = Arc::new(NodeSchedulingManager::new(kv_backend.clone())); + + // The only place a node scheduling decorator is built. Every consumer that reads the + // datanode list (region placement, failover exclusions) must go through + // `Metasrv::selector_ctx().peer_discovery`, and building a `PeerDiscoveryRef` from the raw + // client is impossible by construction. Bypassing this decorator silently drops the + // drain fence, so do not build a second instance. + 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: Arc::new(SchedulingAwarePeerDiscovery::new( + node_scheduling_manager.clone(), + meta_peer_client.clone(), + )), + }; + // Frontends only; the decorator passes them through unchanged. + let database_operator = + Arc::new(DatabaseOperator::new(selector_ctx.peer_discovery.clone())); let event_inserter = Box::new(InsertForwarder::new( database_operator.clone(), @@ -275,18 +301,6 @@ impl MetasrvBuilder { leader_cached_kv_backend.clone() as _, )); - 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(), - }; - let wal_provider = build_wal_provider(&options.wal, kv_backend.clone()) .await .context(BuildWalProviderSnafu)?; @@ -679,6 +693,7 @@ impl MetasrvBuilder { resource_stat: Arc::new(resource_stat), gc_ticker, database_operator, + node_scheduling_manager, }) } } @@ -917,4 +932,44 @@ mod tests { &resolver )); } + + #[tokio::test] + async fn test_builder_wires_node_scheduling_discovery() { + use common_meta::peer::Peer; + + use crate::key::NodeSchedulingRole; + use crate::test_util::put_datanodes; + + let metasrv = MetasrvBuilder::new() + .options(MetasrvOptions { + enable_telemetry: false, + ..Default::default() + }) + .build() + .await + .unwrap(); + + put_datanodes( + metasrv.meta_peer_client(), + vec![Peer::empty(1), Peer::empty(2)], + ) + .await; + metasrv + .node_scheduling_manager() + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + let datanodes = metasrv + .selector_ctx() + .peer_discovery + .active_datanodes(None) + .await + .unwrap(); + assert_eq!(2, datanodes.len()); + for datanode in datanodes { + let expected = datanode.peer.id != 2; + assert_eq!(Some(expected), datanode.schedulable); + } + } } diff --git a/src/meta-srv/src/node_scheduling.rs b/src/meta-srv/src/node_scheduling.rs new file mode 100644 index 00000000000..cb7c34875ec --- /dev/null +++ b/src/meta-srv/src/node_scheduling.rs @@ -0,0 +1,401 @@ +// 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. + +use std::collections::HashSet; +use std::sync::Arc; + +use common_meta::kv_backend::KvBackendRef; +use common_meta::kv_backend::txn::{Compare, CompareOp, Txn, TxnOp}; +use common_meta::rpc::store::RangeRequest; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; + +use crate::error::{self, Result}; +use crate::key::{NodeSchedulingKey, NodeSchedulingRole}; + +pub type NodeSchedulingManagerRef = Arc; + +/// The scheduling state persisted for a node. +/// +/// Only `schedulable: false` is written today; a missing key already means schedulable. +/// New metadata fields must use `#[serde(default)]`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct NodeSchedulingValue { + /// Whether the node may receive new regions. + pub schedulable: bool, +} + +/// Persists which nodes may receive new regions. +/// +/// A missing key means the node is schedulable; a present key stores a +/// [`NodeSchedulingValue`] whose `schedulable` field decides placement. +pub struct NodeSchedulingManager { + kv_backend: KvBackendRef, +} + +impl NodeSchedulingManager { + pub fn new(kv_backend: KvBackendRef) -> Self { + Self { kv_backend } + } + + /// Marks the node as unschedulable. Idempotent. + /// + /// Returns whether the stored state changed. + pub async fn mark_unschedulable(&self, role: NodeSchedulingRole, node_id: u64) -> Result { + let key: Vec = NodeSchedulingKey::new(role, node_id).into(); + let value = serde_json::to_vec(&NodeSchedulingValue { schedulable: false }).context( + error::SerializeToJsonSnafu { + input: String::from_utf8_lossy(&key).into_owned(), + }, + )?; + + let txn = match self + .kv_backend + .get(&key) + .await + .context(error::KvBackendSnafu)? + { + None => Txn::put_if_not_exists(key, value), + Some(kv) if kv.value == value => return Ok(false), + Some(kv) => Txn::compare_and_put(key, kv.value, value), + }; + let response = self + .kv_backend + .txn(txn) + .await + .context(error::KvBackendSnafu)?; + Ok(response.succeeded) + } + + /// Marks the node as schedulable by deleting its state. Idempotent. + /// + /// Returns whether a state was removed; a node without state is already schedulable, + /// so this returns `false` and writes nothing. + pub async fn mark_schedulable(&self, role: NodeSchedulingRole, node_id: u64) -> Result { + let key: Vec = NodeSchedulingKey::new(role, node_id).into(); + // `NotEqual` against a missing target means "the key exists". + let txn = Txn::new() + .when(vec![Compare::new(key.clone(), CompareOp::NotEqual, None)]) + .and_then(vec![TxnOp::Delete(key)]); + let response = self + .kv_backend + .txn(txn) + .await + .context(error::KvBackendSnafu)?; + Ok(response.succeeded) + } + + /// Returns `false` iff the node is marked unschedulable. + pub async fn is_schedulable(&self, role: NodeSchedulingRole, node_id: u64) -> Result { + let key = NodeSchedulingKey::new(role, node_id); + Ok(self + .get_value(&key) + .await? + .is_none_or(|value| value.schedulable)) + } + + /// Returns the ids of nodes of `role` that must not receive new regions. + /// + /// Returns a set because callers test membership per node. + pub async fn list_unschedulable_node_ids( + &self, + role: NodeSchedulingRole, + ) -> Result> { + let request = RangeRequest::new().with_prefix(NodeSchedulingKey::prefix_key(role)); + let response = self + .kv_backend + .range(request) + .await + .context(error::KvBackendSnafu)?; + + let mut node_ids = HashSet::new(); + for kv in response.kvs { + let key = NodeSchedulingKey::try_from(kv.key)?; + if !parse_value(&kv.value)?.schedulable { + node_ids.insert(key.node_id); + } + } + Ok(node_ids) + } + + async fn get_value(&self, key: &NodeSchedulingKey) -> Result> { + let key: Vec = (*key).into(); + let Some(kv) = self + .kv_backend + .get(&key) + .await + .context(error::KvBackendSnafu)? + else { + return Ok(None); + }; + + parse_value(&kv.value).map(Some) + } +} + +fn parse_value(value: &[u8]) -> Result { + serde_json::from_slice(value).context(error::DeserializeFromJsonSnafu { + input_len: value.len(), + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_meta::kv_backend::KvBackend; + use common_meta::kv_backend::memory::MemoryKvBackend; + + use super::*; + use crate::metasrv::SelectorContext; + + fn new_manager() -> ( + Arc>, + NodeSchedulingManager, + ) { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let manager = NodeSchedulingManager::new(kv_backend.clone()); + (kv_backend, manager) + } + + #[tokio::test] + async fn test_mark_unschedulable_is_idempotent() { + let (kv_backend, manager) = new_manager(); + + assert!( + manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + !manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + + assert!( + !manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert_eq!( + HashSet::from([1]), + manager + .list_unschedulable_node_ids(NodeSchedulingRole::Datanode) + .await + .unwrap() + ); + + let key: Vec = NodeSchedulingKey::new(NodeSchedulingRole::Datanode, 1).into(); + let kv = kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!(br#"{"schedulable":false}"#, kv.value.as_slice()); + + assert!( + manager + .mark_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + !manager + .mark_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + + assert!( + manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + manager + .list_unschedulable_node_ids(NodeSchedulingRole::Datanode) + .await + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn test_mark_schedulable_without_state_writes_nothing() { + let (kv_backend, manager) = new_manager(); + + assert!( + !manager + .mark_schedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap() + ); + + let key: Vec = NodeSchedulingKey::new(NodeSchedulingRole::Datanode, 2).into(); + assert!(kv_backend.get(&key).await.unwrap().is_none()); + } + + #[tokio::test] + async fn test_scheduling_state_is_role_scoped() { + let (_, manager) = new_manager(); + + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 7) + .await + .unwrap(); + + assert!( + manager + .is_schedulable(NodeSchedulingRole::Flownode, 7) + .await + .unwrap() + ); + assert!( + manager + .list_unschedulable_node_ids(NodeSchedulingRole::Flownode) + .await + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn test_scheduling_state_is_visible_to_other_manager() { + let (kv_backend, manager_a) = new_manager(); + let manager_b = NodeSchedulingManager::new(kv_backend); + + manager_a + .mark_unschedulable(NodeSchedulingRole::Datanode, 3) + .await + .unwrap(); + + assert!( + !manager_b + .is_schedulable(NodeSchedulingRole::Datanode, 3) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_concurrent_mark_unschedulable_changes_once() { + let (_, manager) = new_manager(); + + let (a, b) = tokio::join!( + manager.mark_unschedulable(NodeSchedulingRole::Datanode, 2), + manager.mark_unschedulable(NodeSchedulingRole::Datanode, 2) + ); + let changed = [a.unwrap(), b.unwrap()]; + assert_eq!(1, changed.iter().filter(|changed| **changed).count()); + } + + /// Returns a selector context that stamps scheduling marks of datanodes 1, 2, 3 + /// where datanode 2 is unschedulable. + async fn new_selector_context() -> SelectorContext { + use common_meta::peer::Peer; + + use crate::test_util::{ + create_meta_peer_client, new_node_scheduling_discovery, put_datanodes, + }; + + let client = create_meta_peer_client(); + put_datanodes( + &client, + vec![Peer::empty(1), Peer::empty(2), Peer::empty(3)], + ) + .await; + + let (manager, discovery) = new_node_scheduling_discovery(&client); + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(common_meta::affinity::NoopAffinity), + peer_discovery: Arc::new(discovery), + } + } + + #[tokio::test] + async fn test_lease_based_selector_skips_unschedulable() { + use crate::error; + use crate::selector::lease_based::LeaseBasedSelector; + use crate::selector::{Selector, SelectorOptions}; + + let ctx = new_selector_context().await; + let selector = LeaseBasedSelector; + + let peers = selector + .select( + &ctx, + SelectorOptions { + min_required_items: 2, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(2, peers.len()); + for peer in peers { + assert!([1, 3].contains(&peer.id), "unexpected peer: {peer:?}"); + } + + let err = selector + .select( + &ctx, + SelectorOptions { + min_required_items: 3, + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!( + matches!(err, error::Error::NoEnoughAvailableNode { .. }), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn test_round_robin_selector_skips_unschedulable() { + use crate::metasrv::SelectTarget; + use crate::selector::round_robin::RoundRobinSelector; + use crate::selector::{Selector, SelectorOptions}; + + let ctx = new_selector_context().await; + let selector = RoundRobinSelector::new(SelectTarget::Datanode); + + let peers = selector + .select( + &ctx, + SelectorOptions { + min_required_items: 2, + ..Default::default() + }, + ) + .await + .unwrap(); + let ids = peers.iter().map(|peer| peer.id).collect::>(); + assert_eq!(vec![1, 3], ids); + } +} diff --git a/src/meta-srv/src/peer.rs b/src/meta-srv/src/peer.rs index 8843d58ad87..e0a0589ab37 100644 --- a/src/meta-srv/src/peer.rs +++ b/src/meta-srv/src/peer.rs @@ -112,7 +112,7 @@ mod tests { use super::*; use crate::selector::Selector; use crate::selector::round_robin::RoundRobinSelector; - use crate::test_util::{create_meta_peer_client, put_datanodes}; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; struct OnlySecondNode; @@ -244,7 +244,7 @@ mod tests { let allocator = MetasrvPeerAllocator::new( SelectorContext { affinity_fallback: fallback, - peer_discovery: discovery.clone(), + peer_discovery: Arc::new(new_node_scheduling_discovery(&discovery).1), affinity, }, selector.clone(), diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 5cf3f8e7fa8..c7c5cf7dab7 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -702,7 +702,19 @@ impl RegionSupervisor { } else { None }; - let exclude_peer_ids = HashSet::from_iter(failure_datanodes.iter().copied()); + let mut exclude_peer_ids: HashSet = + HashSet::from_iter(failure_datanodes.iter().copied()); + let unschedulable_datanodes = self + .selector_context + .peer_discovery + .active_datanodes(None) + .await + .map_err(common_error::ext::BoxedError::new) + .context(error::OtherSnafu)? + .into_iter() + .filter(|datanode| !datanode.is_schedulable()) + .map(|datanode| datanode.peer.id); + exclude_peer_ids.extend(unschedulable_datanodes); match &self.selector { RegionSupervisorSelector::NaiveSelector(selector) => { let opts = SelectorOptions { @@ -944,6 +956,7 @@ pub(crate) mod tests { use std::sync::{Arc, Mutex}; use std::time::Duration; + use common_meta::affinity::NoopAffinity; use common_meta::ddl::RegionFailureDetectorController; use common_meta::ddl::test_util::{ test_create_logical_table_task, test_create_physical_table_task, @@ -964,6 +977,7 @@ pub(crate) mod tests { use tokio::time::sleep; use super::RegionSupervisorSelector; + use crate::metasrv::SelectorContext; use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::RegionMigrationTriggerReason; use crate::procedure::region_migration::manager::{ @@ -977,9 +991,19 @@ pub(crate) mod tests { use crate::selector::test_utils::{RandomNodeSelector, new_test_selector_context}; pub(crate) fn new_test_supervisor() -> (RegionSupervisor, Sender) { + new_test_supervisor_with( + new_test_selector_context(), + RegionSupervisorSelector::NaiveSelector(Arc::new(RandomNodeSelector::new(vec![ + Peer::empty(1), + ]))), + ) + } + + pub(crate) fn new_test_supervisor_with( + selector_context: SelectorContext, + selector: RegionSupervisorSelector, + ) -> (RegionSupervisor, Sender) { let env = TestingEnv::new(); - let selector_context = new_test_selector_context(); - let selector = Arc::new(RandomNodeSelector::new(vec![Peer::empty(1)])); let context_factory = env.context_factory(); let region_migration_manager = Arc::new(RegionMigrationManager::new( env.procedure_manager().clone(), @@ -996,7 +1020,7 @@ pub(crate) mod tests { rx, Default::default(), selector_context, - RegionSupervisorSelector::NaiveSelector(selector), + selector, region_migration_manager, runtime_switch_manager, peer_resolver, @@ -1204,6 +1228,61 @@ pub(crate) mod tests { } } + #[tokio::test] + async fn test_select_peers_excludes_unschedulable_datanodes() { + use std::collections::HashSet as StdHashSet; + + use crate::key::NodeSchedulingRole; + use crate::selector::{Selector, SelectorOptions}; + use crate::test_util::{ + create_meta_peer_client, new_node_scheduling_discovery, put_datanodes, + }; + + /// Selector that records the excluded peers and returns a fixed peer. + struct ExcludeAssertSelector; + + #[async_trait::async_trait] + impl Selector for ExcludeAssertSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + _ctx: &Self::Context, + opts: SelectorOptions<'_>, + ) -> crate::error::Result { + assert_eq!(StdHashSet::from([7, 8]), opts.exclude_peer_ids); + Ok(vec![Peer::empty(9)]) + } + } + + // Datanode 7 is marked unschedulable; datanode 8 failed. + let client = create_meta_peer_client(); + put_datanodes(&client, vec![Peer::empty(7)]).await; + let (manager, discovery) = new_node_scheduling_discovery(&client); + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 7) + .await + .unwrap(); + + let selector_context = SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), + peer_discovery: Arc::new(discovery), + }; + let (supervisor, _tx) = new_test_supervisor_with( + selector_context, + RegionSupervisorSelector::NaiveSelector(Arc::new(ExcludeAssertSelector)), + ); + + let region_peers = supervisor + .select_peers(7, &[RegionId::new(1, 1)], &[8]) + .await + .unwrap(); + assert_eq!(1, region_peers.len()); + assert_eq!(Peer::empty(9), region_peers[0].1); + } + #[tokio::test] async fn test_heartbeat() { let (mut supervisor, sender) = new_test_supervisor(); diff --git a/src/meta-srv/src/selector/common.rs b/src/meta-srv/src/selector/common.rs index 00e72b23096..0fec9a4ec66 100644 --- a/src/meta-srv/src/selector/common.rs +++ b/src/meta-srv/src/selector/common.rs @@ -14,15 +14,33 @@ use std::collections::HashSet; +use api::v1::meta::heartbeat_request::NodeWorkloads; +use common_meta::cluster::NodeInfo; use common_meta::peer::Peer; -use snafu::ensure; +use snafu::{ResultExt, ensure}; use crate::error; use crate::error::Result; -use crate::metasrv::SelectTarget; +use crate::metasrv::{SelectTarget, SelectorContext}; use crate::selector::SelectorOptions; use crate::selector::weighted_choose::{WeightedChoose, WeightedItem}; +/// Returns active datanodes that may receive new regions. +pub(crate) async fn active_schedulable_datanodes( + ctx: &SelectorContext, + filter: Option fn(&'a NodeWorkloads) -> bool>, +) -> Result> { + let datanodes = ctx + .peer_discovery + .active_datanodes(filter) + .await + .context(error::ListActiveDatanodesSnafu)?; + Ok(datanodes + .into_iter() + .filter(NodeInfo::is_schedulable) + .collect()) +} + /// Filter out the excluded peers from the `weight_array`. pub fn filter_out_excluded_peers( weight_array: &mut Vec>, @@ -268,6 +286,7 @@ mod tests { env_vars: pool .map(|pool| HashMap::from([("pool".into(), pool.into())])) .unwrap_or_default(), + schedulable: None, } } diff --git a/src/meta-srv/src/selector/lease_based.rs b/src/meta-srv/src/selector/lease_based.rs index b9ddd852dc4..83363929a57 100644 --- a/src/meta-srv/src/selector/lease_based.rs +++ b/src/meta-srv/src/selector/lease_based.rs @@ -14,11 +14,10 @@ use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; -use snafu::ResultExt; -use crate::error::{self, ListActiveDatanodesSnafu, Result}; +use crate::error::{self, Result}; use crate::metasrv::SelectorContext; -use crate::selector::common::{choose_items, filter_out_excluded_peers}; +use crate::selector::common::{self, choose_items, filter_out_excluded_peers}; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; use crate::selector::{Selector, SelectorOptions}; @@ -33,11 +32,8 @@ impl LeaseBasedSelector { opts: &SelectorOptions<'_>, ) -> Result> { // 1. get alive datanodes. - let alive_datanodes = ctx - .peer_discovery - .active_datanodes(opts.workload_filter) - .await - .context(ListActiveDatanodesSnafu)? + let alive_datanodes = common::active_schedulable_datanodes(ctx, opts.workload_filter) + .await? .into_iter() .filter(|node| { opts.affinity diff --git a/src/meta-srv/src/selector/load_based.rs b/src/meta-srv/src/selector/load_based.rs index 84196f3973e..175c65b4696 100644 --- a/src/meta-srv/src/selector/load_based.rs +++ b/src/meta-srv/src/selector/load_based.rs @@ -18,12 +18,11 @@ 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::{self, ListActiveDatanodesSnafu, Result}; +use crate::error::{self, Result}; use crate::metasrv::SelectorContext; -use crate::selector::common::{choose_items, filter_out_excluded_peers}; +use crate::selector::common::{self, choose_items, filter_out_excluded_peers}; use crate::selector::weight_compute::WeightCompute; use crate::selector::weighted_choose::RandomWeightedChoose; use crate::selector::{Selector, SelectorOptions}; @@ -52,11 +51,8 @@ where opts: &SelectorOptions<'_>, ) -> Result> { // 1. get alive datanodes. - let alive_datanodes = ctx - .peer_discovery - .active_datanodes(opts.workload_filter) - .await - .context(ListActiveDatanodesSnafu)? + let alive_datanodes = common::active_schedulable_datanodes(ctx, opts.workload_filter) + .await? .into_iter() .filter(|node| { opts.affinity diff --git a/src/meta-srv/src/selector/round_robin.rs b/src/meta-srv/src/selector/round_robin.rs index 05921921145..45e2f4f296b 100644 --- a/src/meta-srv/src/selector/round_robin.rs +++ b/src/meta-srv/src/selector/round_robin.rs @@ -18,11 +18,9 @@ use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; use snafu::{ResultExt, ensure}; -use crate::error::{ - self, ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result, -}; +use crate::error::{self, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result}; use crate::metasrv::{SelectTarget, SelectorContext}; -use crate::selector::{Selector, SelectorOptions}; +use crate::selector::{Selector, SelectorOptions, common}; /// Round-robin selector that returns the next peer in the list in sequence. /// Datanodes are ordered by their node_id. @@ -59,19 +57,22 @@ impl RoundRobinSelector { ctx: &SelectorContext, ) -> Result> { let mut peers = match self.select_target { - 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::Datanode => { + // 1. get alive, schedulable datanodes. + let alive_datanodes = + common::active_schedulable_datanodes(ctx, opts.workload_filter).await?; + + // 2. filter by affinity, then drop the excluded datanodes. + alive_datanodes + .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) @@ -157,7 +158,7 @@ mod test { use common_meta::affinity::NoopAffinity; use super::*; - use crate::test_util::{create_meta_peer_client, put_datanodes}; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; #[tokio::test] async fn test_round_robin_selector() { @@ -166,7 +167,7 @@ mod test { let ctx = SelectorContext { affinity_fallback: Default::default(), affinity: Arc::new(NoopAffinity), - peer_discovery: meta_peer_client.clone(), + peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1), }; // add three nodes let peer1 = Peer { @@ -227,7 +228,7 @@ mod test { let ctx = SelectorContext { affinity_fallback: Default::default(), affinity: Arc::new(NoopAffinity), - peer_discovery: meta_peer_client.clone(), + peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1), }; // add three nodes let peer1 = Peer { diff --git a/src/meta-srv/src/selector/test_utils.rs b/src/meta-srv/src/selector/test_utils.rs index 9b79a435e90..8073e41344e 100644 --- a/src/meta-srv/src/selector/test_utils.rs +++ b/src/meta-srv/src/selector/test_utils.rs @@ -23,13 +23,19 @@ use crate::cluster::MetaPeerClientBuilder; use crate::error::Result; use crate::metasrv::SelectorContext; use crate::selector::{Selector, SelectorOptions}; +use crate::test_util::new_node_scheduling_discovery; /// Returns [SelectorContext] for test purpose. +/// +/// Its `peer_discovery` is decorated with +/// [`SchedulingAwarePeerDiscovery`](crate::discovery::node_scheduling::SchedulingAwarePeerDiscovery) +/// backed by an empty scheduling state, so every datanode is schedulable until a test marks it +/// otherwise in the same in-memory backend. pub fn new_test_selector_context() -> SelectorContext { let kv_backend = Arc::new(MemoryKvBackend::new()); let meta_peer_client = MetaPeerClientBuilder::default() .election(None) - .in_memory(kv_backend.clone()) + .in_memory(kv_backend) .build() .map(Arc::new) .unwrap(); @@ -37,7 +43,7 @@ pub fn new_test_selector_context() -> SelectorContext { SelectorContext { affinity_fallback: Default::default(), affinity: Arc::new(NoopAffinity), - peer_discovery: meta_peer_client, + peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1), } } diff --git a/src/meta-srv/src/test_util.rs b/src/meta-srv/src/test_util.rs index 1169a93309d..4696db380c7 100644 --- a/src/meta-srv/src/test_util.rs +++ b/src/meta-srv/src/test_util.rs @@ -24,7 +24,9 @@ use common_time::util as time_util; use common_workload::DatanodeWorkloadType; use crate::cluster::{MetaPeerClientBuilder, MetaPeerClientRef}; +use crate::discovery::node_scheduling::SchedulingAwarePeerDiscovery; use crate::key::{DatanodeLeaseKey, LeaseValue}; +use crate::node_scheduling::{NodeSchedulingManager, NodeSchedulingManagerRef}; pub(crate) fn new_region_route(region_id: u64, peers: &[Peer], leader_node: u64) -> RegionRoute { let region = Region { @@ -55,6 +57,20 @@ pub(crate) fn create_meta_peer_client() -> MetaPeerClientRef { .unwrap() } +/// Wraps `meta_peer_client` into the node scheduling discovery. +/// +/// The manager shares the client's in-memory backend, so marks written through the returned +/// manager are visible to the discovery. +pub(crate) fn new_node_scheduling_discovery( + meta_peer_client: &MetaPeerClientRef, +) -> (NodeSchedulingManagerRef, SchedulingAwarePeerDiscovery) { + let manager = Arc::new(NodeSchedulingManager::new( + meta_peer_client.memory_backend().as_kv_backend_ref(), + )); + let discovery = SchedulingAwarePeerDiscovery::new(manager.clone(), meta_peer_client.clone()); + (manager, discovery) +} + pub(crate) async fn put_datanodes(meta_peer_client: &MetaPeerClientRef, datanodes: Vec) { let backend = meta_peer_client.memory_backend(); for datanode in datanodes { @@ -102,6 +118,7 @@ pub(crate) async fn put_datanodes(meta_peer_client: &MetaPeerClientRef, datanode memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, }; let node_info_key = NodeInfoKey { role: Role::Datanode, diff --git a/src/meta-srv/src/utils/database.rs b/src/meta-srv/src/utils/database.rs index 796d63009d8..a67cdb14aab 100644 --- a/src/meta-srv/src/utils/database.rs +++ b/src/meta-srv/src/utils/database.rs @@ -182,6 +182,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, }]) } diff --git a/src/standalone/src/information_extension.rs b/src/standalone/src/information_extension.rs index 8d52eb12ccc..f42d7a84f7f 100644 --- a/src/standalone/src/information_extension.rs +++ b/src/standalone/src/information_extension.rs @@ -89,6 +89,7 @@ impl InformationExtension for StandaloneInformationExtension { .to_string_lossy() .to_string(), env_vars: Default::default(), + schedulable: None, }; Ok(vec![node_info]) }