From 0fcf61c3fbfbdf436640d77cd8d8bd696f5ef82f Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 14 Jul 2026 14:50:59 +0800 Subject: [PATCH] feat: skip oversized compaction tasks (#8466) fix: skip oversized compaction tasks Signed-off-by: jeremyhi --- src/mito2/src/compaction.rs | 87 ++++++++++++++++++++++++- src/mito2/src/test_util/version_util.rs | 11 +++- 2 files changed, 95 insertions(+), 3 deletions(-) diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 10900b5540..9c0c052045 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -64,7 +64,9 @@ use crate::error::{ RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, RemoteCompactionSnafu, Result, TimeRangePredicateOverflowSnafu, TimeoutSnafu, }; -use crate::metrics::{COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT}; +use crate::metrics::{ + COMPACTION_MEMORY_REJECTED, COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT, +}; use crate::read::FlatSource; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; @@ -613,8 +615,21 @@ impl CompactionScheduler { waiters }; - // Create a local compaction task. + // Check whether this local compaction can ever fit before submitting it. let estimated_bytes = estimate_compaction_bytes(&picker_output); + if let Some(limit_bytes) = self.exceeds_compaction_memory_limit(estimated_bytes) { + COMPACTION_MEMORY_REJECTED + .with_label_values(&["oversized"]) + .inc(); + warn!( + "Skip compaction for region {} because estimated memory {} bytes exceeds compaction memory limit {} bytes", + region_id, estimated_bytes, limit_bytes, + ); + for waiter in waiters { + waiter.send(Ok(0)); + } + return Ok(None); + } let cancel_handle = Arc::new(CancellationHandle::default()); let state = LocalCompactionState::new(cancel_handle.clone()); @@ -652,6 +667,15 @@ impl CompactionScheduler { ) } + fn exceeds_compaction_memory_limit(&self, estimated_bytes: u64) -> Option { + let limit_bytes = self.memory_manager.limit_bytes(); + if limit_bytes > 0 && estimated_bytes > limit_bytes { + Some(limit_bytes) + } else { + None + } + } + 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 { @@ -1551,6 +1575,65 @@ mod tests { assert!(scheduler.region_status.contains_key(®ion_id)); } + #[tokio::test] + async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler.memory_manager = Arc::new(new_compaction_memory_manager(1024 * 1024)); + + let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let end = 1000 * 1000; + let version_control = Arc::new( + builder + .push_l0_file_with_max_row_group_size(0, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(10, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(50, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(80, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(90, end, 1024 * 1024) + .build(), + ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + let (output_tx, output_rx) = oneshot::channel(); + let rejected = COMPACTION_MEMORY_REJECTED.with_label_values(&["oversized"]); + let rejected_before = rejected.get(); + + let scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(output_tx), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .await + .unwrap(); + + assert!(!scheduled); + 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)); + } + #[tokio::test] async fn test_schedule_on_finished() { common_telemetry::init_default_ut_logging(); diff --git a/src/mito2/src/test_util/version_util.rs b/src/mito2/src/test_util/version_util.rs index b14eea0958..1c0d7bfebc 100644 --- a/src/mito2/src/test_util/version_util.rs +++ b/src/mito2/src/test_util/version_util.rs @@ -89,6 +89,15 @@ impl VersionControlBuilder { } pub(crate) fn push_l0_file(&mut self, start_ms: i64, end_ms: i64) -> &mut Self { + self.push_l0_file_with_max_row_group_size(start_ms, end_ms, 0) + } + + pub(crate) fn push_l0_file_with_max_row_group_size( + &mut self, + start_ms: i64, + end_ms: i64, + max_row_group_uncompressed_size: u64, + ) -> &mut Self { let file_id = FileId::random(); self.files.insert( file_id, @@ -101,7 +110,7 @@ impl VersionControlBuilder { ), level: 0, file_size: 0, // We don't care file size. - max_row_group_uncompressed_size: 0, + max_row_group_uncompressed_size, available_indexes: Default::default(), indexes: Default::default(), index_file_size: 0,