1use std::collections::{HashMap, VecDeque};
16use std::future::Future;
17use std::sync::Arc;
18#[cfg(test)]
19use std::sync::atomic::AtomicUsize;
20use std::sync::atomic::{AtomicBool, Ordering};
21
22use common_meta::key::SchemaMetadataManagerRef;
23use prometheus::IntGauge;
24use store_api::compaction_task::{self, CompactionTaskEntry};
25use store_api::storage::RegionId;
26
27use crate::compaction::compactor::{CompactionContext, CompactionRegion};
28use crate::compaction::scheduler::state::{CompactingFiles, CompactionExecution, CompactionPhase};
29use crate::compaction::scheduler::{
30 CompactionCoordinator, CompactionScheduler, CompactionTransition, ReadyUnit,
31};
32use crate::compaction::task::CompactionTask;
33use crate::compaction::unit::CompactionUnit;
34use crate::error::{self, Error};
35use crate::metrics::INFLIGHT_COMPACTION_COUNT;
36use crate::region::ManifestContextRef;
37use crate::request::{
38 BackgroundNotify, CompactionUnitNotification, WorkerRequest, WorkerRequestWithTime,
39};
40use crate::schedule::{CancellableTaskState, RequestCancelResult};
41use crate::sst::file::UncommittedSsts;
42
43#[derive(Debug)]
45struct UnitExecution {
46 state: CancellableTaskState,
47 finished: AtomicBool,
50 released: AtomicBool,
51}
52
53struct InflightCompaction(IntGauge);
56
57impl InflightCompaction {
58 async fn run(job: impl Future<Output = ()>, gauge: IntGauge) {
60 let _inflight = Self::start(gauge);
61 job.await;
62 }
63
64 fn start(gauge: IntGauge) -> Self {
65 gauge.inc();
66 Self(gauge)
67 }
68}
69
70impl Drop for InflightCompaction {
71 fn drop(&mut self) {
72 self.0.dec();
73 }
74}
75
76#[derive(Debug)]
78pub(crate) struct LocalCompaction {
79 coordinator: Arc<CompactionCoordinator>,
80 region_id: RegionId,
81 units: HashMap<u64, Arc<UnitExecution>>,
85 unit_states: Arc<[Arc<UnitExecution>]>,
88 #[cfg(test)]
89 publication_count: Arc<AtomicUsize>,
90 pending: VecDeque<CompactionTask>,
91 parallelism: usize,
92 error: Option<Arc<Error>>,
93 made_progress: bool,
94 cancellation_requested: bool,
96}
97
98#[derive(Debug, Clone)]
100pub(crate) struct LocalCompactionState {
101 units: Arc<[Arc<UnitExecution>]>,
102 cancellation_requested: bool,
103}
104
105impl LocalCompactionState {
106 pub(crate) fn inspect(&self, entry: &mut CompactionTaskEntry) {
110 entry.total_units = Some(self.units.len() as u64);
111 let mut completed = 0;
112 let mut cancelling = self.cancellation_requested;
113 let mut committing = false;
114 let mut running = false;
115 for unit in self.units.iter() {
116 completed += u64::from(unit.finished.load(Ordering::Relaxed));
117 if unit.released.load(Ordering::Relaxed) {
118 continue;
119 }
120 cancelling |= unit.state.is_cancelled();
121 committing |= unit.state.commit_started();
122 running |= unit.state.started();
123 }
124 entry.completed_units = Some(completed);
125 running |= committing;
126 entry.state = if cancelling {
127 compaction_task::STATE_CANCELLING
128 } else if running {
129 compaction_task::STATE_RUNNING
130 } else {
131 compaction_task::STATE_QUEUED
132 }
133 .into();
134 entry.phase = if committing {
135 compaction_task::PHASE_COMMITTING
136 } else if running {
137 compaction_task::PHASE_EXECUTING
138 } else {
139 compaction_task::PHASE_WAITING
140 }
141 .into();
142 }
143}
144
145impl LocalCompaction {
146 pub(crate) fn current_state(&self) -> LocalCompactionState {
148 #[cfg(test)]
149 self.publication_count.fetch_add(1, Ordering::Relaxed);
150 LocalCompactionState {
151 units: self.unit_states.clone(),
152 cancellation_requested: self.cancellation_requested,
153 }
154 }
155
156 #[cfg(test)]
157 pub(crate) fn publication_count_for_test(&self) -> Arc<AtomicUsize> {
158 self.publication_count.clone()
159 }
160
161 #[cfg(test)]
162 pub(crate) fn inspection_stats_for_test(&self) -> (usize, usize, usize) {
163 (
164 Arc::as_ptr(&self.unit_states) as *const () as usize,
165 self.unit_states.len(),
166 Arc::strong_count(&self.unit_states),
167 )
168 }
169
170 #[cfg(test)]
171 pub(crate) fn released_units_for_test(&self) -> usize {
172 self.unit_states
173 .iter()
174 .filter(|unit| unit.released.load(Ordering::Relaxed))
175 .count()
176 }
177
178 #[cfg(test)]
179 pub(crate) fn unit_ids_for_test(&self) -> Vec<u64> {
180 self.units.keys().copied().collect()
181 }
182
183 #[cfg(test)]
184 pub(crate) fn inspect(&self, entry: &mut CompactionTaskEntry) {
185 self.current_state().inspect(entry);
186 }
187
188 pub(crate) fn request_cancel(&mut self) -> RequestCancelResult {
190 let mut result = RequestCancelResult::TooLateToCancel;
191 for unit in self.units.values() {
192 match unit.state.request_cancel() {
193 RequestCancelResult::CancelIssued => result = RequestCancelResult::CancelIssued,
194 RequestCancelResult::AlreadyCancelling
195 if result != RequestCancelResult::CancelIssued =>
196 {
197 result = RequestCancelResult::AlreadyCancelling;
198 }
199 _ => {}
200 }
201 }
202 self.cancellation_requested |= result != RequestCancelResult::TooLateToCancel;
203 self.coordinator.cancel_pending(self.region_id);
204 result
205 }
206
207 fn num_running(&self) -> usize {
209 self.units.len() - self.pending.len()
210 }
211
212 fn can_apply(&self, id: u64) -> bool {
214 self.units.get(&id).is_some_and(|unit| {
215 !unit.finished.load(Ordering::Relaxed) && unit.state.commit_started()
216 })
217 }
218
219 fn release_unit(&mut self, id: u64) {
221 if let Some(unit) = self.units.remove(&id) {
222 unit.released.store(true, Ordering::Relaxed);
223 }
224 }
225
226 fn finish(&mut self, id: u64, result: std::result::Result<bool, Arc<Error>>) {
228 let Some(unit) = self.units.get(&id) else {
229 return;
230 };
231 if unit.finished.load(Ordering::Relaxed) {
232 return;
233 }
234 unit.finished.store(true, Ordering::Relaxed);
235 match result {
236 Ok(progress) => self.made_progress |= progress,
237 Err(err) => {
238 if self.error.as_ref().is_none_or(|current| {
240 matches!(current.as_ref(), Error::CompactionCancelled { .. })
241 }) {
242 self.error = Some(err);
243 }
244 }
245 }
246 }
247}
248
249impl Drop for LocalCompaction {
250 fn drop(&mut self) {
252 self.request_cancel();
253 }
254}
255
256impl CompactionScheduler {
257 pub(crate) fn prepare_local_units(
259 &mut self,
260 region: CompactionRegion,
261 units: Vec<CompactionUnit>,
262 reserved: CompactingFiles,
263 ) -> CompactionPhase {
264 if let Some(info) = self
267 .region_status
268 .get_mut(®ion.region_id)
269 .and_then(|status| status.active.task_state.as_mut())
270 {
271 info.set_plan_stats(
272 units.iter().flat_map(|unit| &unit.inputs),
273 units.iter().flat_map(|unit| &unit.outputs),
274 );
275 }
276 let mut local = LocalCompaction {
277 coordinator: self.coordinator.clone(),
278 region_id: region.region_id,
279 units: HashMap::new(),
280 unit_states: Arc::from([]),
281 #[cfg(test)]
282 publication_count: Arc::new(AtomicUsize::new(0)),
283 pending: VecDeque::new(),
284 parallelism: region.max_parallelism.max(1),
285 error: None,
286 made_progress: false,
287 cancellation_requested: false,
288 };
289 let context = CompactionContext::from(®ion);
290 drop(region);
292 for unit in units {
293 let files = reserved.for_inputs(&unit.inputs);
294 let id = Self::next_plan_id(&mut self.next_plan_id);
295 let execution = CompactionExecution::new(id, files);
296 let state = CancellableTaskState::new();
297 local.units.insert(
298 id,
299 Arc::new(UnitExecution {
300 state: state.clone(),
301 finished: AtomicBool::new(false),
302 released: AtomicBool::new(false),
303 }),
304 );
305 local.pending.push_back(CompactionTask {
306 state,
307 execution,
308 compaction_context: context.clone(),
309 request_sender: self.request_sender.clone(),
310 listener: self.listener.clone(),
311 estimated_memory_bytes: unit.estimated_memory_bytes(),
312 unit,
313 memory_manager: self.memory_manager.clone(),
314 memory_policy: self.memory_policy,
315 uncommitted: UncommittedSsts::new(
316 context.region_id,
317 context.access_layer.clone(),
318 Some(context.cache_manager.clone()),
319 ),
320 });
321 }
322 local.unit_states = local.units.values().cloned().collect();
323 CompactionPhase::Units(local)
324 }
325
326 fn local_units_mut(&mut self, region_id: RegionId) -> Option<&mut LocalCompaction> {
328 match &mut self.region_status.get_mut(®ion_id)?.active.phase {
329 CompactionPhase::Units(local) => Some(local),
330 _ => None,
331 }
332 }
333
334 pub(crate) fn dispatch_local_units(&mut self, region_id: RegionId) {
336 self.submit_local_units(region_id);
337 if let Some(status) = self.region_status.get(®ion_id) {
338 status.publish_state();
339 }
340 }
341
342 fn submit_local_units(&mut self, region_id: RegionId) {
343 loop {
344 let Some(local) = self.local_units_mut(region_id) else {
345 return;
346 };
347 if local.num_running() >= local.parallelism {
348 return;
349 }
350 let Some(mut task) = local.pending.pop_front() else {
351 return;
352 };
353 let id = task.execution.plan_id();
354 let state = task.state.clone();
355 let input_bytes = task.unit.estimated_input_bytes();
356 let sender = self.request_sender.clone();
357 let completion = Box::pin(async move {
358 let _ = sender
359 .send(WorkerRequestWithTime::new(WorkerRequest::Background {
360 region_id,
361 notify: BackgroundNotify::CompactionUnit(
362 CompactionUnitNotification::Released { plan_id: id },
363 ),
364 }))
365 .await;
366 });
367 let job = Box::pin(InflightCompaction::run(
368 async move {
369 task.run().await;
370 drop(task);
371 },
372 INFLIGHT_COMPACTION_COUNT.clone(),
373 ));
374 if let Err(err) = self.coordinator.schedule(
375 region_id,
376 ReadyUnit::new(state, input_bytes, job, completion),
377 ) && let Some(local) = self.local_units_mut(region_id)
378 {
379 local.finish(id, Err(Arc::new(err)));
380 local.release_unit(id);
381 }
382 }
383 }
384
385 pub(crate) fn can_apply_unit(&self, region_id: RegionId, id: u64) -> bool {
387 self.region_status.get(®ion_id).is_some_and(|status| {
388 matches!(&status.active.phase, CompactionPhase::Units(local)
389 if local.can_apply(id))
390 })
391 }
392
393 pub(crate) fn on_unit_finished(
395 &mut self,
396 region_id: RegionId,
397 id: u64,
398 result: std::result::Result<bool, Arc<Error>>,
399 ) {
400 if let Some(local) = self.local_units_mut(region_id) {
401 local.finish(id, result);
402 }
403 if let Some(status) = self.region_status.get(®ion_id) {
404 status.publish_state();
405 }
406 }
407
408 pub(crate) async fn on_unit_released(
410 &mut self,
411 region_id: RegionId,
412 id: u64,
413 manifest: &ManifestContextRef,
414 schemas: SchemaMetadataManagerRef,
415 ) -> CompactionTransition {
416 let Some(local) = self.local_units_mut(region_id) else {
417 return CompactionTransition::NoAction;
418 };
419 if !local.units.contains_key(&id) {
420 return CompactionTransition::NoAction;
421 }
422 let missing_result = if local.units[&id].state.is_cancelled() {
423 error::CompactionCancelledSnafu.build()
424 } else {
425 error::UnexpectedSnafu {
426 reason: "Compaction unit exited without a terminal notification",
427 }
428 .build()
429 };
430 local.finish(id, Err(Arc::new(missing_result)));
431 local.release_unit(id);
432 self.dispatch_local_units(region_id);
433 self.finish_local_units(region_id, manifest, schemas).await
434 }
435
436 pub(crate) async fn finish_local_units(
438 &mut self,
439 region_id: RegionId,
440 manifest: &ManifestContextRef,
441 schemas: SchemaMetadataManagerRef,
442 ) -> CompactionTransition {
443 let Some(local) = self.local_units_mut(region_id) else {
444 return CompactionTransition::NoAction;
445 };
446 if !local.units.is_empty() {
447 return CompactionTransition::NoAction;
448 }
449 if let Some(err) = local.error.take() {
450 if matches!(err.as_ref(), Error::CompactionCancelled { .. }) {
451 return CompactionTransition::from_pending_ddls(
452 self.finish_compaction_on_cancel(region_id),
453 );
454 }
455 self.on_compaction_failed(region_id, err);
456 return CompactionTransition::NoAction;
457 }
458 let progress = local.made_progress;
459 if let Some(status) = self.region_status.get_mut(®ion_id) {
460 for waiter in status.take_waiters() {
461 waiter.send(Ok(0));
462 }
463 }
464 self.on_compaction_finished(region_id, manifest, schemas, progress)
465 .await
466 }
467}
468
469#[cfg(test)]
470mod tests {
471 use futures::FutureExt;
472 use tokio::sync::oneshot;
473
474 use super::*;
475
476 fn local_compaction() -> LocalCompaction {
477 let units: HashMap<_, _> = [1, 2]
478 .into_iter()
479 .map(|id| {
480 (
481 id,
482 Arc::new(UnitExecution {
483 state: CancellableTaskState::new(),
484 finished: AtomicBool::new(false),
485 released: AtomicBool::new(false),
486 }),
487 )
488 })
489 .collect();
490 LocalCompaction {
491 coordinator: Arc::new(CompactionCoordinator::new(
492 Arc::new(crate::test_util::scheduler_util::VecScheduler::default()),
493 2,
494 )),
495 region_id: RegionId::from_u64(1),
496 unit_states: units.values().cloned().collect(),
497 publication_count: Arc::new(AtomicUsize::new(0)),
498 units,
499 pending: VecDeque::new(),
500 parallelism: 2,
501 error: None,
502 made_progress: false,
503 cancellation_requested: false,
504 }
505 }
506
507 #[rstest::rstest]
508 #[case::completed(false)]
509 #[case::panicked(true)]
510 #[tokio::test]
511 async fn test_inflight_job_balances_count_on_exit(#[case] panics: bool) {
512 let gauge = IntGauge::new("test_inflight_compaction", "isolated inflight gauge").unwrap();
513 let (resume, paused) = oneshot::channel();
514 let job = InflightCompaction::run(
515 async move {
516 paused.await.unwrap();
517 if panics {
518 panic!("inflight job panic fixture");
519 }
520 },
521 gauge.clone(),
522 );
523 let mut job = Box::pin(std::panic::AssertUnwindSafe(job).catch_unwind());
524
525 assert_eq!(0, gauge.get());
526 assert!(futures::poll!(job.as_mut()).is_pending());
527 assert_eq!(1, gauge.get());
529 resume.send(()).unwrap();
530 let outcome = job.await;
531 if panics {
532 let payload = outcome.unwrap_err();
533 assert_eq!(
534 Some("inflight job panic fixture"),
535 payload.downcast_ref::<&str>().copied()
536 );
537 } else {
538 assert!(outcome.is_ok());
539 }
540 assert_eq!(0, gauge.get());
541 }
542
543 #[rstest::rstest]
544 #[case::unpolled(false)]
545 #[case::started(true)]
546 #[tokio::test]
547 async fn test_inflight_job_drop_balances_started_count(#[case] started: bool) {
548 let gauge = IntGauge::new("test_inflight_compaction", "isolated inflight gauge").unwrap();
549 let mut job = Box::pin(InflightCompaction::run(
550 std::future::pending::<()>(),
551 gauge.clone(),
552 ));
553 assert_eq!(0, gauge.get());
554 if started {
555 assert!(futures::poll!(job.as_mut()).is_pending());
556 assert_eq!(1, gauge.get());
557 }
558 drop(job);
559 assert_eq!(0, gauge.get());
560 }
561
562 fn task_entry() -> CompactionTaskEntry {
563 use crate::compaction::scheduler::state::{CompactionTaskState, CompactionTrigger};
564
565 CompactionTaskState::new(
566 RegionId::from_u64(1),
567 &api::v1::region::compact_request::Options::Regular(Default::default()),
568 2,
569 None,
570 CompactionTrigger::Manual,
571 )
572 .snapshot(None)
573 }
574
575 #[test]
576 fn test_shared_units_observe_finish_and_release() {
577 let mut local = local_compaction();
578 let published = local.current_state();
579 let mut entry = task_entry();
580 for unit in local.units.values() {
581 assert!(
582 published
583 .units
584 .iter()
585 .any(|shared| Arc::ptr_eq(shared, unit))
586 );
587 }
588 assert!(local.units[&1].state.mark_commit_started());
589 local.finish(1, Ok(true));
590 local.finish(1, Ok(false));
591 published.inspect(&mut entry);
592 assert_eq!(entry.completed_units, Some(1));
593 assert_eq!((&*entry.state, &*entry.phase), ("running", "committing"));
594
595 local.release_unit(1);
596 local.release_unit(1);
597 local.release_unit(99);
598 published.inspect(&mut entry);
599 assert_eq!(entry.completed_units, Some(1));
600 assert_eq!((&*entry.state, &*entry.phase), ("queued", "waiting"));
601 assert!(Arc::ptr_eq(&published.units, &local.current_state().units));
602
603 local.finish(2, Err(Arc::new(error::CompactionCancelledSnafu.build())));
605 local.release_unit(2);
606 published.inspect(&mut entry);
607 assert_eq!(entry.total_units, Some(2));
608 assert_eq!(entry.completed_units, Some(2));
609 assert!(local.units.is_empty());
610 assert!(Arc::ptr_eq(&published.units, &local.current_state().units));
611 }
612
613 #[test]
614 fn test_shared_units_match_live_map_aggregation() {
615 #[derive(Debug, Clone, Copy)]
616 enum Phase {
617 Queued,
618 Running,
619 Committing,
620 Finished,
621 Released,
622 Cancelled,
623 CancelledReleased,
624 }
625 let phases = [
626 Phase::Queued,
627 Phase::Running,
628 Phase::Committing,
629 Phase::Finished,
630 Phase::Released,
631 Phase::Cancelled,
632 Phase::CancelledReleased,
633 ];
634 for first in phases {
635 for second in phases {
636 let mut local = local_compaction();
637 let published = local.current_state();
638 let mut completed = 0;
639 for (id, phase) in [(1, first), (2, second)] {
640 let state = &local.units[&id].state;
641 match phase {
642 Phase::Queued => {}
643 Phase::Running => state.mark_started(),
644 Phase::Committing | Phase::Finished | Phase::Released => {
645 assert!(state.mark_commit_started());
646 }
647 Phase::Cancelled | Phase::CancelledReleased => {
648 assert_eq!(state.request_cancel(), RequestCancelResult::CancelIssued);
649 }
650 }
651 if matches!(
652 phase,
653 Phase::Finished | Phase::Released | Phase::CancelledReleased
654 ) {
655 local.finish(id, Ok(false));
656 completed += 1;
657 }
658 if matches!(phase, Phase::Released | Phase::CancelledReleased) {
659 local.release_unit(id);
660 }
661 }
662 let cancelling = local.units.values().any(|unit| unit.state.is_cancelled());
664 let committing = local.units.values().any(|unit| unit.state.commit_started());
665 let running = committing || local.units.values().any(|unit| unit.state.started());
666 let expected_state = if cancelling {
667 "cancelling"
668 } else if running {
669 "running"
670 } else {
671 "queued"
672 };
673 let expected_phase = if committing {
674 "committing"
675 } else if running {
676 "executing"
677 } else {
678 "waiting"
679 };
680 let mut entry = task_entry();
681 published.inspect(&mut entry);
682 assert_eq!(
683 (
684 &*entry.state,
685 &*entry.phase,
686 entry.total_units,
687 entry.completed_units
688 ),
689 (expected_state, expected_phase, Some(2), Some(completed)),
690 "{first:?}, {second:?}",
691 );
692 }
693 }
694 }
695
696 #[test]
697 fn test_compaction_unit_apply_accepts_independent_executions() {
698 let mut local = local_compaction();
699 assert!(!local.can_apply(1));
700 assert!(!local.can_apply(2));
701 assert!(!local.can_apply(99));
702 for unit in local.units.values() {
703 assert!(unit.state.mark_commit_started());
704 }
705 assert!(local.can_apply(1) && local.can_apply(2));
706 local.finish(2, Ok(true));
707 assert!(local.can_apply(1));
708 assert!(!local.can_apply(2));
709 local.finish(2, Err(Arc::new(error::CompactionCancelledSnafu.build())));
710 local.finish(99, Err(Arc::new(error::CompactionCancelledSnafu.build())));
711 assert!(local.can_apply(1));
712 assert!(local.error.is_none());
713 assert!(local.made_progress);
714 local.finish(1, Ok(false));
715 assert!(!local.can_apply(1));
716 assert_eq!(
717 2,
718 local.units.len(),
719 "terminal units remain tracked until task exit"
720 );
721 }
722
723 #[test]
724 fn test_compaction_cancelling_survives_cancelled_sibling_release() {
725 let mut local = local_compaction();
726 let mut entry = task_entry();
727 let published = local.current_state();
729 published.inspect(&mut entry);
730 assert_eq!((&*entry.state, &*entry.phase), ("queued", "waiting"));
731 local.units[&1].state.mark_started();
732 published.inspect(&mut entry);
733 assert_eq!((&*entry.state, &*entry.phase), ("running", "executing"));
734 assert!(local.units[&1].state.mark_commit_started());
735 published.inspect(&mut entry);
736 assert_eq!((&*entry.state, &*entry.phase), ("running", "committing"));
737 assert_eq!(local.request_cancel(), RequestCancelResult::CancelIssued);
738 published.inspect(&mut entry);
739 assert_eq!((&*entry.state, &*entry.phase), ("cancelling", "committing"));
740
741 local.finish(2, Err(Arc::new(error::CompactionCancelledSnafu.build())));
743 local.release_unit(2);
744 assert!(!local.units[&1].state.is_cancelled());
745 local.inspect(&mut entry);
746 assert_eq!((&*entry.state, &*entry.phase), ("cancelling", "committing"));
747 assert_eq!(entry.completed_units, Some(1));
748 assert_eq!(entry.total_units, Some(2));
749
750 let mut too_late = local_compaction();
752 for unit in too_late.units.values() {
753 assert!(unit.state.mark_commit_started());
754 }
755 assert_eq!(
756 too_late.request_cancel(),
757 RequestCancelResult::TooLateToCancel
758 );
759 too_late.inspect(&mut entry);
760 assert_eq!((&*entry.state, &*entry.phase), ("running", "committing"));
761 }
762}