feat(mito): add an environment switch to disable compaction (#9863)

Signed-off-by: WenyXu <wenymedia@gmail.com>

src/mito2/src/compaction/scheduler.rs

@@ -17,13 +17,13 @@ mod planning;
 mod state;
 
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use api::v1::region::compact_request;
 use common_base::Plugins;
 use common_memory_manager::OnExhaustedPolicy;
 use common_meta::key::SchemaMetadataManagerRef;
-use common_telemetry::{debug, error, info};
+use common_telemetry::{debug, error, info, warn};
 use common_time::range::TimestampRange;
 pub(crate) use planning::CompactionPickFinished;
 pub use planning::CompactionRequest;
@@ -39,7 +39,7 @@ use crate::compaction::task::MAX_PARALLEL_COMPACTION;
 use crate::config::MitoConfig;
 use crate::error::{
     CompactionCancelledSnafu, Error, ManualCompactionAlreadyRunningSnafu, RegionClosedSnafu,
-    RegionDroppedSnafu, RegionTruncatedSnafu, Result,
+    RegionDroppedSnafu, RegionTruncatedSnafu, Result, UnsupportedOperationSnafu,
 };
 use crate::region::version::VersionControlRef;
 use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState};
@@ -49,8 +49,42 @@ use crate::schedule::RequestCancelResult;
 use crate::schedule::scheduler::SchedulerRef;
 use crate::worker::WorkerListener;
 
