mito2/compaction/
compactor.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::num::NonZero;
16use std::sync::Arc;
17use std::time::Duration;
18
19use common_base::Plugins;
20use common_base::cancellation::{CancellableFuture, CancellationHandle};
21use common_meta::key::SchemaMetadataManagerRef;
22use common_telemetry::{debug, info, warn};
23use common_time::TimeToLive;
24use either::Either;
25use itertools::Itertools;
26use object_store::manager::ObjectStoreManagerRef;
27use partition::expr::PartitionExpr;
28use serde::{Deserialize, Serialize};
29use snafu::{OptionExt, ResultExt};
30use store_api::ManifestVersion;
31use store_api::metadata::RegionMetadataRef;
32use store_api::region_request::PathType;
33use store_api::storage::RegionId;
34
35use crate::access_layer::{
36    AccessLayer, AccessLayerRef, Metrics, OperationType, SstWriteRequest, WriteType,
37};
38use crate::cache::{CacheManager, CacheManagerRef};
39use crate::compaction::picker::PickerOutput;
40use crate::compaction::{CompactionOutput, CompactionSstReaderBuilder, find_dynamic_options};
41use crate::config::MitoConfig;
42use crate::engine::region_hook::{RegionHookRef, SstFileInfo};
43use crate::error;
44use crate::error::{
45    EmptyRegionDirSnafu, InvalidPartitionExprSnafu, ObjectStoreNotFoundSnafu, Result,
46};
47use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList};
48use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
49use crate::region::options::RegionOptions;
50use crate::region::version::VersionRef;
51use crate::region::{ManifestContext, RegionLeaderState, RegionRoleState};
52use crate::schedule::scheduler::LocalScheduler;
53use crate::sst::FormatType;
54use crate::sst::file::FileMeta;
55use crate::sst::file_purger::LocalFilePurger;
56use crate::sst::index::intermediate::IntermediateManager;
57use crate::sst::index::puffin_manager::PuffinManagerFactory;
58use crate::sst::location::region_dir_from_table_dir;
59use crate::sst::parquet::metadata::extract_primary_key_range;
60use crate::sst::parquet::{SstInfo, WriteOptions};
61use crate::sst::version::{SstVersion, SstVersionRef};
62
63/// Region version for compaction that does not hold memtables.
64#[derive(Clone)]
65pub struct CompactionVersion {
66    /// Metadata of the region.
67    ///
68    /// Altering metadata isn't frequent, storing metadata in Arc to allow sharing
69    /// metadata and reuse metadata when creating a new `Version`.
70    pub(crate) metadata: RegionMetadataRef,
71    /// Options of the region.
72    pub(crate) options: RegionOptions,
73    /// SSTs of the region.
74    pub(crate) ssts: SstVersionRef,
75    /// Inferred compaction time window.
76    pub(crate) compaction_time_window: Option<Duration>,
77}
78
79impl From<VersionRef> for CompactionVersion {
80    fn from(value: VersionRef) -> Self {
81        Self {
82            metadata: value.metadata.clone(),
83            options: value.options.clone(),
84            ssts: value.ssts.clone(),
85            compaction_time_window: value.compaction_time_window,
86        }
87    }
88}
89
90/// CompactionRegion represents a region that needs to be compacted.
91/// It's the subset of MitoRegion.
92#[derive(Clone)]
93pub struct CompactionRegion {
94    pub region_id: RegionId,
95    pub region_options: RegionOptions,
96
97    pub(crate) engine_config: Arc<MitoConfig>,
98    pub(crate) region_metadata: RegionMetadataRef,
99    pub(crate) cache_manager: CacheManagerRef,
100    /// Access layer to get the table path and path type.
101    pub access_layer: AccessLayerRef,
102    pub(crate) manifest_ctx: Arc<ManifestContext>,
103    pub(crate) current_version: CompactionVersion,
104    pub(crate) file_purger: Option<Arc<LocalFilePurger>>,
105    pub(crate) ttl: Option<TimeToLive>,
106
107    /// Controls the parallelism of this compaction task. Default is 1.
108    ///
109    /// The parallel is inside this compaction task, not across different compaction tasks.
110    /// It can be different windows of the same compaction task or something like this.
111    pub max_parallelism: usize,
112
113    pub(crate) plugins: Plugins,
114}
115
116/// OpenCompactionRegionRequest represents the request to open a compaction region.
117#[derive(Clone)]
118pub struct OpenCompactionRegionRequest {
119    pub region_id: RegionId,
120    pub table_dir: String,
121    pub path_type: PathType,
122    pub region_options: RegionOptions,
123    pub max_parallelism: usize,
124    /// Plugins for the compaction region, used to look up the [`RegionHook`](crate::engine::region_hook::RegionHook).
125    pub plugins: Plugins,
126}
127
128/// Open a compaction region from a compaction request.
129/// It's simple version of RegionOpener::open().
130pub async fn open_compaction_region(
131    req: &OpenCompactionRegionRequest,
132    mito_config: &MitoConfig,
133    object_store_manager: ObjectStoreManagerRef,
134    ttl_provider: Either<TimeToLive, SchemaMetadataManagerRef>,
135) -> Result<CompactionRegion> {
136    let object_store = {
137        let name = &req.region_options.storage;
138        if let Some(name) = name {
139            object_store_manager
140                .find(name)
141                .with_context(|| ObjectStoreNotFoundSnafu {
142                    object_store: name.clone(),
143                })?
144        } else {
145            object_store_manager.default_object_store()
146        }
147    };
148
149    let access_layer = {
150        let puffin_manager_factory = PuffinManagerFactory::new(
151            &mito_config.index.aux_path,
152            mito_config.index.staging_size.as_bytes(),
153            Some(mito_config.index.write_buffer_size.as_bytes() as _),
154            mito_config.index.staging_ttl,
155        )
156        .await?;
157        let intermediate_manager =
158            IntermediateManager::init_fs(mito_config.index.aux_path.clone()).await?;
159
160        Arc::new(AccessLayer::new(
161            &req.table_dir,
162            req.path_type,
163            object_store.clone(),
164            puffin_manager_factory,
165            intermediate_manager,
166        ))
167    };
168
169    let manifest_manager = {
170        let region_dir = region_dir_from_table_dir(&req.table_dir, req.region_id, req.path_type);
171        let region_manifest_options =
172            RegionManifestOptions::new(mito_config, &region_dir, object_store);
173
174        RegionManifestManager::open(region_manifest_options, &Default::default())
175            .await?
176            .with_context(|| EmptyRegionDirSnafu {
177                region_id: req.region_id,
178                region_dir: region_dir_from_table_dir(&req.table_dir, req.region_id, req.path_type),
179            })?
180    };
181
182    let manifest = manifest_manager.manifest();
183    let region_metadata = manifest.metadata.clone();
184    let hook: Option<RegionHookRef> = req.plugins.get();
185    let manifest_ctx = Arc::new(ManifestContext::new(
186        manifest_manager,
187        RegionRoleState::Leader(RegionLeaderState::Writable),
188        hook,
189    ));
190
191    let file_purger = {
192        let purge_scheduler = Arc::new(LocalScheduler::new(mito_config.max_background_purges));
193        Arc::new(LocalFilePurger::new(
194            purge_scheduler.clone(),
195            access_layer.clone(),
196            None,
197        ))
198    };
199
200    let current_version = {
201        let mut ssts = SstVersion::new();
202        ssts.add_files(file_purger.clone(), manifest.files.values().cloned());
203        CompactionVersion {
204            metadata: region_metadata.clone(),
205            options: req.region_options.clone(),
206            ssts: Arc::new(ssts),
207            compaction_time_window: manifest.compaction_time_window,
208        }
209    };
210
211    let ttl = match ttl_provider {
212        // Use the specified ttl.
213        Either::Left(ttl) => ttl,
214        // Get the ttl from the schema metadata manager.
215        Either::Right(schema_metadata_manager) => {
216            let (_, ttl) =
217                find_dynamic_options(req.region_id, &req.region_options, &schema_metadata_manager)
218                    .await
219                    .unwrap_or_else(|e| {
220                        warn!(e; "Failed to get ttl for region: {}", region_metadata.region_id);
221                        (
222                            crate::region::options::CompactionOptions::default(),
223                            TimeToLive::default(),
224                        )
225                    });
226            ttl
227        }
228    };
229
230    Ok(CompactionRegion {
231        region_id: req.region_id,
232        region_options: req.region_options.clone(),
233        engine_config: Arc::new(mito_config.clone()),
234        region_metadata: region_metadata.clone(),
235        cache_manager: Arc::new(CacheManager::default()),
236        access_layer,
237        manifest_ctx,
238        current_version,
239        file_purger: Some(file_purger),
240        ttl: Some(ttl),
241        max_parallelism: req.max_parallelism,
242        plugins: req.plugins.clone(),
243    })
244}
245
246impl CompactionRegion {
247    /// Get the file purger of the compaction region.
248    pub fn file_purger(&self) -> Option<Arc<LocalFilePurger>> {
249        self.file_purger.clone()
250    }
251
252    /// Stop the file purger scheduler of the compaction region.
253    pub async fn stop_purger_scheduler(&self) -> Result<()> {
254        if let Some(file_purger) = &self.file_purger {
255            file_purger.stop_scheduler().await
256        } else {
257            Ok(())
258        }
259    }
260
261    /// Fires [`RegionHook::on_sst_files_written`] for the freshly-merged SST
262    /// files in `merge_output`. Shared by both compaction paths, it must run
263    /// before [`Compactor::update_manifest`], whose `on_manifest_updated`
264    /// drains the per-region state this hook populates.
265    pub async fn invoke_sst_hook(&self, merge_output: &MergeOutput) {
266        let Some(hook) = self.plugins.get::<RegionHookRef>() else {
267            return;
268        };
269
270        // Remote compaction deserializes `MergeOutput` over gRPC, where
271        // `sst_infos` is `#[serde(skip)]`. Rebuild each `SstInfo` from the
272        // matching `FileMeta` so the hook observes real ids/rows/sizes rather
273        // than zeros — see [`sst_info_from_file_meta`] for what stays default.
274        let synthesized: Vec<SstInfo>;
275        let infos: &[SstInfo] = if merge_output.sst_infos.is_empty() {
276            synthesized = merge_output
277                .files_to_add
278                .iter()
279                .map(sst_info_from_file_meta)
280                .collect();
281            &synthesized
282        } else {
283            // `sst_infos` and `files_to_add` are documented as 1:1. If they
284            // ever diverge, `zip` below would silently truncate; warn so the
285            // mismatch is observable rather than dropping files from the hook.
286            if merge_output.sst_infos.len() != merge_output.files_to_add.len() {
287                warn!(
288                    "sst_infos length ({}) does not match files_to_add length ({}) for region {}",
289                    merge_output.sst_infos.len(),
290                    merge_output.files_to_add.len(),
291                    self.region_id
292                );
293            }
294            &merge_output.sst_infos
295        };
296
297        let files: Vec<SstFileInfo<'_>> = merge_output
298            .files_to_add
299            .iter()
300            .zip(infos)
301            .map(|(meta, info)| SstFileInfo {
302                sst_info_ref: info,
303                file_meta: meta,
304            })
305            .collect();
306        hook.on_sst_files_written(self.region_id, &self.region_metadata, &files)
307            .await;
308    }
309}
310
311/// Builds an [`SstInfo`] for the remote-compaction path by copying the scalar
312/// fields that [`FileMeta`] also carries.
313///
314/// Remote compaction runs off-process and ships `MergeOutput` over the wire;
315/// [`MergeOutput::sst_infos`] is `#[serde(skip)]` because the parquet footer
316/// (`ParquetMetaData`) is not serde-serializable. To avoid feeding the hook
317/// default (zero) values for real SSTs, this rebuilds the seven fields `FileMeta`
318/// and `SstInfo` share: `file_id`, `time_range`, `file_size`,
319/// `max_row_group_uncompressed_size`, `num_rows`, `num_row_groups`, `num_series`.
320///
321/// Two fields stay default and **cannot be recovered on the datanode**:
322/// - `file_metadata` — the parquet footer, whose column min/max/null-count
323///   statistics are required by hooks building richer artifacts (e.g. an Iceberg
324///   manifest). That data is produced by the compactor's writer and exists only
325///   in the freshly-written SST on object storage.
326/// - `index_metadata` — `FileMeta` tracks indexes via `available_indexes` /
327///   `indexes`, not as an [`IndexOutput`].
328///
329/// A hook that needs column stats must fetch the footer from object storage
330/// (the [`CompactionRegion`]'s `access_layer` reaches the store), or the footer
331/// must be shipped over the wire (revisit the `#[serde(skip)]` on `sst_infos`).
332/// `num_rows` may read `0` for legacy `FileMeta`s where the count is unknown.
333fn sst_info_from_file_meta(meta: &FileMeta) -> SstInfo {
334    SstInfo {
335        file_id: meta.file_id,
336        time_range: meta.time_range,
337        file_size: meta.file_size,
338        max_row_group_uncompressed_size: meta.max_row_group_uncompressed_size,
339        num_rows: meta.num_rows as usize,
340        num_row_groups: meta.num_row_groups,
341        num_series: meta.num_series,
342        ..Default::default()
343    }
344}
345
346/// `[MergeOutput]` represents the output of merging SST files.
347#[derive(Default, Debug, Serialize, Deserialize)]
348pub struct MergeOutput {
349    pub files_to_add: Vec<FileMeta>,
350    pub files_to_remove: Vec<FileMeta>,
351    pub compaction_time_window: Option<i64>,
352    #[serde(skip)]
353    pub sst_infos: Vec<SstInfo>,
354}
355
356impl MergeOutput {
357    pub fn is_empty(&self) -> bool {
358        self.files_to_add.is_empty() && self.files_to_remove.is_empty()
359    }
360
361    pub fn input_file_size(&self) -> u64 {
362        self.files_to_remove.iter().map(|f| f.file_size).sum()
363    }
364
365    pub fn output_file_size(&self) -> u64 {
366        self.files_to_add.iter().map(|f| f.file_size).sum()
367    }
368}
369
370/// Compactor is the trait that defines the compaction logic.
371#[async_trait::async_trait]
372pub trait Compactor: Send + Sync + 'static {
373    /// Merge SST files for a region.
374    async fn merge_ssts(
375        &self,
376        compaction_region: &CompactionRegion,
377        picker_output: PickerOutput,
378    ) -> Result<MergeOutput>;
379
380    /// Update the manifest after merging SST files.
381    async fn update_manifest(
382        &self,
383        compaction_region: &CompactionRegion,
384        merge_output: MergeOutput,
385    ) -> Result<(RegionEdit, ManifestVersion)>;
386}
387
388/// Trait for merging a single compaction output into SST files.
389///
390/// This is extracted from `DefaultCompactor` to allow injecting mock
391/// implementations in tests.
392#[async_trait::async_trait]
393pub trait SstMerger: Send + Sync + 'static {
394    async fn merge_single_output(
395        &self,
396        compaction_region: CompactionRegion,
397        output: CompactionOutput,
398        write_opts: WriteOptions,
399    ) -> Result<(Vec<FileMeta>, Vec<SstInfo>)>;
400}
401
402/// The production [`SstMerger`] that reads, merges, and writes SST files.
403#[derive(Clone)]
404pub struct DefaultSstMerger;
405
406#[async_trait::async_trait]
407impl SstMerger for DefaultSstMerger {
408    async fn merge_single_output(
409        &self,
410        compaction_region: CompactionRegion,
411        output: CompactionOutput,
412        write_opts: WriteOptions,
413    ) -> Result<(Vec<FileMeta>, Vec<SstInfo>)> {
414        let region_id = compaction_region.region_id;
415        let storage = compaction_region.region_options.storage.clone();
416        let index_options = compaction_region
417            .current_version
418            .options
419            .index_options
420            .clone();
421        let append_mode = compaction_region.current_version.options.append_mode;
422        let merge_mode = compaction_region.current_version.options.merge_mode();
423        let flat_format = compaction_region
424            .region_options
425            .sst_format
426            .map(|format| format == FormatType::Flat)
427            .unwrap_or(compaction_region.engine_config.default_flat_format);
428
429        let index_config = compaction_region.engine_config.index.clone();
430        let inverted_index_config = compaction_region.engine_config.inverted_index.clone();
431        let fulltext_index_config = compaction_region.engine_config.fulltext_index.clone();
432        let bloom_filter_index_config = compaction_region.engine_config.bloom_filter_index.clone();
433        #[cfg(feature = "vector_index")]
434        let vector_index_config = compaction_region.engine_config.vector_index.clone();
435
436        let input_file_names = output
437            .inputs
438            .iter()
439            .map(|f| f.file_id().to_string())
440            .join(",");
441        let max_sequence = output
442            .inputs
443            .iter()
444            .map(|f| f.meta_ref().sequence)
445            .max()
446            .flatten();
447        let builder = CompactionSstReaderBuilder {
448            metadata: compaction_region.region_metadata.clone(),
449            sst_layer: compaction_region.access_layer.clone(),
450            cache: compaction_region.cache_manager.clone(),
451            inputs: &output.inputs,
452            append_mode,
453            filter_deleted: output.filter_deleted,
454            time_range: output.output_time_range,
455            merge_mode,
456        };
457        let source = builder.build_flat_sst_reader().await?;
458
459        let mut metrics = Metrics::new(WriteType::Compaction);
460        let region_metadata = compaction_region.region_metadata.clone();
461        let sst_infos = compaction_region
462            .access_layer
463            .write_sst(
464                SstWriteRequest {
465                    op_type: OperationType::Compact,
466                    metadata: region_metadata.clone(),
467                    source,
468                    cache_manager: compaction_region.cache_manager.clone(),
469                    storage,
470                    max_sequence: max_sequence.map(NonZero::get),
471                    sst_write_format: if flat_format {
472                        FormatType::Flat
473                    } else {
474                        FormatType::PrimaryKey
475                    },
476                    index_options,
477                    index_config,
478                    inverted_index_config,
479                    fulltext_index_config,
480                    bloom_filter_index_config,
481                    #[cfg(feature = "vector_index")]
482                    vector_index_config,
483                },
484                &write_opts,
485                &mut metrics,
486            )
487            .await?;
488        // Convert partition expression once outside the map
489        let partition_expr = match &region_metadata.partition_expr {
490            None => None,
491            Some(json_str) if json_str.is_empty() => None,
492            Some(json_str) => PartitionExpr::from_json_str(json_str).with_context(|_| {
493                InvalidPartitionExprSnafu {
494                    expr: json_str.clone(),
495                }
496            })?,
497        };
498
499        let output_files = sst_infos
500            .iter()
501            .map(|sst_info| {
502                let pk_range = sst_info
503                    .file_metadata
504                    .as_ref()
505                    .and_then(|meta| extract_primary_key_range(meta, &region_metadata));
506                let (primary_key_min, primary_key_max) = match pk_range {
507                    Some((min, max)) => (Some(min), Some(max)),
508                    None => (None, None),
509                };
510
511                FileMeta {
512                    region_id,
513                    file_id: sst_info.file_id,
514                    time_range: sst_info.time_range,
515                    level: output.output_level,
516                    file_size: sst_info.file_size,
517                    max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size,
518                    available_indexes: sst_info.index_metadata.build_available_indexes(),
519                    indexes: sst_info.index_metadata.build_indexes(),
520                    index_file_size: sst_info.index_metadata.file_size,
521                    index_version: 0,
522                    num_rows: sst_info.num_rows as u64,
523                    num_row_groups: sst_info.num_row_groups,
524                    sequence: max_sequence,
525                    partition_expr: partition_expr.clone(),
526                    num_series: sst_info.num_series,
527                    primary_key_min,
528                    primary_key_max,
529                }
530            })
531            .collect::<Vec<_>>();
532        let output_file_names = output_files.iter().map(|f| f.file_id.to_string()).join(",");
533        info!(
534            "Region {} compaction inputs: [{}], outputs: [{}], flat_format: {}, metrics: {:?}",
535            region_id, input_file_names, output_file_names, flat_format, metrics
536        );
537        metrics.observe();
538        Ok((output_files, sst_infos.into_iter().collect()))
539    }
540}
541
542/// DefaultCompactor is the default implementation of Compactor.
543///
544/// It is parameterized by an [`SstMerger`] to allow injecting mock
545/// implementations in tests.
546pub struct DefaultCompactor<M = DefaultSstMerger> {
547    merger: M,
548    cancel_handle: Arc<CancellationHandle>,
549}
550
551#[cfg(test)]
552impl<M: SstMerger> DefaultCompactor<M> {
553    pub fn with_merger(merger: M) -> Self {
554        Self {
555            merger,
556            cancel_handle: Arc::new(CancellationHandle::default()),
557        }
558    }
559}
560
561impl DefaultCompactor {
562    pub fn with_cancel_handle(cancel_handle: Arc<CancellationHandle>) -> Self {
563        Self {
564            merger: DefaultSstMerger,
565            cancel_handle,
566        }
567    }
568}
569
570#[async_trait::async_trait]
571impl<M: SstMerger> Compactor for DefaultCompactor<M>
572where
573    M: Clone,
574{
575    async fn merge_ssts(
576        &self,
577        compaction_region: &CompactionRegion,
578        mut picker_output: PickerOutput,
579    ) -> Result<MergeOutput> {
580        let internal_parallelism = compaction_region.max_parallelism.max(1);
581        let compaction_time_window = picker_output.time_window_size;
582        let region_id = compaction_region.region_id;
583
584        // Build tasks along with their input file metas so we can track which
585        // inputs correspond to each task.
586        let mut tasks: Vec<(Vec<FileMeta>, _)> = Vec::with_capacity(picker_output.outputs.len());
587
588        for output in picker_output.outputs.drain(..) {
589            let inputs_to_remove: Vec<_> =
590                output.inputs.iter().map(|f| f.meta_ref().clone()).collect();
591            let write_opts = WriteOptions {
592                write_buffer_size: compaction_region.engine_config.sst_write_buffer_size,
593                max_file_size: picker_output.max_file_size,
594                ..Default::default()
595            };
596            let merger = self.merger.clone();
597            let compaction_region = compaction_region.clone();
598            let fut = async move {
599                merger
600                    .merge_single_output(compaction_region, output, write_opts)
601                    .await
602            };
603            tasks.push((inputs_to_remove, fut));
604        }
605
606        let hook: Option<RegionHookRef> = compaction_region.plugins.get();
607        let mut output_files = Vec::with_capacity(tasks.len());
608        let mut all_sst_infos: Vec<SstInfo> = Vec::new();
609        let mut compacted_inputs = Vec::with_capacity(
610            tasks.iter().map(|(inputs, _)| inputs.len()).sum::<usize>()
611                + picker_output.expired_ssts.len(),
612        );
613
614        while !tasks.is_empty() {
615            let mut chunk: Vec<(Vec<FileMeta>, _)> = Vec::with_capacity(internal_parallelism);
616            for _ in 0..internal_parallelism {
617                if let Some(task) = tasks.pop() {
618                    chunk.push(task);
619                }
620            }
621            let mut spawned: Vec<_> = chunk
622                .into_iter()
623                .map(|(inputs, fut)| {
624                    let handle = common_runtime::spawn_compact(fut);
625                    (inputs, handle)
626                })
627                .collect();
628
629            while let Some((inputs, handle)) = spawned.pop() {
630                let abort_handle = handle.abort_handle();
631                match CancellableFuture::new(handle, self.cancel_handle.clone()).await {
632                    Ok(Ok(Ok((files, infos)))) => {
633                        output_files.extend(files);
634                        if hook.is_some() {
635                            all_sst_infos.extend(infos);
636                        }
637                        compacted_inputs.extend(inputs);
638                    }
639                    Ok(Ok(Err(e))) => {
640                        warn!(
641                            e; "Failed to merge compaction output for region: {}, inputs: [{}]",
642                            region_id,
643                            inputs.iter().map(|f| f.file_id.to_string()).join(",")
644                        );
645                    }
646                    Ok(Err(e)) => {
647                        warn!(
648                            "Region {} compaction task join error for inputs: [{}], skipping: {}",
649                            region_id,
650                            inputs.iter().map(|f| f.file_id.to_string()).join(","),
651                            e
652                        );
653                        // If the cancel handle is cancelled,
654                        // cancel the remaining tasks before returns the error.
655                        if self.cancel_handle.is_cancelled() {
656                            abort_handle.abort();
657                            for (_, handle) in spawned {
658                                handle.abort();
659                            }
660                        }
661                        return Err(e).context(error::JoinSnafu);
662                    }
663                    Err(_) => {
664                        debug!(
665                            "Compaction merge cancelled for region: {}, aborting remaining {} spawned tasks",
666                            region_id,
667                            spawned.len(),
668                        );
669                        abort_handle.abort();
670                        for (_, handle) in spawned {
671                            handle.abort();
672                        }
673                        break;
674                    }
675                }
676            }
677
678            if self.cancel_handle.is_cancelled() {
679                info!("Compaction merge cancelled for region: {}", region_id);
680                break;
681            }
682        }
683
684        // Include expired SSTs in removals — these don't depend on merge success.
685        compacted_inputs.extend(
686            picker_output
687                .expired_ssts
688                .iter()
689                .map(|f| f.meta_ref().clone()),
690        );
691
692        Ok(MergeOutput {
693            files_to_add: output_files,
694            files_to_remove: compacted_inputs,
695            compaction_time_window: Some(compaction_time_window),
696            sst_infos: all_sst_infos,
697        })
698    }
699
700    async fn update_manifest(
701        &self,
702        compaction_region: &CompactionRegion,
703        merge_output: MergeOutput,
704    ) -> Result<(RegionEdit, ManifestVersion)> {
705        // Write region edit to manifest.
706        let edit = RegionEdit {
707            files_to_add: merge_output.files_to_add,
708            files_to_remove: merge_output.files_to_remove,
709            // Use current timestamp as the edit timestamp.
710            timestamp_ms: Some(chrono::Utc::now().timestamp_millis()),
711            compaction_time_window: merge_output
712                .compaction_time_window
713                .map(|seconds| Duration::from_secs(seconds as u64)),
714            flushed_entry_id: None,
715            flushed_sequence: None,
716            committed_sequence: None,
717        };
718
719        let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone()));
720        // TODO: We might leak files if we fail to update manifest. We can add a cleanup task to remove them later.
721        let manifest_version = compaction_region
722            .manifest_ctx
723            .update_manifest_for_compaction(action_list)
724            .await?;
725
726        Ok((edit, manifest_version))
727    }
728}
729
730#[cfg(test)]
731mod tests {
732    use std::sync::atomic::{AtomicUsize, Ordering};
733    use std::sync::{Arc, Mutex};
734    use std::time::Duration;
735
736    use store_api::storage::{FileId, RegionId};
737    use tokio::time::sleep;
738
739    use super::{DefaultCompactor, *};
740    use crate::cache::CacheManager;
741    use crate::compaction::picker::PickerOutput;
742    use crate::error::Result;
743    use crate::sst::file::FileHandle;
744    use crate::sst::file_purger::NoopFilePurger;
745    use crate::sst::version::SstVersion;
746    use crate::test_util::memtable_util::metadata_for_test;
747    use crate::test_util::scheduler_util::SchedulerEnv;
748
749    fn dummy_file_meta() -> FileMeta {
750        FileMeta {
751            region_id: RegionId::new(1, 1),
752            file_id: FileId::random(),
753            file_size: 100,
754            ..Default::default()
755        }
756    }
757
758    fn new_file_handle(meta: FileMeta) -> FileHandle {
759        FileHandle::new(meta, Arc::new(NoopFilePurger))
760    }
761
762    /// Build a minimal [`CompactionRegion`] suitable for tests where the
763    /// [`SstMerger`] is mocked and never touches the access layer.
764    async fn new_test_compaction_region() -> CompactionRegion {
765        let env = SchedulerEnv::new().await;
766        let metadata = metadata_for_test();
767        let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
768        CompactionRegion {
769            region_id: RegionId::new(1, 1),
770            region_options: RegionOptions::default(),
771            engine_config: Arc::new(MitoConfig::default()),
772            region_metadata: metadata.clone(),
773            cache_manager: Arc::new(CacheManager::default()),
774            access_layer: env.access_layer.clone(),
775            manifest_ctx,
776            current_version: CompactionVersion {
777                metadata,
778                options: RegionOptions::default(),
779                ssts: Arc::new(SstVersion::new()),
780                compaction_time_window: None,
781            },
782            file_purger: None,
783            ttl: None,
784            max_parallelism: 1,
785            plugins: Plugins::new(),
786        }
787    }
788
789    /// An [`SstMerger`] that returns pre-configured results per call index.
790    ///
791    /// Call 0 gets `results[0]`, call 1 gets `results[1]`, etc.
792    #[derive(Clone)]
793    struct MockMerger {
794        results: Arc<Mutex<Vec<Result<Vec<FileMeta>>>>>,
795        call_idx: Arc<AtomicUsize>,
796    }
797
798    impl MockMerger {
799        fn new(results: Vec<Result<Vec<FileMeta>>>) -> Self {
800            Self {
801                results: Arc::new(Mutex::new(results)),
802                call_idx: Arc::new(AtomicUsize::new(0)),
803            }
804        }
805    }
806
807    #[async_trait::async_trait]
808    impl SstMerger for MockMerger {
809        async fn merge_single_output(
810            &self,
811            _compaction_region: CompactionRegion,
812            _output: CompactionOutput,
813            _write_opts: WriteOptions,
814        ) -> Result<(Vec<FileMeta>, Vec<SstInfo>)> {
815            let idx = self.call_idx.fetch_add(1, Ordering::SeqCst);
816            match self.results.lock().unwrap().get(idx) {
817                Some(Ok(files)) => Ok((files.clone(), Vec::new())),
818                Some(Err(_)) => error::InvalidMetaSnafu {
819                    reason: format!("simulated failure at index {idx}"),
820                }
821                .fail(),
822                None => panic!("MockMerger: no result configured for call index {idx}"),
823            }
824        }
825    }
826
827    #[tokio::test]
828    async fn test_partial_merge_failure_collects_only_successful_outputs() {
829        common_telemetry::init_default_ut_logging();
830
831        let compaction_region = new_test_compaction_region().await;
832
833        // Prepare 3 compaction outputs: output 0 and 2 succeed, output 1 fails.
834        let input_meta_0 = dummy_file_meta();
835        let input_meta_1 = dummy_file_meta();
836        let input_meta_2 = dummy_file_meta();
837
838        let output_meta_0 = vec![dummy_file_meta()];
839        let output_meta_2 = vec![dummy_file_meta(), dummy_file_meta()];
840
841        let merger = MockMerger::new(vec![
842            Ok(output_meta_0.clone()),
843            Err(error::InvalidMetaSnafu {
844                reason: "boom".to_string(),
845            }
846            .build()),
847            Ok(output_meta_2.clone()),
848        ]);
849        let compactor = DefaultCompactor::with_merger(merger);
850
851        let picker_output = PickerOutput {
852            outputs: vec![
853                CompactionOutput {
854                    output_level: 1,
855                    inputs: vec![new_file_handle(input_meta_0.clone())],
856                    filter_deleted: false,
857                    output_time_range: None,
858                },
859                CompactionOutput {
860                    output_level: 1,
861                    inputs: vec![new_file_handle(input_meta_1.clone())],
862                    filter_deleted: false,
863                    output_time_range: None,
864                },
865                CompactionOutput {
866                    output_level: 1,
867                    inputs: vec![new_file_handle(input_meta_2.clone())],
868                    filter_deleted: false,
869                    output_time_range: None,
870                },
871            ],
872            expired_ssts: vec![],
873            time_window_size: 3600,
874            max_file_size: None,
875        };
876
877        let merge_output = compactor
878            .merge_ssts(&compaction_region, picker_output)
879            .await
880            .unwrap();
881
882        // Outputs 0 and 2 succeeded (1 + 2 = 3 files added).
883        assert_eq!(merge_output.files_to_add.len(), 3);
884        // Only inputs from successful merges should be removed.
885        assert_eq!(merge_output.files_to_remove.len(), 2);
886
887        let removed_ids: Vec<_> = merge_output
888            .files_to_remove
889            .iter()
890            .map(|f| f.file_id)
891            .collect();
892        assert!(removed_ids.contains(&input_meta_0.file_id));
893        assert!(removed_ids.contains(&input_meta_2.file_id));
894        // The failed output's input must NOT be removed.
895        assert!(!removed_ids.contains(&input_meta_1.file_id));
896    }
897
898    #[tokio::test]
899    async fn test_all_outputs_succeed() {
900        common_telemetry::init_default_ut_logging();
901
902        let compaction_region = new_test_compaction_region().await;
903        let input_meta = dummy_file_meta();
904        let output_meta = vec![dummy_file_meta()];
905
906        let merger = MockMerger::new(vec![Ok(output_meta.clone())]);
907        let compactor = DefaultCompactor::with_merger(merger);
908
909        let picker_output = PickerOutput {
910            outputs: vec![CompactionOutput {
911                output_level: 1,
912                inputs: vec![new_file_handle(input_meta.clone())],
913                filter_deleted: false,
914                output_time_range: None,
915            }],
916            expired_ssts: vec![],
917            time_window_size: 3600,
918            max_file_size: None,
919        };
920
921        let merge_output = compactor
922            .merge_ssts(&compaction_region, picker_output)
923            .await
924            .unwrap();
925
926        assert_eq!(merge_output.files_to_add.len(), 1);
927        assert_eq!(merge_output.files_to_add[0].file_id, output_meta[0].file_id);
928        assert_eq!(merge_output.files_to_remove.len(), 1);
929        assert_eq!(merge_output.files_to_remove[0].file_id, input_meta.file_id);
930    }
931
932    #[tokio::test]
933    async fn test_expired_ssts_always_removed() {
934        common_telemetry::init_default_ut_logging();
935
936        let compaction_region = new_test_compaction_region().await;
937        let input_meta = dummy_file_meta();
938        let expired_meta = dummy_file_meta();
939
940        // The single merge output fails, but expired SSTs should still be removed.
941        let merger = MockMerger::new(vec![Err(error::InvalidMetaSnafu {
942            reason: "fail".to_string(),
943        }
944        .build())]);
945        let compactor = DefaultCompactor::with_merger(merger);
946
947        let picker_output = PickerOutput {
948            outputs: vec![CompactionOutput {
949                output_level: 1,
950                inputs: vec![new_file_handle(input_meta.clone())],
951                filter_deleted: false,
952                output_time_range: None,
953            }],
954            expired_ssts: vec![new_file_handle(expired_meta.clone())],
955            time_window_size: 3600,
956            max_file_size: None,
957        };
958
959        let merge_output = compactor
960            .merge_ssts(&compaction_region, picker_output)
961            .await
962            .unwrap();
963
964        // No files added (merge failed).
965        assert!(merge_output.files_to_add.is_empty());
966        // Only the expired SST should be in files_to_remove (not the failed merge's input).
967        assert_eq!(merge_output.files_to_remove.len(), 1);
968        assert_eq!(
969            merge_output.files_to_remove[0].file_id,
970            expired_meta.file_id
971        );
972    }
973
974    #[derive(Clone)]
975    struct BlockingMerger {
976        call_idx: Arc<AtomicUsize>,
977    }
978
979    #[async_trait::async_trait]
980    impl SstMerger for BlockingMerger {
981        async fn merge_single_output(
982            &self,
983            _compaction_region: CompactionRegion,
984            _output: CompactionOutput,
985            _write_opts: WriteOptions,
986        ) -> Result<(Vec<FileMeta>, Vec<SstInfo>)> {
987            self.call_idx.fetch_add(1, Ordering::SeqCst);
988            std::future::pending().await
989        }
990    }
991
992    #[tokio::test(flavor = "multi_thread")]
993    async fn test_merge_ssts_cancels_spawned_tasks() {
994        common_telemetry::init_default_ut_logging();
995
996        let mut compaction_region = new_test_compaction_region().await;
997        compaction_region.max_parallelism = 2;
998
999        let cancel_handle = Arc::new(CancellationHandle::default());
1000        let call_idx = Arc::new(AtomicUsize::new(0));
1001        let compactor = DefaultCompactor {
1002            merger: BlockingMerger {
1003                call_idx: call_idx.clone(),
1004            },
1005            cancel_handle: cancel_handle.clone(),
1006        };
1007
1008        let picker_output = PickerOutput {
1009            outputs: vec![
1010                CompactionOutput {
1011                    output_level: 1,
1012                    inputs: vec![new_file_handle(dummy_file_meta())],
1013                    filter_deleted: false,
1014                    output_time_range: None,
1015                },
1016                CompactionOutput {
1017                    output_level: 1,
1018                    inputs: vec![new_file_handle(dummy_file_meta())],
1019                    filter_deleted: false,
1020                    output_time_range: None,
1021                },
1022                CompactionOutput {
1023                    output_level: 1,
1024                    inputs: vec![new_file_handle(dummy_file_meta())],
1025                    filter_deleted: false,
1026                    output_time_range: None,
1027                },
1028            ],
1029            expired_ssts: vec![],
1030            time_window_size: 3600,
1031            max_file_size: None,
1032        };
1033
1034        let task = tokio::spawn(async move {
1035            compactor
1036                .merge_ssts(&compaction_region, picker_output)
1037                .await
1038        });
1039
1040        sleep(Duration::from_millis(100)).await;
1041        cancel_handle.cancel();
1042
1043        let merge_output = task
1044            .await
1045            .expect("merge_ssts should stop after cancellation")
1046            .unwrap();
1047
1048        let started = call_idx.load(Ordering::SeqCst);
1049
1050        assert!(merge_output.files_to_add.is_empty());
1051        assert!(merge_output.files_to_remove.is_empty());
1052        assert_eq!(started, 2);
1053    }
1054}
mito2/compaction/compactor.rs

mito2/compaction/
compactor.rs
