diff --git a/grafana/dashboards/metrics/cluster/dashboard.json b/grafana/dashboards/metrics/cluster/dashboard.json index daa0f018836..2fffd27b25f 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.json +++ b/grafana/dashboards/metrics/cluster/dashboard.json @@ -10249,7 +10249,7 @@ "type": "prometheus", "uid": "${metrics}" }, - "description": "Compaction OPS per Instance.", + "description": "Published compaction units per second per Instance. A local unit is one input dependency group; a remote job counts as one unit.", "fieldConfig": { "defaults": { "color": { @@ -10335,9 +10335,9 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count{instance=~\"$datanode\"}[$__rate_interval]))", + "expr": "sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count{instance=~\"$datanode\"}[$__rate_interval]))", "instant": false, - "legendFormat": "[{{ instance }}]-[{{pod}}]", + "legendFormat": "[{{ instance }}]-[{{pod}}]-[{{type}}]", "range": true, "refId": "A" } @@ -10350,7 +10350,7 @@ "type": "prometheus", "uid": "${metrics}" }, - "description": "Ongoing compaction task count", + "description": "Ongoing compaction count: local units and remote jobs", "fieldConfig": { "defaults": { "color": { @@ -10451,7 +10451,7 @@ "type": "prometheus", "uid": "${metrics}" }, - "description": "Compaction P99 and average per Instance.", + "description": "Compaction unit P99 and average elapsed time per Instance. Local units are timed from task start through worker apply; remote jobs from request start.", "fieldConfig": { "defaults": { "color": { @@ -10537,9 +10537,9 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "histogram_quantile(0.99, sum by(instance, pod, le) (rate(greptime_mito_compaction_total_elapsed_bucket{instance=~\"$datanode\"}[$__rate_interval])))", + "expr": "histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_compaction_total_elapsed_bucket{instance=~\"$datanode\"}[$__rate_interval])))", "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p99", + "legendFormat": "[{{instance}}]-[{{pod}}]-[{{type}}]-p99", "range": true, "refId": "A" }, @@ -10549,10 +10549,10 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_sum{instance=~\"$datanode\"}[$__rate_interval])) / sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count{instance=~\"$datanode\"}[$__rate_interval]))", + "expr": "sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_sum{instance=~\"$datanode\"}[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count{instance=~\"$datanode\"}[$__rate_interval]))", "hide": false, "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-avg", + "legendFormat": "[{{instance}}]-[{{pod}}]-[{{type}}]-avg", "range": true, "refId": "B" } diff --git a/grafana/dashboards/metrics/cluster/dashboard.md b/grafana/dashboards/metrics/cluster/dashboard.md index 3278db52bfd..455268cdd43 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.md +++ b/grafana/dashboards/metrics/cluster/dashboard.md @@ -117,9 +117,9 @@ | Flush Elapsed Time | `histogram_quantile(0.95, sum by (instance, pod, le, type) (rate(greptime_mito_flush_elapsed_bucket{instance=~"$datanode"}[$__rate_interval])))`
`histogram_quantile(0.99, sum by (instance, pod, le, type) (rate(greptime_mito_flush_elapsed_bucket{instance=~"$datanode"}[$__rate_interval])))`
`sum by (instance, pod, type) (rate(greptime_mito_flush_elapsed_sum{instance=~"$datanode"}[$__rate_interval])) / sum by (instance, pod, type) (rate(greptime_mito_flush_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Mito flush p95 and p99 elapsed time by datanode and flush type. Use this to identify slow flush jobs. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{type}}]-p95` | | Flush Throughput | `sum by (instance, pod) (rate(greptime_mito_flush_bytes_total{instance=~"$datanode"}[$__rate_interval]))`
`sum by (instance, pod) (rate(greptime_mito_flush_file_total{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Mito flushed bytes and flushed file rates. Use this with flush elapsed time to distinguish slow jobs from large jobs. | `prometheus` | `Bps` | `[{{instance}}]-[{{pod}}]-bytes` | | Inflight Flush | `greptime_mito_inflight_flush_count` | `timeseries` | Ongoing flush task count | `prometheus` | `none` | `[{{instance}}]-[{{pod}}]` | -| Compaction OPS per Instance | `sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Compaction OPS per Instance. | `prometheus` | `ops` | `[{{ instance }}]-[{{pod}}]` | -| Inflight Compaction | `greptime_mito_inflight_compaction_count` | `timeseries` | Ongoing compaction task count | `prometheus` | `none` | `[{{instance}}]-[{{pod}}]` | -| Compaction P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, le) (rate(greptime_mito_compaction_total_elapsed_bucket{instance=~"$datanode"}[$__rate_interval])))`
`sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Compaction P99 and average per Instance. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-p99` | +| Compaction OPS per Instance | `sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Published compaction units per second per Instance. A local unit is one input dependency group; a remote job counts as one unit. | `prometheus` | `ops` | `[{{ instance }}]-[{{pod}}]-[{{type}}]` | +| Inflight Compaction | `greptime_mito_inflight_compaction_count` | `timeseries` | Ongoing compaction count: local units and remote jobs | `prometheus` | `none` | `[{{instance}}]-[{{pod}}]` | +| Compaction P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_compaction_total_elapsed_bucket{instance=~"$datanode"}[$__rate_interval])))`
`sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Compaction unit P99 and average elapsed time per Instance. Local units are timed from task start through worker apply; remote jobs from request start. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{type}}]-p99` | | Compaction Elapsed Time per Instance by Stage | `histogram_quantile(0.99, sum by(instance, pod, le, stage) (rate(greptime_mito_compaction_stage_elapsed_bucket{instance=~"$datanode"}[$__rate_interval])))`
`sum by(instance, pod, stage) (rate(greptime_mito_compaction_stage_elapsed_sum{instance=~"$datanode"}[$__rate_interval]))/sum by(instance, pod, stage) (rate(greptime_mito_compaction_stage_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Compaction latency by stage | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]-p99` | | Compaction Input/Output Bytes | `sum by(instance, pod) (rate(greptime_mito_compaction_input_bytes[$__rate_interval]))`
`sum by(instance, pod) (rate(greptime_mito_compaction_output_bytes[$__rate_interval]))` | `timeseries` | Compaction input and output bytes by datanode. Use this to correlate compaction latency with rewritten data volume. | `prometheus` | `Bps` | `[{{instance}}]-[{{pod}}]-input` | # Index diff --git a/grafana/dashboards/metrics/cluster/dashboard.yaml b/grafana/dashboards/metrics/cluster/dashboard.yaml index bc659b401b5..7513c5e50e1 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.yaml +++ b/grafana/dashboards/metrics/cluster/dashboard.yaml @@ -1417,17 +1417,17 @@ groups: legendFormat: '[{{instance}}]-[{{pod}}]' - title: Compaction OPS per Instance type: timeseries - description: Compaction OPS per Instance. + description: Published compaction units per second per Instance. A local unit is one input dependency group; a remote job counts as one unit. unit: ops queries: - - expr: sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval])) + - expr: sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} - legendFormat: '[{{ instance }}]-[{{pod}}]' + legendFormat: '[{{ instance }}]-[{{pod}}]-[{{type}}]' - title: Inflight Compaction type: timeseries - description: Ongoing compaction task count + description: 'Ongoing compaction count: local units and remote jobs' unit: none queries: - expr: greptime_mito_inflight_compaction_count @@ -1437,19 +1437,19 @@ groups: legendFormat: '[{{instance}}]-[{{pod}}]' - title: Compaction P99 and Avg per Instance type: timeseries - description: Compaction P99 and average per Instance. + description: Compaction unit P99 and average elapsed time per Instance. Local units are timed from task start through worker apply; remote jobs from request start. unit: s queries: - - expr: histogram_quantile(0.99, sum by(instance, pod, le) (rate(greptime_mito_compaction_total_elapsed_bucket{instance=~"$datanode"}[$__rate_interval]))) + - expr: histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_compaction_total_elapsed_bucket{instance=~"$datanode"}[$__rate_interval]))) datasource: type: prometheus uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p99' - - expr: sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval])) + legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-p99' + - expr: sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count{instance=~"$datanode"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-avg' + legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-avg' - title: Compaction Elapsed Time per Instance by Stage type: timeseries description: Compaction latency by stage diff --git a/grafana/dashboards/metrics/standalone/dashboard.json b/grafana/dashboards/metrics/standalone/dashboard.json index c1fda955d90..f17d40b4f96 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.json +++ b/grafana/dashboards/metrics/standalone/dashboard.json @@ -10249,7 +10249,7 @@ "type": "prometheus", "uid": "${metrics}" }, - "description": "Compaction OPS per Instance.", + "description": "Published compaction units per second per Instance. A local unit is one input dependency group; a remote job counts as one unit.", "fieldConfig": { "defaults": { "color": { @@ -10335,9 +10335,9 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))", + "expr": "sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))", "instant": false, - "legendFormat": "[{{ instance }}]-[{{pod}}]", + "legendFormat": "[{{ instance }}]-[{{pod}}]-[{{type}}]", "range": true, "refId": "A" } @@ -10350,7 +10350,7 @@ "type": "prometheus", "uid": "${metrics}" }, - "description": "Ongoing compaction task count", + "description": "Ongoing compaction count: local units and remote jobs", "fieldConfig": { "defaults": { "color": { @@ -10451,7 +10451,7 @@ "type": "prometheus", "uid": "${metrics}" }, - "description": "Compaction P99 and average per Instance.", + "description": "Compaction unit P99 and average elapsed time per Instance. Local units are timed from task start through worker apply; remote jobs from request start.", "fieldConfig": { "defaults": { "color": { @@ -10537,9 +10537,9 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "histogram_quantile(0.99, sum by(instance, pod, le) (rate(greptime_mito_compaction_total_elapsed_bucket[$__rate_interval])))", + "expr": "histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_compaction_total_elapsed_bucket[$__rate_interval])))", "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p99", + "legendFormat": "[{{instance}}]-[{{pod}}]-[{{type}}]-p99", "range": true, "refId": "A" }, @@ -10549,10 +10549,10 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))", + "expr": "sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_sum[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))", "hide": false, "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-avg", + "legendFormat": "[{{instance}}]-[{{pod}}]-[{{type}}]-avg", "range": true, "refId": "B" } diff --git a/grafana/dashboards/metrics/standalone/dashboard.md b/grafana/dashboards/metrics/standalone/dashboard.md index 5dd6f1600e1..5c7a4042444 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.md +++ b/grafana/dashboards/metrics/standalone/dashboard.md @@ -117,9 +117,9 @@ | Flush Elapsed Time | `histogram_quantile(0.95, sum by (instance, pod, le, type) (rate(greptime_mito_flush_elapsed_bucket[$__rate_interval])))`
`histogram_quantile(0.99, sum by (instance, pod, le, type) (rate(greptime_mito_flush_elapsed_bucket[$__rate_interval])))`
`sum by (instance, pod, type) (rate(greptime_mito_flush_elapsed_sum[$__rate_interval])) / sum by (instance, pod, type) (rate(greptime_mito_flush_elapsed_count[$__rate_interval]))` | `timeseries` | Mito flush p95 and p99 elapsed time by datanode and flush type. Use this to identify slow flush jobs. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{type}}]-p95` | | Flush Throughput | `sum by (instance, pod) (rate(greptime_mito_flush_bytes_total[$__rate_interval]))`
`sum by (instance, pod) (rate(greptime_mito_flush_file_total[$__rate_interval]))` | `timeseries` | Mito flushed bytes and flushed file rates. Use this with flush elapsed time to distinguish slow jobs from large jobs. | `prometheus` | `Bps` | `[{{instance}}]-[{{pod}}]-bytes` | | Inflight Flush | `greptime_mito_inflight_flush_count` | `timeseries` | Ongoing flush task count | `prometheus` | `none` | `[{{instance}}]-[{{pod}}]` | -| Compaction OPS per Instance | `sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))` | `timeseries` | Compaction OPS per Instance. | `prometheus` | `ops` | `[{{ instance }}]-[{{pod}}]` | -| Inflight Compaction | `greptime_mito_inflight_compaction_count` | `timeseries` | Ongoing compaction task count | `prometheus` | `none` | `[{{instance}}]-[{{pod}}]` | -| Compaction P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, le) (rate(greptime_mito_compaction_total_elapsed_bucket[$__rate_interval])))`
`sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))` | `timeseries` | Compaction P99 and average per Instance. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-p99` | +| Compaction OPS per Instance | `sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))` | `timeseries` | Published compaction units per second per Instance. A local unit is one input dependency group; a remote job counts as one unit. | `prometheus` | `ops` | `[{{ instance }}]-[{{pod}}]-[{{type}}]` | +| Inflight Compaction | `greptime_mito_inflight_compaction_count` | `timeseries` | Ongoing compaction count: local units and remote jobs | `prometheus` | `none` | `[{{instance}}]-[{{pod}}]` | +| Compaction P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_compaction_total_elapsed_bucket[$__rate_interval])))`
`sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_sum[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval]))` | `timeseries` | Compaction unit P99 and average elapsed time per Instance. Local units are timed from task start through worker apply; remote jobs from request start. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{type}}]-p99` | | Compaction Elapsed Time per Instance by Stage | `histogram_quantile(0.99, sum by(instance, pod, le, stage) (rate(greptime_mito_compaction_stage_elapsed_bucket[$__rate_interval])))`
`sum by(instance, pod, stage) (rate(greptime_mito_compaction_stage_elapsed_sum[$__rate_interval]))/sum by(instance, pod, stage) (rate(greptime_mito_compaction_stage_elapsed_count[$__rate_interval]))` | `timeseries` | Compaction latency by stage | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]-p99` | | Compaction Input/Output Bytes | `sum by(instance, pod) (rate(greptime_mito_compaction_input_bytes[$__rate_interval]))`
`sum by(instance, pod) (rate(greptime_mito_compaction_output_bytes[$__rate_interval]))` | `timeseries` | Compaction input and output bytes by datanode. Use this to correlate compaction latency with rewritten data volume. | `prometheus` | `Bps` | `[{{instance}}]-[{{pod}}]-input` | # Index diff --git a/grafana/dashboards/metrics/standalone/dashboard.yaml b/grafana/dashboards/metrics/standalone/dashboard.yaml index beadb9c3215..627cd8b5d63 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.yaml +++ b/grafana/dashboards/metrics/standalone/dashboard.yaml @@ -1417,17 +1417,17 @@ groups: legendFormat: '[{{instance}}]-[{{pod}}]' - title: Compaction OPS per Instance type: timeseries - description: Compaction OPS per Instance. + description: Published compaction units per second per Instance. A local unit is one input dependency group; a remote job counts as one unit. unit: ops queries: - - expr: sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval])) + - expr: sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval])) datasource: type: prometheus uid: ${metrics} - legendFormat: '[{{ instance }}]-[{{pod}}]' + legendFormat: '[{{ instance }}]-[{{pod}}]-[{{type}}]' - title: Inflight Compaction type: timeseries - description: Ongoing compaction task count + description: 'Ongoing compaction count: local units and remote jobs' unit: none queries: - expr: greptime_mito_inflight_compaction_count @@ -1437,19 +1437,19 @@ groups: legendFormat: '[{{instance}}]-[{{pod}}]' - title: Compaction P99 and Avg per Instance type: timeseries - description: Compaction P99 and average per Instance. + description: Compaction unit P99 and average elapsed time per Instance. Local units are timed from task start through worker apply; remote jobs from request start. unit: s queries: - - expr: histogram_quantile(0.99, sum by(instance, pod, le) (rate(greptime_mito_compaction_total_elapsed_bucket[$__rate_interval]))) + - expr: histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_compaction_total_elapsed_bucket[$__rate_interval]))) datasource: type: prometheus uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p99' - - expr: sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval])) + legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-p99' + - expr: sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_sum[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_compaction_total_elapsed_count[$__rate_interval])) datasource: type: prometheus uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-avg' + legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-avg' - title: Compaction Elapsed Time per Instance by Stage type: timeseries description: Compaction latency by stage diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index efee5f6cca1..f70aa622a27 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -91,6 +91,8 @@ Tests live next to the code as `*_test.rs` (e.g. `src/mito2/src/engine/flush_tes - Loaded regions have resident compaction state, including `Idle` and handed-off DDLs. Publish create/open/catchup results through the worker's `register_region`; only close/drop unregisters them. See `compaction/scheduler/state.rs`. +- Local compaction units use the snapshot-free `CompactionContext`. Retaining a + full `CompactionRegion` in a unit pins other units' inputs and delays SST purge. ## Maintenance contract diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 1fbbfad5613..0f1b637e096 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -26,8 +26,10 @@ mod task; #[cfg(test)] mod test_util; mod twcs; +mod unit; mod window; +use std::any::Any; use std::collections::HashMap; use common_meta::key::SchemaMetadataManagerRef; @@ -155,3 +157,14 @@ pub struct SerializedCompactionOutput { filter_deleted: bool, output_time_range: Option, } + +/// Returns the printable reason carried by a caught panic payload. +pub(crate) fn format_panic_reason(payload: &(dyn Any + Send + 'static)) -> String { + if let Some(message) = payload.downcast_ref::<&str>() { + message.to_string() + } else if let Some(message) = payload.downcast_ref::() { + message.clone() + } else { + "unknown panic".to_string() + } +} diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 1a025bcaeec..527750bf33e 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -38,6 +38,7 @@ use crate::access_layer::{ use crate::cache::{CacheManager, CacheManagerRef}; use crate::compaction::picker::PickerOutput; use crate::compaction::reader::CompactionSstReaderBuilder; +use crate::compaction::unit::CompactionUnit; use crate::compaction::{CompactionOutput, find_dynamic_options}; use crate::config::MitoConfig; use crate::engine::region_hook::{RegionHookRef, SstFileInfo}; @@ -123,6 +124,70 @@ pub struct CompactionRegion { pub(crate) plugins: Plugins, } +/// Execution resources for compaction, without a region-wide SST snapshot. +/// +/// Inputs are owned by each output or unit so a slow sibling cannot retain +/// files that another unit has already replaced. +#[derive(Clone)] +pub struct CompactionContext { + /// Region whose files are being compacted. + pub region_id: RegionId, + /// Region options including resolved dynamic compaction options. + pub region_options: RegionOptions, + /// Options from the version used to plan the merge. + version_options: RegionOptions, + pub(crate) engine_config: Arc, + region_metadata: RegionMetadataRef, + pub(crate) cache_manager: CacheManagerRef, + pub(crate) access_layer: AccessLayerRef, + manifest_ctx: Arc, + plugins: Plugins, +} + +impl From<&CompactionRegion> for CompactionContext { + fn from(region: &CompactionRegion) -> Self { + Self { + region_id: region.region_id, + region_options: region.region_options.clone(), + version_options: region.current_version.options.clone(), + engine_config: region.engine_config.clone(), + region_metadata: region.region_metadata.clone(), + cache_manager: region.cache_manager.clone(), + access_layer: region.access_layer.clone(), + manifest_ctx: region.manifest_ctx.clone(), + plugins: region.plugins.clone(), + } + } +} + +/// Builds a minimal [`CompactionRegion`] for tests that do not touch the access layer. +#[cfg(test)] +pub(crate) async fn new_test_compaction_region() -> CompactionRegion { + let env = crate::test_util::scheduler_util::SchedulerEnv::new().await; + let metadata = crate::test_util::memtable_util::metadata_for_test(); + let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; + CompactionRegion { + region_id: RegionId::new(1, 1), + region_options: RegionOptions::default(), + engine_config: Arc::new(MitoConfig::default()), + region_metadata: metadata.clone(), + cache_manager: Arc::new(CacheManager::default()), + access_layer: env.access_layer.clone(), + manifest_ctx, + current_version: CompactionVersion { + metadata: metadata.clone(), + options: RegionOptions::default(), + ssts: Arc::new(SstVersion::new(metadata)), + memtable_min_sequence: None, + compaction_time_window: None, + }, + file_purger: None, + ttl: None, + max_parallelism: 1, + plugins: Plugins::new(), + } +} + /// OpenCompactionRegionRequest represents the request to open a compaction region. #[derive(Clone)] pub struct OpenCompactionRegionRequest { @@ -273,9 +338,18 @@ impl CompactionRegion { /// Fires [`RegionHook::on_sst_files_written`] for the freshly-merged SST /// files in `merge_output`. Shared by both compaction paths, it must run - /// before [`Compactor::update_manifest`], whose `on_manifest_updated` - /// drains the per-region state this hook populates. + /// before this unit's [`Compactor::update_manifest`] so its committed edit + /// can be associated with the SST metadata reported here. pub async fn invoke_sst_hook(&self, merge_output: &MergeOutput) { + CompactionContext::from(self) + .invoke_sst_hook(merge_output) + .await; + } +} + +impl CompactionContext { + /// Reports this unit's output metadata before publishing its manifest edit. + pub(crate) async fn invoke_sst_hook(&self, merge_output: &MergeOutput) { let Some(hook) = self.plugins.get::() else { return; }; @@ -319,6 +393,32 @@ impl CompactionRegion { hook.on_sst_files_written(self.region_id, &self.region_metadata, &files) .await; } + + /// Publishes outputs through the existing manifest validation and hook path. + pub(crate) async fn update_manifest( + &self, + merge_output: MergeOutput, + ) -> Result<(RegionEdit, ManifestVersion)> { + let edit = RegionEdit { + files_to_add: merge_output.files_to_add, + files_to_remove: merge_output.files_to_remove, + timestamp_ms: Some(chrono::Utc::now().timestamp_millis()), + compaction_time_window: merge_output + .compaction_time_window + .map(|seconds| Duration::from_secs(seconds as u64)), + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }; + + let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())); + let manifest_version = self + .manifest_ctx + .update_manifest_for_compaction(action_list) + .await?; + + Ok((edit, manifest_version)) + } } /// Builds an [`SstInfo`] for the remote-compaction path by copying the scalar @@ -406,7 +506,7 @@ pub trait Compactor: Send + Sync + 'static { pub trait SstMerger: Send + Sync + 'static { async fn merge_single_output( &self, - compaction_region: CompactionRegion, + context: CompactionContext, output: CompactionOutput, write_opts: WriteOptions, ) -> Result<(Vec, Vec)>; @@ -426,13 +526,13 @@ struct OutputSequenceMetadata { /// without allocating a new admission marker. Retaining physical row sequences /// does not by itself restore exact-read capability for untrusted inputs. fn output_sequence_metadata( - region: &CompactionRegion, + context: &CompactionContext, inputs: &[FileHandle], ) -> OutputSequenceMetadata { - let exact_sequence_trusted = region.region_options.preserve_row_sequence + let exact_sequence_trusted = context.region_options.preserve_row_sequence && inputs .iter() - .all(|f| f.is_effective_target_sequence_trusted(region.region_id)); + .all(|f| f.is_effective_target_sequence_trusted(context.region_id)); OutputSequenceMetadata { file_sequence: known_max_input_sequence(inputs), exact_sequence_trusted, @@ -456,39 +556,35 @@ fn known_max_input_sequence(inputs: &[FileHandle]) -> Option { impl SstMerger for DefaultSstMerger { async fn merge_single_output( &self, - compaction_region: CompactionRegion, + context: CompactionContext, output: CompactionOutput, write_opts: WriteOptions, ) -> Result<(Vec, Vec)> { - let region_id = compaction_region.region_id; - let index_options = compaction_region - .current_version - .options - .index_options - .clone(); - let append_mode = compaction_region.current_version.options.append_mode; - let merge_mode = compaction_region.current_version.options.merge_mode(); - let flat_format = compaction_region + let region_id = context.region_id; + let index_options = context.version_options.index_options.clone(); + let append_mode = context.version_options.append_mode; + let merge_mode = context.version_options.merge_mode(); + let flat_format = context .region_options .sst_format .map(|format| format == FormatType::Flat) - .unwrap_or(compaction_region.engine_config.default_flat_format); + .unwrap_or(context.engine_config.default_flat_format); - let index_config = compaction_region.engine_config.index.clone(); - let inverted_index_config = compaction_region.engine_config.inverted_index.clone(); - let fulltext_index_config = compaction_region.engine_config.fulltext_index.clone(); - let bloom_filter_index_config = compaction_region.engine_config.bloom_filter_index.clone(); + let index_config = context.engine_config.index.clone(); + let inverted_index_config = context.engine_config.inverted_index.clone(); + let fulltext_index_config = context.engine_config.fulltext_index.clone(); + let bloom_filter_index_config = context.engine_config.bloom_filter_index.clone(); let input_file_names = output .inputs .iter() .map(|f| f.file_id().to_string()) .join(","); - let sequence_metadata = output_sequence_metadata(&compaction_region, &output.inputs); + let sequence_metadata = output_sequence_metadata(&context, &output.inputs); let builder = CompactionSstReaderBuilder { - metadata: compaction_region.region_metadata.clone(), - sst_layer: compaction_region.access_layer.clone(), - cache: compaction_region.cache_manager.clone(), + metadata: context.region_metadata.clone(), + sst_layer: context.access_layer.clone(), + cache: context.cache_manager.clone(), inputs: &output.inputs, append_mode, filter_deleted: output.filter_deleted, @@ -498,15 +594,15 @@ impl SstMerger for DefaultSstMerger { let source = builder.build_flat_sst_reader().await?; let mut metrics = Metrics::new(WriteType::Compaction); - let region_metadata = compaction_region.region_metadata.clone(); - let sst_infos = compaction_region + let region_metadata = context.region_metadata.clone(); + let sst_infos = context .access_layer .write_sst( SstWriteRequest { op_type: OperationType::Compact, metadata: region_metadata.clone(), source, - cache_manager: compaction_region.cache_manager.clone(), + cache_manager: context.cache_manager.clone(), // Readers resolve file overrides before merge/dedup. Replacing // their effective sequences here could promote old rows above // versions in SSTs that were not part of this merge. @@ -629,6 +725,50 @@ impl DefaultCompactor { } } +impl DefaultCompactor { + /// Merges a complete atomic unit. No partial input replacement is returned. + pub(crate) async fn merge_unit( + &self, + context: &CompactionContext, + unit: &CompactionUnit, + ) -> Result { + let mut result = MergeOutput { + files_to_remove: unit.inputs.iter().map(|f| f.meta_ref().clone()).collect(), + compaction_time_window: Some(unit.time_window_size), + ..Default::default() + }; + for output in &unit.outputs { + if self.cancel_handle.is_cancelled() { + return error::CompactionCancelledSnafu.fail(); + } + // Let the writer finish and register its files before observing cancellation. + // Dropping it mid-write can lose ownership of already finalized SSTs. + let (files, infos) = self + .merger + .merge_single_output( + context.clone(), + output.clone(), + WriteOptions { + write_buffer_size: context.engine_config.sst_write_buffer_size, + max_file_size: unit.max_file_size, + row_group_size: context.region_options.row_group_size(), + float_field_encoding: context.region_options.float_field_encoding, + }, + ) + .await?; + if let Some(uncommitted) = &self.uncommitted { + uncommitted.track(&infos); + } + result.files_to_add.extend(files); + result.sst_infos.extend(infos); + } + if self.cancel_handle.is_cancelled() { + return error::CompactionCancelledSnafu.fail(); + } + Ok(result) + } +} + #[async_trait::async_trait] impl Compactor for DefaultCompactor where @@ -642,6 +782,7 @@ where let internal_parallelism = compaction_region.max_parallelism.max(1); let compaction_time_window = picker_output.time_window_size; let region_id = compaction_region.region_id; + let context = CompactionContext::from(compaction_region); // Build tasks along with their input file metas so we can track which // inputs correspond to each task. @@ -657,11 +798,11 @@ where float_field_encoding: compaction_region.region_options.float_field_encoding, }; let merger = self.merger.clone(); - let compaction_region = compaction_region.clone(); + let context = context.clone(); let uncommitted = self.uncommitted.clone(); let fut = async move { let result = merger - .merge_single_output(compaction_region, output, write_opts) + .merge_single_output(context, output, write_opts) .await; if let (Some(uncommitted), Ok((_, infos))) = (&uncommitted, &result) { uncommitted.track(infos); @@ -777,28 +918,9 @@ where compaction_region: &CompactionRegion, merge_output: MergeOutput, ) -> Result<(RegionEdit, ManifestVersion)> { - // Write region edit to manifest. - let edit = RegionEdit { - files_to_add: merge_output.files_to_add, - files_to_remove: merge_output.files_to_remove, - // Use current timestamp as the edit timestamp. - timestamp_ms: Some(chrono::Utc::now().timestamp_millis()), - compaction_time_window: merge_output - .compaction_time_window - .map(|seconds| Duration::from_secs(seconds as u64)), - flushed_entry_id: None, - flushed_sequence: None, - committed_sequence: None, - }; - - let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())); - // TODO: We might leak files if we fail to update manifest. We can add a cleanup task to remove them later. - let manifest_version = compaction_region - .manifest_ctx - .update_manifest_for_compaction(action_list) - .await?; - - Ok((edit, manifest_version)) + CompactionContext::from(compaction_region) + .update_manifest(merge_output) + .await } } @@ -812,14 +934,10 @@ mod tests { use tokio::time::sleep; use super::{DefaultCompactor, *}; - use crate::cache::CacheManager; use crate::compaction::picker::PickerOutput; use crate::error::Result; use crate::sst::file::FileHandle; use crate::sst::file_purger::NoopFilePurger; - use crate::sst::version::SstVersion; - use crate::test_util::memtable_util::metadata_for_test; - use crate::test_util::scheduler_util::SchedulerEnv; fn dummy_file_meta() -> FileMeta { FileMeta { @@ -963,7 +1081,8 @@ mod tests { for preserve in [false, true] { region.region_options.preserve_row_sequence = preserve; for (inputs, known_bound, inputs_trusted) in &cases { - let metadata = output_sequence_metadata(®ion, inputs); + let metadata = + output_sequence_metadata(&CompactionContext::from(®ion), inputs); let trusted = preserve && *inputs_trusted; assert_eq!(trusted, metadata.exact_sequence_trusted); assert_eq!( @@ -1149,34 +1268,6 @@ mod tests { } } - /// Build a minimal [`CompactionRegion`] suitable for tests where the - /// [`SstMerger`] is mocked and never touches the access layer. - async fn new_test_compaction_region() -> CompactionRegion { - let env = SchedulerEnv::new().await; - let metadata = metadata_for_test(); - let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; - CompactionRegion { - region_id: RegionId::new(1, 1), - region_options: RegionOptions::default(), - engine_config: Arc::new(MitoConfig::default()), - region_metadata: metadata.clone(), - cache_manager: Arc::new(CacheManager::default()), - access_layer: env.access_layer.clone(), - manifest_ctx, - current_version: CompactionVersion { - metadata: metadata.clone(), - options: RegionOptions::default(), - ssts: Arc::new(SstVersion::new(metadata)), - memtable_min_sequence: None, - compaction_time_window: None, - }, - file_purger: None, - ttl: None, - max_parallelism: 1, - plugins: Plugins::new(), - } - } - /// An [`SstMerger`] that returns pre-configured results per call index. /// /// Call 0 gets `results[0]`, call 1 gets `results[1]`, etc. @@ -1199,7 +1290,7 @@ mod tests { impl SstMerger for MockMerger { async fn merge_single_output( &self, - _compaction_region: CompactionRegion, + _context: CompactionContext, _output: CompactionOutput, _write_opts: WriteOptions, ) -> Result<(Vec, Vec)> { @@ -1215,6 +1306,56 @@ mod tests { } } + #[tokio::test] + async fn test_compaction_unit_group_failure_is_atomic() { + let context = CompactionContext::from(&new_test_compaction_region().await); + let input = new_file_handle(dummy_file_meta()); + let unit = CompactionUnit::from_picker(PickerOutput { + outputs: (0..2) + .map(|_| CompactionOutput { + output_level: 1, + inputs: vec![input.clone()], + filter_deleted: false, + output_time_range: None, + }) + .collect(), + ..Default::default() + }) + .unwrap() + .pop() + .unwrap(); + for fail_second in [false, true] { + let second = if fail_second { + error::InvalidMetaSnafu { + reason: "second output failed", + } + .fail() + } else { + Ok(vec![dummy_file_meta(), dummy_file_meta()]) + }; + let compactor = + DefaultCompactor::with_merger(MockMerger::new(vec![Ok(Vec::new()), second])); + let result = compactor.merge_unit(&context, &unit).await; + if fail_second { + assert!( + result.is_err(), + "cannot publish a successful subset of a dependency group" + ); + } else { + let result = result.unwrap(); + assert_eq!(2, result.files_to_add.len()); + assert_eq!( + vec![input.meta_ref().file_id], + result + .files_to_remove + .iter() + .map(|f| f.file_id) + .collect::>() + ); + } + } + } + #[tokio::test] async fn test_partial_merge_failure_collects_only_successful_outputs() { common_telemetry::init_default_ut_logging(); @@ -1371,7 +1512,7 @@ mod tests { impl SstMerger for BlockingMerger { async fn merge_single_output( &self, - _compaction_region: CompactionRegion, + _context: CompactionContext, _output: CompactionOutput, _write_opts: WriteOptions, ) -> Result<(Vec, Vec)> { diff --git a/src/mito2/src/compaction/picker.rs b/src/mito2/src/compaction/picker.rs index ec9fa91f4c8..3628953cc31 100644 --- a/src/mito2/src/compaction/picker.rs +++ b/src/mito2/src/compaction/picker.rs @@ -31,11 +31,6 @@ use crate::sst::file::{FileHandle, FileMeta}; use crate::sst::file_purger::FilePurger; use crate::sst::version::LevelMeta; -#[async_trait::async_trait] -pub(crate) trait CompactionTask: Debug + Send + Sync + 'static { - async fn run(&mut self); -} - /// Picker picks input SST files for compaction. /// Different compaction strategy may implement different pickers. #[async_trait::async_trait] diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index e9da3002254..76459abf3bc 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod local; mod planning; mod state; @@ -324,28 +325,6 @@ impl CompactionScheduler { .await } - /// Completes a cooperatively canceled execution. - /// - /// # Effects - /// - /// Resets the matching region to idle, notifies compaction waiters, and - /// returns DDLs that were waiting for cancellation. - /// - /// # Constraints - /// - /// The owning worker must execute every returned DDL request. A stale - /// cancellation is ignored. - pub(crate) async fn on_execution_cancelled( - &mut self, - region_id: RegionId, - execution: &CompactionExecution, - ) -> Vec { - if !self.is_current_execution(region_id, execution) { - return Vec::new(); - } - self.on_compaction_cancelled(region_id).await - } - /// Records failure of the installed execution. /// /// # Effects @@ -684,6 +663,7 @@ impl CompactionScheduler { } /// Notifies the scheduler that the compaction job is cancelled cooperatively. + #[cfg(test)] async fn on_compaction_cancelled(&mut self, region_id: RegionId) -> Vec { self.finish_compaction_on_cancel(region_id) } diff --git a/src/mito2/src/compaction/scheduler/local.rs b/src/mito2/src/compaction/scheduler/local.rs new file mode 100644 index 00000000000..8ff2a6cb150 --- /dev/null +++ b/src/mito2/src/compaction/scheduler/local.rs @@ -0,0 +1,341 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; + +use common_meta::key::SchemaMetadataManagerRef; +use store_api::storage::RegionId; + +use crate::compaction::compactor::{CompactionContext, CompactionRegion}; +use crate::compaction::scheduler::state::{CompactingFiles, CompactionExecution, CompactionPhase}; +use crate::compaction::scheduler::{CompactionScheduler, CompactionTransition}; +use crate::compaction::task::CompactionTask; +use crate::compaction::unit::CompactionUnit; +use crate::error::{self, Error}; +use crate::metrics::INFLIGHT_COMPACTION_COUNT; +use crate::region::ManifestContextRef; +use crate::request::{ + BackgroundNotify, CompactionUnitNotification, WorkerRequest, WorkerRequestWithTime, +}; +use crate::schedule::{CancellableTaskState, RequestCancelResult}; +use crate::sst::file::UncommittedSsts; + +/// Tracks one unit through its terminal result and subsequent task exit. +#[derive(Debug)] +struct UnitExecution { + state: CancellableTaskState, + finished: bool, +} + +/// A request's local units. All transitions run on the owning region worker. +#[derive(Debug)] +pub(crate) struct LocalCompaction { + /// All pending and submitted units, retained until task exit. Tasks own their input leases. + /// Keyed by each unit's execution plan ID, allocated from the scheduler's `next_plan_id`; + /// it equals `CompactionExecution::plan_id()` and the `plan_id` in unit notifications. + units: HashMap, + pending: VecDeque, + parallelism: usize, + error: Option>, + made_progress: bool, +} + +impl LocalCompaction { + /// Requests cancellation for all units that have not started committing. + pub(crate) fn request_cancel(&self) -> RequestCancelResult { + let mut result = RequestCancelResult::TooLateToCancel; + for unit in self.units.values() { + match unit.state.request_cancel() { + RequestCancelResult::CancelIssued => result = RequestCancelResult::CancelIssued, + RequestCancelResult::AlreadyCancelling + if result != RequestCancelResult::CancelIssued => + { + result = RequestCancelResult::AlreadyCancelling; + } + _ => {} + } + } + result + } + + /// Counts submitted units that have not exited, including queued jobs. + fn num_running(&self) -> usize { + self.units.len() - self.pending.len() + } + + /// Checks this attempt's phase without depending on any sibling's apply status. + fn can_apply(&self, id: u64) -> bool { + self.units + .get(&id) + .is_some_and(|unit| !unit.finished && unit.state.commit_started()) + } + + /// Records a terminal result once, but still awaits task exit. + fn finish(&mut self, id: u64, result: std::result::Result>) { + let Some(unit) = self.units.get_mut(&id) else { + return; + }; + if unit.finished { + return; + } + unit.finished = true; + match result { + Ok(progress) => self.made_progress |= progress, + Err(err) => { + // A real failure must not be hidden by a sibling's cooperative cancellation. + if self.error.as_ref().is_none_or(|current| { + matches!(current.as_ref(), Error::CompactionCancelled { .. }) + }) { + self.error = Some(err); + } + } + } + } +} + +impl Drop for LocalCompaction { + /// Stops cancellable work when its region or compaction cycle is removed. + fn drop(&mut self) { + self.request_cancel(); + } +} + +impl CompactionScheduler { + /// Builds unit tasks with fresh execution IDs and shares their already-reserved inputs. + pub(crate) fn prepare_local_units( + &mut self, + region: CompactionRegion, + units: Vec, + reserved: CompactingFiles, + ) -> CompactionPhase { + let mut local = LocalCompaction { + units: HashMap::new(), + pending: VecDeque::new(), + parallelism: region.max_parallelism.max(1), + error: None, + made_progress: false, + }; + let context = CompactionContext::from(®ion); + // Planning is complete; unit tasks must not pin the other units' inputs. + drop(region); + for unit in units { + let files = reserved.for_inputs(&unit.inputs); + let id = Self::next_plan_id(&mut self.next_plan_id); + let execution = CompactionExecution::new(id, files); + let state = CancellableTaskState::new(); + local.units.insert( + id, + UnitExecution { + state: state.clone(), + finished: false, + }, + ); + local.pending.push_back(CompactionTask { + state, + execution, + compaction_context: context.clone(), + request_sender: self.request_sender.clone(), + listener: self.listener.clone(), + estimated_memory_bytes: unit.estimated_memory_bytes(), + unit, + memory_manager: self.memory_manager.clone(), + memory_policy: self.memory_policy, + uncommitted: UncommittedSsts::new( + context.region_id, + context.access_layer.clone(), + Some(context.cache_manager.clone()), + ), + }); + } + CompactionPhase::Units(local) + } + + /// Returns local unit state only while the region is in a local execution phase. + fn local_units_mut(&mut self, region_id: RegionId) -> Option<&mut LocalCompaction> { + match &mut self.region_status.get_mut(®ion_id)?.active.phase { + CompactionPhase::Units(local) => Some(local), + _ => None, + } + } + + /// Submits pending units up to the request limit and arranges task-exit notifications. + pub(crate) fn dispatch_local_units(&mut self, region_id: RegionId) { + loop { + let Some(local) = self.local_units_mut(region_id) else { + return; + }; + if local.num_running() >= local.parallelism { + return; + } + let Some(mut task) = local.pending.pop_front() else { + return; + }; + let id = task.execution.plan_id(); + let sender = self.request_sender.clone(); + let completion = Box::pin(async move { + let _ = sender + .send(WorkerRequestWithTime::new(WorkerRequest::Background { + region_id, + notify: BackgroundNotify::CompactionUnit( + CompactionUnitNotification::Released { plan_id: id }, + ), + })) + .await; + }); + let job = Box::pin(async move { + INFLIGHT_COMPACTION_COUNT.inc(); + task.run().await; + drop(task); + INFLIGHT_COMPACTION_COUNT.dec(); + }); + if let Err(err) = self.scheduler.schedule_with_completion(job, completion) + && let Some(local) = self.local_units_mut(region_id) + { + local.finish(id, Err(Arc::new(err))); + local.units.remove(&id); + } + } + } + + /// Accepts an edit from any current unit that has started committing and is not finished. + pub(crate) fn can_apply_unit(&self, region_id: RegionId, id: u64) -> bool { + self.region_status.get(®ion_id).is_some_and(|status| { + matches!(&status.active.phase, CompactionPhase::Units(local) + if local.can_apply(id)) + }) + } + + /// Records apply progress or failure without ending sibling executions. + pub(crate) fn on_unit_finished( + &mut self, + region_id: RegionId, + id: u64, + result: std::result::Result>, + ) { + if let Some(local) = self.local_units_mut(region_id) { + local.finish(id, result); + } + } + + /// Accounts for task exit, starts pending units, and checks request completion. + pub(crate) async fn on_unit_released( + &mut self, + region_id: RegionId, + id: u64, + manifest: &ManifestContextRef, + schemas: SchemaMetadataManagerRef, + ) -> CompactionTransition { + let Some(local) = self.local_units_mut(region_id) else { + return CompactionTransition::NoAction; + }; + if !local.units.contains_key(&id) { + return CompactionTransition::NoAction; + } + local.finish( + id, + Err(Arc::new( + error::UnexpectedSnafu { + reason: "Compaction unit exited without a terminal notification", + } + .build(), + )), + ); + local.units.remove(&id); + self.dispatch_local_units(region_id); + self.finish_local_units(region_id, manifest, schemas).await + } + + /// Completes the request only after all units exit, using the existing terminal paths. + pub(crate) async fn finish_local_units( + &mut self, + region_id: RegionId, + manifest: &ManifestContextRef, + schemas: SchemaMetadataManagerRef, + ) -> CompactionTransition { + let Some(local) = self.local_units_mut(region_id) else { + return CompactionTransition::NoAction; + }; + if !local.units.is_empty() { + return CompactionTransition::NoAction; + } + if let Some(err) = local.error.take() { + if matches!(err.as_ref(), Error::CompactionCancelled { .. }) { + return CompactionTransition::from_pending_ddls( + self.finish_compaction_on_cancel(region_id), + ); + } + self.on_compaction_failed(region_id, err); + return CompactionTransition::NoAction; + } + let progress = local.made_progress; + if let Some(status) = self.region_status.get_mut(®ion_id) { + for waiter in status.take_waiters() { + waiter.send(Ok(0)); + } + } + self.on_compaction_finished(region_id, manifest, schemas, progress) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_compaction_unit_apply_accepts_independent_executions() { + let units = [1, 2] + .into_iter() + .map(|id| { + ( + id, + UnitExecution { + state: CancellableTaskState::new(), + finished: false, + }, + ) + }) + .collect(); + let mut local = LocalCompaction { + units, + pending: VecDeque::new(), + parallelism: 2, + error: None, + made_progress: false, + }; + assert!(!local.can_apply(1)); + assert!(!local.can_apply(2)); + assert!(!local.can_apply(99)); + for unit in local.units.values() { + assert!(unit.state.mark_commit_started()); + } + assert!(local.can_apply(1) && local.can_apply(2)); + local.finish(2, Ok(true)); + assert!(local.can_apply(1)); + assert!(!local.can_apply(2)); + local.finish(2, Err(Arc::new(error::CompactionCancelledSnafu.build()))); + local.finish(99, Err(Arc::new(error::CompactionCancelledSnafu.build()))); + assert!(local.can_apply(1)); + assert!(local.error.is_none()); + assert!(local.made_progress); + local.finish(1, Ok(false)); + assert!(!local.can_apply(1)); + assert_eq!( + 2, + local.units.len(), + "terminal units remain tracked until task exit" + ); + } +} diff --git a/src/mito2/src/compaction/scheduler/planning.rs b/src/mito2/src/compaction/scheduler/planning.rs index 5553d6435da..d427044b955 100644 --- a/src/mito2/src/compaction/scheduler/planning.rs +++ b/src/mito2/src/compaction/scheduler/planning.rs @@ -14,7 +14,7 @@ use std::fmt; use std::future::Future; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use std::time::Instant; use api::v1::region::compact_request; @@ -30,12 +30,12 @@ use tokio::sync::mpsc::{self, Sender}; use crate::access_layer::AccessLayerRef; use crate::cache::CacheManagerRef; -use crate::compaction::compactor::{CompactionRegion, CompactionVersion, DefaultCompactor}; -use crate::compaction::picker::{CompactionTask, PickerOutput, new_picker}; +use crate::compaction::compactor::{CompactionRegion, CompactionVersion}; +use crate::compaction::picker::{PickerOutput, new_picker}; use crate::compaction::scheduler::state::{CompactingFiles, CompactionExecution, CompactionPhase}; use crate::compaction::scheduler::{CompactionScheduler, CompactionTransition}; -use crate::compaction::task::CompactionTaskImpl; -use crate::compaction::{CompactionOutput, find_dynamic_options}; +use crate::compaction::unit::CompactionUnit; +use crate::compaction::{CompactionOutput, find_dynamic_options, format_panic_reason}; use crate::config::MitoConfig; use crate::error::{CompactRegionSnafu, Error, RemoteCompactionSnafu, Result, UnexpectedSnafu}; use crate::metrics::{ @@ -44,11 +44,10 @@ use crate::metrics::{ use crate::region::ManifestContextRef; use crate::region::options::RegionOptions; use crate::request::{BackgroundNotify, OutputTx, WorkerRequest, WorkerRequestWithTime}; -use crate::schedule::CancellableTaskState; use crate::schedule::remote_job_scheduler::{ CompactionJob, DefaultNotifier, RemoteJob, RemoteJobSchedulerRef, }; -use crate::sst::file::{FileHandle, UncommittedSsts}; +use crate::sst::file::FileHandle; use crate::sst::version::SstVersion; use crate::worker::WorkerListener; @@ -149,13 +148,7 @@ impl CompactionScheduler { ) { // The idiomatic way to handle a panic result. let result = std::panic::AssertUnwindSafe(planning).catch_unwind().await.unwrap_or_else(|payload| { - let reason = if let Some(message) = payload.as_ref().downcast_ref::<&str>() { - message.to_string() - } else if let Some(message) = payload.as_ref().downcast_ref::() { - message.clone() - } else { - "unknown panic".to_string() - }; + let reason = format_panic_reason(payload.as_ref()); CompactionPlanningResult::Error(Arc::new( UnexpectedSnafu { reason: format!( @@ -333,16 +326,18 @@ impl CompactionScheduler { }; let waiters = status.take_waiters(); match self - .submit_prepared_compaction(prepared, files, waiters, plan_id) + .prepare_or_submit_compaction(prepared, files, waiters, plan_id) .await { Ok(Some(phase)) => { if let Some(status) = self.region_status.get_mut(®ion_id) { status.set_phase(phase); } + self.dispatch_local_units(region_id); // The execution now owns the terminal transition. Any // fenced DDLs remain queued until its callback arrives. - CompactionTransition::NoAction + self.finish_local_units(region_id, manifest_ctx, schema_metadata_manager) + .await } Ok(None) => { self.finish_compaction_planning( @@ -427,12 +422,14 @@ impl CompactionScheduler { CompactionTransition::NoAction } - async fn submit_prepared_compaction( + /// Submits a remote job, or prepares local unit tasks without submitting them. + /// The caller dispatches local units after installing the returned phase. + async fn prepare_or_submit_compaction( &mut self, prepared: PreparedCompaction, files: CompactingFiles, waiters: Vec, - mut plan_id: u64, + plan_id: u64, ) -> Result> { let PreparedCompaction { compaction_region, @@ -492,7 +489,6 @@ impl CompactionScheduler { error!(e; "Failed to schedule remote compaction job for region {}, fallback to local compaction", region_id); // An error may be ambiguous after the remote scheduler consumed // the notifier. Fence a delayed remote callback from the local fallback. - plan_id = Self::next_plan_id(&mut self.next_plan_id); e.waiters } } @@ -507,9 +503,19 @@ impl CompactionScheduler { waiters }; - // 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) { + let units = CompactionUnit::from_picker(picker_output); + if let Some(status) = self.region_status.get_mut(®ion_id) { + status.extend_waiters(waiters); + } + let Some(mut units) = units else { + return UnexpectedSnafu { + reason: "Compaction plan expires a merge input", + } + .fail(); + }; + units.retain(|unit| { + let estimated_bytes = unit.estimated_memory_bytes(); + let Some(limit_bytes) = self.exceeds_compaction_memory_limit(estimated_bytes) else { return true }; COMPACTION_MEMORY_REJECTED .with_label_values(&["oversized"]) .inc(); @@ -517,82 +523,16 @@ impl CompactionScheduler { "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)); - } + false + }); + if units.is_empty() { return Ok(None); } - - let state = CancellableTaskState::new(); - let cancel_handle = state.cancel_handle(); - let execution = CompactionExecution::new(plan_id, files); - let uncommitted = UncommittedSsts::new( - region_id, - compaction_region.access_layer.clone(), - Some(compaction_region.cache_manager.clone()), - ); - let local_compaction_task = Box::new(CompactionTaskImpl { - state: state.clone(), - execution: execution.clone(), - request_sender: self.request_sender.clone(), - waiters, - start_time, - listener: self.listener.clone(), - picker_output, + Ok(Some(self.prepare_local_units( compaction_region, - compactor: Arc::new(DefaultCompactor::with_cancel_handle( - cancel_handle.clone(), - uncommitted.clone(), - )), - memory_manager: self.memory_manager.clone(), - memory_policy: self.memory_policy, - estimated_memory_bytes: estimated_bytes, - uncommitted, - }); - - match self.submit_compaction_task(local_compaction_task, region_id) { - Ok(()) => Ok(Some(CompactionPhase::Local { state, execution })), - Err((err, task)) => { - if let (Some(status), Some(mut task)) = - (self.region_status.get_mut(®ion_id), task) - { - status.append_waiters(&mut task.waiters); - } - Err(err) - } - } - } - - fn submit_compaction_task( - &mut self, - task: Box, - region_id: RegionId, - ) -> std::result::Result<(), (Error, Option>)> { - let task = Arc::new(Mutex::new(Some(task))); - let task_to_run = task.clone(); - match self.scheduler.schedule(Box::pin(async move { - let task = task_to_run - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .take(); - if let Some(mut task) = task { - INFLIGHT_COMPACTION_COUNT.inc(); - task.run().await; - INFLIGHT_COMPACTION_COUNT.dec(); - } else { - error!("Compaction task was missing when the scheduled job started"); - } - })) { - Ok(()) => Ok(()), - Err(err) => { - error!(err; "Failed to submit compaction request for region {}", region_id); - let task = task - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .take(); - Err((err, task)) - } - } + units, + files, + ))) } fn exceeds_compaction_memory_limit(&self, estimated_bytes: u64) -> Option { @@ -605,20 +545,6 @@ impl CompactionScheduler { } } -/// Estimates compaction memory as the sum of all input files' maximum row-group -/// uncompressed sizes. -fn estimate_compaction_bytes(picker_output: &PickerOutput) -> u64 { - picker_output - .outputs - .iter() - .flat_map(|output| output.inputs.iter()) - .map(|file: &FileHandle| { - let meta = file.meta_ref(); - meta.max_row_group_uncompressed_size - }) - .sum() -} - /// Rebuilds picker output with current SST handles while preserving the picker's grouping. /// /// Picking runs in background on a version snapshot that may be stale by the diff --git a/src/mito2/src/compaction/scheduler/state.rs b/src/mito2/src/compaction/scheduler/state.rs index 4fe4a686751..02423f8e32b 100644 --- a/src/mito2/src/compaction/scheduler/state.rs +++ b/src/mito2/src/compaction/scheduler/state.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet, hash_map}; use std::sync::Arc; use std::time::Instant; @@ -28,6 +28,7 @@ use crate::access_layer::AccessLayerRef; use crate::cache::CacheManagerRef; use crate::compaction::compactor::CompactionVersion; use crate::compaction::picker::PickerOutput; +use crate::compaction::scheduler::local::LocalCompaction; use crate::compaction::scheduler::planning::CompactionRequest; use crate::config::MitoConfig; use crate::error::{ @@ -36,8 +37,10 @@ use crate::error::{ use crate::region::ManifestContextRef; use crate::region::version::VersionControlRef; use crate::request::{OptionOutputTx, OutputTx, SenderDdlRequest, WorkerRequestWithTime}; -use crate::schedule::{CancellableTaskState, RequestCancelResult}; -use crate::sst::file::FileHandle; +#[cfg(test)] +use crate::schedule::CancellableTaskState; +use crate::schedule::RequestCancelResult; +use crate::sst::file::{FileHandle, RegionFileId}; use crate::worker::WorkerListener; /// Identifies an accepted compaction attempt and keeps its SST reservations alive. @@ -49,6 +52,10 @@ pub(crate) struct CompactionExecution { } impl CompactionExecution { + /// Returns the attempt ID used to reject stale execution notifications. + pub(crate) fn plan_id(&self) -> u64 { + self.plan_id + } pub(super) fn new(plan_id: u64, files: CompactingFiles) -> Self { Self { plan_id, @@ -67,13 +74,14 @@ impl CompactionExecution { } #[derive(Debug)] -pub(super) enum CompactionPhase { +pub(crate) enum CompactionPhase { /// The region is registered but has no compaction in progress. Idle, Picking { plan_id: u64, cancelled: bool, }, + #[cfg(test)] Local { state: CancellableTaskState, execution: CompactionExecution, @@ -81,6 +89,7 @@ pub(super) enum CompactionPhase { Remote { execution: CompactionExecution, }, + Units(LocalCompaction), } /// Describes how the current compaction cycle was triggered. @@ -179,9 +188,10 @@ impl ActiveCompaction { pub(super) fn matches_execution(&self, execution: &CompactionExecution) -> bool { match &self.phase { CompactionPhase::Idle | CompactionPhase::Picking { .. } => None, - CompactionPhase::Local { execution, .. } | CompactionPhase::Remote { execution } => { - Some(execution) - } + #[cfg(test)] + CompactionPhase::Local { execution, .. } => Some(execution), + CompactionPhase::Units(_) => None, + CompactionPhase::Remote { execution } => Some(execution), } .is_some_and(|current| current.matches(execution)) } @@ -197,7 +207,9 @@ impl ActiveCompaction { RequestCancelResult::CancelIssued } } + #[cfg(test)] CompactionPhase::Local { state, .. } => state.request_cancel(), + CompactionPhase::Units(units) => units.request_cancel(), CompactionPhase::Remote { .. } => RequestCancelResult::TooLateToCancel, } } @@ -211,56 +223,73 @@ impl ActiveCompaction { /// Owns atomic reservations for every SST selected by a compaction plan. #[derive(Debug, Clone)] -pub(super) struct CompactingFiles { - _inner: Arc, +pub(crate) struct CompactingFiles { + files: HashMap>, } +/// Keeps an SST reserved until the last shared lease is dropped. #[derive(Debug)] -struct CompactingFilesInner { - files: Vec, +struct CompactingFile { + file: FileHandle, } impl CompactingFiles { + /// Reserves both merge inputs and expired files from an accepted picker result. pub(super) fn try_new(output: &PickerOutput) -> Option { - let mut seen = HashSet::new(); - let mut files: Vec = Vec::new(); - let selected_files = output - .outputs - .iter() - .flat_map(|output| output.inputs.iter()) - .chain(output.expired_ssts.iter()); + Self::try_reserve( + output + .outputs + .iter() + .flat_map(|o| &o.inputs) + .chain(&output.expired_ssts), + ) + } + + /// Reserves distinct inputs, releasing acquired leases if any input is already busy. + fn try_reserve<'a>(selected_files: impl IntoIterator) -> Option { + let mut files = HashMap::new(); for file in selected_files { - if !seen.insert(file.file_id()) { + let hash_map::Entry::Vacant(entry) = files.entry(file.file_id()) else { continue; - } + }; if !file.try_set_compacting() { - for reserved in &files { - reserved.set_compacting(false); - } return None; } - files.push(file.clone()); + entry.insert(Arc::new(CompactingFile { file: file.clone() })); } - Some(Self { - _inner: Arc::new(CompactingFilesInner { files }), - }) + Some(Self { files }) + } + + /// Clones the shared leases of a dependency group's inputs without briefly releasing any input. + /// `self` retains its leases, so an input is released only after every owner, including + /// the original collection, drops its lease. + /// A delayed remote notifier may retain the same leases across local fallback. + /// Indexing keeps splitting proportional to the group's input count. + pub(crate) fn for_inputs(&self, inputs: &[FileHandle]) -> Self { + Self { + files: inputs + .iter() + .filter_map(|file| { + let id = file.file_id(); + self.files.get(&id).map(|lease| (id, lease.clone())) + }) + .collect(), + } } #[cfg(test)] pub(super) fn empty() -> Self { Self { - _inner: Arc::new(CompactingFilesInner { files: Vec::new() }), + files: HashMap::new(), } } } -impl Drop for CompactingFilesInner { +impl Drop for CompactingFile { fn drop(&mut self) { - for file in &self.files { - file.set_compacting(false); - } + self.file.set_compacting(false); } } @@ -426,10 +455,6 @@ impl CompactionStatus { self.active.waiters.extend(waiters); } - pub(super) fn append_waiters(&mut self, waiters: &mut Vec) { - self.active.waiters.append(waiters); - } - pub(super) fn set_phase(&mut self, phase: CompactionPhase) { self.active.phase = phase; } @@ -529,3 +554,34 @@ pub(super) struct PendingCompaction { /// Optional time range that constrains candidate compaction windows. pub(crate) time_range: Option, } + +#[cfg(test)] +mod tests { + use store_api::storage::FileId; + + use super::*; + use crate::compaction::test_util::new_file_handle; + + #[test] + fn test_compaction_unit_input_leases_release_independently() { + let x = new_file_handle(FileId::random(), 0, 100, 0); + let y = new_file_handle(FileId::random(), 100, 200, 0); + let unused = new_file_handle(FileId::random(), 200, 300, 0); + let reserved = CompactingFiles::try_reserve([&x, &y, &x, &unused]).unwrap(); + let first = reserved.for_inputs(&[x.clone(), x.clone()]); + let second = reserved.for_inputs(std::slice::from_ref(&y)); + drop(reserved); + assert!(x.compacting() && y.compacting()); + assert!( + !unused.compacting(), + "unit leases must not retain unrelated inputs" + ); + drop(first); + assert!(!x.compacting() && y.compacting()); + let notification = second.clone(); + drop(second); + assert!(y.compacting()); + drop(notification); + assert!(!y.compacting()); + } +} diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index fe624316357..d8c7b3944a5 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -1033,7 +1033,7 @@ async fn test_remote_fallback_uses_new_execution_plan_id() { assert_eq!(job_scheduler.num_jobs(), 1); assert!(matches!( &scheduler.region_status[®ion_id].active.phase, - CompactionPhase::Local { .. } + CompactionPhase::Units(_) )); assert!( !scheduler.region_status[®ion_id] diff --git a/src/mito2/src/compaction/task.rs b/src/mito2/src/compaction/task.rs index 6f87c37ddf8..cc6c3c9b859 100644 --- a/src/mito2/src/compaction/task.rs +++ b/src/mito2/src/compaction/task.rs @@ -14,57 +14,42 @@ use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use std::time::Instant; use common_base::cancellation::CancellableFuture; use common_memory_manager::OnExhaustedPolicy; -use common_telemetry::{error, info, warn}; -use itertools::Itertools; +use common_telemetry::error; +use futures::FutureExt; use snafu::ResultExt; -use store_api::ManifestVersion; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, oneshot}; -use crate::compaction::CompactionExecution; -use crate::compaction::compactor::{CompactionRegion, Compactor, MergeOutput}; -use crate::compaction::memory_manager::{CompactionMemoryGuard, CompactionMemoryManager}; -use crate::compaction::picker::{CompactionTask, PickerOutput}; -use crate::error::{CompactRegionSnafu, CompactionMemoryExhaustedSnafu}; -use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; -use crate::metrics::{COMPACTION_FAILURE_COUNT, COMPACTION_MEMORY_WAIT, COMPACTION_STAGE_ELAPSED}; -use crate::region::RegionRoleState; +use crate::compaction::compactor::{CompactionContext, DefaultCompactor, MergeOutput}; +use crate::compaction::memory_manager::CompactionMemoryManager; +use crate::compaction::unit::CompactionUnit; +use crate::compaction::{CompactionExecution, format_panic_reason}; +use crate::error::{self, CompactionCancelledSnafu, CompactionMemoryExhaustedSnafu, Result}; +use crate::manifest::action::RegionEdit; +use crate::metrics::{ + COMPACTION_FAILURE_COUNT, COMPACTION_INPUT_BYTES, COMPACTION_OUTPUT_BYTES, + COMPACTION_STAGE_ELAPSED, +}; use crate::request::{ - BackgroundNotify, CompactionCancelled, CompactionFailed, CompactionFinished, OutputTx, - RegionEditResult, Waiters, WorkerRequest, WorkerRequestWithTime, + BackgroundNotify, CompactionUnitNotification, WorkerRequest, WorkerRequestWithTime, }; use crate::schedule::CancellableTaskState; -use crate::sst::file::{FileMeta, UncommittedSsts}; +use crate::sst::file::UncommittedSsts; use crate::worker::WorkerListener; -use crate::{error, metrics}; -/// Maximum number of compaction tasks in parallel. pub const MAX_PARALLEL_COMPACTION: usize = 1; -pub(crate) struct CompactionTaskImpl { - /// Shared local-compaction state for cooperative cancellation. +/// One local atomic unit. The scheduler owns request waiters and sibling units. +pub(crate) struct CompactionTask { pub(crate) state: CancellableTaskState, - /// Identity and reservation lease of this accepted execution. pub(crate) execution: CompactionExecution, - pub compaction_region: CompactionRegion, - /// Request sender to notify the worker. + pub(crate) compaction_context: CompactionContext, pub(crate) request_sender: mpsc::Sender, - /// Senders that are used to notify waiters waiting for pending compaction tasks. - pub waiters: Vec, - /// Start time of compaction task - pub start_time: Instant, - /// Event listener. pub(crate) listener: WorkerListener, - /// Compactor to handle compaction. - pub(crate) compactor: Arc, - /// Output of the picker. - pub(crate) picker_output: PickerOutput, - /// Memory manager to acquire memory budget. + pub(crate) unit: CompactionUnit, pub(crate) memory_manager: Arc, - /// Policy when memory is exhausted. pub(crate) memory_policy: OnExhaustedPolicy, /// Estimated memory bytes needed for this compaction. pub(crate) estimated_memory_bytes: u64, @@ -72,397 +57,229 @@ pub(crate) struct CompactionTaskImpl { pub(crate) uncommitted: UncommittedSsts, } -impl Debug for CompactionTaskImpl { +impl Debug for CompactionTask { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TwcsCompactionTask") - .field("region_id", &self.compaction_region.region_id) - .field("picker_output", &self.picker_output) + f.debug_struct("CompactionTask") + .field("region_id", &self.compaction_context.region_id) + .field("unit", &self.unit) .field( "append_mode", - &self.compaction_region.region_options.append_mode, + &self.compaction_context.region_options.append_mode, ) .finish() } } -impl CompactionTaskImpl { - /// Acquires memory budget based on the configured policy. - /// - /// Returns an error if memory cannot be acquired according to the policy. - async fn acquire_memory_with_policy(&self) -> error::Result { - let region_id = self.compaction_region.region_id; - let requested_bytes = self.estimated_memory_bytes; - let policy = self.memory_policy; - - let _timer = COMPACTION_MEMORY_WAIT.start_timer(); - self.memory_manager - .acquire_with_policy(requested_bytes, policy) +impl CompactionTask { + /// Executes one unit and retains its resources until worker apply or failure acknowledgement. + pub(crate) async fn run(&mut self) { + let start = std::time::Instant::now(); + let result = std::panic::AssertUnwindSafe(self.run_once()) + .catch_unwind() .await - .context(CompactionMemoryExhaustedSnafu { - region_id, - policy: format!("{policy:?}"), - }) + .unwrap_or_else(|payload| Err(self.panic_to_error(payload.as_ref()))); + self.complete(result, start).await; } - fn cancelled_notify(&mut self) -> BackgroundNotify { - let senders = std::mem::take(&mut self.waiters); - BackgroundNotify::CompactionCancelled(CompactionCancelled { - region_id: self.compaction_region.region_id, - execution: self.execution.clone(), - senders, - }) - } - - /// Remove expired ssts files, update manifest immediately - /// and apply the edit to region version. - /// - /// This function logs errors but does not stop the compaction process if removal fails. - async fn remove_expired( - &self, - compaction_region: &CompactionRegion, - expired_files: Vec, - ) { - let region_id = compaction_region.region_id; - let expired_files_str = expired_files.iter().map(|f| f.file_id).join(","); - let (expire_delete_sender, expire_delete_listener) = tokio::sync::oneshot::channel(); - // Update manifest to remove expired SSTs - let edit = RegionEdit { - files_to_add: Vec::new(), - files_to_remove: expired_files, - timestamp_ms: Some(chrono::Utc::now().timestamp_millis()), - compaction_time_window: None, - flushed_entry_id: None, - flushed_sequence: None, - committed_sequence: None, + /// Acquires the unit's memory budget and merges on the compaction runtime. + async fn merge(&self) -> Result { + let region_id = self.compaction_context.region_id; + self.listener + .on_compaction_unit_merge_begin(region_id, self.execution.plan_id()) + .await; + let bytes = self.estimated_memory_bytes; + let _memory_guard = if bytes == 0 { + None + } else { + let _wait_timer = crate::metrics::COMPACTION_MEMORY_WAIT.start_timer(); + Some( + CancellableFuture::new( + self.memory_manager + .acquire_with_policy(bytes, self.memory_policy), + self.state.cancel_handle(), + ) + .await + .map_err(|_| CompactionCancelledSnafu.build())? + .context(CompactionMemoryExhaustedSnafu { + region_id, + policy: format!("{:?}", self.memory_policy), + })?, + ) }; - - // 1. Update manifest - let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())); - let RegionRoleState::Leader(current_region_state) = - compaction_region.manifest_ctx.current_state() - else { - warn!( - "Region {} not in leader state, skip removing expired files", - region_id - ); - return; - }; - if let Err(e) = compaction_region - .manifest_ctx - .update_manifest(current_region_state, action_list, false) - .await - { - warn!( - e; - "Failed to update manifest for expired files removal, region: {region_id}, files: [{expired_files_str}]. Compaction will continue." - ); - return; - } - - // 2. Notify region worker loop to remove expired files from region version. - self.send_to_worker(WorkerRequest::Background { - region_id, - notify: BackgroundNotify::RegionEdit(RegionEditResult { - region_id, - waiters: Waiters::one(expire_delete_sender), - edit, - result: Ok(()), - update_region_state: false, - is_staging: false, - }), - }) - .await; - - if let Err(e) = expire_delete_listener - .await - .context(error::RecvSnafu) - .flatten() - { - warn!( - e; - "Failed to remove expired files from region version, region: {region_id}, files: [{expired_files_str}]. Compaction will continue." - ); - return; - } - - info!( - "Successfully removed expired files, region: {region_id}, files: [{expired_files_str}]" - ); - } - - async fn handle_expiration(&mut self) { - // 1. In case of local compaction, we can delete expired ssts in advance. - if !self.picker_output.expired_ssts.is_empty() { - let remove_timer = COMPACTION_STAGE_ELAPSED - .with_label_values(&["remove_expired"]) - .start_timer(); - let expired_ssts = self - .picker_output - .expired_ssts - .drain(..) - .map(|f| f.meta_ref().clone()) - .collect(); - // remove_expired logs errors but doesn't stop compaction - self.remove_expired(&self.compaction_region, expired_ssts) - .await; - remove_timer.observe_duration(); - } - } - - async fn handle_compaction(&mut self) -> error::Result { - // 2. Merge inputs - let merge_timer = COMPACTION_STAGE_ELAPSED + let _timer = COMPACTION_STAGE_ELAPSED .with_label_values(&["merge"]) .start_timer(); - - let compaction_result = match self - .compactor - .merge_ssts(&self.compaction_region, self.picker_output.clone()) - .await - { - Ok(v) => v, - Err(e) => { - error!(e; "Failed to compact region: {}", self.compaction_region.region_id); - merge_timer.stop_and_discard(); - return Err(e); - } - }; - let merge_time = merge_timer.stop_and_record(); - - metrics::COMPACTION_INPUT_BYTES.inc_by(compaction_result.input_file_size() as f64); - metrics::COMPACTION_OUTPUT_BYTES.inc_by(compaction_result.output_file_size() as f64); - info!( - "Compacted SST files, region_id: {}, input: {:?}, output: {:?}, window: {:?}, waiter_num: {}, merge_time: {}s", - self.compaction_region.region_id, - compaction_result.files_to_remove, - compaction_result.files_to_add, - compaction_result.compaction_time_window, - self.waiters.len(), - merge_time, + let compactor = DefaultCompactor::with_cancel_handle( + self.state.cancel_handle(), + self.uncommitted.clone(), ); - - self.listener - .on_merge_ssts_finished(self.compaction_region.region_id) - .await; - - Ok(compaction_result) + let context = self.compaction_context.clone(); + let unit = self.unit.clone(); + let output = + common_runtime::spawn_compact( + async move { compactor.merge_unit(&context, &unit).await }, + ) + .await + .context(error::JoinSnafu)??; + COMPACTION_INPUT_BYTES.inc_by(output.input_file_size() as f64); + COMPACTION_OUTPUT_BYTES.inc_by(output.output_file_size() as f64); + self.listener.on_merge_ssts_finished(region_id).await; + Ok(output) } - async fn update_manifest( - &self, - compaction_result: crate::compaction::compactor::MergeOutput, - ) -> error::Result<(RegionEdit, ManifestVersion)> { - let _manifest_timer = COMPACTION_STAGE_ELAPSED + /// Commits this unit through the existing manifest path, independently of sibling applies. + async fn publish(&self, output: MergeOutput) -> Result { + if !self.state.mark_commit_started() { + return CompactionCancelledSnafu.fail(); + } + // Report this unit's SSTs before its manifest update; sibling units may progress independently. + self.compaction_context.invoke_sst_hook(&output).await; + self.listener + .on_compaction_commit_begin(self.compaction_context.region_id) + .await; + let _timer = COMPACTION_STAGE_ELAPSED .with_label_values(&["write_manifest"]) .start_timer(); - - self.compactor - .update_manifest(&self.compaction_region, compaction_result) - .await + let (edit, _manifest_version) = self.compaction_context.update_manifest(output).await?; + self.uncommitted.disarm_cleanup(); + self.listener + .on_compaction_unit_committed( + self.compaction_context.region_id, + self.execution.plan_id(), + ) + .await; + Ok(edit) } - /// Handles compaction failure, notifies all waiters. - pub(crate) fn on_failure(&mut self, err: Arc) { - COMPACTION_FAILURE_COUNT.inc(); - for waiter in self.waiters.drain(..) { - waiter.send(Err(err.clone()).context(CompactRegionSnafu { - region_id: self.compaction_region.region_id, - })); + /// Converts a caught panic into the unit's terminal error. + fn panic_to_error(&self, payload: &(dyn std::any::Any + Send + 'static)) -> error::Error { + let region_id = self.compaction_context.region_id; + let plan_id = self.execution.plan_id(); + error::UnexpectedSnafu { + reason: format!( + "Compaction unit panicked for region {region_id}, plan_id {plan_id}: {}", + format_panic_reason(payload) + ), } + .build() } - /// Notifies region worker to handle post-compaction tasks. - async fn send_to_worker(&self, request: WorkerRequest) { - if let Err(e) = self - .request_sender - .send(WorkerRequestWithTime::new(request)) - .await - { - error!( - "Failed to notify compaction job status for region {}, request: {:?}", - self.compaction_region.region_id, e.0 - ); - } + /// Runs merge and publication. + async fn run_once(&mut self) -> Result { + let output = self.merge().await?; + self.publish(output).await } - async fn invoke_sst_hook(&self, merge_output: &MergeOutput) { - self.compaction_region.invoke_sst_hook(merge_output).await; - } -} - -#[async_trait::async_trait] -impl CompactionTask for CompactionTaskImpl { - async fn run(&mut self) { - // Acquire memory budget before starting compaction - let cancel_handle = self.state.cancel_handle(); - let _memory_guard = match CancellableFuture::new( - self.acquire_memory_with_policy(), - cancel_handle, - ) - .await - { - Ok(Ok(guard)) => guard, - Ok(Err(e)) => { - error!(e; "Failed to acquire memory for compaction, region id: {}", self.compaction_region.region_id); - let err = Arc::new(e); - self.on_failure(err.clone()); - let notify = BackgroundNotify::CompactionFailed(CompactionFailed { - region_id: self.compaction_region.region_id, - execution: self.execution.clone(), - err, - }); - self.send_to_worker(WorkerRequest::Background { - region_id: self.compaction_region.region_id, - notify, - }) - .await; - return; + /// Cleans up and reports a terminal result before task exit. + async fn complete(&self, result: Result, start: std::time::Instant) { + if let Err(err) = &result { + if !matches!(err, error::Error::CompactionCancelled { .. }) { + COMPACTION_FAILURE_COUNT.inc(); + error!(err; "Compaction unit failed, region: {}, plan_id: {}", + self.compaction_context.region_id, self.execution.plan_id()); } - Err(_) => { - info!( - "Compaction cancelled while waiting for memory, region id: {}", - self.compaction_region.region_id - ); - let notify = self.cancelled_notify(); - self.send_to_worker(WorkerRequest::Background { - region_id: self.compaction_region.region_id, - notify, - }) - .await; - return; - } - }; - - self.handle_expiration().await; - - // The local compactor owns cancellation of its spawned merge tasks. Waiting for it to - // return ensures all finalized outputs are tracked before cleanup starts. - let notify = match self.handle_compaction().await { - Ok(merge_output) => { - self.invoke_sst_hook(&merge_output).await; - // Stop accepting cancellation once we are about to publish the compaction edit. - if !self.state.mark_commit_started() { - self.uncommitted.cleanup().await; - self.cancelled_notify() - } else { - self.listener - .on_compaction_commit_begin(self.compaction_region.region_id) - .await; - match self.update_manifest(merge_output).await { - Ok((edit, _manifest_version)) => { - self.uncommitted.disarm_cleanup(); - let senders = std::mem::take(&mut self.waiters); - BackgroundNotify::CompactionFinished(CompactionFinished { - region_id: self.compaction_region.region_id, - execution: self.execution.clone(), - senders, - start_time: self.start_time, - edit, - }) - } - Err(e) => { - if e.may_have_persisted_manifest_update() { - self.uncommitted.disarm_cleanup(); - } else { - info!( - "Cleaning uncommitted SSTs because the manifest update was not persisted, region: {}, job: compaction, error: {:?}", - self.compaction_region.region_id, e - ); - self.uncommitted.cleanup().await; - } - error!(e; "Failed to compact region, region id: {}", self.compaction_region.region_id); - let err = Arc::new(e); - self.on_failure(err.clone()); - BackgroundNotify::CompactionFailed(CompactionFailed { - region_id: self.compaction_region.region_id, - execution: self.execution.clone(), - err, - }) - } - } - } - } - Err(e) => { - error!(e; "Failed to compact region, region id: {}", self.compaction_region.region_id); + if self.state.commit_started() && err.may_have_persisted_manifest_update() { + self.uncommitted.disarm_cleanup(); + } else { self.uncommitted.cleanup().await; - let err = Arc::new(e); - // notify compaction waiters - self.on_failure(err.clone()); - BackgroundNotify::CompactionFailed(CompactionFailed { - region_id: self.compaction_region.region_id, - execution: self.execution.clone(), - err, - }) } + } + let succeeded = result.is_ok(); + let (applied, receiver) = oneshot::channel(); + let notify = CompactionUnitNotification::Finished { + plan_id: self.execution.plan_id(), + result: result.map_err(Arc::new), + applied, }; + if let Err(err) = self.notify_result(notify).await { + error!(err; "Failed to notify compaction unit completion, region: {}", self.compaction_context.region_id); + return; + } + // This keeps the execution slot and inputs alive through worker apply. + if receiver.await.is_ok() && succeeded { + crate::metrics::COMPACTION_ELAPSED_TOTAL + .with_label_values(&["local"]) + .observe(start.elapsed().as_secs_f64()); + } + } - self.send_to_worker(WorkerRequest::Background { - region_id: self.compaction_region.region_id, - notify, - }) - .await; + /// Sends a unit lifecycle event to the owning region worker. + async fn notify_result(&self, notify: CompactionUnitNotification) -> Result<()> { + self.request_sender + .send(WorkerRequestWithTime::new(WorkerRequest::Background { + region_id: self.compaction_context.region_id, + notify: BackgroundNotify::CompactionUnit(notify), + })) + .await + .map_err(|_| error::InvalidSenderSnafu.build()) } } #[cfg(test)] mod tests { - use store_api::storage::FileId; + use std::time::Duration; - use crate::compaction::picker::PickerOutput; - use crate::compaction::test_util::new_file_handle; + use store_api::storage::RegionId; - #[test] - fn test_picker_output_with_expired_ssts() { - // Test that PickerOutput correctly includes expired_ssts - // This verifies that expired SSTs are properly identified and included - // in the picker output, which is then handled by handle_expiration() + use super::*; + use crate::compaction::compactor::new_test_compaction_region; + use crate::compaction::memory_manager::new_compaction_memory_manager; + use crate::engine::listener::EventListener; + use crate::error::Error; - let file_ids = (0..3).map(|_| FileId::random()).collect::>(); - let expired_ssts = vec![ - new_file_handle(file_ids[0], 0, 999, 0), - new_file_handle(file_ids[1], 1000, 1999, 0), - ]; + struct PanicOnUnitMergeBegin; - let picker_output = PickerOutput { - outputs: vec![], - expired_ssts: expired_ssts.clone(), - time_window_size: 3600, - max_file_size: None, - }; - - // Verify expired_ssts are included - assert_eq!(picker_output.expired_ssts.len(), 2); - assert_eq!( - picker_output.expired_ssts[0].file_id(), - expired_ssts[0].file_id() - ); - assert_eq!( - picker_output.expired_ssts[1].file_id(), - expired_ssts[1].file_id() - ); + #[async_trait::async_trait] + impl EventListener for PanicOnUnitMergeBegin { + async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) { + panic!("unit merge listener boom"); + } } - #[test] - fn test_picker_output_without_expired_ssts() { - // Test that PickerOutput works correctly when there are no expired SSTs - let picker_output = PickerOutput { - outputs: vec![], - expired_ssts: vec![], - time_window_size: 3600, - max_file_size: None, + #[tokio::test] + async fn test_compaction_unit_panic_reports_terminal_error() { + let region = new_test_compaction_region().await; + let (request_sender, mut requests) = mpsc::channel(1); + let mut task = CompactionTask { + state: CancellableTaskState::new(), + execution: CompactionExecution::for_test(7), + compaction_context: CompactionContext::from(®ion), + request_sender, + listener: WorkerListener::new(Some(Arc::new(PanicOnUnitMergeBegin))), + unit: CompactionUnit { + inputs: Vec::new(), + outputs: Vec::new(), + time_window_size: 0, + max_file_size: None, + }, + memory_manager: Arc::new(new_compaction_memory_manager(0)), + memory_policy: OnExhaustedPolicy::default(), + estimated_memory_bytes: 0, + uncommitted: UncommittedSsts::new(region.region_id, region.access_layer.clone(), None), }; + let task_handle = tokio::spawn(async move { task.run().await }); - // Verify empty expired_ssts - assert!(picker_output.expired_ssts.is_empty()); + let request = tokio::time::timeout(Duration::from_secs(5), requests.recv()) + .await + .unwrap() + .unwrap(); + let WorkerRequest::Background { + notify: + BackgroundNotify::CompactionUnit(CompactionUnitNotification::Finished { + plan_id, + result, + applied, + }), + .. + } = request.request + else { + panic!("expected unit terminal notification"); + }; + assert_eq!(7, plan_id); + let err = result.unwrap_err(); + assert!(matches!(err.as_ref(), Error::Unexpected { .. })); + assert!(err.to_string().contains("unit merge listener boom")); + applied.send(()).unwrap(); + task_handle.await.unwrap(); } - - // Note: Testing remove_expired() directly requires extensive mocking of: - // - manifest_ctx (ManifestContext) - // - request_sender (mpsc::Sender) - // - WorkerRequest handling - // - // The behavior is tested indirectly through integration tests: - // - remove_expired() logs errors but doesn't stop compaction - // - handle_expiration() continues even if remove_expired() encounters errors - // - The expiration stage is designed to be non-blocking for compaction } diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index c170845be3d..a1f51b64091 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -3489,6 +3489,25 @@ mod tests { .collect::>(); assert_eq!(vec![(0, 1), (0, 0), (1, 3), (1, 2), (1, 5)], pop_priority); + let units = crate::compaction::unit::CompactionUnit::from_picker(PickerOutput { + outputs: output, + time_window_size: 1, + ..Default::default() + }) + .unwrap(); + let unit_priority = units + .iter() + .map(|unit| { + ( + unit.outputs[0].inputs[0].level(), + output_window(&unit.outputs[0]), + ) + }) + .collect::>(); + assert_eq!( + pop_priority, unit_priority, + "local FIFO must follow the picker's execution priority" + ); } #[tokio::test] diff --git a/src/mito2/src/compaction/unit.rs b/src/mito2/src/compaction/unit.rs new file mode 100644 index 00000000000..19a2224a9d3 --- /dev/null +++ b/src/mito2/src/compaction/unit.rs @@ -0,0 +1,251 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{BTreeMap, HashMap, HashSet}; + +use crate::compaction::CompactionOutput; +use crate::compaction::picker::PickerOutput; +use crate::sst::file::FileHandle; + +/// A complete input dependency group, merged and published atomically. +/// +/// Logical outputs retain their original input sets and time bounds. In particular, +/// splitting physical SSTs must never split a strict-window input dependency group. +#[derive(Debug, Clone)] +pub(crate) struct CompactionUnit { + /// Distinct input SSTs reserved and removed together. + pub(crate) inputs: Vec, + /// Original logical outputs; empty for expiration-only maintenance. + pub(crate) outputs: Vec, + /// Compaction window size in seconds. + pub(crate) time_window_size: i64, + /// Maximum physical output SST size in bytes. + pub(crate) max_file_size: Option, +} + +impl CompactionUnit { + /// Groups outputs by the transitive closure of shared input FileIds. + /// Returns `None` for an inconsistent plan that both expires and merges a file. + /// Expiration maintenance runs before merge units, matching the legacy pre-merge removal. + /// Merge units follow the legacy compactor's tail-first execution order; a shared-input + /// group inherits the position of its highest-priority output. + pub(crate) fn from_picker(output: PickerOutput) -> Option> { + let PickerOutput { + outputs, + expired_ssts, + time_window_size, + max_file_size, + } = output; + let mut parents: Vec<_> = (0..outputs.len()).collect(); + let mut owners = HashMap::new(); + for (index, output) in outputs.iter().enumerate() { + for file in &output.inputs { + if let Some(&previous) = owners.get(&file.file_id()) { + let a = root(&mut parents, previous); + let b = root(&mut parents, index); + // Higher output indices are consumed first by the legacy compactor. + parents[a.min(b)] = a.max(b); + } else { + owners.insert(file.file_id(), index); + } + } + } + if expired_ssts + .iter() + .any(|f| owners.contains_key(&f.file_id())) + { + return None; + } + + let mut groups: BTreeMap> = BTreeMap::new(); + for (index, output) in outputs.into_iter().enumerate() { + groups + .entry(root(&mut parents, index)) + .or_default() + .push(output); + } + let mut units = Vec::with_capacity(groups.len() + 1); + if !expired_ssts.is_empty() { + let mut seen = HashSet::new(); + units.push(Self { + inputs: expired_ssts + .into_iter() + .filter(|f| seen.insert(f.file_id())) + .collect(), + outputs: Vec::new(), + time_window_size, + max_file_size, + }); + } + for outputs in groups.into_values().rev() { + let mut seen = HashSet::new(); + let inputs: Vec<_> = outputs + .iter() + .flat_map(|o| &o.inputs) + .filter(|f| seen.insert(f.file_id())) + .cloned() + .collect(); + if inputs.is_empty() { + continue; + } + units.push(Self { + inputs, + outputs, + time_window_size, + max_file_size, + }); + } + Some(units) + } + + /// Estimates merge memory once per input, even when several outputs read it. + pub(crate) fn estimated_memory_bytes(&self) -> u64 { + if self.outputs.is_empty() { + return 0; + } + self.inputs.iter().fold(0u64, |total, file| { + total.saturating_add(file.meta_ref().max_row_group_uncompressed_size) + }) + } +} + +/// Finds a dependency group's representative while shortening the parent chain. +fn root(parents: &mut [usize], mut index: usize) -> usize { + while parents[index] != index { + parents[index] = parents[parents[index]]; + index = parents[index]; + } + index +} + +#[cfg(test)] +mod tests { + use store_api::storage::FileId; + + use super::*; + use crate::compaction::test_util::new_file_handle; + + /// Builds an L1 output that preserves deletion markers for grouping tests. + fn output(inputs: Vec) -> CompactionOutput { + CompactionOutput { + inputs, + output_level: 1, + filter_deleted: false, + output_time_range: None, + } + } + + #[test] + fn test_compaction_unit_transitive_inputs() { + let files: Vec<_> = (0..3) + .map(|_| new_file_handle(FileId::random(), 0, 100, 0)) + .collect(); + let units = CompactionUnit::from_picker(PickerOutput { + outputs: vec![ + output(vec![files[0].clone()]), + output(vec![files[2].clone()]), + output(vec![files[1].clone()]), + output(vec![files[0].clone(), files[1].clone()]), + ], + time_window_size: 100, + max_file_size: Some(1024), + ..Default::default() + }) + .unwrap(); + assert_eq!(2, units.len()); + assert_eq!(3, units[0].outputs.len()); + assert_eq!(2, units[0].inputs.len()); + assert_eq!(files[2].file_id(), units[1].inputs[0].file_id()); + assert_eq!(Some(1024), units[0].max_file_size); + assert_eq!(100, units[1].time_window_size); + } + + #[test] + fn test_compaction_unit_group_inherits_highest_output_priority() { + let files: Vec<_> = (0..3) + .map(|_| new_file_handle(FileId::random(), 0, 100, 0)) + .collect(); + let units = CompactionUnit::from_picker(PickerOutput { + outputs: [0, 1, 1, 2, 0] + .into_iter() + .map(|index| output(vec![files[index].clone()])) + .collect(), + ..Default::default() + }) + .unwrap(); + assert_eq!( + vec![files[0].file_id(), files[2].file_id(), files[1].file_id()], + units + .iter() + .map(|unit| unit.inputs[0].file_id()) + .collect::>() + ); + assert_eq!( + vec![2, 1, 2], + units + .iter() + .map(|unit| unit.outputs.len()) + .collect::>() + ); + } + + #[test] + fn test_compaction_unit_expiration_precedes_merge_units() { + let files: Vec<_> = (0..3) + .map(|_| new_file_handle(FileId::random(), 0, 100, 0)) + .collect(); + let units = CompactionUnit::from_picker(PickerOutput { + outputs: vec![ + output(vec![files[0].clone()]), + output(vec![files[1].clone()]), + ], + expired_ssts: vec![files[2].clone()], + ..Default::default() + }) + .unwrap(); + + assert_eq!(3, units.len()); + assert!(units[0].outputs.is_empty()); + assert_eq!(files[2].file_id(), units[0].inputs[0].file_id()); + assert_eq!(files[1].file_id(), units[1].inputs[0].file_id()); + assert_eq!(files[0].file_id(), units[2].inputs[0].file_id()); + } + + #[test] + fn test_compaction_unit_expiration_inputs() { + let file = new_file_handle(FileId::random(), 0, 100, 0); + assert!( + CompactionUnit::from_picker(PickerOutput { + outputs: vec![output(vec![file.clone()])], + expired_ssts: vec![file.clone()], + ..Default::default() + }) + .is_none() + ); + let units = CompactionUnit::from_picker(PickerOutput { + expired_ssts: vec![file.clone(), file], + ..Default::default() + }) + .unwrap(); + assert_eq!(1, units.len()); + assert_eq!(1, units[0].inputs.len()); + assert!(units[0].outputs.is_empty()); + assert_eq!(0, units[0].estimated_memory_bytes()); + assert!( + CompactionUnit::from_picker(PickerOutput::default()) + .unwrap() + .is_empty() + ); + } +} diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index f20eef75460..82f8bcaecc8 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -18,9 +18,10 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; +use api::region::RegionResponse; use api::v1::{ColumnSchema, Rows}; use async_trait::async_trait; -use common_error::ext::ErrorExt; +use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; use common_time::Timestamp; @@ -40,6 +41,7 @@ use tokio::sync::{Notify, Semaphore}; use crate::config::MitoConfig; use crate::engine::MitoEngine; use crate::engine::listener::{CompactionListener, EventListener}; +use crate::sst::file::RegionFileId; use crate::test_util::batch_util::sort_batches_and_print; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows_for_key, column_metadata_to_column_schema, put_rows, @@ -542,6 +544,979 @@ async fn assert_partial_compaction_preserves_delete_order(flat_format: bool) { ); } +/// Region fixture containing only identifiers, so it cannot pin an old SST snapshot. +struct CompactionTestRegion { + region_id: RegionId, + table_dir: String, + path_type: PathType, +} + +/// Scan observations from one version, without retaining the scanner or its SST handles. +struct CompactionScanResult { + file_ids: HashSet, + timestamps: Vec, +} + +impl CompactionTestRegion { + async fn create(engine: &MitoEngine, region_id: RegionId, ranges: &[Range]) -> Self { + let create = CreateRequestBuilder::new().build(); + let fixture = Self { + region_id, + table_dir: create.table_dir.clone(), + path_type: create.path_type, + }; + let columns = crate::test_util::rows_schema(&create); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + for rows in ranges { + put_and_flush(engine, region_id, &columns, rows.clone()).await; + } + fixture + } + + fn start_compaction( + &self, + engine: &MitoEngine, + ) -> tokio::task::JoinHandle> { + let engine = engine.clone(); + let region_id = self.region_id; + tokio::spawn(async move { + engine + .handle_request(region_id, strict_window_request(60, Some(2))) + .await + }) + } + + async fn scanner(&self, engine: &MitoEngine) -> crate::read::scan_region::Scanner { + engine + .scanner(self.region_id, ScanRequest::default()) + .await + .unwrap() + } + + async fn file_ids(&self, engine: &MitoEngine) -> HashSet { + self.scanner(engine).await.file_ids().into_iter().collect() + } + + async fn scan(&self, engine: &MitoEngine) -> CompactionScanResult { + let scanner = self.scanner(engine).await; + let file_ids = scanner.file_ids().into_iter().collect(); + let timestamps = collect_stream_ts(scanner.scan().await.unwrap()).await; + CompactionScanResult { + file_ids, + timestamps, + } + } + + async fn open(&self, engine: &MitoEngine) { + engine + .handle_request( + self.region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: self.table_dir.clone(), + path_type: self.path_type, + options: Default::default(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + } + + async fn reopen(&self, engine: &MitoEngine) { + engine + .handle_request( + self.region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + self.open(engine).await; + } +} + +fn strict_window_request(window_seconds: i64, parallelism: Option) -> RegionRequest { + RegionRequest::Compact(RegionCompactRequest { + options: api::v1::region::compact_request::Options::StrictWindow( + api::v1::region::StrictWindow { window_seconds }, + ), + parallelism, + ..Default::default() + }) +} + +/// Pauses the first unit at merge or commit while observing sibling progress. +struct UnitVisibilityGate { + first: std::sync::atomic::AtomicU64, + at_commit: bool, + entered: Notify, + resume: Semaphore, + applied: Notify, + cancel_requested: Notify, +} + +/// Occupies one pool slot and holds the target's first commit to expose dispatch order. +struct WindowPriorityGate { + blocker: RegionId, + target: RegionId, + blocker_started: Notify, + release_blocker: Semaphore, + target_committed: Notify, + release_target: Semaphore, + target_started: std::sync::atomic::AtomicUsize, + first_target_commit: AtomicBool, +} + +/// Releases both background tasks if an assertion aborts the test. +struct WindowPriorityGuard(Arc); + +impl Drop for WindowPriorityGuard { + fn drop(&mut self) { + self.0.release_blocker.add_permits(1); + self.0.release_target.add_permits(1); + } +} + +#[async_trait] +impl EventListener for WindowPriorityGate { + async fn on_compaction_unit_merge_begin(&self, region_id: RegionId, _plan_id: u64) { + if region_id == self.blocker { + self.blocker_started.notify_one(); + self.release_blocker.acquire().await.unwrap().forget(); + } else if region_id == self.target { + self.target_started.fetch_add(1, Ordering::SeqCst); + } + } + + async fn on_compaction_unit_committed(&self, region_id: RegionId, _plan_id: u64) { + if region_id == self.target && self.first_target_commit.swap(false, Ordering::SeqCst) { + self.target_committed.notify_one(); + self.release_target.acquire().await.unwrap().forget(); + } + } +} + +#[tokio::test] +async fn test_compaction_unit_newest_window_starts_first_with_one_available_slot() { + let blocker = RegionId::new(46, 1); + let target = RegionId::new(47, 1); + let gate = Arc::new(WindowPriorityGate { + blocker, + target, + blocker_started: Notify::new(), + release_blocker: Semaphore::new(0), + target_committed: Notify::new(), + release_target: Semaphore::new(0), + target_started: std::sync::atomic::AtomicUsize::new(0), + first_target_commit: AtomicBool::new(true), + }); + let guard = WindowPriorityGuard(gate.clone()); + let mut env = TestEnv::new().await; + let engine = env + .create_engine_with_time( + MitoConfig { + max_background_compactions: 2, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }, + None, + Some(gate.clone()), + Arc::new(MockTimeProvider::new(1000)), + ) + .await; + for region_id in [blocker, target] { + let create = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.time_window", "1h") + .insert_option("compaction.twcs.trigger_file_num", "4") + .insert_option("compaction.twcs.inactive_window_trigger_file_num", "4") + .build(); + let columns = crate::test_util::rows_schema(&create); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + if region_id == blocker { + put_and_flush(&engine, region_id, &columns, 0..2).await; + } else { + for window_start in [0, 3600] { + for _ in 0..4 { + put_and_flush(&engine, region_id, &columns, window_start..window_start + 2) + .await; + } + } + } + } + let region = engine.get_region(target).unwrap(); + let version = region.version(); + let (newest, older): (HashSet<_>, HashSet<_>) = version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .map(|file| (file.meta_ref().file_id, file.time_range().0)) + .partition(|(_, start)| *start >= Timestamp::new_second(3600)); + assert_eq!(4, newest.len()); + assert_eq!(4, older.len()); + let blocker_engine = engine.clone(); + let blocked = tokio::spawn(async move { + blocker_engine + .handle_request(blocker, strict_window_request(3600, None)) + .await + }); + tokio::time::timeout(Duration::from_secs(10), gate.blocker_started.notified()) + .await + .unwrap(); + let target_engine = engine.clone(); + let compact = tokio::spawn(async move { + target_engine + .handle_request( + target, + RegionRequest::Compact(RegionCompactRequest { + parallelism: Some(2), + ..Default::default() + }), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(10), gate.target_committed.notified()) + .await + .unwrap(); + assert_eq!(1, gate.target_started.load(Ordering::SeqCst)); + let manifest = region.manifest_ctx.manifest_manager.read().await.manifest(); + assert!( + newest + .iter() + .all(|(id, _)| !manifest.files.contains_key(id)), + "the newest window must consume the only available execution slot" + ); + assert!(older.iter().all(|(id, _)| manifest.files.contains_key(id))); + drop(guard); + tokio::time::timeout(Duration::from_secs(10), async { + compact.await.unwrap().unwrap(); + blocked.await.unwrap().unwrap(); + }) + .await + .unwrap(); +} + +/// Fails a selected SST close and records earlier outputs for cleanup assertions. +struct FailingUnitSstWriter { + inner: object_store::layers::mock::Writer, + path: String, + remaining: Arc, + finalized: Arc>>, +} + +impl object_store::layers::mock::Write for FailingUnitSstWriter { + async fn write( + &mut self, + bytes: object_store::layers::mock::Buffer, + ) -> object_store::layers::mock::Result<()> { + self.inner.write(bytes).await + } + /// Injects failure at the configured close without publishing that SST. + async fn close( + &mut self, + ) -> object_store::layers::mock::Result { + let previous = self + .remaining + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1)); + if previous == Ok(1) { + return Err(object_store::layers::mock::Error::new( + object_store::layers::mock::ErrorKind::Unexpected, + "second unit SST failed", + )); + } + let metadata = self.inner.close().await?; + if previous.is_ok() { + self.finalized.lock().unwrap().push(self.path.clone()); + } + Ok(metadata) + } + async fn abort(&mut self) -> object_store::layers::mock::Result<()> { + self.inner.abort().await + } +} + +fn failing_unit_sst_layer( + remaining: Arc, + finalized: Arc>>, +) -> object_store::layers::mock::MockLayer { + object_store::layers::mock::MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + if path.ends_with(".parquet") { + Box::new(FailingUnitSstWriter { + inner, + path: path.to_string(), + remaining: remaining.clone(), + finalized: finalized.clone(), + }) + } else { + inner + } + })) + .build() + .unwrap() +} + +#[tokio::test] +async fn test_compaction_unit_failed_dependency_group_keeps_input_and_cleans_output() { + let remaining = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let finalized = Arc::new(std::sync::Mutex::new(Vec::new())); + let layer = failing_unit_sst_layer(remaining.clone(), finalized.clone()); + let mut env = TestEnv::new().await.with_mock_layer(layer); + let config = MitoConfig { + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env.create_engine(config.clone()).await; + let region_id = RegionId::new(43, 1); + let test_region = + CompactionTestRegion::create(&engine, region_id, std::slice::from_ref(&(0..122))).await; + let original = test_region.file_ids(&engine).await; + remaining.store(2, Ordering::SeqCst); + let error = test_region + .start_compaction(&engine) + .await + .unwrap() + .unwrap_err(); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + let current = test_region.scan(&engine).await; + assert_eq!(original, current.file_ids); + let expected = (0..122).map(|n| n * 1000).collect::>(); + assert_eq!(expected, current.timestamps); + let paths = finalized.lock().unwrap().clone(); + assert_eq!( + 1, + paths.len(), + "one output must have finished before the failing output" + ); + let store = env.get_object_store().unwrap(); + for path in paths { + assert!( + !store.exists(&path).await.unwrap(), + "abandoned output {path}" + ); + } + let engine = env.reopen_engine(engine, config).await; + test_region.open(&engine).await; + assert_eq!(expected, test_region.scan(&engine).await.timestamps); +} + +#[rstest::rstest] +#[tokio::test] +async fn test_compaction_unit_applied_sibling_survives_failure_and_reopen( + #[values(false, true)] flat_format: bool, +) { + let remaining = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let finalized = Arc::new(std::sync::Mutex::new(Vec::new())); + let layer = failing_unit_sst_layer(remaining.clone(), finalized.clone()); + let gate = Arc::new(UnitVisibilityGate::before_merge()); + let guard = UnitVisibilityGuard(gate.clone()); + let mut env = TestEnv::new().await.with_mock_layer(layer); + let config = MitoConfig { + default_flat_format: flat_format, + max_background_compactions: 2, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env + .create_engine_with(config.clone(), None, Some(gate.clone()), None) + .await; + let region_id = RegionId::new(48, 1); + // Disjoint input SSTs produce two independent units, not one dependency group. + let test_region = CompactionTestRegion::create(&engine, region_id, &[0..2, 120..122]).await; + let original = test_region.file_ids(&engine).await; + assert_eq!(2, original.len()); + remaining.store(2, Ordering::SeqCst); + let compact = test_region.start_compaction(&engine); + // B is held before merge until A's manifest edit has been committed and applied. + tokio::time::timeout(Duration::from_secs(10), gate.entered.notified()) + .await + .expect("unit B did not reach the merge gate"); + tokio::time::timeout(Duration::from_secs(10), gate.applied.notified()) + .await + .expect("unit A did not apply independently of B"); + assert!(!compact.is_finished(), "request must still wait for B"); + assert_eq!(1, remaining.load(Ordering::SeqCst)); + let applied = test_region.scan(&engine).await; + let applied_ids = applied.file_ids; + assert_eq!(2, applied_ids.len()); + assert_eq!(1, applied_ids.intersection(&original).count()); + assert_eq!(1, applied_ids.difference(&original).count()); + let expected = vec![0, 1000, 120_000, 121_000]; + assert_eq!(expected, applied.timestamps); + + drop(guard); + let error = tokio::time::timeout(Duration::from_secs(10), compact) + .await + .expect("compaction did not finish after B failed") + .unwrap() + .unwrap_err(); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + assert_eq!(0, remaining.load(Ordering::SeqCst)); + let failed = test_region.scan(&engine).await; + assert_eq!( + applied_ids, failed.file_ids, + "B's failure must retain A's output and B's original input" + ); + assert_eq!(expected, failed.timestamps); + assert!( + engine + .get_region(region_id) + .unwrap() + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .all(|file| !file.compacting()) + ); + let paths = finalized.lock().unwrap().clone(); + assert_eq!(1, paths.len()); + assert!( + env.get_object_store() + .unwrap() + .exists(&paths[0]) + .await + .unwrap() + ); + + let engine = tokio::time::timeout(Duration::from_secs(10), env.reopen_engine(engine, config)) + .await + .expect("engine shutdown must not wait for a failed unit"); + test_region.reopen(&engine).await; + let reopened = test_region.scan(&engine).await; + assert_eq!( + applied_ids, reopened.file_ids, + "recovery must preserve the partial-success file set" + ); + assert_eq!(expected, reopened.timestamps); +} + +/// Holds one unit before merge and observes its sibling's task exit. +struct UnitReleaseListener { + merge: Arc, + released: Semaphore, +} + +#[async_trait] +impl EventListener for UnitReleaseListener { + async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, plan_id: u64) { + self.merge.block_first(plan_id).await; + } + + async fn on_compaction_result_notified(&self, _region_id: RegionId) { + self.released.add_permits(1); + } +} + +/// Regression for #9067: a slow sibling must not pin a completed unit's old SSTs. +#[tokio::test] +async fn test_compaction_unit_releases_inputs_before_sibling_finishes() { + let gate = Arc::new(UnitVisibilityGate::before_merge()); + let guard = UnitVisibilityGuard(gate.clone()); + let listener = Arc::new(UnitReleaseListener { + merge: gate.clone(), + released: Semaphore::new(0), + }); + let mut env = TestEnv::new().await; + let engine = env + .create_engine_with( + MitoConfig { + max_background_compactions: 2, + // A queued barrier can drain every preceding purge job. + max_background_purges: 1, + // TestEnv constructs its purge scheduler with flush concurrency. + max_background_flushes: 1, + min_compaction_interval: Duration::from_secs(3600), + gc: crate::gc::GcConfig { + enable: false, + ..Default::default() + }, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + let region_id = RegionId::new(49, 1); + // The third unit starts only after a submitted unit exits; one sibling stays blocked. + let test_region = + CompactionTestRegion::create(&engine, region_id, &[0..2, 120..122, 240..242]).await; + + // Keep only IDs and paths, never a scanner, Version, or input FileHandle. + let original = test_region.file_ids(&engine).await; + assert_eq!(3, original.len()); + let paths: HashMap<_, _> = original + .iter() + .map(|id| { + ( + *id, + crate::sst::location::sst_file_path( + &test_region.table_dir, + *id, + test_region.path_type, + ), + ) + }) + .collect(); + let store = env.get_object_store().unwrap(); + for path in paths.values() { + assert!(store.exists(path).await.unwrap()); + } + + let compact = test_region.start_compaction(&engine); + tokio::time::timeout(Duration::from_secs(10), gate.entered.notified()) + .await + .expect("unit B did not reach the merge gate"); + tokio::time::timeout(Duration::from_secs(10), listener.released.acquire_many(2)) + .await + .expect("submitted and pending units did not exit while their sibling was blocked") + .unwrap() + .forget(); + assert!(!compact.is_finished()); + let current = test_region.file_ids(&engine).await; + assert_eq!(3, current.len()); + assert_eq!(1, original.intersection(¤t).count()); + assert_eq!(2, original.difference(¤t).count()); + let blocked_input = *original.intersection(¤t).next().unwrap(); + + let drain_purges = || async { + let (sender, receiver) = tokio::sync::oneshot::channel(); + engine + .purge_scheduler() + .schedule(Box::pin(async move { + let _ = sender.send(()); + })) + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), receiver) + .await + .expect("purge queue did not drain") + .unwrap(); + }; + drain_purges().await; + let mut retained_before_sibling_exit = Vec::new(); + for input in original.difference(¤t) { + if store.exists(&paths[input]).await.unwrap() { + retained_before_sibling_exit.push(*input); + } + } + assert!(store.exists(&paths[&blocked_input]).await.unwrap()); + + // Finish B before asserting the regression, also checking that purging works at all. + drop(guard); + tokio::time::timeout(Duration::from_secs(10), compact) + .await + .expect("compaction did not finish after releasing B") + .unwrap() + .unwrap(); + drain_purges().await; + for path in paths.values() { + assert!( + !store.exists(path).await.unwrap(), + "input not purged: {path}" + ); + } + engine.stop().await.unwrap(); + assert!( + retained_before_sibling_exit.is_empty(), + "completed units exited and the purge queue drained, but their old inputs remained until the blocked sibling exited: {retained_before_sibling_exit:?}", + ); +} + +impl UnitVisibilityGate { + fn before_merge() -> Self { + Self { + first: std::sync::atomic::AtomicU64::new(u64::MAX), + at_commit: false, + entered: Notify::new(), + resume: Semaphore::new(0), + applied: Notify::new(), + cancel_requested: Notify::new(), + } + } + + fn after_commit() -> Self { + Self { + at_commit: true, + ..Self::before_merge() + } + } + + /// Blocks only the first arriving unit, leaving its siblings free to progress. + async fn block_first(&self, id: u64) { + if self + .first + .compare_exchange(u64::MAX, id, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + self.entered.notify_one(); + self.resume.acquire().await.unwrap().forget(); + } + } +} + +/// Releases the visibility gate even when a test exits through a panic. +struct UnitVisibilityGuard(Arc); + +/// Holds one unit before merge and its sibling at commit for deterministic failure ordering. +struct UnitFailureGate { + merge: Arc, + commit: Arc, +} + +#[async_trait] +impl EventListener for UnitFailureGate { + async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, id: u64) { + self.merge.block_first(id).await; + } + + async fn on_compaction_commit_begin(&self, region_id: RegionId) { + self.commit.on_compaction_commit_begin(region_id).await; + } + + fn on_compaction_cancel_requested(&self, region_id: RegionId) { + self.commit.on_compaction_cancel_requested(region_id); + } + + async fn on_compaction_result_notified(&self, region_id: RegionId) { + self.commit.on_compaction_result_notified(region_id).await; + } +} + +/// Distinguishes a rejected write from a lost reply after persistence. +#[derive(Clone, Copy)] +enum ManifestWriteFailure { + BeforeWrite, + AfterWrite, +} + +/// Injects manifest failures around a real storage write to verify output ownership. +struct FailingManifestWriter { + inner: object_store::layers::mock::Writer, + failure: ManifestWriteFailure, +} + +impl object_store::layers::mock::Write for FailingManifestWriter { + /// Rejects the pre-write case before any bytes reach the real writer. + async fn write( + &mut self, + bytes: object_store::layers::mock::Buffer, + ) -> object_store::layers::mock::Result<()> { + if matches!(self.failure, ManifestWriteFailure::BeforeWrite) { + return Err(object_store::layers::mock::Error::new( + object_store::layers::mock::ErrorKind::PermissionDenied, + "manifest write rejected before persistence", + )); + } + self.inner.write(bytes).await + } + /// Persists the manifest before simulating a lost success reply. + async fn close( + &mut self, + ) -> object_store::layers::mock::Result { + self.inner.close().await?; + Err(object_store::layers::mock::Error::new( + object_store::layers::mock::ErrorKind::Unexpected, + "lost manifest response after persistence", + )) + } + async fn abort(&mut self) -> object_store::layers::mock::Result<()> { + self.inner.abort().await + } +} + +#[rstest::rstest] +#[case::before_write(ManifestWriteFailure::BeforeWrite)] +#[case::after_write(ManifestWriteFailure::AfterWrite)] +#[tokio::test] +async fn test_compaction_unit_manifest_failure_preserves_outputs( + #[case] failure: ManifestWriteFailure, +) { + let armed = Arc::new(AtomicBool::new(false)); + let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let ssts = Arc::new(std::sync::Mutex::new(Vec::new())); + let (inject, writes, paths) = (armed.clone(), attempts.clone(), ssts.clone()); + let layer = object_store::layers::mock::MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + if path.ends_with(".parquet") { + paths.lock().unwrap().push(path.to_string()); + } + if path.contains("/manifest/") && path.contains(".json") { + writes.fetch_add(1, Ordering::SeqCst); + if inject.swap(false, Ordering::SeqCst) { + return Box::new(FailingManifestWriter { inner, failure }); + } + } + inner + })) + .build() + .unwrap(); + let region_id = RegionId::new(44, 1); + let gate = Arc::new(CompactionPlanningGate::new(region_id)); + let merge_gate = Arc::new(UnitVisibilityGate::before_merge()); + let merge_guard = UnitVisibilityGuard(merge_gate.clone()); + let mut env = TestEnv::new().await.with_mock_layer(layer); + let engine = env + .create_engine_with( + MitoConfig { + max_background_compactions: 2, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }, + None, + Some(Arc::new(UnitFailureGate { + merge: merge_gate, + commit: gate.clone(), + })), + None, + ) + .await; + let test_region = CompactionTestRegion::create(&engine, region_id, &[0..2, 120..122]).await; + let original = test_region.file_ids(&engine).await; + assert_eq!(2, original.len()); + let region = engine.get_region(region_id).unwrap(); + let version = region + .manifest_ctx + .manifest_manager + .read() + .await + .manifest() + .manifest_version; + ssts.lock().unwrap().clear(); + attempts.store(0, Ordering::SeqCst); + armed.store(true, Ordering::SeqCst); + let guard = gate.arm_commit(); + let compact = test_region.start_compaction(&engine); + tokio::time::timeout(Duration::from_secs(10), gate.wait_until_commit_entered()) + .await + .unwrap(); + let sibling_guard = gate.arm_pending_ddl_dispatch(); + let ddl_engine = engine.clone(); + let truncate = tokio::spawn(async move { + ddl_engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::All), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(10), gate.wait_until_cancel_requested()) + .await + .unwrap(); + drop(merge_guard); + // Observe the cancelled sibling's exit before releasing the failing commit. + // Its cancellation must neither finish the request nor mask the later I/O error. + tokio::time::timeout( + Duration::from_secs(10), + gate.wait_until_pending_ddl_dispatch(), + ) + .await + .unwrap(); + assert!(!compact.is_finished()); + assert!(!truncate.is_finished()); + sibling_guard.release(); + guard.release(); + let error = tokio::time::timeout(Duration::from_secs(10), compact) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + let error = tokio::time::timeout(Duration::from_secs(10), truncate) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + assert_eq!( + 1, + attempts.load(Ordering::SeqCst), + "one failed publication, without automatic retry" + ); + assert!(!armed.load(Ordering::SeqCst)); + assert_eq!( + version, + region + .manifest_ctx + .manifest_manager + .read() + .await + .manifest() + .manifest_version + ); + let scan = test_region.scan(&engine).await; + assert_eq!(original, scan.file_ids, "failed edits must not be applied"); + assert_eq!(vec![0, 1000, 120_000, 121_000], scan.timestamps); + assert!( + region + .version() + .ssts + .levels() + .iter() + .flat_map(|l| l.files()) + .all(|f| !f.compacting()) + ); + let paths = ssts.lock().unwrap().clone(); + assert!(!paths.is_empty()); + let store = env.get_object_store().unwrap(); + let mut retained = 0; + for path in paths { + retained += usize::from(store.exists(&path).await.unwrap()); + } + assert_eq!( + usize::from(matches!(failure, ManifestWriteFailure::AfterWrite)), + retained + ); +} + +impl Drop for UnitVisibilityGuard { + fn drop(&mut self) { + self.0.resume.add_permits(1); + } +} + +#[async_trait] +impl EventListener for UnitVisibilityGate { + fn on_compaction_cancel_requested(&self, _region_id: RegionId) { + self.cancel_requested.notify_one(); + } + async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, id: u64) { + if !self.at_commit { + self.block_first(id).await; + } + } + async fn on_compaction_unit_committed(&self, _region_id: RegionId, id: u64) { + if self.at_commit { + self.block_first(id).await; + } + } + fn on_compaction_unit_applied(&self, _region_id: RegionId, _id: u64) { + self.applied.notify_one(); + } +} + +enum UnitInterleaving { + SlowMerge, + PendingApply, + Truncate, + Reopen, +} + +#[rstest::rstest] +#[case::slow_merge(UnitInterleaving::SlowMerge)] +#[case::pending_apply(UnitInterleaving::PendingApply)] +#[case::truncate(UnitInterleaving::Truncate)] +#[case::reopen(UnitInterleaving::Reopen)] +#[tokio::test] +async fn test_compaction_unit_commit_and_apply_can_interleave(#[case] scenario: UnitInterleaving) { + let gate = Arc::new(match scenario { + UnitInterleaving::PendingApply | UnitInterleaving::Reopen => { + UnitVisibilityGate::after_commit() + } + UnitInterleaving::SlowMerge | UnitInterleaving::Truncate => { + UnitVisibilityGate::before_merge() + } + }); + let guard = UnitVisibilityGuard(gate.clone()); + let mut env = TestEnv::new().await; + let config = MitoConfig { + max_background_compactions: 2, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env + .create_engine_with(config.clone(), None, Some(gate.clone()), None) + .await; + let region_id = RegionId::new(42, 1); + let test_region = CompactionTestRegion::create(&engine, region_id, &[0..2, 120..122]).await; + // Keep this scanner explicitly: this scenario also verifies old query snapshots. + let old = test_region.scanner(&engine).await; + let old_ids = old.file_ids(); + assert_eq!(2, old.num_files()); + let task = test_region.start_compaction(&engine); + tokio::time::timeout(Duration::from_secs(10), gate.entered.notified()) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), gate.applied.notified()) + .await + .expect("sibling unit waited for another unit's merge or apply"); + let current = test_region.scan(&engine).await; + let unchanged = current + .file_ids + .iter() + .filter(|id| old_ids.contains(id)) + .count(); + assert_eq!(1, unchanged, "the ungated unit must already be visible"); + assert!( + !task.is_finished(), + "manual request finished before all units" + ); + let mut expected = vec![0, 1000, 120_000, 121_000]; + assert_eq!(expected, current.timestamps); + assert_eq!(expected, collect_stream_ts(old.scan().await.unwrap()).await); + let truncate_task = if matches!(scenario, UnitInterleaving::Truncate) { + let engine = engine.clone(); + let ddl = tokio::spawn(async move { + engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::All), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(10), gate.cancel_requested.notified()) + .await + .unwrap(); + assert!(!ddl.is_finished()); + assert!(!task.is_finished()); + Some(ddl) + } else { + None + }; + if matches!(scenario, UnitInterleaving::Reopen) { + test_region.reopen(&engine).await; + assert_eq!(expected, test_region.scan(&engine).await.timestamps); + } + drop(guard); + let result = tokio::time::timeout(Duration::from_secs(10), task) + .await + .unwrap() + .unwrap(); + if let Some(ddl) = truncate_task { + assert_eq!(StatusCode::Cancelled, result.unwrap_err().status_code()); + tokio::time::timeout(Duration::from_secs(10), ddl) + .await + .unwrap() + .unwrap() + .unwrap(); + expected.clear(); + } else if matches!(scenario, UnitInterleaving::Reopen) { + assert_eq!(StatusCode::Cancelled, result.unwrap_err().status_code()); + } else { + result.unwrap(); + } + let completed = test_region.scan(&engine).await; + assert!(completed.file_ids.iter().all(|id| !old_ids.contains(id))); + assert_eq!( + expected, completed.timestamps, + "applying the delayed edit must preserve its sibling's result" + ); + let engine = tokio::time::timeout(Duration::from_secs(10), env.reopen_engine(engine, config)) + .await + .unwrap(); + test_region.open(&engine).await; + assert_eq!(expected, test_region.scan(&engine).await.timestamps); +} + /// Real-engine regression for the tombstone resurrection chain described in /// https://greptime.feishu.cn/wiki/L2WYwIM40iG9tckowrnctyw8nQc. /// No synthetic file sizes, PK ranges, or manually selected compaction inputs. @@ -885,6 +1860,7 @@ struct CompactionPlanningGate { pending_ddl_armed: AtomicBool, pending_ddl_entered: Notify, pending_ddl_permits: Semaphore, + flush_completed: Notify, } /// Releases an armed [`CompactionPlanningGate`] when a test exits unexpectedly. @@ -980,6 +1956,7 @@ impl CompactionPlanningGate { pending_ddl_armed: AtomicBool::new(false), pending_ddl_entered: Notify::new(), pending_ddl_permits: Semaphore::new(0), + flush_completed: Notify::new(), } } @@ -1050,6 +2027,12 @@ impl CompactionPlanningGate { #[async_trait] impl EventListener for CompactionPlanningGate { + fn on_flush_success(&self, region_id: RegionId) { + if region_id == self.region_id { + self.flush_completed.notify_one(); + } + } + async fn on_compaction_pick_begin(&self, region_id: RegionId) { if region_id != self.region_id { return; @@ -1211,6 +2194,9 @@ async fn assert_automatic_followup_updates_schedule_time(preexisting_flushes: us offset * 10..offset * 10 + 10, ) .await; + tokio::time::timeout(Duration::from_secs(5), gate.flush_completed.notified()) + .await + .unwrap(); } let first_schedule_time = initial_time + interval.as_millis() as i64; @@ -1227,6 +2213,9 @@ async fn assert_automatic_followup_updates_schedule_time(preexisting_flushes: us tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered()) .await .expect("initial automatic planning did not reach the gate"); + tokio::time::timeout(Duration::from_secs(5), gate.flush_completed.notified()) + .await + .unwrap(); let trigger_time = first_schedule_time + interval.as_millis() as i64; time_provider.set_now(trigger_time); @@ -1237,6 +2226,11 @@ async fn assert_automatic_followup_updates_schedule_time(preexisting_flushes: us first_start + 10..first_start + 20, ) .await; + // The flush reply precedes the worker's compaction trigger; wait for that trigger + // before advancing the mock clock and releasing the picker. + tokio::time::timeout(Duration::from_secs(5), gate.flush_completed.notified()) + .await + .unwrap(); let followup_schedule_time = trigger_time + 1; time_provider.set_now(followup_schedule_time); gate_guard.release(); diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index bbfdc97afd6..d5e4b8e8854 100644 --- a/src/mito2/src/engine/listener.rs +++ b/src/mito2/src/engine/listener.rs @@ -86,6 +86,15 @@ pub trait EventListener: Send + Sync { /// Notifies the listener after local compaction becomes non-cancellable and before commit. async fn on_compaction_commit_begin(&self, _region_id: RegionId) {} + /// Gates a particular unit before it reads or writes SSTs. + async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) {} + + /// Gates a persisted unit before its notification reaches the region worker. + async fn on_compaction_unit_committed(&self, _region_id: RegionId, _plan_id: u64) {} + + /// Observes independent query visibility, before request-level completion. + fn on_compaction_unit_applied(&self, _region_id: RegionId, _plan_id: u64) {} + /// Notifies the listener after compaction results are sent and before pending DDL dispatch. async fn on_compaction_result_notified(&self, _region_id: RegionId) {} diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index 4c275c248c9..cd1116c4fbd 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -136,11 +136,14 @@ lazy_static! { exponential_buckets(1.0, 10.0, 6).unwrap(), ) .unwrap(); - /// Timer of whole compaction task. - pub static ref COMPACTION_ELAPSED_TOTAL: Histogram = - register_histogram!( + /// Elapsed time of each successfully published compaction unit, labeled by `type`. + /// - local: one dependency unit, from task start through worker apply acknowledgement + /// - remote: one remote job, from request start through publication + pub static ref COMPACTION_ELAPSED_TOTAL: HistogramVec = + register_histogram_vec!( "greptime_mito_compaction_total_elapsed", - "mito compaction total elapsed", + "mito compaction elapsed per published unit", + &[TYPE_LABEL], // 1 ~ 100000 exponential_buckets(1.0, 10.0, 6).unwrap(), ).unwrap(); @@ -151,7 +154,7 @@ lazy_static! { pub static ref COMPACTION_FAILURE_COUNT: IntCounter = register_int_counter!("greptime_mito_compaction_failure_total", "mito compaction failure total").unwrap(); - /// Gauge for inflight compaction tasks. + /// Gauge for inflight compaction tasks: local units and remote jobs. pub static ref INFLIGHT_COMPACTION_COUNT: IntGauge = register_int_gauge!( "greptime_mito_inflight_compaction_count", diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index cd2902f68dd..c7ec359d6b0 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -49,9 +49,9 @@ use tokio::sync::oneshot::{self, Receiver, Sender}; use crate::compaction::{CompactionExecution, CompactionPickFinished}; use crate::error::{ - CompactRegionSnafu, CompactionCancelledSnafu, ConvertColumnDataTypeSnafu, CreateDefaultSnafu, - Error, FillDefaultSnafu, FlushRegionSnafu, InvalidPartitionExprSnafu, InvalidRequestSnafu, - MissingPartitionExprSnafu, Result, UnexpectedSnafu, + CompactRegionSnafu, ConvertColumnDataTypeSnafu, CreateDefaultSnafu, Error, FillDefaultSnafu, + FlushRegionSnafu, InvalidPartitionExprSnafu, InvalidRequestSnafu, MissingPartitionExprSnafu, + Result, UnexpectedSnafu, }; use crate::flush::FlushReason; use crate::manifest::action::{RegionEdit, TruncateKind}; @@ -905,6 +905,8 @@ pub(crate) enum BackgroundNotify { region: MitoRegionRef, registered: Sender<()>, }, + /// Local unit publication and resource completion events. + CompactionUnit(CompactionUnitNotification), /// A DDL released by compaction has actually completed, including failure. CompactionDdlComplete { generation: u64, @@ -926,8 +928,6 @@ pub(crate) enum BackgroundNotify { IndexBuildRetry(BuildIndexRequest), /// Compaction has finished. CompactionFinished(CompactionFinished), - /// Compaction has been cancelled cooperatively. - CompactionCancelled(CompactionCancelled), /// Compaction has failed. CompactionFailed(CompactionFailed), /// Truncate result. @@ -1017,6 +1017,22 @@ pub(crate) struct IndexBuildFailed { pub(crate) err: Arc, } +/// A local unit event carries only identity, never an input reservation lease. +#[derive(Debug)] +pub(crate) enum CompactionUnitNotification { + /// A persisted edit or error awaits worker apply or failure acknowledgement. + Finished { + plan_id: u64, + result: std::result::Result>, + applied: tokio::sync::oneshot::Sender<()>, + }, + /// The task has exited and released its owned resources. + Released { + /// Identifies the exited attempt, not the whole compaction request. + plan_id: u64, + }, +} + /// Notifies a compaction job has finished. #[derive(Debug)] pub(crate) struct CompactionFinished { @@ -1032,30 +1048,12 @@ pub(crate) struct CompactionFinished { pub(crate) edit: RegionEdit, } -/// Notifies a compaction job has been cancelled cooperatively. -#[derive(Debug)] -pub(crate) struct CompactionCancelled { - /// Region id. - pub(crate) region_id: RegionId, - /// Identity and reservation lease of the accepted execution. - pub(crate) execution: CompactionExecution, - /// Waiters to wake once the cancellation has been observed by the worker. - pub(crate) senders: Vec, -} - -impl CompactionCancelled { - pub(crate) fn on_success(self) { - for sender in self.senders { - sender.send(CompactionCancelledSnafu {}.fail()); - } - info!("Compaction cancelled for region: {}", self.region_id); - } -} - impl CompactionFinished { pub fn on_success(self) { // only update compaction time on success - COMPACTION_ELAPSED_TOTAL.observe(self.start_time.elapsed().as_secs_f64()); + COMPACTION_ELAPSED_TOTAL + .with_label_values(&["remote"]) + .observe(self.start_time.elapsed().as_secs_f64()); for sender in self.senders { sender.send(Ok(0)); @@ -1270,8 +1268,6 @@ pub(crate) struct CopyRegionFromRequest { mod tests { use api::v1::value::ValueData; use api::v1::{Row, SemanticType}; - use common_error::ext::ErrorExt; - use common_error::status_code::StatusCode; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnDefaultConstraint; use mito_codec::test_util::i64_value; @@ -1389,22 +1385,6 @@ mod tests { assert_eq!(None, request.column_index_by_name("c2")); } - #[test] - fn test_compaction_cancelled_sends_cancelled_error() { - let (tx, rx) = oneshot::channel(); - let request = CompactionCancelled { - region_id: RegionId::new(1, 1), - execution: crate::compaction::CompactionExecution::for_test(0), - senders: vec![OutputTx::new(tx)], - }; - - request.on_success(); - - let err = rx.blocking_recv().unwrap().unwrap_err(); - assert!(matches!(err, Error::CompactionCancelled { .. })); - assert_eq!(err.status_code(), StatusCode::Cancelled); - } - #[test] fn test_write_request_column_num() { let rows = Rows { diff --git a/src/mito2/src/schedule.rs b/src/mito2/src/schedule.rs index ae3145f3b95..a4f9cb60522 100644 --- a/src/mito2/src/schedule.rs +++ b/src/mito2/src/schedule.rs @@ -42,6 +42,11 @@ impl CancellableTaskState { self.cancel_handle.is_cancelled() } + /// Returns whether the task has crossed the non-cancellable publication boundary. + pub(crate) fn commit_started(&self) -> bool { + *self.commit_started.lock().unwrap() + } + /// Starts the non-cancellable commit phase. /// /// Returns false if cancellation was requested first. diff --git a/src/mito2/src/schedule/scheduler.rs b/src/mito2/src/schedule/scheduler.rs index ae25815a2a7..b93f5f705b9 100644 --- a/src/mito2/src/schedule/scheduler.rs +++ b/src/mito2/src/schedule/scheduler.rs @@ -38,16 +38,41 @@ pub trait Scheduler: Send + Sync { /// Schedules a Job fn schedule(&self, job: Job) -> Result<()>; + /// Runs completion only after the job and its owned resources have been dropped. + fn schedule_with_completion(&self, job: Job, completion: Job) -> Result<()> { + self.schedule(Box::pin(async move { + job.await; + completion.await; + })) + } + /// Stops scheduler. If `await_termination` is set to true, the scheduler will wait until all tasks are processed. async fn stop(&self, await_termination: bool) -> Result<()>; } pub type SchedulerRef = Arc; +/// A queued job with an optional completion future dispatched after the job exits. +struct ScheduledJob { + job: Job, + completion: Option, +} + +impl ScheduledJob { + /// Finishes the job before dispatching completion outside its runner. + async fn run(self) { + self.job.await; + if let Some(completion) = self.completion { + // The original job has exited. Notification must not occupy its runner. + common_runtime::spawn_global(completion); + } + } +} + /// Request scheduler based on local state. pub struct LocalScheduler { /// Sends jobs to flume bounded channel - sender: RwLock>>, + sender: RwLock>>, /// Task handles handles: Mutex>>, /// Token used to halt the scheduler @@ -61,7 +86,7 @@ impl LocalScheduler { /// /// concurrency: the number of bounded receiver pub fn new(concurrency: usize) -> Self { - let (tx, rx) = async_channel::unbounded(); + let (tx, rx) = async_channel::unbounded::(); let token = CancellationToken::new(); let state = Arc::new(AtomicU8::new(STATE_RUNNING)); @@ -79,7 +104,7 @@ impl LocalScheduler { } req_opt = receiver.recv() =>{ if let Ok(job) = req_opt { - job.await; + job.run().await; } } } @@ -88,7 +113,7 @@ impl LocalScheduler { if state_clone.load(Ordering::Relaxed) == STATE_AWAIT_TERMINATION { // recv_async waits until all sender's been dropped. while let Ok(job) = receiver.recv().await { - job.await; + job.run().await; } state_clone.store(STATE_STOP, Ordering::Relaxed); } @@ -108,13 +133,10 @@ impl LocalScheduler { fn is_running(&self) -> bool { self.state.load(Ordering::Relaxed) == STATE_RUNNING } -} -#[async_trait::async_trait] -impl Scheduler for LocalScheduler { - fn schedule(&self, job: Job) -> Result<()> { + /// Enqueues work only while the scheduler and its sender remain available. + fn submit(&self, job: ScheduledJob) -> Result<()> { ensure!(self.is_running(), InvalidSchedulerStateSnafu); - self.sender .read() .unwrap() @@ -123,6 +145,24 @@ impl Scheduler for LocalScheduler { .try_send(job) .map_err(|_| InvalidSenderSnafu {}.build()) } +} + +#[async_trait::async_trait] +impl Scheduler for LocalScheduler { + fn schedule(&self, job: Job) -> Result<()> { + self.submit(ScheduledJob { + job, + completion: None, + }) + } + + /// Queues the job and its completion together so completion cannot precede task exit. + fn schedule_with_completion(&self, job: Job, completion: Job) -> Result<()> { + self.submit(ScheduledJob { + job, + completion: Some(completion), + }) + } /// if await_termination is true, scheduler will wait all tasks finished before stopping async fn stop(&self, await_termination: bool) -> Result<()> { diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 7fca91cf3a3..4bf9d4b244a 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1360,6 +1360,9 @@ impl RegionWorkerLoop { sender.send(Ok(0)); } } + BackgroundNotify::CompactionUnit(req) => { + self.handle_compaction_unit(region_id, req).await; + } BackgroundNotify::CompactionDdlComplete { generation, request_index, @@ -1390,9 +1393,6 @@ impl RegionWorkerLoop { BackgroundNotify::CompactionFinished(req) => { self.handle_compaction_finished(region_id, req).await } - BackgroundNotify::CompactionCancelled(req) => { - self.handle_compaction_cancelled(region_id, req).await - } BackgroundNotify::CompactionFailed(req) => self.handle_compaction_failure(req).await, BackgroundNotify::Truncate(req) => self.handle_truncate_result(req).await, BackgroundNotify::DiscardUnflushed(req) => { @@ -1593,6 +1593,34 @@ impl WorkerListener { } } + /// Exposes a per-unit test gate before merge starts. + pub(crate) async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener + .on_compaction_unit_merge_begin(_region_id, _plan_id) + .await; + } + } + + /// Exposes a per-unit test gate between manifest commit and worker notification. + pub(crate) async fn on_compaction_unit_committed(&self, _region_id: RegionId, _plan_id: u64) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener + .on_compaction_unit_committed(_region_id, _plan_id) + .await; + } + } + + /// Reports unit visibility to the test listener independently of request completion. + pub(crate) fn on_compaction_unit_applied(&self, _region_id: RegionId, _plan_id: u64) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_compaction_unit_applied(_region_id, _plan_id); + } + } + pub(crate) async fn on_compaction_result_notified(&self, _region_id: RegionId) { #[cfg(any(test, feature = "test"))] if let Some(listener) = &self.listener { diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index 238c6c216da..50147584bcd 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -24,7 +24,7 @@ use crate::error::{RegionNotFoundSnafu, StaleCompactionExecutionSnafu}; use crate::metrics::COMPACTION_REQUEST_COUNT; use crate::region::MitoRegionRef; use crate::request::{ - BuildIndexRequest, CompactionCancelled, CompactionFailed, CompactionFinished, OnFailure, + BuildIndexRequest, CompactionFailed, CompactionFinished, CompactionUnitNotification, OnFailure, OptionOutputTx, }; use crate::sst::index::IndexBuildType; @@ -35,6 +35,90 @@ fn made_progress(files_to_add: usize, files_to_remove: usize) -> bool { } impl RegionWorkerLoop { + /// Drives local unit publication and applies successful edits before acknowledging completion. + pub(crate) async fn handle_compaction_unit( + &mut self, + region_id: RegionId, + notify: CompactionUnitNotification, + ) where + S: LogStore, + { + match notify { + CompactionUnitNotification::Finished { + plan_id, + result, + applied, + } => { + let Some(region) = self.regions.get_region(region_id) else { + return; + }; + match result { + Ok(edit) => { + if !self.compaction_scheduler.can_apply_unit(region_id, plan_id) { + let _ = applied.send(()); + return; + } + let progress = + made_progress(edit.files_to_add.len(), edit.files_to_remove.len()); + region.version_control.apply_edit( + Some(edit.clone()), + &[], + region.file_purger.clone(), + ); + self.listener.on_compaction_unit_applied(region_id, plan_id); + self.compaction_scheduler.on_unit_finished( + region_id, + plan_id, + Ok(progress), + ); + if self.config.index.build_mode == IndexBuildMode::Async + && !edit.files_to_add.is_empty() + { + self.handle_rebuild_index( + BuildIndexRequest { + region_id, + build_type: IndexBuildType::Compact, + file_metas: edit.files_to_add, + }, + OptionOutputTx::new(None), + ) + .await; + } + } + Err(err) => { + self.compaction_scheduler + .on_unit_finished(region_id, plan_id, Err(err)); + } + } + let _ = applied.send(()); + } + CompactionUnitNotification::Released { plan_id } => { + let Some(region) = self.regions.get_region(region_id) else { + return; + }; + let transition = self + .compaction_scheduler + .on_unit_released( + region_id, + plan_id, + ®ion.manifest_ctx, + self.schema_metadata_manager.clone(), + ) + .await; + self.listener.on_compaction_result_notified(region_id).await; + match transition { + CompactionTransition::AutomaticFollowupScheduled => { + region.update_schedule_compaction_millis() + } + CompactionTransition::DdlReady(mut ddls) => { + self.handle_ddl_requests(&mut ddls).await + } + CompactionTransition::NoAction => {} + } + } + } + } + pub(crate) async fn handle_compaction_pick_finished( &mut self, region_id: RegionId, @@ -189,36 +273,6 @@ impl RegionWorkerLoop { } } - pub(crate) async fn handle_compaction_cancelled( - &mut self, - region_id: RegionId, - request: CompactionCancelled, - ) where - S: LogStore, - { - let execution = request.execution.clone(); - let is_current = self.regions.get_region(region_id).is_some_and(|_| { - self.compaction_scheduler - .is_current_execution(region_id, &execution) - }); - request.on_success(); - - if !is_current { - return; - } - - // Reuse the scheduler's finish path to wake pending DDLs after a cooperative stop. - let mut pending_ddls = self - .compaction_scheduler - .on_execution_cancelled(region_id, &execution) - .await; - if !pending_ddls.is_empty() { - self.listener.on_compaction_result_notified(region_id).await; - } - - self.handle_ddl_requests(&mut pending_ddls).await; - } - /// When compaction fails, we simply log the error. pub(crate) async fn handle_compaction_failure(&mut self, req: CompactionFailed) { if self.regions.get_region(req.region_id).is_none() {