diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index b8fd2f8892c..efee5f6cca1 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -88,6 +88,9 @@ Tests live next to the code as `*_test.rs` (e.g. `src/mito2/src/engine/flush_tes - Lock ordering: take the manifest lock before updating `version_control`; the reverse deadlocks against concurrent flush/compaction. - All region I/O runs on tokio workers; never `block_on` inside a worker. +- Loaded regions have resident compaction state, including `Idle` and handed-off + DDLs. Publish create/open/catchup results through the worker's `register_region`; + only close/drop unregisters them. See `compaction/scheduler/state.rs`. ## Maintenance contract diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index ff92f8a41c3..94efddc2345 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -27,7 +27,7 @@ use common_time::range::TimestampRange; pub(crate) use planning::CompactionPickFinished; pub use planning::CompactionRequest; pub(crate) use state::CompactionExecution; -use state::{CompactionStatus, CompactionTrigger, PendingCompaction}; +use state::{CompactionStatus, CompactionTrigger, DdlExecution, PendingCompaction}; use store_api::storage::RegionId; use tokio::sync::mpsc::Sender; @@ -51,7 +51,7 @@ use crate::worker::WorkerListener; /// Compaction scheduler tracks and manages compaction tasks. pub(crate) struct CompactionScheduler { scheduler: SchedulerRef, - /// Compacting regions. + /// All loaded regions, including idle regions and regions executing DDLs. region_status: HashMap, /// Request sender of the worker that this scheduler belongs to. request_sender: Sender, @@ -139,11 +139,38 @@ impl CompactionScheduler { } } + /// Registers a loaded region; replacing its VersionControl invalidates the old incarnation. + pub(crate) fn register_region( + &mut self, + version_control: &VersionControlRef, + access_layer: &AccessLayerRef, + ) { + let region_id = version_control.region_id(); + if let Some(status) = self.region_status.get_mut(®ion_id) { + if Arc::ptr_eq(&status.version_control, version_control) { + return; + } + status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build())); + } + self.region_status.insert( + region_id, + CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()), + ); + } + + /// Reports active work rather than the existence of a resident region record. + #[cfg(test)] + fn is_compacting(&self, region_id: RegionId) -> bool { + self.region_status + .get(®ion_id) + .is_some_and(CompactionStatus::is_compacting) + } + /// Accepts an automatic compaction trigger for a region. /// /// # Effects /// - /// Starts planning when no scheduler status exists for the region. + /// Starts planning when the registered region is idle. /// Otherwise, an active status coalesces the trigger into one unrestricted /// follow-up cycle. Returns `true` only when this call dispatches planning. /// @@ -178,7 +205,7 @@ impl CompactionScheduler { /// /// # Effects /// - /// Starts planning when no scheduler status exists for the region. During + /// Starts planning when the registered region is idle. During /// an automatic compaction, queues the request and replaces any older /// pending manual request; during a manual compaction, rejects the new /// request through `waiter`. Returns `true` only when this call dispatches @@ -265,7 +292,7 @@ impl CompactionScheduler { /// # Effects /// /// Notifies waiters, schedules a pending manual or automatic follow-up, or - /// removes the region status. A follow-up is scheduled whenever the cycle + /// returns the region to idle. A follow-up is scheduled whenever the cycle /// latched an automatic trigger, or when `made_progress` is true (the /// execution produced output or reduced the file count): successful compaction /// keeps draining the region until the picker returns no plan. The returned @@ -301,7 +328,7 @@ impl CompactionScheduler { /// /// # Effects /// - /// Removes the matching region status, notifies compaction waiters, and + /// Resets the matching region to idle, notifies compaction waiters, and /// returns DDLs that were waiting for cancellation. /// /// # Constraints @@ -323,7 +350,7 @@ impl CompactionScheduler { /// /// # Effects /// - /// Removes the matching region status and fails all compaction waiters and + /// Resets the matching region to idle and fails all compaction waiters and /// dependent DDL requests with the supplied error. /// /// # Constraints @@ -352,7 +379,7 @@ impl CompactionScheduler { /// /// The owning worker must invoke this as part of serialized region teardown. pub(crate) fn on_region_dropped(&mut self, region_id: RegionId) { - self.remove_region_on_failure( + self.remove_region( region_id, Arc::new(RegionDroppedSnafu { region_id }.build()), ); @@ -368,10 +395,10 @@ impl CompactionScheduler { /// /// The owning worker must invoke this as part of serialized region teardown. pub(crate) fn on_region_closed(&mut self, region_id: RegionId) { - self.remove_region_on_failure(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); + self.remove_region(region_id, Arc::new(RegionClosedSnafu { region_id }.build())); } - /// Removes compaction state because the region was truncated. + /// Resets the cycle after truncate, retaining registration and outstanding DDL replies. /// /// # Effects /// @@ -381,7 +408,7 @@ impl CompactionScheduler { /// /// The owning worker must invoke this after truncate completes. pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) { - self.remove_region_on_failure( + self.fail_compaction( region_id, Arc::new(RegionTruncatedSnafu { region_id }.build()), ); @@ -412,6 +439,9 @@ impl CompactionScheduler { let Some(status) = self.region_status.get_mut(®ion_id) else { return Err((sender, request)); }; + if !status.is_compacting() { + return Err((sender, request)); + } status.request_cancel(); let request = SenderDdlRequest { @@ -454,6 +484,13 @@ impl CompactionScheduler { time_range: Option, ) -> Result { let region_id = version_control.region_id(); + let status = self.region_status.entry(region_id).or_insert_with(|| { + CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()) + }); + if status.has_ddl() { + waiter.send(CompactionCancelledSnafu.fail()); + return Ok(false); + } let current_state = manifest_ctx.current_state(); if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) { info!( @@ -464,18 +501,7 @@ impl CompactionScheduler { return Ok(false); } - if let Some(status) = self.region_status.get_mut(®ion_id) { - // Pending Truncate/EnterStaging requests form a scheduling fence. Any later - // manual request receives CompactionCancelled; automatic triggers are ignored. - if !status.pending_ddl_requests.is_empty() { - waiter.send(CompactionCancelledSnafu.fail()); - info!( - "Region {} has pending DDL requests, ignoring compaction: {:?}", - region_id, compact_options - ); - return Ok(false); - } - + if status.is_compacting() { match trigger { CompactionTrigger::Automatic => status.mark_automatic_trigger(), CompactionTrigger::Manual if status.is_manual_compaction() => { @@ -503,13 +529,7 @@ impl CompactionScheduler { // Publish the picking phase before dispatching background planning. let plan_id = Self::next_plan_id(&mut self.next_plan_id); - let mut status = CompactionStatus::new( - region_id, - version_control.clone(), - access_layer.clone(), - plan_id, - trigger, - ); + status.start_picking_with_trigger(plan_id, trigger); let request = status.new_compaction_request( self.request_sender.clone(), self.engine_config.clone(), @@ -520,7 +540,6 @@ impl CompactionScheduler { max_parallelism, ); status.merge_waiter(waiter); - self.region_status.insert(region_id, status); self.dispatch_compaction_planning(plan_id, request, compact_options, time_range); self.listener.on_compaction_scheduled(region_id); Ok(true) @@ -587,6 +606,10 @@ impl CompactionScheduler { return CompactionTransition::NoAction; } + if let Some(ddls) = self.take_ready_ddls(region_id) { + return CompactionTransition::DdlReady(ddls); + } + if self.handle_pending_compaction_request( region_id, manifest_ctx, @@ -595,8 +618,6 @@ impl CompactionScheduler { return CompactionTransition::NoAction; } - // The region status might be removed by the previous steps. - // So we return empty DDL requests. let Some(status) = self.region_status.get_mut(®ion_id) else { return CompactionTransition::NoAction; }; @@ -607,15 +628,6 @@ impl CompactionScheduler { // A queued DDL was waiting for the current task to terminate; chaining // another compaction ahead of it would delay the DDL by a whole extra // plan/execution cycle, so dispatch the DDLs first. - let pending_ddl_requests = std::mem::take(&mut status.pending_ddl_requests); - if !pending_ddl_requests.is_empty() { - // The DDL supersedes any retained automatic follow-up. - self.region_status.remove(®ion_id); - // If there are pending DDL requests, we should return them to the caller. - // And skip try to schedule next compaction task. - return CompactionTransition::DdlReady(pending_ddl_requests); - } - // Keep draining when the cycle latched an automatic trigger, or when the // execution produced output or reduced the file count. An empty edit stops // here; the next flush trigger resumes. @@ -625,7 +637,9 @@ impl CompactionScheduler { { return CompactionTransition::AutomaticFollowupScheduled; } - self.region_status.remove(®ion_id); + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.become_idle(); + } CompactionTransition::NoAction } @@ -668,13 +682,13 @@ impl CompactionScheduler { /// Notifies the scheduler that the compaction job is cancelled cooperatively. async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec { - self.remove_region_on_cancel(region_id) + self.finish_compaction_on_cancel(region_id) } /// Notifies the scheduler that the compaction job is failed. fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc) { error!(err; "Region {} failed to compact, cancel all pending tasks", region_id); - self.remove_region_on_failure(region_id, err); + self.fail_compaction(region_id, err); } #[cfg(test)] @@ -708,28 +722,104 @@ impl CompactionScheduler { .request_cancel() } - fn remove_region_on_failure(&mut self, region_id: RegionId, err: Arc) { - // Remove this region. - let Some(status) = self.region_status.remove(®ion_id) else { - return; - }; - - // Notifies all pending tasks. - status.on_failure(err); + /// Unregisters a closed or dropped region and fails its queued requests. + fn remove_region(&mut self, region_id: RegionId, err: Arc) { + if let Some(mut status) = self.region_status.remove(®ion_id) { + status.on_failure(err); + } } - fn remove_region_on_cancel(&mut self, region_id: RegionId) -> Vec { - let Some(status) = self.region_status.remove(®ion_id) else { + /// Fails a cycle without discarding the loaded region's registration. + fn fail_compaction(&mut self, region_id: RegionId, err: Arc) { + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.on_failure(err); + } + } + + /// Returns the cycle to idle and records any DDL handoff in the same resident entry. + fn finish_compaction_on_cancel(&mut self, region_id: RegionId) -> Vec { + let Some(status) = self.region_status.get_mut(®ion_id) else { return Vec::new(); }; - status.on_cancel() + let ddls = status.on_cancel(); + self.fence_ddl_replies(region_id, ddls) + } + + /// Hands off pending DDLs after compaction ends and cancels remaining compaction waiters. + fn take_ready_ddls(&mut self, region_id: RegionId) -> Option> { + if self + .region_status + .get(®ion_id)? + .pending_ddl_requests + .is_empty() + { + return None; + } + Some(self.finish_compaction_on_cancel(region_id)) + } + + /// Wraps DDL replies so compaction stays paused until every handed-off request completes. + fn fence_ddl_replies( + &mut self, + region_id: RegionId, + mut ddls: Vec, + ) -> Vec { + if ddls.is_empty() { + return ddls; + } + let generation = Self::next_plan_id(&mut self.next_plan_id); + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.executing_ddl = Some(DdlExecution::new(generation, ddls.len())); + } + for (request_id, ddl) in ddls.iter_mut().enumerate() { + let (tx, rx) = tokio::sync::oneshot::channel(); + let reply = std::mem::replace(&mut ddl.sender, tx.into()); + let sender = self.request_sender.clone(); + common_runtime::spawn_global(async move { + use snafu::ResultExt; + let result = rx + .await + .context(crate::error::RecvSnafu) + .and_then(|result| result); + let _ = sender + .send(WorkerRequestWithTime::new( + crate::request::WorkerRequest::Background { + region_id, + notify: crate::request::BackgroundNotify::CompactionDdlComplete { + generation, + request_id, + }, + }, + )) + .await; + reply.send(result); + }); + } + ddls + } + + /// Releases the DDL fence after its last reply, ignoring replies from an older generation. + pub(crate) fn on_ddl_complete( + &mut self, + region_id: RegionId, + generation: u64, + request_id: usize, + ) { + let Some(status) = self.region_status.get_mut(®ion_id) else { + return; + }; + if let Some(ddl) = &mut status.executing_ddl + && ddl.complete(generation, request_id) + { + status.executing_ddl = None; + } } } impl Drop for CompactionScheduler { fn drop(&mut self) { - for (region_id, status) in self.region_status.drain() { + for (region_id, mut status) in self.region_status.drain() { // We are shutting down so notify all pending tasks. status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build())); } diff --git a/src/mito2/src/compaction/scheduler/planning.rs b/src/mito2/src/compaction/scheduler/planning.rs index 23318ad7669..5553d6435da 100644 --- a/src/mito2/src/compaction/scheduler/planning.rs +++ b/src/mito2/src/compaction/scheduler/planning.rs @@ -298,7 +298,7 @@ impl CompactionScheduler { // execution callback will follow, so return DDLs released by the fence. if !status.accept_plan(finished.plan_id) { return CompactionTransition::from_pending_ddls( - self.remove_region_on_cancel(region_id), + self.finish_compaction_on_cancel(region_id), ); } @@ -354,7 +354,7 @@ impl CompactionScheduler { .await } Err(err) => { - self.remove_region_on_failure(region_id, Arc::new(err)); + self.fail_compaction(region_id, Arc::new(err)); CompactionTransition::NoAction } } @@ -400,6 +400,9 @@ impl CompactionScheduler { } } + if let Some(ddls) = self.take_ready_ddls(region_id) { + return CompactionTransition::DdlReady(ddls); + } if self.handle_pending_compaction_request( region_id, manifest_ctx, @@ -412,21 +415,15 @@ impl CompactionScheduler { return CompactionTransition::NoAction; }; - // A queued DDL supersedes a retained automatic follow-up, matching the - // execution terminal path in `on_compaction_finished`. - let pending_ddls = std::mem::take(&mut status.pending_ddl_requests); - if !pending_ddls.is_empty() { - self.region_status.remove(®ion_id); - return CompactionTransition::DdlReady(pending_ddls); - } - if status.active.reset_automatic_followup() && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager) { return CompactionTransition::AutomaticFollowupScheduled; } - self.region_status.remove(®ion_id); + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.become_idle(); + } CompactionTransition::NoAction } diff --git a/src/mito2/src/compaction/scheduler/state.rs b/src/mito2/src/compaction/scheduler/state.rs index c5aa0816c55..3357b09119b 100644 --- a/src/mito2/src/compaction/scheduler/state.rs +++ b/src/mito2/src/compaction/scheduler/state.rs @@ -68,6 +68,8 @@ impl CompactionExecution { #[derive(Debug)] pub(super) enum CompactionPhase { + /// The region is registered but has no compaction in progress. + Idle, Picking { plan_id: u64, cancelled: bool, @@ -102,6 +104,17 @@ pub(super) struct ActiveCompaction { } impl ActiveCompaction { + /// Creates an empty cycle while retaining the region's scheduling record. + pub(super) fn idle() -> Self { + Self { + phase: CompactionPhase::Idle, + trigger: CompactionTrigger::Automatic, + waiters: Vec::new(), + automatic_followup_required: false, + } + } + + #[cfg(test)] pub(super) fn picking( plan_id: u64, waiters: Vec, @@ -165,7 +178,7 @@ impl ActiveCompaction { pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool { match &self.phase { - CompactionPhase::Picking { .. } => None, + CompactionPhase::Idle | CompactionPhase::Picking { .. } => None, CompactionPhase::Local { execution, .. } | CompactionPhase::Remote { execution } => { Some(execution) } @@ -175,6 +188,7 @@ impl ActiveCompaction { pub(super) fn request_cancel(&mut self) -> RequestCancelResult { match &mut self.phase { + CompactionPhase::Idle => RequestCancelResult::TooLateToCancel, CompactionPhase::Picking { cancelled, .. } => { if *cancelled { RequestCancelResult::AlreadyCancelling @@ -250,7 +264,31 @@ impl Drop for CompactingFilesInner { } } -/// Status of running and pending region compaction tasks. +/// A handed-off DDL batch, retained until each reply has been observed once. +#[derive(Debug)] +pub(super) struct DdlExecution { + generation: u64, + pending_replies: HashSet, +} + +impl DdlExecution { + /// Tracks a batch's individual replies so duplicate notifications are harmless. + pub(super) fn new(generation: u64, count: usize) -> Self { + Self { + generation, + pending_replies: (0..count).collect(), + } + } + + /// Returns true only for the last distinct reply of this batch. + pub(super) fn complete(&mut self, generation: u64, request_id: usize) -> bool { + self.generation == generation + && self.pending_replies.remove(&request_id) + && self.pending_replies.is_empty() + } +} + +/// Scheduling state retained for the lifetime of a loaded region. pub(super) struct CompactionStatus { /// Id of the region. pub(super) region_id: RegionId, @@ -271,24 +309,25 @@ pub(super) struct CompactionStatus { /// [`crate::request::DdlRequest::Truncate`] and [`crate::request::DdlRequest::EnterStaging`] here. Both must serialize with /// compaction so they observe the version after compaction terminates. pub(super) pending_ddl_requests: Vec, + /// DDLs already handed to the worker keep new compaction paused until their replies. + pub(super) executing_ddl: Option, } impl CompactionStatus { - /// Creates a new picking [CompactionStatus]. + /// Registers an idle region without starting a compaction cycle. pub(super) fn new( region_id: RegionId, version_control: VersionControlRef, access_layer: AccessLayerRef, - plan_id: u64, - trigger: CompactionTrigger, ) -> CompactionStatus { CompactionStatus { region_id, version_control, access_layer, - active: ActiveCompaction::picking(plan_id, Vec::new(), trigger), + active: ActiveCompaction::idle(), pending_request: None, pending_ddl_requests: Vec::new(), + executing_ddl: None, } } @@ -298,13 +337,24 @@ impl CompactionStatus { version_control: VersionControlRef, access_layer: AccessLayerRef, ) -> Self { - Self::new( - region_id, - version_control, - access_layer, - 0, - CompactionTrigger::Automatic, - ) + let mut status = Self::new(region_id, version_control, access_layer); + status.active = ActiveCompaction::picking(0, Vec::new(), CompactionTrigger::Automatic); + status + } + + /// Registration alone does not indicate active work. + pub(super) fn is_compacting(&self) -> bool { + !matches!(self.active.phase, CompactionPhase::Idle) + } + + /// Both queued and handed-off DDLs prevent new compaction admission. + pub(super) fn has_ddl(&self) -> bool { + !self.pending_ddl_requests.is_empty() || self.executing_ddl.is_some() + } + + /// Ends the current cycle without losing region or DDL execution state. + pub(super) fn become_idle(&mut self) { + self.active = ActiveCompaction::idle(); } #[cfg(test)] @@ -395,14 +445,15 @@ impl CompactionStatus { } } - pub(super) fn on_failure(mut self, err: Arc) { - for waiter in self.active.waiters.drain(..) { + pub(super) fn on_failure(&mut self, err: Arc) { + let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle()); + for waiter in active.waiters.drain(..) { waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id: self.region_id, })); } - if let Some(pending_compaction) = self.pending_request { + if let Some(pending_compaction) = self.pending_request.take() { pending_compaction .waiter .send(Err(err.clone()).context(CompactRegionSnafu { @@ -410,7 +461,7 @@ impl CompactionStatus { })); } - for pending_ddl in self.pending_ddl_requests { + for pending_ddl in self.pending_ddl_requests.drain(..) { pending_ddl .sender .send(Err(err.clone()).context(CompactRegionSnafu { @@ -420,12 +471,13 @@ impl CompactionStatus { } #[must_use] - pub(super) fn on_cancel(mut self) -> Vec { - for waiter in self.active.waiters.drain(..) { + pub(super) fn on_cancel(&mut self) -> Vec { + let mut active = std::mem::replace(&mut self.active, ActiveCompaction::idle()); + for waiter in active.waiters.drain(..) { waiter.send(CompactionCancelledSnafu.fail()); } - if let Some(pending_compaction) = self.pending_request { + if let Some(pending_compaction) = self.pending_request.take() { pending_compaction.waiter.send( Err(Arc::new(CompactionCancelledSnafu.build())).context(CompactRegionSnafu { region_id: self.region_id, diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index 42c1df23f07..c561aa0cfa5 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -358,12 +358,214 @@ async fn test_find_compaction_options_priority() { } } +#[tokio::test] +async fn test_resident_region_survives_terminal_cycles_and_truncate() { + let env = SchedulerEnv::new().await; + let (tx, _rx) = mpsc::channel(8); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version = Arc::new(VersionControlBuilder::new().build()); + let id = version.region_id(); + scheduler.register_region(&version, &env.access_layer); + assert!(!scheduler.is_compacting(id)); + assert!(scheduler.region_status.contains_key(&id)); + for cancelled in [false, true] { + let status = scheduler.region_status.get_mut(&id).unwrap(); + status.start_picking(7); + let (reply, response) = oneshot::channel(); + status.merge_waiter(reply.into()); + scheduler.register_region(&version, &env.access_layer); + assert!( + scheduler.is_compacting(id), + "duplicate registration must retain the active cycle" + ); + if cancelled { + assert!(scheduler.on_compaction_cancelled(id).await.is_empty()); + } else { + scheduler.on_compaction_failed(id, Arc::new(InvalidSchedulerStateSnafu.build())); + } + assert!(response.await.unwrap().is_err()); + assert!(!scheduler.is_compacting(id)); + assert!(Arc::ptr_eq( + &scheduler.region_status[&id].version_control, + &version + )); + } + scheduler.on_region_truncated(id); + assert!(scheduler.region_status.contains_key(&id)); + assert!(!scheduler.is_compacting(id)); + scheduler.on_region_closed(id); + assert!(!scheduler.region_status.contains_key(&id)); + scheduler.register_region(&version, &env.access_layer); + let (reply, response) = oneshot::channel(); + let status = scheduler.region_status.get_mut(&id).unwrap(); + status.start_picking(8); + status.merge_waiter(reply.into()); + let replacement = Arc::new(VersionControlBuilder::new().build()); + scheduler.register_region(&replacement, &env.access_layer); + assert!(response.await.unwrap().is_err()); + assert!(Arc::ptr_eq( + &scheduler.region_status[&id].version_control, + &replacement + )); + assert!(!scheduler.is_compacting(id)); + scheduler.on_region_dropped(id); + assert!(!scheduler.region_status.contains_key(&id)); +} + +/// Reads the real completion event produced by a wrapped DDL reply. +async fn recv_ddl_completion(rx: &mut mpsc::Receiver) -> (u64, usize) { + let request = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + let WorkerRequest::Background { + notify: + BackgroundNotify::CompactionDdlComplete { + generation, + request_id, + }, + .. + } = request.request + else { + panic!("expected DDL completion"); + }; + (generation, request_id) +} + +#[tokio::test] +async fn test_resident_region_ddl_handoff_and_late_replies() { + let env = SchedulerEnv::new().await; + let (tx, mut rx) = mpsc::channel(8); + let mut scheduler = env.mock_compaction_scheduler(tx); + let version = Arc::new(VersionControlBuilder::new().build()); + let id = version.region_id(); + let manifest = env + .mock_manifest_context(version.current().version.metadata.clone()) + .await; + let (schemas, _) = mock_schema_metadata_manager(); + scheduler.register_region(&version, &env.access_layer); + scheduler + .region_status + .get_mut(&id) + .unwrap() + .start_picking(7); + let mut responses = Vec::new(); + for _ in 0..2 { + let (reply, response) = oneshot::channel(); + scheduler + .try_cancel_and_add_ddl(id, reply.into(), (), |_| { + DdlRequest::Truncate(store_api::region_request::RegionTruncateRequest::All) + }) + .unwrap(); + responses.push(response); + } + let mut ddls = scheduler.on_compaction_cancelled(id).await; + assert_eq!(2, ddls.len()); + assert!(!scheduler.is_compacting(id)); + assert!(scheduler.region_status[&id].has_ddl()); + assert!(scheduler.region_status[&id].pending_ddl_requests.is_empty()); + assert!( + !scheduler + .schedule_automatic_compaction( + Options::Regular(Default::default()), + &version, + &env.access_layer, + &manifest, + schemas.clone() + ) + .unwrap() + ); + let (reply, response) = oneshot::channel(); + assert!( + !scheduler + .schedule_manual_compaction( + Options::Regular(Default::default()), + &version, + &env.access_layer, + reply.into(), + &manifest, + schemas.clone(), + 1, + None + ) + .unwrap() + ); + assert_matches!( + response.await.unwrap(), + Err(Error::CompactionCancelled { .. }) + ); + + ddls.remove(0).sender.send(Ok(0)); + let (old_generation, first) = recv_ddl_completion(&mut rx).await; + scheduler.on_ddl_complete(id, old_generation, first); + scheduler.on_ddl_complete(id, old_generation, first); + scheduler.on_region_truncated(id); + scheduler.register_region(&version, &env.access_layer); + assert!( + scheduler.region_status[&id].has_ddl(), + "duplicate replies and truncate must not end the batch" + ); + ddls.pop().unwrap().sender.send(Ok(0)); + let (_, late_reply) = recv_ddl_completion(&mut rx).await; + for response in responses { + assert_eq!(0, response.await.unwrap().unwrap()); + } + + scheduler.on_region_closed(id); + let reopened = Arc::new(VersionControlBuilder::new().build()); + scheduler.register_region(&reopened, &env.access_layer); + scheduler + .region_status + .get_mut(&id) + .unwrap() + .start_picking(8); + let (reply, response) = oneshot::channel(); + scheduler + .try_cancel_and_add_ddl(id, reply.into(), (), |_| { + DdlRequest::Truncate(store_api::region_request::RegionTruncateRequest::All) + }) + .unwrap(); + let mut ddls = scheduler.on_compaction_cancelled(id).await; + scheduler.on_ddl_complete(id, old_generation, late_reply); + assert!( + scheduler.region_status[&id].has_ddl(), + "old incarnation must not clear the new batch" + ); + ddls.pop() + .unwrap() + .sender + .send(InvalidSchedulerStateSnafu.fail()); + let (generation, request_id) = recv_ddl_completion(&mut rx).await; + scheduler.on_ddl_complete(id, generation, request_id); + scheduler.on_ddl_complete(id, generation, request_id); + assert!(response.await.unwrap().is_err()); + assert!(!scheduler.region_status[&id].has_ddl()); + assert!( + scheduler + .schedule_automatic_compaction( + Options::Regular(Default::default()), + &reopened, + &env.access_layer, + &manifest, + schemas.clone() + ) + .unwrap() + ); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest, schemas) + .await; + assert!(!scheduler.is_compacting(id)); + assert!(scheduler.region_status.contains_key(&id)); +} + #[tokio::test] async fn test_schedule_empty() { let env = SchedulerEnv::new().await; let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); schema_metadata_manager .register_region_table_info( @@ -402,10 +604,12 @@ async fn test_schedule_empty() { .await; let output = output_rx.await.unwrap().unwrap(); assert_eq!(output, 0); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); // Only one file, picker won't compact it. let version_control = Arc::new(builder.push_l0_file(0, 1000).build()); + scheduler.register_region(&version_control, &env.access_layer); let (output_tx, output_rx) = oneshot::channel(); let waiter = OptionOutputTx::from(output_tx); let scheduled = scheduler @@ -428,7 +632,8 @@ async fn test_schedule_empty() { .await; let output = output_rx.await.unwrap().unwrap(); assert_eq!(output, 0); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -484,7 +689,7 @@ async fn test_schedule_compaction_returns_true_when_task_scheduled() { .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) .await; assert_eq!(1, job_scheduler.num_jobs()); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); } #[tokio::test] @@ -521,7 +726,7 @@ async fn test_planning_followup_reports_automatic_schedule() { } #[tokio::test] -async fn test_planning_panic_notifies_and_clears_status() { +async fn test_planning_panic_notifies_and_returns_region_to_idle() { let env = SchedulerEnv::new().await; let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx.clone()); @@ -555,7 +760,8 @@ async fn test_planning_panic_notifies_and_clears_status() { assert!(pending_ddls.is_empty()); assert!(waiter_rx.await.unwrap().is_err()); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -613,14 +819,11 @@ async fn test_ddl_fence_prevents_repeated_regular_followups() { schema_metadata_manager.clone(), ) .await; - assert!(pending_ddls.is_empty()); - - let mut followup_finished = tokio::time::timeout( - Duration::from_secs(5), - recv_compaction_pick_finished(&mut rx), - ) - .await - .expect("pre-fence regular follow-up was not planned"); + assert_eq!(pending_ddls.len(), 1); + assert!( + rx.try_recv().is_err(), + "DDL must precede even an earlier pending manual request" + ); let (post_fence_tx, post_fence_rx) = oneshot::channel(); assert!( !scheduler @@ -641,13 +844,9 @@ async fn test_ddl_fence_prevents_repeated_regular_followups() { Error::CompactionCancelled { .. } ); - followup_finished.result = CompactionPlanningResult::NoPlan; - let pending_ddls = scheduler - .handle_compaction_pick_finished(followup_finished, &manifest_ctx, schema_metadata_manager) - .await; - assert_eq!(pending_ddls.len(), 1); - assert_eq!(pre_fence_rx.await.unwrap().unwrap(), 0); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(pre_fence_rx.await.unwrap().is_err()); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert!(rx.try_recv().is_err()); } @@ -776,7 +975,8 @@ async fn test_pick_result_local_submission_failure_releases_and_notifies_once() assert!(waiter_rx.await.unwrap().is_err()); assert!(selected.iter().all(|file| !file.compacting())); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -806,7 +1006,8 @@ async fn test_pick_result_remote_submission_failure_releases_and_notifies_once() assert!(waiter_rx.await.unwrap().is_err()); assert!(selected.iter().all(|file| !file.compacting())); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); } #[tokio::test] @@ -882,13 +1083,13 @@ async fn test_stale_plan_execution_does_not_affect_replacement_status() { ) .await; assert!(pending_ddls.is_empty()); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); scheduler.on_execution_failed( region_id, &stale_execution, Arc::new(InvalidSchedulerStateSnafu.build()), ); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); assert_matches!( waiter_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty) @@ -955,7 +1156,8 @@ async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { assert_eq!(output_rx.await.unwrap().unwrap(), 0); assert_eq!(rejected_before + 1, rejected.get()); assert_eq!(0, job_scheduler.num_jobs()); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert!(selected.iter().all(|file| !file.compacting())); } @@ -1039,7 +1241,7 @@ async fn test_execution_finished_drains_until_no_plan() { ) .await; assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); - assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.is_compacting(region_id)); let finished = recv_compaction_pick_finished(&mut rx).await; scheduler .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) @@ -1069,12 +1271,13 @@ async fn test_execution_finished_drains_until_no_plan() { .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) .await; assert_matches!(transition, CompactionTransition::NoAction); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] -async fn test_execution_finished_without_progress_removes_status() { +async fn test_execution_finished_without_progress_keeps_idle_status() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); @@ -1136,7 +1339,8 @@ async fn test_execution_finished_without_progress_removes_status() { ) .await; assert_matches!(transition, CompactionTransition::NoAction); - assert!(scheduler.region_status.is_empty()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(1, job_scheduler.num_jobs()); } @@ -1672,7 +1876,8 @@ async fn test_compaction_bypass_in_staging_mode() { let result = rx.await.unwrap(); assert_eq!(result.unwrap(), 0); - assert_eq!(0, scheduler.region_status.len()); + assert_eq!(1, scheduler.region_status.len()); + assert!(!scheduler.is_compacting(version_control.region_id())); } #[tokio::test] @@ -1851,6 +2056,10 @@ async fn test_try_cancel_and_add_ddl_returns_request_when_not_running() { let (tx, _rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let region_id = RegionId::new(1, 1); + let version = Arc::new(VersionControlBuilder::new().build()); + scheduler.register_region(&version, &env.access_layer); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); let (ddl_tx, ddl_rx) = oneshot::channel(); let result = @@ -1940,7 +2149,8 @@ async fn test_on_compaction_cancelled_returns_pending_ddl_requests() { assert_eq!(pending_ddls.len(), 1); assert!(!scheduler.has_pending_ddls(region_id)); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(job_scheduler.num_jobs(), 0); } @@ -1987,7 +2197,8 @@ async fn test_on_compaction_cancelled_prioritizes_pending_ddls_over_pending_comp let pending_ddls = scheduler.on_compaction_cancelled(region_id).await; assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(job_scheduler.num_jobs(), 0); assert_matches!(manual_rx.await.unwrap(), Err(_)); } @@ -2173,7 +2384,8 @@ async fn test_on_compaction_finished_returns_pending_ddl_requests() { assert_eq!(pending_ddls.len(), 1); assert!(!scheduler.has_pending_ddls(region_id)); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_eq!(job_scheduler.num_jobs(), 0); } @@ -2225,14 +2437,11 @@ async fn test_planning_terminal_prioritizes_pending_ddl_over_automatic_followup( ) .await; - assert!(pending_ddls.is_empty()); - let finished = recv_compaction_pick_finished(&mut rx).await; - let pending_ddls = scheduler - .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) - .await; assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert_eq!(manual_rx.await.unwrap().unwrap(), 0); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); + assert!(rx.try_recv().is_err()); + assert!(manual_rx.await.unwrap().is_err()); } #[tokio::test] @@ -2275,7 +2484,8 @@ async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regul .await; assert_eq!(pending_ddls.len(), 1); - assert!(!scheduler.region_status.contains_key(®ion_id)); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert!(rx.try_recv().is_err()); } @@ -2300,7 +2510,7 @@ async fn test_on_compaction_finished_returns_empty_when_region_absent() { } #[tokio::test] -async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { +async fn test_on_compaction_finished_ddl_bypasses_failing_manual_scheduler() { let env = SchedulerEnv::new() .await .scheduler(Arc::new(FailingScheduler)); @@ -2356,15 +2566,16 @@ async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { ) .await; - assert!(pending_ddls.is_empty()); - let finished = recv_compaction_pick_finished(&mut rx).await; - let pending_ddls = scheduler - .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager) - .await; - assert!(pending_ddls.is_empty()); - assert!(!scheduler.region_status.contains_key(®ion_id)); + let super::CompactionTransition::DdlReady(mut ddls) = pending_ddls else { + panic!("DDL not released") + }; + assert_eq!(ddls.len(), 1); + assert!(rx.try_recv().is_err()); + assert!(scheduler.region_status.contains_key(®ion_id)); + assert!(!scheduler.is_compacting(region_id)); assert_matches!(manual_rx.await.unwrap(), Err(_)); - assert_matches!(ddl_rx.await.unwrap(), Err(_)); + ddls.pop().unwrap().sender.send(Ok(0)); + assert_eq!(ddl_rx.await.unwrap().unwrap(), 0); } #[tokio::test] diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 092abc296a6..f20eef75460 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -1097,6 +1097,61 @@ impl EventListener for CompactionPlanningGate { } } +#[tokio::test] +async fn test_idle_region_accepts_ddl_after_create_and_reopen() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(45, 1); + let create = CreateRequestBuilder::new().build(); + let table_dir = create.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + for reopen in [false, true] { + if reopen { + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: table_dir.clone(), + path_type: PathType::Bare, + options: Default::default(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + } + for request in [ + RegionRequest::Compact(RegionCompactRequest::default()), + RegionRequest::Truncate(RegionTruncateRequest::All), + RegionRequest::Compact(RegionCompactRequest::default()), + ] { + tokio::time::timeout( + Duration::from_secs(5), + engine.handle_request(region_id, request), + ) + .await + .expect("idle registration must not stall DDL or the next compaction") + .unwrap(); + } + } +} + #[tokio::test] async fn test_planning_followup_updates_schedule_time() { assert_automatic_followup_updates_schedule_time(0).await; diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 10002570008..81993a49afe 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -58,6 +58,7 @@ use crate::manifest::action::{RegionEdit, TruncateKind}; use crate::memtable::MemtableId; use crate::memtable::bulk::part::BulkPart; use crate::metrics::COMPACTION_ELAPSED_TOTAL; +use crate::region::MitoRegionRef; use crate::region::options::RegionOptions; use crate::sst::file::FileMeta; use crate::sst::index::IndexBuildType; @@ -899,6 +900,13 @@ pub(crate) struct SenderDdlRequest { /// Notification from a background job. #[derive(Debug)] pub(crate) enum BackgroundNotify { + /// Registers a loaded region and its scheduling state before acknowledging open. + RegionOpened { + region: MitoRegionRef, + registered: Sender<()>, + }, + /// A DDL released by compaction has actually completed, including failure. + CompactionDdlComplete { generation: u64, request_id: usize }, /// Compaction planning has finished. CompactionPickFinished(CompactionPickFinished), /// Flush has finished. diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 2eda4fff49f..6de814385b4 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1010,6 +1010,13 @@ struct RegionWorkerLoop { } impl RegionWorkerLoop { + /// Publishes a loaded region together with its resident compaction state. + fn register_region(&mut self, region: MitoRegionRef) { + self.compaction_scheduler + .register_region(®ion.version_control, ®ion.access_layer); + self.regions.insert_region(region); + } + /// Starts the worker loop. async fn run(&mut self) { let init_check_delay = worker_init_check_delay(); @@ -1338,6 +1345,18 @@ impl RegionWorkerLoop { /// Handles region background request async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) { match notify { + BackgroundNotify::RegionOpened { region, registered } => { + self.register_region(region); + self.region_count.inc(); + let _ = registered.send(()); + } + BackgroundNotify::CompactionDdlComplete { + generation, + request_id, + } => { + self.compaction_scheduler + .on_ddl_complete(region_id, generation, request_id); + } BackgroundNotify::CompactionPickFinished(req) => { self.handle_compaction_pick_finished(region_id, req).await } diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index 576ffed7ba2..81838060385 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -133,7 +133,7 @@ impl RegionWorkerLoop { .open(&self.config, &self.wal) .await?; debug_assert!(!reopened_region.is_writable()); - self.regions.insert_region(reopened_region.clone()); + self.register_region(reopened_region.clone()); if let Some(state) = &self.series_index_task_state { state.wake(); } diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index ad7872b7237..8d1861a041f 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -45,8 +45,8 @@ impl RegionWorkerLoop { let Some(region) = self.regions.get_region(region_id) else { return; }; - // A terminal pick (canceled, no plan, or failed) may remove the - // compaction status and release DDLs fenced behind its picking phase. + // A terminal pick (canceled, no plan, or failed) may return the region + // to idle and release DDLs fenced behind its picking phase. // Such a pick never produces an execution callback, so the worker must // execute the returned DDLs from this notification. let transition = self diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 2263cfa2ffb..cff77b50f5f 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -95,7 +95,7 @@ impl RegionWorkerLoop { } // Insert the MitoRegion into the RegionMap. - self.regions.insert_region(region); + self.register_region(region); if let Some(state) = &self.series_index_task_state { state.wake(); } diff --git a/src/mito2/src/worker/handle_enter_staging.rs b/src/mito2/src/worker/handle_enter_staging.rs index f9e24541241..e5b0674aaac 100644 --- a/src/mito2/src/worker/handle_enter_staging.rs +++ b/src/mito2/src/worker/handle_enter_staging.rs @@ -64,6 +64,24 @@ impl RegionWorkerLoop { return; } + // Retain the compaction fence while the DDL later waits for flush. + let (sender, partition_directive) = match self.compaction_scheduler.try_cancel_and_add_ddl( + region_id, + sender, + partition_directive, + |partition_directive| { + DdlRequest::EnterStaging(EnterStagingRequest { + partition_directive, + }) + }, + ) { + Ok(()) => { + self.listener.on_compaction_cancel_requested(region_id); + return; + } + Err(request) => request, + }; + let version = region.version(); if !version.memtables.is_empty() { // If memtable is not empty, we can't enter staging directly and need to flush @@ -98,23 +116,6 @@ impl RegionWorkerLoop { return; } - let (sender, partition_directive) = match self.compaction_scheduler.try_cancel_and_add_ddl( - region_id, - sender, - partition_directive, - |partition_directive| { - DdlRequest::EnterStaging(EnterStagingRequest { - partition_directive, - }) - }, - ) { - Ok(()) => { - self.listener.on_compaction_cancel_requested(region_id); - return; - } - Err(request) => request, - }; - self.handle_enter_staging(region, partition_directive, sender); } diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index 52b902ffe92..ae56a366d46 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -35,7 +35,7 @@ use crate::region::opener::{ RegionOpener, get_object_store, provider_from_wal_options, sanitize_open_request_options, }; use crate::region::options::RegionOptions; -use crate::request::OptionOutputTx; +use crate::request::{BackgroundNotify, OptionOutputTx, WorkerRequest, WorkerRequestWithTime}; use crate::sst::location::region_dir_from_table_dir; use crate::wal::entry_distributor::WalEntryReceiver; use crate::worker::handle_drop::{ @@ -221,11 +221,10 @@ impl RegionWorkerLoop { } let now = Instant::now(); - let regions = self.regions.clone(); + let request_sender = self.sender.clone(); let wal = self.wal.clone(); let config = self.config.clone(); let opening_regions = self.opening_regions.clone(); - let region_count = self.region_count.clone(); let worker_id = self.id; let series_index_task_state = self.series_index_task_state.clone(); opening_regions.insert_sender(region_id, sender); @@ -239,8 +238,6 @@ impl RegionWorkerLoop { worker_id, now.elapsed() ); - region_count.inc(); - // Notify the region hook that the region has been opened. // Fires before registration; allocates nothing when no hook // is registered. @@ -248,15 +245,31 @@ impl RegionWorkerLoop { hook.on_region_opened(region_id, ®ion.metadata()).await; } - // Insert the Region into the RegionMap. - regions.insert_region(region); - if let Some(state) = &series_index_task_state { + // The worker installs both the region and its resident scheduler state. + let (registered, receiver) = tokio::sync::oneshot::channel(); + let sent = request_sender + .send(WorkerRequestWithTime::new(WorkerRequest::Background { + region_id, + notify: BackgroundNotify::RegionOpened { + region: region.clone(), + registered, + }, + })) + .await + .is_ok(); + let accepted = sent && receiver.await.is_ok(); + if !accepted { + region.stop().await; + } else if let Some(state) = &series_index_task_state { state.wake(); } - let senders = opening_regions.remove_sender(region_id); for sender in senders { - sender.send(Ok(0)); + if accepted { + sender.send(Ok(0)); + } else { + sender.send(crate::error::RegionClosedSnafu { region_id }.fail()); + } } } Err(err) => { diff --git a/src/mito2/src/worker/handle_truncate.rs b/src/mito2/src/worker/handle_truncate.rs index d62bf8369a4..409998592d1 100644 --- a/src/mito2/src/worker/handle_truncate.rs +++ b/src/mito2/src/worker/handle_truncate.rs @@ -40,19 +40,7 @@ impl RegionWorkerLoop { } }; - let (sender, req) = match self.flush_scheduler.try_cancel_and_add_ddl( - region_id, - sender, - req, - DdlRequest::Truncate, - ) { - Ok(()) => { - self.listener.on_flush_cancel_requested(region_id); - return; - } - Err(request) => request, - }; - + // Fence compaction before any later flush wait so handed-off DDLs keep admission paused. let (sender, req) = match self.compaction_scheduler.try_cancel_and_add_ddl( region_id, sender, @@ -66,6 +54,19 @@ impl RegionWorkerLoop { Err(request) => request, }; + let (sender, req) = match self.flush_scheduler.try_cancel_and_add_ddl( + region_id, + sender, + req, + DdlRequest::Truncate, + ) { + Ok(()) => { + self.listener.on_flush_cancel_requested(region_id); + return; + } + Err(request) => request, + }; + let version_data = region.version_control.current(); match req {