1mod local;
16mod planning;
17mod state;
18
19use std::collections::HashMap;
20use std::sync::Arc;
21
22use api::v1::region::compact_request;
23use common_base::Plugins;
24use common_memory_manager::OnExhaustedPolicy;
25use common_meta::key::SchemaMetadataManagerRef;
26use common_telemetry::{debug, error, info};
27use common_time::range::TimestampRange;
28pub(crate) use planning::CompactionPickFinished;
29pub use planning::CompactionRequest;
30pub(crate) use state::CompactionExecution;
31use state::{CompactionStatus, CompactionTrigger, DdlExecution, PendingCompaction};
32use store_api::storage::RegionId;
33use tokio::sync::mpsc::Sender;
34
35use crate::access_layer::AccessLayerRef;
36use crate::cache::CacheManagerRef;
37use crate::compaction::memory_manager::CompactionMemoryManager;
38use crate::compaction::task::MAX_PARALLEL_COMPACTION;
39use crate::config::MitoConfig;
40use crate::error::{
41 CompactionCancelledSnafu, Error, ManualCompactionAlreadyRunningSnafu, RegionClosedSnafu,
42 RegionDroppedSnafu, RegionTruncatedSnafu, Result,
43};
44use crate::region::version::VersionControlRef;
45use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState};
46use crate::request::{DdlRequest, OptionOutputTx, SenderDdlRequest, WorkerRequestWithTime};
47#[cfg(test)]
48use crate::schedule::RequestCancelResult;
49use crate::schedule::scheduler::SchedulerRef;
50use crate::worker::WorkerListener;
51
52pub(crate) struct CompactionScheduler {
54 scheduler: SchedulerRef,
55 region_status: HashMap<RegionId, CompactionStatus>,
57 request_sender: Sender<WorkerRequestWithTime>,
59 cache_manager: CacheManagerRef,
60 engine_config: Arc<MitoConfig>,
61 memory_manager: Arc<CompactionMemoryManager>,
62 memory_policy: OnExhaustedPolicy,
63 listener: WorkerListener,
64 plugins: Plugins,
66 next_plan_id: u64,
69}
70
71#[derive(Debug)]
73pub(crate) enum CompactionTransition {
74 NoAction,
76 AutomaticFollowupScheduled,
78 DdlReady(Vec<SenderDdlRequest>),
80}
81
82impl CompactionTransition {
83 fn from_pending_ddls(pending_ddls: Vec<SenderDdlRequest>) -> Self {
84 if pending_ddls.is_empty() {
85 Self::NoAction
86 } else {
87 Self::DdlReady(pending_ddls)
88 }
89 }
90
91 #[cfg(test)]
92 fn is_empty(&self) -> bool {
93 !matches!(self, Self::DdlReady(pending_ddls) if !pending_ddls.is_empty())
94 }
95
96 #[cfg(test)]
97 fn len(&self) -> usize {
98 match self {
99 Self::DdlReady(pending_ddls) => pending_ddls.len(),
100 Self::NoAction | Self::AutomaticFollowupScheduled => 0,
101 }
102 }
103}
104
105impl CompactionScheduler {
107 #[allow(clippy::too_many_arguments)]
119 pub(crate) fn new(
120 scheduler: SchedulerRef,
121 request_sender: Sender<WorkerRequestWithTime>,
122 cache_manager: CacheManagerRef,
123 engine_config: Arc<MitoConfig>,
124 listener: WorkerListener,
125 plugins: Plugins,
126 memory_manager: Arc<CompactionMemoryManager>,
127 memory_policy: OnExhaustedPolicy,
128 ) -> Self {
129 Self {
130 scheduler,
131 region_status: HashMap::new(),
132 request_sender,
133 cache_manager,
134 engine_config,
135 memory_manager,
136 memory_policy,
137 listener,
138 plugins,
139 next_plan_id: 0,
140 }
141 }
142
143 pub(crate) fn register_region(
145 &mut self,
146 version_control: &VersionControlRef,
147 access_layer: &AccessLayerRef,
148 ) {
149 let region_id = version_control.region_id();
150 if let Some(status) = self.region_status.get_mut(®ion_id) {
151 if Arc::ptr_eq(&status.version_control, version_control) {
152 return;
153 }
154 status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
155 }
156 self.region_status.insert(
157 region_id,
158 CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()),
159 );
160 }
161
162 #[cfg(test)]
164 fn is_compacting(&self, region_id: RegionId) -> bool {
165 self.region_status
166 .get(®ion_id)
167 .is_some_and(CompactionStatus::is_compacting)
168 }
169
170 pub(crate) fn schedule_automatic_compaction(
185 &mut self,
186 compact_options: compact_request::Options,
187 version_control: &VersionControlRef,
188 access_layer: &AccessLayerRef,
189 manifest_ctx: &ManifestContextRef,
190 schema_metadata_manager: SchemaMetadataManagerRef,
191 ) -> Result<bool> {
192 self.schedule_compaction(
193 CompactionTrigger::Automatic,
194 compact_options,
195 version_control,
196 access_layer,
197 OptionOutputTx::none(),
198 manifest_ctx,
199 schema_metadata_manager,
200 1, None,
202 )
203 }
204
205 #[allow(clippy::too_many_arguments)]
223 pub(crate) fn schedule_manual_compaction(
224 &mut self,
225 compact_options: compact_request::Options,
226 version_control: &VersionControlRef,
227 access_layer: &AccessLayerRef,
228 waiter: OptionOutputTx,
229 manifest_ctx: &ManifestContextRef,
230 schema_metadata_manager: SchemaMetadataManagerRef,
231 max_parallelism: usize,
232 time_range: Option<TimestampRange>,
233 ) -> Result<bool> {
234 self.schedule_compaction(
235 CompactionTrigger::Manual,
236 compact_options,
237 version_control,
238 access_layer,
239 waiter,
240 manifest_ctx,
241 schema_metadata_manager,
242 max_parallelism,
243 time_range,
244 )
245 }
246
247 pub(crate) async fn handle_compaction_pick_finished(
262 &mut self,
263 finished: CompactionPickFinished,
264 manifest_ctx: &ManifestContextRef,
265 schema_metadata_manager: SchemaMetadataManagerRef,
266 ) -> CompactionTransition {
267 self.handle_compaction_pick_finished_inner(finished, manifest_ctx, schema_metadata_manager)
268 .await
269 }
270
271 pub(crate) fn is_current_execution(
282 &self,
283 region_id: RegionId,
284 execution: &CompactionExecution,
285 ) -> bool {
286 self.region_status
287 .get(®ion_id)
288 .is_some_and(|status| status.matches_execution(execution))
289 }
290
291 pub(crate) async fn on_execution_finished(
309 &mut self,
310 region_id: RegionId,
311 execution: &CompactionExecution,
312 manifest_ctx: &ManifestContextRef,
313 schema_metadata_manager: SchemaMetadataManagerRef,
314 made_progress: bool,
315 ) -> CompactionTransition {
316 if !self.is_current_execution(region_id, execution) {
317 return CompactionTransition::NoAction;
318 }
319 self.on_compaction_finished(
320 region_id,
321 manifest_ctx,
322 schema_metadata_manager,
323 made_progress,
324 )
325 .await
326 }
327
328 pub(crate) fn on_execution_failed(
340 &mut self,
341 region_id: RegionId,
342 execution: &CompactionExecution,
343 err: Arc<Error>,
344 ) {
345 if !self.is_current_execution(region_id, execution) {
346 return;
347 }
348 self.on_compaction_failed(region_id, err);
349 }
350
351 pub(crate) fn on_region_dropped(&mut self, region_id: RegionId) {
361 self.remove_region(
362 region_id,
363 Arc::new(RegionDroppedSnafu { region_id }.build()),
364 );
365 }
366
367 pub(crate) fn on_region_closed(&mut self, region_id: RegionId) {
377 self.remove_region(region_id, Arc::new(RegionClosedSnafu { region_id }.build()));
378 }
379
380 pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) {
390 self.fail_compaction(
391 region_id,
392 Arc::new(RegionTruncatedSnafu { region_id }.build()),
393 );
394 }
395
396 pub(crate) fn try_cancel_and_add_ddl<T>(
412 &mut self,
413 region_id: RegionId,
414 sender: OptionOutputTx,
415 request: T,
416 into_ddl_request: impl FnOnce(T) -> DdlRequest,
417 ) -> std::result::Result<(), (OptionOutputTx, T)> {
418 let Some(status) = self.region_status.get_mut(®ion_id) else {
419 return Err((sender, request));
420 };
421 if !status.is_compacting() {
422 return Err((sender, request));
423 }
424 status.request_cancel();
425
426 let request = SenderDdlRequest {
427 region_id,
428 sender,
429 request: into_ddl_request(request),
430 };
431 debug!(
432 "Added pending DDL request for region: {}, ddl: {:?}",
433 request.region_id, request.request
434 );
435 status.pending_ddl_requests.push(request);
436 Ok(())
437 }
438}
439
440impl CompactionScheduler {
442 fn next_plan_id(counter: &mut u64) -> u64 {
447 let plan_id = *counter;
448 *counter = counter.wrapping_add(1);
449 plan_id
450 }
451
452 #[allow(clippy::too_many_arguments)]
453 fn schedule_compaction(
454 &mut self,
455 trigger: CompactionTrigger,
456 compact_options: compact_request::Options,
457 version_control: &VersionControlRef,
458 access_layer: &AccessLayerRef,
459 waiter: OptionOutputTx,
460 manifest_ctx: &ManifestContextRef,
461 schema_metadata_manager: SchemaMetadataManagerRef,
462 max_parallelism: usize,
463 time_range: Option<TimestampRange>,
464 ) -> Result<bool> {
465 let region_id = version_control.region_id();
466 let status = self.region_status.entry(region_id).or_insert_with(|| {
467 CompactionStatus::new(region_id, version_control.clone(), access_layer.clone())
468 });
469 if status.has_ddl() {
470 info!(
471 "Region {} has pending DDL requests, ignoring compaction: queued={}, executing={}",
472 region_id,
473 status.pending_ddl_requests.len(),
474 status.executing_ddl.is_some(),
475 );
476 waiter.send(CompactionCancelledSnafu.fail());
477 return Ok(false);
478 }
479 let current_state = manifest_ctx.current_state();
480 if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) {
481 info!(
482 "Skipping compaction for region {} in staging mode, options: {:?}",
483 region_id, compact_options
484 );
485 waiter.send(Ok(0));
486 return Ok(false);
487 }
488
489 if status.is_compacting() {
490 match trigger {
491 CompactionTrigger::Automatic => status.mark_automatic_trigger(),
492 CompactionTrigger::Manual if status.is_manual_compaction() => {
493 waiter.send(ManualCompactionAlreadyRunningSnafu { region_id }.fail());
494 info!(
495 "Region {} already has a manually triggered compaction running",
496 region_id
497 );
498 }
499 CompactionTrigger::Manual => {
500 status.set_pending_request(PendingCompaction {
501 options: compact_options,
502 waiter,
503 max_parallelism,
504 time_range,
505 });
506 info!(
507 "Region {} is running an automatic compaction; manual compaction will be re-scheduled",
508 region_id
509 );
510 }
511 }
512 return Ok(false);
513 }
514
515 let plan_id = Self::next_plan_id(&mut self.next_plan_id);
517 status.start_picking_with_trigger(plan_id, trigger);
518 let request = status.new_compaction_request(
519 self.request_sender.clone(),
520 self.engine_config.clone(),
521 self.cache_manager.clone(),
522 manifest_ctx,
523 self.listener.clone(),
524 schema_metadata_manager,
525 max_parallelism,
526 );
527 status.merge_waiter(waiter);
528 self.dispatch_compaction_planning(plan_id, request, compact_options, time_range);
529 self.listener.on_compaction_scheduled(region_id);
530 Ok(true)
531 }
532
533 fn handle_pending_compaction_request(
537 &mut self,
538 region_id: RegionId,
539 manifest_ctx: &ManifestContextRef,
540 schema_metadata_manager: SchemaMetadataManagerRef,
541 ) -> bool {
542 let Some(status) = self.region_status.get_mut(®ion_id) else {
543 return true;
544 };
545
546 let Some(pending_request) = std::mem::take(&mut status.pending_request) else {
549 return false;
550 };
551
552 let PendingCompaction {
553 options,
554 waiter,
555 max_parallelism,
556 time_range,
557 } = pending_request;
558
559 let request = status.new_compaction_request(
560 self.request_sender.clone(),
561 self.engine_config.clone(),
562 self.cache_manager.clone(),
563 manifest_ctx,
564 self.listener.clone(),
565 schema_metadata_manager,
566 max_parallelism,
567 );
568 status.merge_waiter(waiter);
569 let plan_id = Self::next_plan_id(&mut self.next_plan_id);
573 status.start_picking_with_trigger(plan_id, CompactionTrigger::Manual);
574 self.dispatch_compaction_planning(plan_id, request, options, time_range);
575 debug!(
576 "Successfully scheduled manual compaction planning for region id: {}",
577 region_id
578 );
579 true
580 }
581
582 async fn on_compaction_finished(
584 &mut self,
585 region_id: RegionId,
586 manifest_ctx: &ManifestContextRef,
587 schema_metadata_manager: SchemaMetadataManagerRef,
588 made_progress: bool,
589 ) -> CompactionTransition {
590 if !self.region_status.contains_key(®ion_id) {
591 return CompactionTransition::NoAction;
592 }
593
594 if let Some(ddls) = self.take_ready_ddls(region_id) {
595 return CompactionTransition::DdlReady(ddls);
596 }
597
598 if self.handle_pending_compaction_request(
599 region_id,
600 manifest_ctx,
601 schema_metadata_manager.clone(),
602 ) {
603 return CompactionTransition::NoAction;
604 }
605
606 let Some(status) = self.region_status.get_mut(®ion_id) else {
607 return CompactionTransition::NoAction;
608 };
609 for waiter in status.take_waiters() {
610 waiter.send(Ok(0));
611 }
612
613 let should_continue = status.active.reset_automatic_followup() || made_progress;
617 if should_continue
618 && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager)
619 {
620 return CompactionTransition::AutomaticFollowupScheduled;
621 }
622 if let Some(status) = self.region_status.get_mut(®ion_id) {
623 status.become_idle();
624 }
625 CompactionTransition::NoAction
626 }
627
628 fn schedule_automatic_followup(
629 &mut self,
630 region_id: RegionId,
631 manifest_ctx: &ManifestContextRef,
632 schema_metadata_manager: SchemaMetadataManagerRef,
633 ) -> bool {
634 let Some(status) = self.region_status.get_mut(®ion_id) else {
635 return false;
636 };
637 let request = status.new_compaction_request(
639 self.request_sender.clone(),
640 self.engine_config.clone(),
641 self.cache_manager.clone(),
642 manifest_ctx,
643 self.listener.clone(),
644 schema_metadata_manager,
645 MAX_PARALLEL_COMPACTION,
646 );
647 let plan_id = Self::next_plan_id(&mut self.next_plan_id);
651 status.start_regular_picking(plan_id);
652 self.dispatch_compaction_planning(
653 plan_id,
654 request,
655 compact_request::Options::Regular(Default::default()),
656 None,
657 );
658 debug!(
659 "Successfully scheduled next compaction planning for region id: {}",
660 region_id
661 );
662 true
663 }
664
665 #[cfg(test)]
667 async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
668 self.finish_compaction_on_cancel(region_id)
669 }
670
671 fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc<Error>) {
673 error!(err; "Region {} failed to compact, cancel all pending tasks", region_id);
674 self.fail_compaction(region_id, err);
675 }
676
677 #[cfg(test)]
678 fn add_ddl_request_to_pending(&mut self, request: SenderDdlRequest) {
679 self.region_status
680 .get_mut(&request.region_id)
681 .unwrap()
682 .pending_ddl_requests
683 .push(request);
684 }
685
686 #[cfg(test)]
687 fn has_pending_ddls(&self, region_id: RegionId) -> bool {
688 let has_pending = self
689 .region_status
690 .get(®ion_id)
691 .map(|status| !status.pending_ddl_requests.is_empty())
692 .unwrap_or(false);
693 debug!(
694 "Checked pending DDL requests for region: {}, has_pending: {}",
695 region_id, has_pending
696 );
697 has_pending
698 }
699
700 #[cfg(test)]
701 fn request_cancel(&mut self, region_id: RegionId) -> RequestCancelResult {
702 self.region_status
703 .get_mut(®ion_id)
704 .unwrap()
705 .request_cancel()
706 }
707
708 fn remove_region(&mut self, region_id: RegionId, err: Arc<Error>) {
710 if let Some(mut status) = self.region_status.remove(®ion_id) {
711 status.on_failure(err);
712 }
713 }
714
715 fn fail_compaction(&mut self, region_id: RegionId, err: Arc<Error>) {
717 if let Some(status) = self.region_status.get_mut(®ion_id) {
718 status.on_failure(err);
719 }
720 }
721
722 fn finish_compaction_on_cancel(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
724 let Some(status) = self.region_status.get_mut(®ion_id) else {
725 return Vec::new();
726 };
727
728 let ddls = status.on_cancel();
729 self.fence_ddl_replies(region_id, ddls)
730 }
731
732 fn take_ready_ddls(&mut self, region_id: RegionId) -> Option<Vec<SenderDdlRequest>> {
734 if self
735 .region_status
736 .get(®ion_id)?
737 .pending_ddl_requests
738 .is_empty()
739 {
740 return None;
741 }
742 Some(self.finish_compaction_on_cancel(region_id))
743 }
744
745 fn fence_ddl_replies(
754 &mut self,
755 region_id: RegionId,
756 mut ddls: Vec<SenderDdlRequest>,
757 ) -> Vec<SenderDdlRequest> {
758 if ddls.is_empty() {
759 return ddls;
760 }
761 let generation = Self::next_plan_id(&mut self.next_plan_id);
762 if let Some(status) = self.region_status.get_mut(®ion_id) {
763 status.executing_ddl = Some(DdlExecution::new(generation, ddls.len()));
765 }
766
767 for (request_index, ddl) in ddls.iter_mut().enumerate() {
768 let (tx, rx) = tokio::sync::oneshot::channel();
772 let reply = std::mem::replace(&mut ddl.sender, tx.into());
773 let sender = self.request_sender.clone();
774 common_runtime::spawn_global(async move {
775 use snafu::ResultExt;
776 let result = rx
777 .await
778 .context(crate::error::RecvSnafu)
779 .and_then(|result| result);
780 let _ = sender
791 .send(WorkerRequestWithTime::new(
792 crate::request::WorkerRequest::Background {
793 region_id,
794 notify: crate::request::BackgroundNotify::CompactionDdlComplete {
795 generation,
796 request_index,
797 },
798 },
799 ))
800 .await;
801 reply.send(result);
802 });
803 }
804 ddls
805 }
806
807 pub(crate) fn on_ddl_complete(
809 &mut self,
810 region_id: RegionId,
811 generation: u64,
812 request_index: usize,
813 ) {
814 let Some(status) = self.region_status.get_mut(®ion_id) else {
815 return;
816 };
817 if let Some(ddl) = &mut status.executing_ddl
821 && ddl.complete(generation, request_index)
822 {
823 status.executing_ddl = None;
824 }
825 }
826}
827
828impl Drop for CompactionScheduler {
829 fn drop(&mut self) {
830 for (region_id, mut status) in self.region_status.drain() {
831 status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
833 }
834 }
835}
836
837#[cfg(test)]
838#[path = "scheduler_test.rs"]
839mod tests;