From cecb2d93a26daea3b67712cbdda90d5520aade68 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 29 Jul 2026 20:04:16 +0800 Subject: [PATCH] fix(repartition): enforce GC across lifecycle (#8678) * fix(repartition): enforce GC across lifecycle Signed-off-by: Dennis Zhuang * refactor(repartition): narrow GC recovery enforcement Signed-off-by: Dennis Zhuang * fix(metasrv): avoid repeated GC preflight during bootstrap Signed-off-by: Dennis Zhuang * docs(metasrv): clarify legacy GC reconciliation Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/meta/src/ddl_manager.rs | 21 +- src/common/meta/src/error.rs | 11 +- src/common/meta/src/key.rs | 1 + src/common/procedure/src/local.rs | 9 + src/common/procedure/src/procedure.rs | 5 + src/meta-srv/src/bootstrap.rs | 31 +- src/meta-srv/src/error.rs | 32 +- src/meta-srv/src/metasrv.rs | 13 + src/meta-srv/src/metasrv/builder.rs | 6 + src/meta-srv/src/procedure/repartition.rs | 63 +++- .../procedure/repartition/gc_requirement.rs | 310 ++++++++++++++++++ src/standalone/src/procedure.rs | 5 + 12 files changed, 494 insertions(+), 13 deletions(-) create mode 100644 src/meta-srv/src/procedure/repartition/gc_requirement.rs diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 22fd67581e..36d6b0e08b 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -48,10 +48,11 @@ use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::undrop_table::UndropTableProcedure; use crate::ddl::{DdlContext, utils}; use crate::error::{ - self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, ProcedureOutputSnafu, - RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result, - SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu, - UnexpectedLogicalRouteTableSnafu, WaitProcedureSnafu, + self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, + PersistRepartitionGcRequirementSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, + RegisterRepartitionProcedureLoaderSnafu, Result, SubmitProcedureSnafu, TableInfoNotFoundSnafu, + TableNotFoundSnafu, TableRouteNotFoundSnafu, UnexpectedLogicalRouteTableSnafu, + WaitProcedureSnafu, }; use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; @@ -170,6 +171,7 @@ pub enum RepartitionSource { }, } +#[async_trait::async_trait] pub trait RepartitionProcedureFactory: Send + Sync { fn create( &self, @@ -186,6 +188,9 @@ pub trait RepartitionProcedureFactory: Send + Sync { ddl_ctx: &DdlContext, procedure_manager: &ProcedureManagerRef, ) -> std::result::Result<(), BoxedError>; + + /// Persists the cluster-level GC requirement before submitting a repartition. + async fn ensure_gc_requirement(&self) -> std::result::Result<(), BoxedError>; } /// The options for DDL tasks. @@ -356,6 +361,10 @@ impl DdlManager { Some(timeout), ) .context(CreateRepartitionProcedureSnafu)?; + self.repartition_procedure_factory + .ensure_gc_requirement() + .await + .context(PersistRepartitionGcRequirementSnafu)?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); if wait { self.execute_procedure_and_wait(procedure_with_id).await @@ -1328,6 +1337,10 @@ mod tests { ) -> std::result::Result<(), BoxedError> { Ok(()) } + + async fn ensure_gc_requirement(&self) -> std::result::Result<(), BoxedError> { + Ok(()) + } } struct LoaderCommitter; diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 447caf911a..aa27227c7c 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -127,6 +127,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to persist the repartition GC requirement"))] + PersistRepartitionGcRequirement { + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to submit procedure"))] SubmitProcedure { #[snafu(implicit)] @@ -1264,6 +1271,7 @@ impl ErrorExt for Error { ProcedureStateReceiver { source, .. } => source.status_code(), RegisterRepartitionProcedureLoader { source, .. } => source.status_code(), CreateRepartitionProcedure { source, .. } => source.status_code(), + PersistRepartitionGcRequirement { source, .. } => source.status_code(), ParseProcedureId { .. } | InvalidNumTopics { .. } @@ -1336,7 +1344,8 @@ impl ErrorExt for Error { | OperateDatanode { source, .. } | AbortProcedure { source, .. } | RegisterRepartitionProcedureLoader { source, .. } - | CreateRepartitionProcedure { source, .. } => source.retry_hint(), + | CreateRepartitionProcedure { source, .. } + | PersistRepartitionGcRequirement { source, .. } => source.retry_hint(), Table { source, .. } => source.retry_hint(), ConvertAlterTableRequest { source, .. } => source.retry_hint(), ConvertColumnDef { source, .. } => source.retry_hint(), diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index c9005e20d8..644ba47223 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -174,6 +174,7 @@ use crate::wal_provider::RegionWalOptions; pub const TOPIC_NAME_PATTERN: &str = r"[a-zA-Z0-9_:-][a-zA-Z0-9_:\-\.@#]*"; pub const LEGACY_MAINTENANCE_KEY: &str = "__maintenance"; pub const MAINTENANCE_KEY: &str = "__switches/maintenance"; +pub const REPARTITION_GC_REQUIRED_KEY: &str = "__requirements/gc/repartition"; pub const PAUSE_PROCEDURE_KEY: &str = "__switches/pause_procedure"; pub const RECOVERY_MODE_KEY: &str = "__switches/recovery"; diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index f70065e01e..fd89cd9c7b 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -973,6 +973,15 @@ impl ProcedureManager for LocalManager { async fn list_procedures(&self) -> Result> { Ok(self.manager_ctx.list_procedure()) } + + async fn has_unfinished_procedure(&self, type_names: &[&str]) -> Result { + let messages = self.procedure_store.load_messages().await?; + Ok(messages + .messages + .values() + .chain(messages.rollback_messages.values()) + .any(|message| type_names.contains(&message.type_name.as_str()))) + } } struct RemoveOutdatedMetaFunction { diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index fac945453b..615fb834e2 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -701,6 +701,11 @@ pub trait ProcedureManager: Send + Sync + 'static { /// Returns the details of the procedure. async fn list_procedures(&self) -> Result>; + + /// Returns whether persisted unfinished procedures contain any requested type. + /// + /// This inspects durable state without submitting procedures for execution. + async fn has_unfinished_procedure(&self, type_names: &[&str]) -> Result; } /// Ref-counted pointer to the [ProcedureManager]. diff --git a/src/meta-srv/src/bootstrap.rs b/src/meta-srv/src/bootstrap.rs index f94c0bfe06..700cc5d8ce 100644 --- a/src/meta-srv/src/bootstrap.rs +++ b/src/meta-srv/src/bootstrap.rs @@ -105,6 +105,8 @@ impl MetasrvInstance { } pub async fn start(&mut self) -> Result<()> { + self.metasrv.ensure_repartition_gc_enabled().await?; + if let Some(builder) = self.http_server.as_mut().left() && let Some(mut builder) = builder.take() { @@ -129,7 +131,7 @@ impl MetasrvInstance { return Ok(()); }; - self.metasrv.try_start().await?; + self.metasrv.try_start_after_gc_check().await?; let (tx, rx) = mpsc::channel::<()>(1); @@ -472,6 +474,7 @@ mod tests { use super::*; use crate::metasrv::{SelectorFactory, SelectorFactoryContext}; + use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManager; struct RecordingSelectorFactory { called: Arc, @@ -507,6 +510,32 @@ mod tests { assert!(called.load(Ordering::Relaxed)); } + #[tokio::test] + async fn gc_requirement_is_checked_before_http_server_start() { + let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::new()); + RepartitionGcRequirementManager::new(kv_backend.clone()) + .require_gc() + .await + .unwrap(); + + let opts = MetasrvOptions { + enable_telemetry: false, + ..Default::default() + }; + let metasrv = MetasrvBuilder::new() + .options(opts) + .kv_backend(kv_backend) + .build() + .await + .unwrap(); + let mut instance = MetasrvInstance::new(metasrv).await.unwrap(); + + let err = instance.start().await.unwrap_err(); + assert!(matches!(err, error::Error::RepartitionGcRequired { .. })); + assert!(instance.http_server().is_none()); + assert!(matches!(instance.mut_http_server(), Either::Left(Some(_)))); + } + struct TestExtraHttpRouterProvider; impl ExtraHttpRouterProvider for TestExtraHttpRouterProvider { diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index aa689080cd..53872d4817 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -791,6 +791,28 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to manage the repartition GC requirement"))] + RepartitionGcRequirement { + source: common_meta::error::Error, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Failed to inspect persisted repartition procedures"))] + InspectRepartitionProcedures { + source: common_procedure::Error, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Metasrv GC must be enabled because the cluster has durable repartition state" + ))] + RepartitionGcRequired { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "Source partition expression '{}' does not match any existing region", expr @@ -1261,7 +1283,9 @@ impl ErrorExt for Error { Error::ListActiveFrontends { source, .. } | Error::ListActiveDatanodes { source, .. } | Error::ListActiveFlownodes { source, .. } => source.status_code(), - Error::NoAvailableFrontend { .. } => StatusCode::IllegalState, + Error::NoAvailableFrontend { .. } | Error::RepartitionGcRequired { .. } => { + StatusCode::IllegalState + } Error::InitMetadata { source, .. } | Error::InitDdlManager { source, .. } @@ -1270,6 +1294,8 @@ impl ErrorExt for Error { Error::BuildTlsOptions { source, .. } => source.status_code(), Error::Other { source, .. } => source.status_code(), Error::RepartitionCreateSubtasks { source, .. } => source.status_code(), + Error::RepartitionGcRequirement { source, .. } => source.status_code(), + Error::InspectRepartitionProcedures { source, .. } => source.status_code(), Error::RepartitionSubprocedureStateReceiver { source, .. } => source.status_code(), Error::AllocateRegions { source, .. } => source.status_code(), Error::DeallocateRegions { source, .. } => source.status_code(), @@ -1360,7 +1386,8 @@ impl ErrorExt for Error { | Error::DeallocateRegions { source, .. } | Error::BuildCreateRequest { source, .. } | Error::AllocateRegionRoutes { source, .. } - | Error::AllocateWalOptions { source, .. } => source.retry_hint(), + | Error::AllocateWalOptions { source, .. } + | Error::RepartitionGcRequirement { source, .. } => source.retry_hint(), Error::Other { source, .. } | Error::ListCatalogs { source, .. } @@ -1374,6 +1401,7 @@ impl ErrorExt for Error { | Error::StartProcedureManager { source, .. } | Error::StopProcedureManager { source, .. } | Error::RegisterProcedureLoader { source, .. } + | Error::InspectRepartitionProcedures { source, .. } | Error::RepartitionSubprocedureStateReceiver { source, .. } => source.retry_hint(), Error::ShutdownServer { source, .. } | Error::StartHttp { source, .. } => { diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index bbf6c64ca5..65a079bec1 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -75,6 +75,7 @@ use crate::gc::{GcSchedulerOptions, GcTickerRef}; use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatHandlerGroupRef}; use crate::procedure::ProcedureManagerListenerAdapter; use crate::procedure::region_migration::manager::RegionMigrationManagerRef; +use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManagerRef; use crate::procedure::wal_prune::manager::WalPruneTickerRef; use crate::pubsub::{PublisherRef, SubscriptionManagerRef}; use crate::region::flush_trigger::RegionFlushTickerRef; @@ -582,6 +583,7 @@ pub struct Metasrv { wal_provider: WalProviderRef, table_metadata_manager: TableMetadataManagerRef, runtime_switch_manager: RuntimeSwitchManagerRef, + repartition_gc_requirement_manager: RepartitionGcRequirementManagerRef, memory_region_keeper: MemoryRegionKeeperRef, greptimedb_telemetry_task: Arc, region_migration_manager: RegionMigrationManagerRef, @@ -601,7 +603,18 @@ pub struct Metasrv { } impl Metasrv { + pub(crate) async fn ensure_repartition_gc_enabled(&self) -> Result<()> { + self.repartition_gc_requirement_manager + .ensure_gc_enabled(self.options.gc.enable, &self.procedure_manager) + .await + } + pub async fn try_start(&self) -> Result<()> { + self.ensure_repartition_gc_enabled().await?; + self.try_start_after_gc_check().await + } + + pub(crate) async fn try_start_after_gc_check(&self) -> Result<()> { if self .started .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index a528583ba6..25a179c3d4 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -71,6 +71,7 @@ use crate::metasrv::{ use crate::peer::MetasrvPeerAllocator; use crate::procedure::region_migration::DefaultContextFactory; use crate::procedure::region_migration::manager::RegionMigrationManager; +use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManager; use crate::procedure::repartition::{ DefaultRepartitionProcedureFactory, GcDisabledRepartitionProcedureFactory, }; @@ -255,6 +256,8 @@ impl MetasrvBuilder { let pushers = Pushers::default(); let mailbox = build_mailbox(&kv_backend, &pushers); let runtime_switch_manager = Arc::new(RuntimeSwitchManager::new(kv_backend.clone())); + let repartition_gc_requirement_manager = + Arc::new(RepartitionGcRequirementManager::new(kv_backend.clone())); let procedure_manager = build_procedure_manager( &options, &kv_backend, @@ -436,11 +439,13 @@ impl MetasrvBuilder { Arc::new(DefaultRepartitionProcedureFactory::new( mailbox.clone(), options.grpc.server_addr.clone(), + repartition_gc_requirement_manager.clone(), )) } else { Arc::new(GcDisabledRepartitionProcedureFactory::new( mailbox.clone(), options.grpc.server_addr.clone(), + repartition_gc_requirement_manager.clone(), )) }; let ddl_manager = DdlManager::new( @@ -628,6 +633,7 @@ impl MetasrvBuilder { wal_provider, table_metadata_manager, runtime_switch_manager, + repartition_gc_requirement_manager, greptimedb_telemetry_task: get_greptimedb_telemetry_task( Some(metasrv_home), meta_peer_client, diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index bc3757319b..4577d65255 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -16,6 +16,7 @@ pub mod allocate_region; pub mod collect; pub mod deallocate_region; pub mod dispatch; +pub mod gc_requirement; pub mod group; pub mod plan; pub mod repartition_end; @@ -60,6 +61,7 @@ use table::table_name::TableName; use crate::error::{self, Result}; use crate::procedure::repartition::collect::ProcedureMeta; use crate::procedure::repartition::deallocate_region::DeallocateRegion; +use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManagerRef; use crate::procedure::repartition::group::{ Context as RepartitionGroupContext, RepartitionGroupProcedure, region_routes, }; @@ -792,13 +794,19 @@ impl Procedure for RepartitionProcedure { pub struct DefaultRepartitionProcedureFactory { mailbox: MailboxRef, server_addr: String, + gc_requirement_manager: RepartitionGcRequirementManagerRef, } impl DefaultRepartitionProcedureFactory { - pub fn new(mailbox: MailboxRef, server_addr: String) -> Self { + pub fn new( + mailbox: MailboxRef, + server_addr: String, + gc_requirement_manager: RepartitionGcRequirementManagerRef, + ) -> Self { Self { mailbox, server_addr, + gc_requirement_manager, } } } @@ -806,19 +814,28 @@ impl DefaultRepartitionProcedureFactory { /// Rejects new repartition requests when metasrv GC is disabled. /// /// Procedure loaders are still delegated to the enabled factory so procedures -/// persisted before a metasrv restart can be recovered. +/// persisted before a metasrv restart remain recoverable after GC is re-enabled. pub struct GcDisabledRepartitionProcedureFactory { enabled_factory: DefaultRepartitionProcedureFactory, } impl GcDisabledRepartitionProcedureFactory { - pub fn new(mailbox: MailboxRef, server_addr: String) -> Self { + pub fn new( + mailbox: MailboxRef, + server_addr: String, + gc_requirement_manager: RepartitionGcRequirementManagerRef, + ) -> Self { Self { - enabled_factory: DefaultRepartitionProcedureFactory::new(mailbox, server_addr), + enabled_factory: DefaultRepartitionProcedureFactory::new( + mailbox, + server_addr, + gc_requirement_manager, + ), } } } +#[async_trait::async_trait] impl RepartitionProcedureFactory for GcDisabledRepartitionProcedureFactory { fn create( &self, @@ -845,8 +862,18 @@ impl RepartitionProcedureFactory for GcDisabledRepartitionProcedureFactory { self.enabled_factory .register_loaders(ddl_ctx, procedure_manager) } + + async fn ensure_gc_requirement(&self) -> std::result::Result<(), BoxedError> { + Err(BoxedError::new( + error::InvalidArgumentsSnafu { + err_msg: "Repartition requires metasrv GC to be enabled".to_string(), + } + .build(), + )) + } } +#[async_trait::async_trait] impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { fn create( &self, @@ -950,6 +977,13 @@ impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { Ok(()) } + + async fn ensure_gc_requirement(&self) -> std::result::Result<(), BoxedError> { + self.gc_requirement_manager + .require_gc() + .await + .map_err(BoxedError::new) + } } #[cfg(test)] @@ -968,7 +1002,9 @@ mod tests { use common_meta::peer::Peer; use common_meta::region_keeper::MemoryRegionKeeper; use common_meta::rpc::router::{LeaderState, Region, RegionRoute}; + use common_meta::state_store::KvStateStore; use common_meta::test_util::MockDatanodeManager; + use common_procedure::local::{LocalManager, ManagerConfig}; use common_procedure::{Error as ProcedureError, Procedure, ProcedureId, ProcedureState}; use store_api::region_engine::RegionRole; use store_api::storage::RegionId; @@ -981,6 +1017,7 @@ mod tests { use crate::procedure::repartition::collect::Collect; use crate::procedure::repartition::deallocate_region::DeallocateRegion; use crate::procedure::repartition::dispatch::Dispatch; + use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManager; use crate::procedure::repartition::group::update_metadata::UpdateMetadata; use crate::procedure::repartition::plan::{SourceRegionDescriptor, TargetRegionDescriptor}; use crate::procedure::repartition::repartition_end::RepartitionEnd; @@ -1090,13 +1127,14 @@ mod tests { } #[test] - fn test_gc_disabled_factory_rejects_repartition() { + fn test_gc_disabled_factory_rejects_repartition_and_registers_loaders() { let env = TestingEnv::new(); let node_manager = Arc::new(MockDatanodeManager::new(UnexpectedErrorDatanodeHandler)); let ddl_ctx = env.ddl_context(node_manager); let factory = GcDisabledRepartitionProcedureFactory::new( env.mailbox_ctx.mailbox().clone(), env.server_addr.clone(), + Arc::new(RepartitionGcRequirementManager::new(env.kv_backend.clone())), ); let err = factory @@ -1118,6 +1156,21 @@ mod tests { "Invalid arguments: Repartition requires metasrv GC to be enabled", err.to_string() ); + + let state_store = Arc::new(KvStateStore::new(env.kv_backend)); + let procedure_manager = Arc::new(LocalManager::new( + ManagerConfig::default(), + state_store.clone(), + state_store, + None, + None, + )); + let procedure_manager_ref: ProcedureManagerRef = procedure_manager.clone(); + factory + .register_loaders(&ddl_ctx, &procedure_manager_ref) + .unwrap(); + assert!(procedure_manager.contains_loader(RepartitionProcedure::TYPE_NAME)); + assert!(procedure_manager.contains_loader(RepartitionGroupProcedure::TYPE_NAME)); } #[test] diff --git a/src/meta-srv/src/procedure/repartition/gc_requirement.rs b/src/meta-srv/src/procedure/repartition/gc_requirement.rs new file mode 100644 index 0000000000..11e82d294a --- /dev/null +++ b/src/meta-srv/src/procedure/repartition/gc_requirement.rs @@ -0,0 +1,310 @@ +// 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::sync::Arc; + +use common_meta::key::REPARTITION_GC_REQUIRED_KEY; +use common_meta::key::table_repart::TableRepartManager; +use common_meta::kv_backend::KvBackendRef; +use common_meta::rpc::store::PutRequest; +use common_procedure::ProcedureManagerRef; +use snafu::{ResultExt, ensure}; + +use crate::error::{self, Result}; +use crate::procedure::repartition::RepartitionProcedure; +use crate::procedure::repartition::group::RepartitionGroupProcedure; + +const REPARTITION_GC_REQUIREMENT_VALUE: &[u8] = b"repartition"; + +pub type RepartitionGcRequirementManagerRef = Arc; + +/// Persists and enforces the cluster-level GC requirement introduced by repartition. +pub struct RepartitionGcRequirementManager { + kv_backend: KvBackendRef, +} + +impl RepartitionGcRequirementManager { + pub fn new(kv_backend: KvBackendRef) -> Self { + Self { kv_backend } + } + + /// Persists the requirement before a repartition procedure can be submitted. + pub async fn require_gc(&self) -> Result<()> { + self.kv_backend + .put( + PutRequest::new() + .with_key(REPARTITION_GC_REQUIRED_KEY.as_bytes().to_vec()) + .with_value(REPARTITION_GC_REQUIREMENT_VALUE.to_vec()), + ) + .await + .context(error::RepartitionGcRequirementSnafu)?; + Ok(()) + } + + pub async fn is_gc_required(&self) -> Result { + self.kv_backend + .get(REPARTITION_GC_REQUIRED_KEY.as_bytes()) + .await + .map(|value| value.is_some()) + .context(error::RepartitionGcRequirementSnafu) + } + + /// Backfills the requirement from durable state written by older versions. + /// + /// An unfinished procedure covers the crash window before repartition mappings + /// are written. A non-empty mapping covers completed repartitions whose + /// manifests can still contain cross-region file references. + pub async fn reconcile_legacy_state( + &self, + procedure_manager: &ProcedureManagerRef, + ) -> Result { + if self.is_gc_required().await? { + return Ok(true); + } + + let table_reparts = TableRepartManager::new(self.kv_backend.clone()) + .table_reparts() + .await + .context(error::RepartitionGcRequirementSnafu)?; + let has_cross_region_references = table_reparts + .iter() + .any(|(_, value)| !value.src_to_dst.is_empty()); + let has_unfinished_repartition = procedure_manager + .has_unfinished_procedure(&[ + RepartitionProcedure::TYPE_NAME, + RepartitionGroupProcedure::TYPE_NAME, + ]) + .await + .context(error::InspectRepartitionProceduresSnafu)?; + + if has_cross_region_references || has_unfinished_repartition { + self.require_gc().await?; + return Ok(true); + } + + Ok(false) + } + + pub async fn ensure_gc_enabled( + &self, + gc_enabled: bool, + procedure_manager: &ProcedureManagerRef, + ) -> Result<()> { + // Reconciliation is also a legacy migration and must run before GC can + // remove the repartition mappings used to backfill the durable marker. + let required = self.reconcile_legacy_state(procedure_manager).await?; + ensure!(!required || gc_enabled, error::RepartitionGcRequiredSnafu); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::time::Duration; + + use async_trait::async_trait; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::state_store::KvStateStore; + use common_procedure::local::{LocalManager, ManagerConfig}; + use common_procedure::{ + Context, LockKey, Procedure, ProcedureId, ProcedureManager, ProcedureWithId, Status, + }; + use store_api::storage::RegionId; + use tokio::sync::oneshot; + use tokio::time::timeout; + + use super::*; + + fn local_procedure_manager(kv_backend: KvBackendRef) -> Arc { + let state_store = Arc::new(KvStateStore::new(kv_backend)); + Arc::new(LocalManager::new( + ManagerConfig::default(), + state_store.clone(), + state_store, + None, + None, + )) + } + + fn procedure_manager(kv_backend: KvBackendRef) -> ProcedureManagerRef { + local_procedure_manager(kv_backend) + } + + #[derive(Debug)] + struct LegacyRepartitionProcedure { + persisted: bool, + block_after_persist: bool, + persisted_tx: Option>, + } + + #[async_trait] + impl Procedure for LegacyRepartitionProcedure { + fn type_name(&self) -> &str { + RepartitionProcedure::TYPE_NAME + } + + async fn execute(&mut self, _ctx: &Context) -> common_procedure::error::Result { + if !self.persisted { + self.persisted = true; + return Ok(Status::executing(true)); + } + + if self.block_after_persist { + if let Some(tx) = self.persisted_tx.take() { + let _ = tx.send(()); + } + return std::future::pending().await; + } + + Ok(Status::done()) + } + + fn dump(&self) -> common_procedure::error::Result { + Ok("{}".to_string()) + } + + fn lock_key(&self) -> LockKey { + LockKey::default() + } + } + + #[tokio::test] + async fn test_completed_repartition_requirement_rejects_gc_disabled_restart() { + let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::new()); + let manager = RepartitionGcRequirementManager::new(kv_backend.clone()); + let procedure_manager = procedure_manager(kv_backend.clone()); + + manager.require_gc().await.unwrap(); + let marker = kv_backend + .get(REPARTITION_GC_REQUIRED_KEY.as_bytes()) + .await + .unwrap() + .unwrap(); + assert_eq!(REPARTITION_GC_REQUIREMENT_VALUE, marker.value.as_slice()); + + let err = manager + .ensure_gc_enabled(false, &procedure_manager) + .await + .unwrap_err(); + assert!(matches!(err, error::Error::RepartitionGcRequired { .. })); + assert!(manager.is_gc_required().await.unwrap()); + manager + .ensure_gc_enabled(true, &procedure_manager) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_reconcile_legacy_repartition_mapping() { + let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::new()); + let manager = RepartitionGcRequirementManager::new(kv_backend.clone()); + let procedure_manager = procedure_manager(kv_backend.clone()); + let table_id = 1024; + let source = RegionId::new(table_id, 1); + let destination = RegionId::new(table_id, 2); + TableRepartManager::new(kv_backend) + .update_mappings(table_id, &HashMap::from([(source, vec![destination])])) + .await + .unwrap(); + + assert!( + manager + .reconcile_legacy_state(&procedure_manager) + .await + .unwrap() + ); + assert!(manager.is_gc_required().await.unwrap()); + } + + #[tokio::test] + async fn test_legacy_unfinished_repartition_fails_closed_and_recovers() { + let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::new()); + let first_manager = local_procedure_manager(kv_backend.clone()); + first_manager.start().await.unwrap(); + + let procedure_id = ProcedureId::random(); + let (persisted_tx, persisted_rx) = oneshot::channel(); + first_manager + .submit(ProcedureWithId { + id: procedure_id, + procedure: Box::new(LegacyRepartitionProcedure { + persisted: false, + block_after_persist: true, + persisted_tx: Some(persisted_tx), + }), + }) + .await + .unwrap(); + timeout(Duration::from_secs(10), persisted_rx) + .await + .unwrap() + .unwrap(); + first_manager.stop().await.unwrap(); + + let requirement_manager = RepartitionGcRequirementManager::new(kv_backend.clone()); + let disabled_manager = local_procedure_manager(kv_backend.clone()); + let disabled_manager_ref: ProcedureManagerRef = disabled_manager.clone(); + let err = requirement_manager + .ensure_gc_enabled(false, &disabled_manager_ref) + .await + .unwrap_err(); + assert!(matches!(err, error::Error::RepartitionGcRequired { .. })); + assert!(requirement_manager.is_gc_required().await.unwrap()); + + assert!( + disabled_manager + .has_unfinished_procedure(&[RepartitionProcedure::TYPE_NAME]) + .await + .unwrap() + ); + assert!( + disabled_manager + .procedure_state(procedure_id) + .await + .unwrap() + .is_none() + ); + + let enabled_manager = local_procedure_manager(kv_backend); + enabled_manager + .register_loader( + RepartitionProcedure::TYPE_NAME, + Box::new(|_| { + Ok(Box::new(LegacyRepartitionProcedure { + persisted: true, + block_after_persist: false, + persisted_tx: None, + }) as _) + }), + ) + .unwrap(); + enabled_manager.start().await.unwrap(); + let mut watcher = enabled_manager.procedure_watcher(procedure_id).unwrap(); + timeout(Duration::from_secs(10), async { + while !watcher.borrow().is_done() { + watcher.changed().await.unwrap(); + } + }) + .await + .unwrap(); + assert!( + !enabled_manager + .has_unfinished_procedure(&[RepartitionProcedure::TYPE_NAME]) + .await + .unwrap() + ); + enabled_manager.stop().await.unwrap(); + } +} diff --git a/src/standalone/src/procedure.rs b/src/standalone/src/procedure.rs index a232fba950..ec63b11703 100644 --- a/src/standalone/src/procedure.rs +++ b/src/standalone/src/procedure.rs @@ -63,6 +63,7 @@ pub fn build_procedure_manager( /// work in [`RepartitionProcedureFactory::register_loaders`]. pub struct StandaloneRepartitionProcedureFactory; +#[async_trait::async_trait] impl RepartitionProcedureFactory for StandaloneRepartitionProcedureFactory { fn create( &self, @@ -83,4 +84,8 @@ impl RepartitionProcedureFactory for StandaloneRepartitionProcedureFactory { ) -> std::result::Result<(), BoxedError> { Ok(()) } + + async fn ensure_gc_requirement(&self) -> std::result::Result<(), BoxedError> { + Err(BoxedError::new(NoSupportRepartitionProcedureSnafu.build())) + } }