+/// Internal testing switch to prevent SST compaction from changing file layouts
+/// during query regression tests. This is not a public configuration option.
+const DISABLE_COMPACTION_ENV: &str = "GREPTIME_DISABLE_COMPACTION";
+
+/// Internal query-regression switch, read once when the first scheduler starts.
+/// Applies to SST compaction only; ingestion and flushing remain enabled.
+static DISABLE_COMPACTION: LazyLock<bool> = LazyLock::new(|| {
+    let value = std::env::var(DISABLE_COMPACTION_ENV).ok();
+    let disabled = parse_disable_compaction(value.as_deref());
+    if disabled {
+        warn!(
+            "SST compaction is disabled by {DISABLE_COMPACTION_ENV}; restart without it to resume"
+        );
+    }
+    disabled
+});
+
+fn parse_disable_compaction(value: Option<&str>) -> bool {
+    match value {
+        None | Some("false") => false,
+        Some("true") => true,
+        Some(_) => {
+            warn!(
+                "Invalid {DISABLE_COMPACTION_ENV}: expected true or false; compaction remains enabled"
+            );
+            false
+        }
+    }
+}
+
 /// Compaction scheduler tracks and manages compaction tasks.
 pub(crate) struct CompactionScheduler {
+    /// Disables SST compaction for internal query regression tests.
+    /// Initialized from [`DISABLE_COMPACTION`] and fixed for the scheduler's lifetime
+    /// in production; unit tests can override it without changing the process environment.
+    disable_compaction: bool,
     scheduler: SchedulerRef,
     /// All loaded regions, including idle regions and regions executing DDLs.
     region_status: HashMap<RegionId, CompactionStatus>,
@@ -127,6 +161,7 @@ impl CompactionScheduler {
         memory_policy: OnExhaustedPolicy,
     ) -> Self {
         Self {
+            disable_compaction: *DISABLE_COMPACTION,
             scheduler,
             region_status: HashMap::new(),
             request_sender,
@@ -462,6 +497,17 @@ impl CompactionScheduler {
         max_parallelism: usize,
         time_range: Option<TimestampRange>,
     ) -> Result<bool> {
+        if self.disable_compaction {
+            if matches!(trigger, CompactionTrigger::Manual) {
+                waiter.send(
+                    UnsupportedOperationSnafu {
+                        err_msg: format!("SST compaction is disabled by {DISABLE_COMPACTION_ENV}"),
+                    }
+                    .fail(),
+                );
+            }
+            return Ok(false);
+        }
         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())
@@ -631,6 +677,9 @@ impl CompactionScheduler {
         manifest_ctx: &ManifestContextRef,
         schema_metadata_manager: SchemaMetadataManagerRef,
     ) -> bool {
+        if self.disable_compaction {
+            return false;
+        }
         let Some(status) = self.region_status.get_mut(&region_id) else {
             return false;
         };

src/mito2/src/compaction/scheduler_test.rs

@@ -2342,12 +2342,16 @@ async fn test_pending_ddl_request_failed_on_region_truncated() {
     assert_matches!(result, Err(_));
 }
 
+#[rstest::rstest]
 #[tokio::test]
-async fn test_on_compaction_finished_returns_pending_ddl_requests() {
+async fn test_on_compaction_finished_returns_pending_ddl_requests(
+    #[values(false, true)] disable_compaction: bool,
+) {
     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.disable_compaction = disable_compaction;
     let builder = VersionControlBuilder::new();
     let version_control = Arc::new(builder.build());
     let region_id = builder.region_id();
@@ -2612,3 +2616,95 @@ async fn test_concurrent_memory_competition() {
     drop(results);
     assert_eq!(manager.used_bytes(), 0);
 }
+
+#[test]
+fn test_parse_disable_compaction() {
+    for (value, expected) in [
+        (None, false),
+        (Some("true"), true),
+        (Some("false"), false),
+        (Some(""), false),
+        (Some("1"), false),
+        (Some("TRUE"), false),
+        (Some(" true "), false),
+    ] {
+        assert_eq!(parse_disable_compaction(value), expected, "{value:?}");
+    }
+}
+
+#[rstest::rstest]
+#[case::regular(Options::Regular(Default::default()))]
+#[case::strict_window(Options::StrictWindow(StrictWindow::default()))]
+#[tokio::test]
+async fn test_disable_compaction_scheduling(#[case] options: Options) {
+    let job_scheduler = Arc::new(VecScheduler::default());
+    let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone());
+    let (tx, mut rx) = mpsc::channel(4);
+    let mut scheduler = env.mock_compaction_scheduler(tx);
+    scheduler.disable_compaction = true;
+    let version_control = compactable_version();
+    let region_id = version_control.region_id();
+    scheduler.register_region(&version_control, &env.access_layer);
+    let manifest_ctx = env
+        .mock_manifest_context(version_control.current().version.metadata.clone())
+        .await;
+    let (schemas, _kv_backend) = mock_schema_metadata_manager();
+
+    assert!(
+        !scheduler
+            .schedule_automatic_compaction(
+                Options::Regular(Default::default()),
+                &version_control,
+                &env.access_layer,
+                &manifest_ctx,
+                schemas.clone(),
+            )
+            .unwrap()
+    );
+    let (output_tx, mut output_rx) = oneshot::channel();
+    assert!(
+        !scheduler
+            .schedule_manual_compaction(
+                options,
+                &version_control,
+                &env.access_layer,
+                OptionOutputTx::from(output_tx),
+                &manifest_ctx,
+                schemas.clone(),
+                1,
+                None,
+            )
+            .unwrap()
+    );
+    let err = output_rx.try_recv().unwrap().unwrap_err();
+    assert_matches!(err, Error::UnsupportedOperation { .. });
+    assert!(err.to_string().contains(DISABLE_COMPACTION_ENV));
+    assert!(!scheduler.schedule_automatic_followup(region_id, &manifest_ctx, schemas.clone()));
+    assert!(!scheduler.is_compacting(region_id));
+    assert_eq!(job_scheduler.num_jobs(), 0);
+    assert!(rx.try_recv().is_err());
+
+    // The same compactable fixture dispatches planning with the switch off.
+    scheduler.disable_compaction = false;
+    assert!(
+        scheduler
+            .schedule_automatic_compaction(
+                Options::Regular(Default::default()),
+                &version_control,
+                &env.access_layer,
+                &manifest_ctx,
+                schemas,
+            )
+            .unwrap()
+    );
+    let finished = tokio::time::timeout(
+        Duration::from_secs(10),
+        recv_compaction_pick_finished(&mut rx),
+    )
+    .await
+    .expect("enabled compaction should finish planning");
+    assert!(matches!(
+        finished.result,
+        CompactionPlanningResult::Prepared(_)
+    ));
+}