1mod local;
16mod planning;
17mod state;
18
19use std::collections::HashMap;
20use std::sync::{Arc, LazyLock};
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, warn};
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, UnsupportedOperationSnafu,
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
52const DISABLE_COMPACTION_ENV: &str = "GREPTIME_DISABLE_COMPACTION";
55
56static DISABLE_COMPACTION: LazyLock<bool> = LazyLock::new(|| {
59 let value = std::env::var(DISABLE_COMPACTION_ENV).ok();
60 let disabled = parse_disable_compaction(value.as_deref());
61 if disabled {
62 warn!(
63 "SST compaction is disabled by {DISABLE_COMPACTION_ENV}; restart without it to resume"
64 );
65 }
66 disabled
67});
68
69fn parse_disable_compaction(value: Option<&str>) -> bool {
70 match value {
71 None | Some("false") => false,
72 Some("true") => true,
73 Some(_) => {
74 warn!(
75 "Invalid {DISABLE_COMPACTION_ENV}: expected true or false; compaction remains enabled"
76 );
77 false
78 }
79 }
80}
81
82pub(crate) struct CompactionScheduler {
84 disable_compaction: bool,
88 scheduler: SchedulerRef,
89 region_status: HashMap<RegionId, CompactionStatus>,
91 request_sender: Sender<WorkerRequestWithTime>,
93 cache_manager: CacheManagerRef,
94 engine_config: Arc<MitoConfig>,
95 memory_manager: Arc<CompactionMemoryManager>,
96 memory_policy: OnExhaustedPolicy,
97 listener: WorkerListener,
98 plugins: Plugins,
100 next_plan_id: u64,
103}
104
105#[derive(Debug)]
107pub(crate) enum CompactionTransition {
108 NoAction,
110 AutomaticFollowupScheduled,
112 DdlReady(Vec<SenderDdlRequest>),
114}
115
116impl CompactionTransition {
117 fn from_pending_ddls(pending_ddls: Vec<SenderDdlRequest>) -> Self {
118 if pending_ddls.is_empty() {
119 Self::NoAction
120 } else {
121 Self::DdlReady(pending_ddls)
122 }
123 }
124
125 #[cfg(test)]
126 fn is_empty(&self) -> bool {
127 !matches!(self, Self::DdlReady(pending_ddls) if !pending_ddls.is_empty())
128 }
129
130 #[cfg(test)]
131 fn len(&self) -> usize {
132 match self {
133 Self::DdlReady(pending_ddls) => pending_ddls.len(),
134 Self::NoAction | Self::AutomaticFollowupScheduled => 0,
135 }
136 }
137}
138
139impl CompactionScheduler {
141 #[allow(clippy::too_many_arguments)]
153 pub(crate) fn new(
154 scheduler: SchedulerRef,
155 request_sender: Sender<WorkerRequestWithTime>,
156 cache_manager: CacheManagerRef,
157 engine_config: Arc<MitoConfig>,
158 listener: WorkerListener,
159 plugins: Plugins,
160 memory_manager: Arc<CompactionMemoryManager>,
161 memory_policy: OnExhaustedPolicy,
162 ) -> Self {
163 Self {
164 disable_compaction: *DISABLE_COMPACTION,
165 scheduler,
166 region_status: HashMap::new(),
167 request_sender,
168 cache_manager,
169 engine_config,
170 memory_manager,
171 memory_policy,
172 listener,
173 plugins,
174 next_plan_id: 0,
175 }
176 }
177
178 pub(crate) fn register_region(
180 &mut self,
181 version_control: &VersionControlRef,
182 access_layer: &AccessLayerRef,
183 ) {
184 let region_id = version_control.region_id();
185 if let Some(status) = self.region_status.get_mut(®ion_id) {
186 if Arc::ptr_eq(&status.version_control, version_control) {
187 return;
188 }
189 status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
190 }
191 self.region_status.insert(
192 region_id,
193 CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()),
194 );
195 }
196
197 #[cfg(test)]
199 fn is_compacting(&self, region_id: RegionId) -> bool {
200 self.region_status
201 .get(®ion_id)
202 .is_some_and(CompactionStatus::is_compacting)
203 }
204
205 pub(crate) fn schedule_automatic_compaction(
220 &mut self,
221 compact_options: compact_request::Options,
222 version_control: &VersionControlRef,
223 access_layer: &AccessLayerRef,
224 manifest_ctx: &ManifestContextRef,
225 schema_metadata_manager: SchemaMetadataManagerRef,
226 ) -> Result<bool> {
227 self.schedule_compaction(
228 CompactionTrigger::Automatic,
229 compact_options,
230 version_control,
231 access_layer,
232 OptionOutputTx::none(),
233 manifest_ctx,
234 schema_metadata_manager,
235 1, None,
237 )
238 }
239
240 #[allow(clippy::too_many_arguments)]
258 pub(crate) fn schedule_manual_compaction(
259 &mut self,
260 compact_options: compact_request::Options,
261 version_control: &VersionControlRef,
262 access_layer: &AccessLayerRef,
263 waiter: OptionOutputTx,
264 manifest_ctx: &ManifestContextRef,
265 schema_metadata_manager: SchemaMetadataManagerRef,
266 max_parallelism: usize,
267 time_range: Option<TimestampRange>,
268 ) -> Result<bool> {
269 self.schedule_compaction(
270 CompactionTrigger::Manual,
271 compact_options,
272 version_control,
273 access_layer,
274 waiter,
275 manifest_ctx,
276 schema_metadata_manager,
277 max_parallelism,
278 time_range,
279 )
280 }
281
282 pub(crate) async fn handle_compaction_pick_finished(
297 &mut self,
298 finished: CompactionPickFinished,
299 manifest_ctx: &ManifestContextRef,
300 schema_metadata_manager: SchemaMetadataManagerRef,
301 ) -> CompactionTransition {
302 self.handle_compaction_pick_finished_inner(finished, manifest_ctx, schema_metadata_manager)
303 .await
304 }
305
306 pub(crate) fn is_current_execution(
317 &self,
318 region_id: RegionId,
319 execution: &CompactionExecution,
320 ) -> bool {
321 self.region_status
322 .get(®ion_id)
323 .is_some_and(|status| status.matches_execution(execution))
324 }
325
326 pub(crate) async fn on_execution_finished(
344 &mut self,
345 region_id: RegionId,
346 execution: &CompactionExecution,
347 manifest_ctx: &ManifestContextRef,
348 schema_metadata_manager: SchemaMetadataManagerRef,
349 made_progress: bool,
350 ) -> CompactionTransition {
351 if !self.is_current_execution(region_id, execution) {
352 return CompactionTransition::NoAction;
353 }
354 self.on_compaction_finished(
355 region_id,
356 manifest_ctx,
357 schema_metadata_manager,
358 made_progress,
359 )
360 .await
361 }
362
363 pub(crate) fn on_execution_failed(
375 &mut self,
376 region_id: RegionId,
377 execution: &CompactionExecution,
378 err: Arc<Error>,
379 ) {
380 if !self.is_current_execution(region_id, execution) {
381 return;
382 }
383 self.on_compaction_failed(region_id, err);
384 }
385
386 pub(crate) fn on_region_dropped(&mut self, region_id: RegionId) {
396 self.remove_region(
397 region_id,
398 Arc::new(RegionDroppedSnafu { region_id }.build()),
399 );
400 }
401
402 pub(crate) fn on_region_closed(&mut self, region_id: RegionId) {
412 self.remove_region(region_id, Arc::new(RegionClosedSnafu { region_id }.build()));
413 }
414
415 pub(crate) fn on_region_truncated(&mut self, region_id: RegionId) {
425 self.fail_compaction(
426 region_id,
427 Arc::new(RegionTruncatedSnafu { region_id }.build()),
428 );
429 }
430
431 pub(crate) fn try_cancel_and_add_ddl<T>(
447 &mut self,
448 region_id: RegionId,
449 sender: OptionOutputTx,
450 request: T,
451 into_ddl_request: impl FnOnce(T) -> DdlRequest,
452 ) -> std::result::Result<(), (OptionOutputTx, T)> {
453 let Some(status) = self.region_status.get_mut(®ion_id) else {
454 return Err((sender, request));
455 };
456 if !status.is_compacting() {
457 return Err((sender, request));
458 }
459 status.request_cancel();
460
461 let request = SenderDdlRequest {
462 region_id,
463 sender,
464 request: into_ddl_request(request),
465 };
466 debug!(
467 "Added pending DDL request for region: {}, ddl: {:?}",
468 request.region_id, request.request
469 );
470 status.pending_ddl_requests.push(request);
471 Ok(())
472 }
473}
474
475impl CompactionScheduler {
477 fn next_plan_id(counter: &mut u64) -> u64 {
482 let plan_id = *counter;
483 *counter = counter.wrapping_add(1);
484 plan_id
485 }
486
487 #[allow(clippy::too_many_arguments)]
488 fn schedule_compaction(
489 &mut self,
490 trigger: CompactionTrigger,
491 compact_options: compact_request::Options,
492 version_control: &VersionControlRef,
493 access_layer: &AccessLayerRef,
494 waiter: OptionOutputTx,
495 manifest_ctx: &ManifestContextRef,
496 schema_metadata_manager: SchemaMetadataManagerRef,
497 max_parallelism: usize,
498 time_range: Option<TimestampRange>,
499 ) -> Result<bool> {
500 if self.disable_compaction {
501 if matches!(trigger, CompactionTrigger::Manual) {
502 waiter.send(
503 UnsupportedOperationSnafu {
504 err_msg: format!("SST compaction is disabled by {DISABLE_COMPACTION_ENV}"),
505 }
506 .fail(),
507 );
508 }
509 return Ok(false);
510 }
511 let region_id = version_control.region_id();
512 let status = self.region_status.entry(region_id).or_insert_with(|| {
513 CompactionStatus::new(region_id, version_control.clone(), access_layer.clone())
514 });
515 if status.has_ddl() {
516 info!(
517 "Region {} has pending DDL requests, ignoring compaction: queued={}, executing={}",
518 region_id,
519 status.pending_ddl_requests.len(),
520 status.executing_ddl.is_some(),
521 );
522 waiter.send(CompactionCancelledSnafu.fail());
523 return Ok(false);
524 }
525 let current_state = manifest_ctx.current_state();
526 if current_state == RegionRoleState::Leader(RegionLeaderState::Staging) {
527 info!(
528 "Skipping compaction for region {} in staging mode, options: {:?}",
529 region_id, compact_options
530 );
531 waiter.send(Ok(0));
532 return Ok(false);
533 }
534
535 if status.is_compacting() {
536 match trigger {
537 CompactionTrigger::Automatic => status.mark_automatic_trigger(),
538 CompactionTrigger::Manual if status.is_manual_compaction() => {
539 waiter.send(ManualCompactionAlreadyRunningSnafu { region_id }.fail());
540 info!(
541 "Region {} already has a manually triggered compaction running",
542 region_id
543 );
544 }
545 CompactionTrigger::Manual => {
546 status.set_pending_request(PendingCompaction {
547 options: compact_options,
548 waiter,
549 max_parallelism,
550 time_range,
551 });
552 info!(
553 "Region {} is running an automatic compaction; manual compaction will be re-scheduled",
554 region_id
555 );
556 }
557 }
558 return Ok(false);
559 }
560
561 let plan_id = Self::next_plan_id(&mut self.next_plan_id);
563 status.start_picking_with_trigger(plan_id, trigger);
564 let request = status.new_compaction_request(
565 self.request_sender.clone(),
566 self.engine_config.clone(),
567 self.cache_manager.clone(),
568 manifest_ctx,
569 self.listener.clone(),
570 schema_metadata_manager,
571 max_parallelism,
572 );
573 status.merge_waiter(waiter);
574 self.dispatch_compaction_planning(plan_id, request, compact_options, time_range, trigger);
575 self.listener.on_compaction_scheduled(region_id);
576 Ok(true)
577 }
578
579 fn handle_pending_compaction_request(
583 &mut self,
584 region_id: RegionId,
585 manifest_ctx: &ManifestContextRef,
586 schema_metadata_manager: SchemaMetadataManagerRef,
587 ) -> bool {
588 let Some(status) = self.region_status.get_mut(®ion_id) else {
589 return true;
590 };
591
592 let Some(pending_request) = std::mem::take(&mut status.pending_request) else {
595 return false;
596 };
597
598 let PendingCompaction {
599 options,
600 waiter,
601 max_parallelism,
602 time_range,
603 } = pending_request;
604
605 let request = status.new_compaction_request(
606 self.request_sender.clone(),
607 self.engine_config.clone(),
608 self.cache_manager.clone(),
609 manifest_ctx,
610 self.listener.clone(),
611 schema_metadata_manager,
612 max_parallelism,
613 );
614 status.merge_waiter(waiter);
615 let plan_id = Self::next_plan_id(&mut self.next_plan_id);
619 status.start_picking_with_trigger(plan_id, CompactionTrigger::Manual);
620 self.dispatch_compaction_planning(
621 plan_id,
622 request,
623 options,
624 time_range,
625 CompactionTrigger::Manual,
626 );
627 debug!(
628 "Successfully scheduled manual compaction planning for region id: {}",
629 region_id
630 );
631 true
632 }
633
634 async fn on_compaction_finished(
636 &mut self,
637 region_id: RegionId,
638 manifest_ctx: &ManifestContextRef,
639 schema_metadata_manager: SchemaMetadataManagerRef,
640 made_progress: bool,
641 ) -> CompactionTransition {
642 if !self.region_status.contains_key(®ion_id) {
643 return CompactionTransition::NoAction;
644 }
645
646 if let Some(ddls) = self.take_ready_ddls(region_id) {
647 return CompactionTransition::DdlReady(ddls);
648 }
649
650 if self.handle_pending_compaction_request(
651 region_id,
652 manifest_ctx,
653 schema_metadata_manager.clone(),
654 ) {
655 return CompactionTransition::NoAction;
656 }
657
658 let Some(status) = self.region_status.get_mut(®ion_id) else {
659 return CompactionTransition::NoAction;
660 };
661 for waiter in status.take_waiters() {
662 waiter.send(Ok(0));
663 }
664
665 let should_continue = status.active.reset_automatic_followup() || made_progress;
669 if should_continue
670 && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager)
671 {
672 return CompactionTransition::AutomaticFollowupScheduled;
673 }
674 if let Some(status) = self.region_status.get_mut(®ion_id) {
675 status.become_idle();
676 }
677 CompactionTransition::NoAction
678 }
679
680 fn schedule_automatic_followup(
681 &mut self,
682 region_id: RegionId,
683 manifest_ctx: &ManifestContextRef,
684 schema_metadata_manager: SchemaMetadataManagerRef,
685 ) -> bool {
686 if self.disable_compaction {
687 return false;
688 }
689 let Some(status) = self.region_status.get_mut(®ion_id) else {
690 return false;
691 };
692 let request = status.new_compaction_request(
694 self.request_sender.clone(),
695 self.engine_config.clone(),
696 self.cache_manager.clone(),
697 manifest_ctx,
698 self.listener.clone(),
699 schema_metadata_manager,
700 MAX_PARALLEL_COMPACTION,
701 );
702 let plan_id = Self::next_plan_id(&mut self.next_plan_id);
706 status.start_regular_picking(plan_id);
707 self.dispatch_compaction_planning(
708 plan_id,
709 request,
710 compact_request::Options::Regular(Default::default()),
711 None,
712 CompactionTrigger::Automatic,
713 );
714 debug!(
715 "Successfully scheduled next compaction planning for region id: {}",
716 region_id
717 );
718 true
719 }
720
721 #[cfg(test)]
723 async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
724 self.finish_compaction_on_cancel(region_id)
725 }
726
727 fn on_compaction_failed(&mut self, region_id: RegionId, err: Arc<Error>) {
729 error!(err; "Region {} failed to compact, cancel all pending tasks", region_id);
730 self.fail_compaction(region_id, err);
731 }
732
733 #[cfg(test)]
734 fn add_ddl_request_to_pending(&mut self, request: SenderDdlRequest) {
735 self.region_status
736 .get_mut(&request.region_id)
737 .unwrap()
738 .pending_ddl_requests
739 .push(request);
740 }
741
742 #[cfg(test)]
743 fn has_pending_ddls(&self, region_id: RegionId) -> bool {
744 let has_pending = self
745 .region_status
746 .get(®ion_id)
747 .map(|status| !status.pending_ddl_requests.is_empty())
748 .unwrap_or(false);
749 debug!(
750 "Checked pending DDL requests for region: {}, has_pending: {}",
751 region_id, has_pending
752 );
753 has_pending
754 }
755
756 #[cfg(test)]
757 fn request_cancel(&mut self, region_id: RegionId) -> RequestCancelResult {
758 self.region_status
759 .get_mut(®ion_id)
760 .unwrap()
761 .request_cancel()
762 }
763
764 fn remove_region(&mut self, region_id: RegionId, err: Arc<Error>) {
766 if let Some(mut status) = self.region_status.remove(®ion_id) {
767 status.on_failure(err);
768 }
769 }
770
771 fn fail_compaction(&mut self, region_id: RegionId, err: Arc<Error>) {
773 if let Some(status) = self.region_status.get_mut(®ion_id) {
774 status.on_failure(err);
775 }
776 }
777
778 fn finish_compaction_on_cancel(&mut self, region_id: RegionId) -> Vec<SenderDdlRequest> {
780 let Some(status) = self.region_status.get_mut(®ion_id) else {
781 return Vec::new();
782 };
783
784 let ddls = status.on_cancel();
785 self.fence_ddl_replies(region_id, ddls)
786 }
787
788 fn take_ready_ddls(&mut self, region_id: RegionId) -> Option<Vec<SenderDdlRequest>> {
790 if self
791 .region_status
792 .get(®ion_id)?
793 .pending_ddl_requests
794 .is_empty()
795 {
796 return None;
797 }
798 Some(self.finish_compaction_on_cancel(region_id))
799 }
800
801 fn fence_ddl_replies(
810 &mut self,
811 region_id: RegionId,
812 mut ddls: Vec<SenderDdlRequest>,
813 ) -> Vec<SenderDdlRequest> {
814 if ddls.is_empty() {
815 return ddls;
816 }
817 let generation = Self::next_plan_id(&mut self.next_plan_id);
818 if let Some(status) = self.region_status.get_mut(®ion_id) {
819 status.executing_ddl = Some(DdlExecution::new(generation, ddls.len()));
821 }
822
823 for (request_index, ddl) in ddls.iter_mut().enumerate() {
824 let (tx, rx) = tokio::sync::oneshot::channel();
828 let reply = std::mem::replace(&mut ddl.sender, tx.into());
829 let sender = self.request_sender.clone();
830 common_runtime::spawn_global(async move {
831 use snafu::ResultExt;
832 let result = rx
833 .await
834 .context(crate::error::RecvSnafu)
835 .and_then(|result| result);
836 let _ = sender
847 .send(WorkerRequestWithTime::new(
848 crate::request::WorkerRequest::Background {
849 region_id,
850 notify: crate::request::BackgroundNotify::CompactionDdlComplete {
851 generation,
852 request_index,
853 },
854 },
855 ))
856 .await;
857 reply.send(result);
858 });
859 }
860 ddls
861 }
862
863 pub(crate) fn on_ddl_complete(
865 &mut self,
866 region_id: RegionId,
867 generation: u64,
868 request_index: usize,
869 ) {
870 let Some(status) = self.region_status.get_mut(®ion_id) else {
871 return;
872 };
873 if let Some(ddl) = &mut status.executing_ddl
877 && ddl.complete(generation, request_index)
878 {
879 status.executing_ddl = None;
880 }
881 }
882}
883
884impl Drop for CompactionScheduler {
885 fn drop(&mut self) {
886 for (region_id, mut status) in self.region_status.drain() {
887 status.on_failure(Arc::new(RegionClosedSnafu { region_id }.build()));
889 }
890 }
891}
892
893#[cfg(test)]
894#[path = "scheduler_test.rs"]
895mod tests;