From 2d00ae25c82e363eaba62f63ec1a21e85a348337 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 26 Oct 2026 11:52:41 +0000 Subject: [PATCH] feat(meta): fence a draining datanode and plan the drain (#9960) Introduce the metasrv side procedure that fences a datanode before it is drained, and collect the regions placed on it: - mark the datanode unschedulable through the node scheduling manager - build the drain plan from the datanode table index and the authoritative table routes, splitting leader regions from follower regions - keep the plan in the volatile context: it is derived from the table routes, so a restarted procedure rebuilds it instead of reading a persisted copy The leader regions are not migrated yet: the migration state machine lands separately. The procedure has no caller yet, so it stays inert. Signed-off-by: WenyXu --- src/meta-srv/src/error.rs | 8 + src/meta-srv/src/metasrv.rs | 7 + src/meta-srv/src/metasrv/builder.rs | 9 + src/meta-srv/src/procedure.rs | 1 + src/meta-srv/src/procedure/drain_datanode.rs | 652 ++++++++++++++++++ .../drain_datanode/drain_datanode_end.rs | 41 ++ .../drain_datanode/drain_datanode_start.rs | 61 ++ .../src/procedure/drain_datanode/manager.rs | 187 +++++ .../src/procedure/drain_datanode/utils.rs | 83 +++ 9 files changed, 1049 insertions(+) create mode 100644 src/meta-srv/src/procedure/drain_datanode.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/manager.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/utils.rs diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 011576f9b27..8ceb39ed56e 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -91,6 +91,13 @@ pub enum Error { region_id: RegionId, }, + #[snafu(display("Another drain procedure is running for datanode: {}", node_id))] + DrainDatanodeRunning { + #[snafu(implicit)] + location: Location, + node_id: DatanodeId, + }, + #[snafu(display( "The region migration procedure is completed for region: {}, target_peer: {}", region_id, @@ -1254,6 +1261,7 @@ impl ErrorExt for Error { | Error::RegionRouteNotFound { .. } | Error::MigrationAbort { .. } | Error::MigrationRunning { .. } + | Error::DrainDatanodeRunning { .. } | Error::RegionMigrated { .. } => StatusCode::Unexpected, Error::TableNotFound { .. } => StatusCode::TableNotFound, Error::SaveClusterInfo { source, .. } diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index f59e15cd5cf..37d8d3ebb2c 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -77,6 +77,7 @@ use crate::gc::{GcSchedulerOptions, GcTickerRef}; use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatHandlerGroupRef}; use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::ProcedureManagerListenerAdapter; +use crate::procedure::drain_datanode::manager::DrainDatanodeManagerRef; use crate::procedure::region_migration::manager::RegionMigrationManagerRef; use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManagerRef; use crate::procedure::wal_prune::manager::WalPruneTickerRef; @@ -606,6 +607,7 @@ pub struct Metasrv { memory_region_keeper: MemoryRegionKeeperRef, greptimedb_telemetry_task: Arc, region_migration_manager: RegionMigrationManagerRef, + drain_datanode_manager: DrainDatanodeManagerRef, region_supervisor_ticker: Option, cache_invalidator: CacheInvalidatorRef, leader_region_registry: LeaderRegionRegistryRef, @@ -916,6 +918,11 @@ impl Metasrv { &self.region_migration_manager } + /// Returns the [`DrainDatanodeManagerRef`]. + pub fn drain_datanode_manager(&self) -> &DrainDatanodeManagerRef { + &self.drain_datanode_manager + } + pub fn publish(&self) -> Option { self.plugins.get::() } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 0a0e1d339e8..bc10a27af85 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -74,6 +74,7 @@ use crate::metasrv::{ }; use crate::node_scheduling::NodeSchedulingManager; use crate::peer::MetasrvPeerAllocator; +use crate::procedure::drain_datanode::manager::DrainDatanodeManager; use crate::procedure::region_migration::DefaultContextFactory; use crate::procedure::region_migration::manager::RegionMigrationManager; use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManager; @@ -410,6 +411,13 @@ impl MetasrvBuilder { ), )); region_migration_manager.try_start()?; + + let drain_datanode_manager = Arc::new(DrainDatanodeManager::new( + procedure_manager.clone(), + table_metadata_manager.clone(), + node_scheduling_manager.clone(), + )); + drain_datanode_manager.try_start()?; let affinity_resolver = plugins .as_ref() .and_then(|plugins| plugins.get::()) @@ -688,6 +696,7 @@ impl MetasrvBuilder { plugins: plugins.unwrap_or_else(Plugins::default), memory_region_keeper, region_migration_manager, + drain_datanode_manager, region_supervisor_ticker, cache_invalidator, leader_region_registry, diff --git a/src/meta-srv/src/procedure.rs b/src/meta-srv/src/procedure.rs index da1a1b00e7a..231671eb683 100644 --- a/src/meta-srv/src/procedure.rs +++ b/src/meta-srv/src/procedure.rs @@ -18,6 +18,7 @@ use common_meta::leadership_notifier::LeadershipChangeListener; use common_procedure::ProcedureManagerRef; use snafu::ResultExt; +pub mod drain_datanode; pub mod region_migration; pub mod repartition; #[cfg(any(test, feature = "testing"))] diff --git a/src/meta-srv/src/procedure/drain_datanode.rs b/src/meta-srv/src/procedure/drain_datanode.rs new file mode 100644 index 00000000000..41bdb85c217 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode.rs @@ -0,0 +1,652 @@ +// 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. + +pub(crate) mod drain_datanode_end; +pub(crate) mod drain_datanode_start; +pub(crate) mod manager; +pub(crate) mod utils; + +use std::any::Any; +use std::fmt::Debug; + +use common_meta::DatanodeId; +use common_meta::key::TableMetadataManagerRef; +use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; +use common_procedure::{ + Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, + Result as ProcedureResult, Status, +}; +use common_telemetry::{error, info}; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; +use store_api::storage::RegionId; + +use crate::error::Result; +use crate::node_scheduling::NodeSchedulingManagerRef; +use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart; +use crate::procedure::drain_datanode::manager::{ + DrainDatanodeProcedureGuard, DrainDatanodeProcedureTracker, +}; + +/// The regions placed on a datanode, split by role. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct DrainDatanodePlan { + /// The leader regions still to migrate, sorted and deduplicated. + pub leader_regions: Vec, + /// The follower regions of the datanode, sorted and deduplicated. + pub follower_regions: Vec, +} + +/// The persistent data of [`DrainDatanodeProcedure`]. +/// +/// Only what cannot be derived from the cluster metadata is persisted: the plan is rebuilt +/// from the authoritative table routes after a restart. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PersistentContext { + /// The datanode to drain. + pub node_id: DatanodeId, + /// The pinned destination of leader regions. + pub target_node_id: Option, +} + +/// The in-memory data of [`DrainDatanodeProcedure`]. +/// +/// It is not persisted: a restarted procedure rebuilds the plan by rescanning the datanode. +#[derive(Debug, Clone, Default)] +pub struct VolatileContext { + /// The drain plan, `None` until the first state builds it. + pub plan: Option, +} + +/// The task of [`DrainDatanodeProcedure`]. +#[derive(Debug, Clone)] +pub struct DrainDatanodeProcedureTask { + /// The datanode to drain. + pub node_id: DatanodeId, + /// The pinned destination of leader regions. + pub target_node_id: Option, +} + +/// The context of [`DrainDatanodeProcedure`]. +#[derive(Clone)] +pub struct Context { + /// The persistent data of the procedure. + pub persistent_ctx: PersistentContext, + /// The in-memory data of the procedure. + pub volatile_ctx: VolatileContext, + /// The table metadata manager. + pub table_metadata_manager: TableMetadataManagerRef, + /// The node scheduling manager, used to fence the drained datanode. + pub node_scheduling_manager: NodeSchedulingManagerRef, +} + +impl Context { + /// Builds the context of a run. + /// + /// The per-run state starts empty: the plan is rebuilt from the table routes when the + /// procedure needs it. + pub(crate) fn new( + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + persistent_ctx: PersistentContext, + ) -> Self { + Self { + persistent_ctx, + volatile_ctx: VolatileContext::default(), + table_metadata_manager, + node_scheduling_manager, + } + } +} + +impl Context { + /// Builds the drain plan when it is not available yet. + /// + /// The plan is derived from the authoritative table routes, so a restarted procedure + /// rebuilds it instead of reading a persisted copy. + pub(crate) async fn ensure_plan(&mut self) -> Result<()> { + if self.volatile_ctx.plan.is_some() { + return Ok(()); + } + + let node_id = self.persistent_ctx.node_id; + let plan = utils::build_datanode_plan(&self.table_metadata_manager, node_id).await?; + self.volatile_ctx.plan = Some(plan); + + Ok(()) + } +} + +impl VolatileContext { + /// Returns the leader regions still to migrate. + pub(crate) fn leader_regions(&self) -> &[RegionId] { + self.plan + .as_ref() + .map_or(&[], |plan| plan.leader_regions.as_slice()) + } + + /// Returns the follower regions collected by the plan. + pub(crate) fn follower_regions(&self) -> &[RegionId] { + self.plan + .as_ref() + .map_or(&[], |plan| plan.follower_regions.as_slice()) + } +} + +/// The persistent data of [`DrainDatanodeProcedure`] as read by `from_json`. +#[derive(Debug, Deserialize)] +pub struct DrainDatanodeDataOwned { + persistent_ctx: PersistentContext, + state: Box, +} + +/// The persistent data of [`DrainDatanodeProcedure`] as written by `dump`. +#[derive(Debug, Serialize)] +pub struct DrainDatanodeData<'a> { + persistent_ctx: &'a PersistentContext, + state: &'a dyn State, +} + +/// A state of [`DrainDatanodeProcedure`]. +#[async_trait::async_trait] +#[typetag::serde(tag = "drain_datanode_state")] +pub(crate) trait State: Sync + Send + Debug { + /// Returns the short name of this state. + fn name(&self) -> &'static str { + let type_name = std::any::type_name::(); + type_name.split("::").last().unwrap_or(type_name) + } + + /// Yields the next [`State`] and [`Status`]. + async fn next( + &mut self, + ctx: &mut Context, + procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)>; + + /// Returns as [`Any`]. + fn as_any(&self) -> &dyn Any; +} + +/// The procedure draining a datanode. +/// +/// It fences the datanode and collects the leader and follower regions placed on it. +/// Migrating the leader regions off the datanode lands in a follow-up. +pub(crate) struct DrainDatanodeProcedure { + state: Box, + context: Context, + /// Keeps the drained datanode in the running set of the manager until the procedure ends, + /// so a second drain of the datanode is rejected at submission instead of being queued. + _guards: Vec, +} + +impl DrainDatanodeProcedure { + const TYPE_NAME: &'static str = "metasrv-procedure::DrainDatanode"; + + /// Returns a new [`DrainDatanodeProcedure`] in its initial state. + pub fn new(context: Context, guards: Vec) -> Self { + Self { + state: Box::new(DrainDatanodeStart), + context, + _guards: guards, + } + } + + /// Returns a new [`DrainDatanodeProcedure`] from the given serialized data. + /// + /// The procedure takes its place back in `tracker`, so a reloaded drain still rejects a + /// second drain of the same datanode. + pub fn from_json( + json: &str, + tracker: &DrainDatanodeProcedureTracker, + ctx_factory: F, + ) -> ProcedureResult + where + F: FnOnce(PersistentContext) -> Context, + { + let DrainDatanodeDataOwned { + persistent_ctx, + state, + } = serde_json::from_str(json).context(FromJsonSnafu)?; + let guards = tracker + .insert_running_procedure(&DrainDatanodeProcedureTask { + node_id: persistent_ctx.node_id, + target_node_id: persistent_ctx.target_node_id, + }) + .into_iter() + .collect(); + + Ok(Self { + state, + context: ctx_factory(persistent_ctx), + _guards: guards, + }) + } +} + +#[async_trait::async_trait] +impl Procedure for DrainDatanodeProcedure { + fn type_name(&self) -> &str { + Self::TYPE_NAME + } + + #[tracing::instrument(skip_all, fields( + state = %self.state.name(), + node_id = %self.context.persistent_ctx.node_id + ))] + async fn execute(&mut self, ctx: &ProcedureContext) -> ProcedureResult { + let state_name = self.state.name(); + info!( + "Drain datanode procedure executing state: {}, node_id: {}", + state_name, self.context.persistent_ctx.node_id + ); + match self.state.next(&mut self.context, ctx).await { + Ok((next, status)) => { + self.state = next; + Ok(status) + } + Err(e) => { + if e.is_retryable() { + Err(ProcedureError::retry_later(e)) + } else { + error!( + e; + "Drain datanode procedure failed, node_id: {}", + self.context.persistent_ctx.node_id + ); + Err(ProcedureError::external(e)) + } + } + } + } + + fn dump(&self) -> ProcedureResult { + let data = DrainDatanodeData { + persistent_ctx: &self.context.persistent_ctx, + state: self.state.as_ref(), + }; + serde_json::to_string(&data).context(ToJsonSnafu) + } + + fn lock_key(&self) -> LockKey { + LockKey::single_exclusive(format!( + "__drain_datanode/{}", + self.context.persistent_ctx.node_id + )) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use common_meta::key::TableMetadataManagerRef; + use common_meta::key::table_route::TableRouteValue; + use common_meta::key::test_utils::new_test_table_info; + use common_meta::peer::Peer; + use common_meta::rpc::router::{Region, RegionRoute}; + + use super::*; + use crate::error; + use crate::key::NodeSchedulingRole; + use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; + use crate::procedure::drain_datanode::manager::{ + DrainDatanodeManager, DrainDatanodeProcedureTracker, + }; + use crate::procedure::region_migration::test_util::TestingEnv; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; + + fn new_persistent_context(target_node_id: Option) -> PersistentContext { + PersistentContext { + node_id: 1, + target_node_id, + } + } + + fn leader_region(region_id: RegionId, leader_id: DatanodeId) -> RegionRoute { + RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(leader_id)), + follower_peers: vec![], + ..Default::default() + } + } + + fn follower_region( + region_id: RegionId, + leader_id: DatanodeId, + follower_id: DatanodeId, + ) -> RegionRoute { + RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(leader_id)), + follower_peers: vec![Peer::empty(follower_id)], + ..Default::default() + } + } + + struct TestEnv { + env: TestingEnv, + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + } + + impl TestEnv { + async fn new() -> Self { + let env = TestingEnv::new(); + // Shares one backend between the scheduling manager and the peer discovery, so + // the selector observes the scheduling marks written by the procedure. + let meta_peer_client = create_meta_peer_client(); + put_datanodes( + &meta_peer_client, + vec![ + Peer { + id: 1, + addr: "dn1:3001".to_string(), + }, + Peer { + id: 2, + addr: "dn2:3001".to_string(), + }, + ], + ) + .await; + let (node_scheduling_manager, _discovery) = + new_node_scheduling_discovery(&meta_peer_client); + let table_metadata_manager = env.table_metadata_manager().clone(); + + Self { + env, + table_metadata_manager, + node_scheduling_manager, + } + } + + fn context(&self, persistent_ctx: PersistentContext) -> Context { + Context::new( + self.table_metadata_manager.clone(), + self.node_scheduling_manager.clone(), + persistent_ctx, + ) + } + + fn table_metadata_manager(&self) -> &TableMetadataManagerRef { + &self.table_metadata_manager + } + + async fn create_table(&self, region_routes: Vec) { + create_table(&self.table_metadata_manager, region_routes).await; + } + } + + async fn create_table( + table_metadata_manager: &TableMetadataManagerRef, + region_routes: Vec, + ) { + table_metadata_manager + .create_table_metadata( + new_test_table_info(1024), + TableRouteValue::physical(region_routes), + HashMap::default(), + ) + .await + .unwrap(); + } + + async fn run_until_done( + state: &mut Box, + ctx: &mut Context, + procedure_ctx: &ProcedureContext, + ) -> Result { + for _ in 0..10 { + let (next, status) = state.next(ctx, procedure_ctx).await?; + *state = next; + if !matches!(status, Status::Executing { .. }) { + return Ok(status); + } + } + + panic!("the drain procedure did not finish in 10 steps"); + } + + #[tokio::test] + async fn test_build_datanode_plan() { + let env = TestEnv::new().await; + let first = RegionId::new(1024, 1); + let second = RegionId::new(1024, 2); + env.create_table(vec![leader_region(first, 1), follower_region(second, 2, 1)]) + .await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert_eq!(vec![first], plan.leader_regions); + assert_eq!(vec![second], plan.follower_regions); + } + + #[tokio::test] + async fn test_build_datanode_plan_of_empty_datanode() { + let env = TestEnv::new().await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert!(plan.leader_regions.is_empty()); + assert!(plan.follower_regions.is_empty()); + } + + #[tokio::test] + async fn test_build_datanode_plan_reads_the_table_route() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + // The datanode is a follower of the region, so it is not a leader even though the + // datanode table index of the datanode contains the table. + env.create_table(vec![follower_region(region_id, 2, 1)]) + .await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert!(plan.leader_regions.is_empty()); + assert_eq!(vec![region_id], plan.follower_regions); + } + + #[tokio::test] + async fn test_drain_datanode_start_transitions_to_end() { + let env = TestEnv::new().await; + env.create_table(vec![leader_region(RegionId::new(1024, 1), 1)]) + .await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeStart + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + // Registering the procedure is the only persisted step of the run. + assert!(matches!(status, Status::Executing { persist: true, .. })); + assert_eq!( + vec![RegionId::new(1024, 1)], + ctx.volatile_ctx.leader_regions() + ); + assert!( + !ctx.node_scheduling_manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_drain_datanode_end_transitions_to_done() { + let env = TestEnv::new().await; + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeEnd + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Done { output: None })); + } + + #[tokio::test] + async fn test_drain_datanode_fences_and_collects_the_plan() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + let follower_region_id = RegionId::new(1024, 2); + env.create_table(vec![ + follower_region(follower_region_id, 2, 1), + leader_region(region_id, 1), + ]) + .await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + let mut state: Box = Box::new(DrainDatanodeStart); + + let status = run_until_done(&mut state, &mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(matches!(status, Status::Done { .. })); + assert_eq!(vec![region_id], ctx.volatile_ctx.leader_regions()); + assert_eq!( + vec![follower_region_id], + ctx.volatile_ctx.follower_regions() + ); + assert!( + !ctx.node_scheduling_manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_drain_datanode_dump_roundtrip() { + let env = TestEnv::new().await; + let persistent_ctx = new_persistent_context(Some(2)); + + let procedure = DrainDatanodeProcedure::new(env.context(persistent_ctx), vec![]); + let json = procedure.dump().unwrap(); + let tracker = DrainDatanodeProcedureTracker::default(); + let reloaded = DrainDatanodeProcedure::from_json(&json, &tracker, |persistent_ctx| { + env.context(persistent_ctx) + }) + .unwrap(); + + assert_eq!("DrainDatanodeStart", reloaded.state.name()); + let ctx = &reloaded.context.persistent_ctx; + assert_eq!(1, ctx.node_id); + assert_eq!(Some(2), ctx.target_node_id); + // The plan is volatile, so it is not dumped. + assert!(!json.contains("leader_regions")); + } + + #[tokio::test] + async fn test_drain_datanode_end_finishes() { + let env = TestEnv::new().await; + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeEnd + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Done { .. })); + } + + #[tokio::test] + async fn test_submit_procedure_rejects_same_target() { + let env = TestEnv::new().await; + let manager = DrainDatanodeManager::new( + env.env.procedure_manager().clone(), + env.table_metadata_manager.clone(), + env.node_scheduling_manager.clone(), + ); + + let error = manager + .submit_procedure( + common_procedure::ProcedureContext::default(), + DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(1), + }, + ) + .await + .unwrap_err(); + + assert!(matches!(error, error::Error::InvalidArguments { .. })); + } + + #[tokio::test] + async fn test_submit_procedure_rejects_a_running_drain() { + let env = TestEnv::new().await; + let manager = DrainDatanodeManager::new( + env.env.procedure_manager().clone(), + env.table_metadata_manager.clone(), + env.node_scheduling_manager.clone(), + ); + let task = DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + }; + + // A drain of datanode 1 is running, for example one recovered after a restart. + let _guard = manager.tracker().insert_running_procedure(&task).unwrap(); + + let error = manager + .submit_procedure(common_procedure::ProcedureContext::default(), task) + .await + .unwrap_err(); + + assert!(matches!(error, error::Error::DrainDatanodeRunning { .. })); + assert!(!error.is_retryable()); + } + + #[tokio::test] + async fn test_drain_datanode_tracker_takes_the_place_back_on_reload() { + let env = TestEnv::new().await; + let procedure = + DrainDatanodeProcedure::new(env.context(new_persistent_context(Some(2))), vec![]); + let json = procedure.dump().unwrap(); + let tracker = DrainDatanodeProcedureTracker::default(); + + // Reloading takes the place back, so a second drain of the datanode is rejected. + let reloaded = DrainDatanodeProcedure::from_json(&json, &tracker, |persistent_ctx| { + env.context(persistent_ctx) + }) + .unwrap(); + let task = DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + }; + assert!(tracker.insert_running_procedure(&task).is_none()); + + // Dropping the reloaded procedure releases the place. + drop(reloaded); + assert!(tracker.insert_running_procedure(&task).is_some()); + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs new file mode 100644 index 00000000000..1ff5e3041e1 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs @@ -0,0 +1,41 @@ +// 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::any::Any; + +use common_procedure::{Context as ProcedureContext, Status}; +use serde::{Deserialize, Serialize}; + +use crate::error::Result; +use crate::procedure::drain_datanode::{Context, State}; + +/// Finishes the procedure. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct DrainDatanodeEnd; + +#[async_trait::async_trait] +#[typetag::serde] +impl State for DrainDatanodeEnd { + async fn next( + &mut self, + _ctx: &mut Context, + _procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)> { + Ok((Box::new(DrainDatanodeEnd), Status::done())) + } + + fn as_any(&self) -> &dyn Any { + self + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs new file mode 100644 index 00000000000..0e657ebbd00 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs @@ -0,0 +1,61 @@ +// 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::any::Any; + +use common_procedure::{Context as ProcedureContext, Status}; +use common_telemetry::info; +use serde::{Deserialize, Serialize}; + +use crate::error::Result; +use crate::key::NodeSchedulingRole; +use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; +use crate::procedure::drain_datanode::{Context, State}; + +/// Fences the datanode and builds the drain plan. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct DrainDatanodeStart; + +#[async_trait::async_trait] +#[typetag::serde] +impl State for DrainDatanodeStart { + async fn next( + &mut self, + ctx: &mut Context, + _procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)> { + let node_id = ctx.persistent_ctx.node_id; + let marked = ctx + .node_scheduling_manager + .mark_unschedulable(NodeSchedulingRole::Datanode, node_id) + .await?; + if marked { + info!("Marked datanode {} as unschedulable", node_id); + } + + ctx.ensure_plan().await?; + info!( + "Planned drain of datanode {}, leader regions: {}, follower regions: {}", + node_id, + ctx.volatile_ctx.leader_regions().len(), + ctx.volatile_ctx.follower_regions().len() + ); + + Ok((Box::new(DrainDatanodeEnd), Status::executing(true))) + } + + fn as_any(&self) -> &dyn Any { + self + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/manager.rs b/src/meta-srv/src/procedure/drain_datanode/manager.rs new file mode 100644 index 00000000000..1211356bea0 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/manager.rs @@ -0,0 +1,187 @@ +// 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::HashMap; +use std::collections::hash_map::Entry; +use std::sync::{Arc, RwLock}; + +use common_meta::DatanodeId; +use common_meta::key::TableMetadataManagerRef; +use common_procedure::{ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId}; +use common_telemetry::info; +use snafu::ResultExt; + +use crate::error::{self, Result}; +use crate::node_scheduling::NodeSchedulingManagerRef; +use crate::procedure::drain_datanode::{ + Context as DrainDatanodeContext, DrainDatanodeProcedure, DrainDatanodeProcedureTask, + PersistentContext, +}; + +pub type DrainDatanodeManagerRef = Arc; + +/// Tracks the running drains, indexed by the drained datanode. +/// +/// Mirrors `RegionMigrationProcedureTracker`: a drain of a datanode that is already being +/// drained is rejected at submission instead of waiting for the procedure lock. +#[derive(Default, Clone)] +pub struct DrainDatanodeProcedureTracker { + running_procedures: Arc>>, +} + +impl DrainDatanodeProcedureTracker { + /// Returns the guard of a drain of `task.node_id`, or `None` if that datanode is already + /// being drained. + pub(crate) fn insert_running_procedure( + &self, + task: &DrainDatanodeProcedureTask, + ) -> Option { + let mut procedures = self.running_procedures.write().unwrap(); + match procedures.entry(task.node_id) { + Entry::Occupied(_) => None, + Entry::Vacant(v) => { + v.insert(task.clone()); + Some(DrainDatanodeProcedureGuard { + node_id: task.node_id, + running_procedures: self.running_procedures.clone(), + }) + } + } + } +} + +/// Removes the drained datanode from [`DrainDatanodeProcedureTracker`] when dropped. +pub(crate) struct DrainDatanodeProcedureGuard { + node_id: DatanodeId, + running_procedures: Arc>>, +} + +impl Drop for DrainDatanodeProcedureGuard { + fn drop(&mut self) { + self.running_procedures + .write() + .unwrap() + .remove(&self.node_id); + } +} + +/// Manages the [`DrainDatanodeProcedure`]s. +/// +/// It registers the procedure loader and submits drain tasks. A datanode is drained by at +/// most one procedure at a time: [`DrainDatanodeManager::submit_procedure`] rejects a second +/// drain through the [`DrainDatanodeProcedureTracker`]. +pub struct DrainDatanodeManager { + procedure_manager: ProcedureManagerRef, + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + tracker: DrainDatanodeProcedureTracker, +} + +impl DrainDatanodeManager { + /// Returns a new [`DrainDatanodeManager`]. + pub fn new( + procedure_manager: ProcedureManagerRef, + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + ) -> Self { + Self { + procedure_manager, + table_metadata_manager, + node_scheduling_manager, + tracker: DrainDatanodeProcedureTracker::default(), + } + } + + /// Returns the tracker of the running drains. + #[cfg(test)] + pub(crate) fn tracker(&self) -> DrainDatanodeProcedureTracker { + self.tracker.clone() + } + + /// Returns the context of a drain run. + fn context(&self, persistent_ctx: PersistentContext) -> DrainDatanodeContext { + DrainDatanodeContext::new( + self.table_metadata_manager.clone(), + self.node_scheduling_manager.clone(), + persistent_ctx, + ) + } + + /// Submits a [`DrainDatanodeProcedure`] and returns its id. + pub async fn submit_procedure( + &self, + procedure_context: ProcedureContext, + task: DrainDatanodeProcedureTask, + ) -> Result { + if task.target_node_id == Some(task.node_id) { + return error::InvalidArgumentsSnafu { + err_msg: "the target datanode must differ from the drained datanode", + } + .fail(); + } + + let Some(guard) = self.tracker.insert_running_procedure(&task) else { + return error::DrainDatanodeRunningSnafu { + node_id: task.node_id, + } + .fail(); + }; + + let persistent_ctx = PersistentContext { + node_id: task.node_id, + target_node_id: task.target_node_id, + }; + let procedure = DrainDatanodeProcedure::new(self.context(persistent_ctx), vec![guard]); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); + let procedure_id = procedure_with_id.id; + info!( + "Starting drain datanode procedure {procedure_id} for datanode {}", + task.node_id + ); + self.procedure_manager + .submit(procedure_with_id) + .await + .context(error::SubmitProcedureSnafu)?; + + Ok(procedure_id) + } + + /// Registers the loader of [`DrainDatanodeProcedure`]. + pub(crate) fn try_start(&self) -> Result<()> { + let table_metadata_manager = self.table_metadata_manager.clone(); + let node_scheduling_manager = self.node_scheduling_manager.clone(); + let tracker = self.tracker.clone(); + self.procedure_manager + .register_loader( + DrainDatanodeProcedure::TYPE_NAME, + Box::new(move |json| { + let table_metadata_manager = table_metadata_manager.clone(); + let node_scheduling_manager = node_scheduling_manager.clone(); + let factory = move |persistent_ctx| { + DrainDatanodeContext::new( + table_metadata_manager, + node_scheduling_manager, + persistent_ctx, + ) + }; + DrainDatanodeProcedure::from_json(json, &tracker, factory) + .map(|procedure| Box::new(procedure) as _) + }), + ) + .context(error::RegisterProcedureLoaderSnafu { + type_name: DrainDatanodeProcedure::TYPE_NAME, + }) + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/utils.rs b/src/meta-srv/src/procedure/drain_datanode/utils.rs new file mode 100644 index 00000000000..a43ba815e18 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/utils.rs @@ -0,0 +1,83 @@ +// 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::BTreeSet; + +use common_meta::DatanodeId; +use common_meta::key::TableMetadataManagerRef; +use common_meta::key::table_route::TableRouteValue; +use common_meta::rpc::router::region_distribution; +use futures::TryStreamExt; +use snafu::ResultExt; +use store_api::storage::RegionId; + +use crate::error::{self, Result}; +use crate::procedure::drain_datanode::DrainDatanodePlan; + +/// Builds the drain plan of `node_id`. +/// +/// Walks the datanode table index and reads the authoritative table routes, so the plan +/// reflects the routes rather than the index. Each table route is split by role with +/// [`region_distribution`], which keys the regions by datanode id. +pub(crate) async fn build_datanode_plan( + table_metadata_manager: &TableMetadataManagerRef, + node_id: DatanodeId, +) -> Result { + let mut leader_regions = BTreeSet::new(); + let mut follower_regions = BTreeSet::new(); + + let mut tables = table_metadata_manager + .datanode_table_manager() + .tables(node_id); + while let Some(value) = tables + .try_next() + .await + .context(error::TableMetadataManagerSnafu)? + { + let Some(table_route) = table_metadata_manager + .table_route_manager() + .table_route_storage() + .get(value.table_id) + .await + .context(error::TableMetadataManagerSnafu)? + else { + continue; + }; + let TableRouteValue::Physical(table_route) = table_route else { + continue; + }; + + let mut distribution = region_distribution(&table_route.region_routes); + let Some(roles) = distribution.remove(&node_id) else { + continue; + }; + leader_regions.extend( + roles + .leader_regions + .into_iter() + .map(|number| RegionId::new(value.table_id, number)), + ); + follower_regions.extend( + roles + .follower_regions + .into_iter() + .map(|number| RegionId::new(value.table_id, number)), + ); + } + + Ok(DrainDatanodePlan { + leader_regions: leader_regions.into_iter().collect(), + follower_regions: follower_regions.into_iter().collect(), + }) +}