diff --git a/licenserc-enterprise.toml b/licenserc-enterprise.toml index 939346f5cba..22b35ded3e8 100644 --- a/licenserc-enterprise.toml +++ b/licenserc-enterprise.toml @@ -26,6 +26,10 @@ includes = [ "src/sql/src/parsers/create_parser/trigger.rs", "src/sql/src/parsers/show_parser/trigger.rs", "src/mito2/src/extension.rs", + "src/mito2/src/engine/table_compaction.rs", + "src/mito2/src/engine/table_compaction/output.rs", + "src/mito2/src/engine/table_compaction/tests.rs", + "src/mito2/src/worker/handle_table_compaction.rs", "src/mito2/src/compaction/multi_region.rs", "src/mito2/src/compaction/multi_region/source.rs", "src/mito2/src/compaction/multi_region/memory.rs", diff --git a/licenserc.toml b/licenserc.toml index e601ca6f7c1..172a6af5d61 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -42,6 +42,10 @@ excludes = [ "src/sql/src/parsers/create_parser/trigger.rs", "src/sql/src/parsers/show_parser/trigger.rs", "src/mito2/src/extension.rs", + "src/mito2/src/engine/table_compaction.rs", + "src/mito2/src/engine/table_compaction/output.rs", + "src/mito2/src/engine/table_compaction/tests.rs", + "src/mito2/src/worker/handle_table_compaction.rs", "src/mito2/src/compaction/multi_region.rs", "src/mito2/src/compaction/multi_region/source.rs", "src/mito2/src/compaction/multi_region/memory.rs", diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index ffd674917a2..9dc57b3f1fc 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -25,6 +25,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `flush` | `src/mito2/src/flush.rs` | `FlushScheduler`, `WriteBufferManager`, memtable → SST | | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `compaction/multi_region` | `src/mito2/src/compaction/multi_region.rs` | Enterprise lossless routing by a captured table partition rule; selected SSTs to per-region outputs, without manifest publication | +| `engine/table_compaction` | `src/mito2/src/engine/table_compaction.rs` | Enterprise conditional publication of table-compaction results: shared-input snapshot, verified-output checks, and the manifest-locked swap of shared inputs for L0 outputs | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | | `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning, catalogs, worker-owned reconciliation (`maintenance.rs`), shared approximate usage (`task.rs`), catalogs, and index lifecycle | diff --git a/src/mito2/src/compaction/multi_region.rs b/src/mito2/src/compaction/multi_region.rs index 2fa6ab1b737..de4549a4c5c 100644 --- a/src/mito2/src/compaction/multi_region.rs +++ b/src/mito2/src/compaction/multi_region.rs @@ -102,10 +102,10 @@ pub struct MultiRegionCompactionStats { /// Complete per-target outputs, including explicit empty partitions. pub struct MultiRegionCompactionOutput { - /// Output files, initially assigned to level 1 (L1). + /// Output files, already assigned to region-local L0. /// - /// The caller owns publication and must set [FileMeta::level] before - /// publishing if its intended output level differs from L1. + /// Window merging and semantic folding stay with ordinary region compaction, + /// which runs after publication; the caller does not adjust the level. pub files: BTreeMap>, /// Complete footer/index statistics for target hook delivery; not fabricated from FileMeta. pub sst_infos: BTreeMap>, diff --git a/src/mito2/src/compaction/multi_region/tests.rs b/src/mito2/src/compaction/multi_region/tests.rs index 3aadc435cb5..67709443086 100644 --- a/src/mito2/src/compaction/multi_region/tests.rs +++ b/src/mito2/src/compaction/multi_region/tests.rs @@ -543,7 +543,7 @@ async fn test_single_scan_routes_all_versions_without_folding( } for file in files { assert!(protected.files.lock().unwrap().contains(&file.file_id())); - assert_eq!(1, file.level); + assert_eq!(0, file.level); assert_eq!(id, file.region_id); assert_eq!(id != fixture.metadata.region_id, file.preserve_row_sequence); let rows = fixture @@ -1665,7 +1665,7 @@ async fn test_streaming_merge_preserves_overlapping_versions_and_field_routing( for (file, info) in files.into_iter().zip(&output.sst_infos[&id]) { assert!(protection.files.lock().unwrap().contains(&file.file_id())); assert_eq!(id, file.region_id); - assert_eq!(1, file.level); + assert_eq!(0, file.level); assert_eq!(info.file_id, file.file_id); assert_eq!(info.num_rows as u64, file.num_rows); assert_eq!(info.file_size, file.file_size); diff --git a/src/mito2/src/compaction/multi_region/writer.rs b/src/mito2/src/compaction/multi_region/writer.rs index 737fae51e44..0d2ce6ea838 100644 --- a/src/mito2/src/compaction/multi_region/writer.rs +++ b/src/mito2/src/compaction/multi_region/writer.rs @@ -183,7 +183,9 @@ impl<'a> TargetWriter<'a> { region_id: self.target.metadata.region_id, file_id: info.file_id, time_range: info.time_range, - level: 1, + // Region-local L0: window merges and semantic folding stay with + // ordinary region compaction. + level: 0, file_size: info.file_size, max_row_group_uncompressed_size: info.max_row_group_uncompressed_size, available_indexes: info.index_metadata.build_available_indexes(), diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 8ff2b0f6ce0..de9e796f8d1 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -87,6 +87,8 @@ mod skip_wal_test; mod staging_test; #[cfg(test)] mod sync_test; +#[cfg(feature = "enterprise")] +pub mod table_compaction; #[cfg(test)] mod truncate_test; diff --git a/src/mito2/src/engine/region_hook.rs b/src/mito2/src/engine/region_hook.rs index b3ecb4a1cdc..9f138405842 100644 --- a/src/mito2/src/engine/region_hook.rs +++ b/src/mito2/src/engine/region_hook.rs @@ -157,6 +157,7 @@ use std::fmt::Debug; use std::sync::Arc; use async_trait::async_trait; +use parquet::file::metadata::ParquetMetaData; use store_api::ManifestVersion; use store_api::metadata::RegionMetadataRef; use store_api::storage::RegionId; @@ -269,6 +270,13 @@ pub struct SstFileInfo<'a> { pub file_meta: &'a FileMeta, } +/// Externally produced SST with its verified footer. Index-builder statistics +/// are deliberately absent: they cannot be reconstructed from FileMeta. +pub struct VerifiedSstFileInfo<'a> { + pub file_meta: &'a FileMeta, + pub footer: &'a Arc, +} + /// Hook for observing region mutations in mito2. /// /// Implementations can be registered via the `Plugins` system: @@ -315,6 +323,21 @@ pub trait RegionHook: Send + Sync + Debug { Ok(HashSet::new()) } + /// Called for verified table-compaction outputs before manifest publication. + /// Unlike `on_sst_files_written`, this carries actual remote footers without + /// fabricating unavailable index-builder statistics. Retries may repeat it; + /// no manifest lock is held and implementations must be idempotent. Publication + /// may still fail after this callback; it is not a publication acknowledgement. + /// Use `on_manifest_updated` to observe a successful manifest commit. + async fn on_sst_files_verified( + &self, + region_id: RegionId, + region_metadata: &RegionMetadataRef, + files: &[VerifiedSstFileInfo<'_>], + ) { + let _ = (region_id, region_metadata, files); + } + /// Called after the region manifest is successfully committed to the **live** /// (normal) manifest directory. /// diff --git a/src/mito2/src/engine/table_compaction.rs b/src/mito2/src/engine/table_compaction.rs new file mode 100644 index 00000000000..d50cbb0944e --- /dev/null +++ b/src/mito2/src/engine/table_compaction.rs @@ -0,0 +1,131 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +//! One-shot publication of already-admitted shared inputs. Task tracking is external. + +pub(crate) mod output; +#[cfg(test)] +mod tests; + +use serde::{Deserialize, Serialize}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::RegionId; +use tokio::sync::oneshot; + +use crate::engine::MitoEngine; +use crate::error::{ + Error, InvalidRequestSnafu, RecvSnafu, RegionBusySnafu, RegionNotFoundSnafu, Result, + SerdeJsonSnafu, +}; +use crate::manifest::action::RegionEdit; +use crate::region::options::RegionOptions; +use crate::region::{RegionLeaderState, RegionRoleState}; +use crate::request::{RegionEditKind, RegionEditRequest, WorkerRequest}; +use crate::sst::file::{FileMeta, SHARED_LEVEL_SLOT}; + +/// A query-visible input snapshot. Callers may select a subset of `inputs` for a job. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TableCompactionSnapshot { + pub metadata: RegionMetadataRef, + pub options: serde_json::Value, + pub inputs: Vec, +} + +/// Transient publication preconditions, never stored in a manifest. +#[derive(Debug)] +pub(crate) struct PublicationContext { + pub metadata: RegionMetadataRef, + pub options: RegionOptions, +} + +/// A result for this publication attempt, not a durable task receipt. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PublicationOutcome { + Applied, + /// The captured inputs/schema/options no longer match. Re-snapshot and replan. + StaleInputs, +} + +impl MitoEngine { + pub fn table_compaction_snapshot( + &self, + region_id: RegionId, + ) -> Result { + let region = self + .find_region(region_id) + .context(RegionNotFoundSnafu { region_id })?; + ensure!( + region.state() == RegionRoleState::Leader(RegionLeaderState::Writable), + RegionBusySnafu { region_id } + ); + let current = region.version_control.current(); + Ok(TableCompactionSnapshot { + metadata: current.version.metadata.clone(), + options: serde_json::to_value(¤t.version.options).context(SerdeJsonSnafu)?, + inputs: current.version.ssts.levels()[SHARED_LEVEL_SLOT] + .files() + .map(|file| file.meta_ref().clone()) + .collect(), + }) + } + + /// Publishes one region independently, using compaction rather than import sequences. + /// + /// Duplicate submissions are allowed. A stale outcome does NOT prove these + /// outputs were never published (an earlier ACK may have been lost). Never + /// blindly delete outputs on stale/error: keep their protection until recovery + /// or manifest-aware GC determines liveness. No input registration is performed. + /// Input admission must be complete/quiescent; a compaction retry must not replay + /// imports. The caller must produce a logically equivalent replacement under this snapshot. + pub async fn publish_table_compaction( + &self, + region_id: RegionId, + snapshot: TableCompactionSnapshot, + outputs: Vec, + ) -> Result { + ensure!( + snapshot.metadata.region_id == region_id && !snapshot.inputs.is_empty(), + InvalidRequestSnafu { + region_id, + reason: "Publication requires this region's nonempty input snapshot" + } + ); + // The public snapshot keeps JSON options; parse them once here so the + // worker compares typed options and reads typed flags without reparsing. + let options = serde_json::from_value(snapshot.options).context(SerdeJsonSnafu)?; + let edit = RegionEdit { + files_to_remove: snapshot.inputs, + files_to_add: outputs, + timestamp_ms: Some(chrono::Utc::now().timestamp_millis()), + ..Default::default() + }; + let (tx, rx) = oneshot::channel(); + let mut request = RegionEditRequest::new(region_id, edit, false, tx); + request.kind = RegionEditKind::Compaction(Box::new(PublicationContext { + metadata: snapshot.metadata, + options, + })); + self.inner + .workers + .submit_to_worker(region_id, WorkerRequest::EditRegion(request)) + .await?; + match rx.await.context(RecvSnafu)? { + Ok(()) => Ok(PublicationOutcome::Applied), + Err(Error::EditRegion { source, .. }) + if matches!(source.as_ref(), Error::StaleCompactionExecution { .. }) => + { + Ok(PublicationOutcome::StaleInputs) + } + Err(error) => Err(error), + } + } +} diff --git a/src/mito2/src/engine/table_compaction/output.rs b/src/mito2/src/engine/table_compaction/output.rs new file mode 100644 index 00000000000..8f85123439b --- /dev/null +++ b/src/mito2/src/engine/table_compaction/output.rs @@ -0,0 +1,207 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +use std::sync::Arc; + +use parquet::file::metadata::ParquetMetaData; +use parquet::file::statistics::Statistics; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metadata::RegionMetadataRef; + +use crate::access_layer::AccessLayer; +use crate::cache::CachedSstMeta; +use crate::engine::region_hook::VerifiedSstFileInfo; +use crate::error::{InvalidRequestSnafu, OpenDalSnafu, Result}; +use crate::region::MitoRegionRef; +use crate::sst::file::FileMeta; +use crate::sst::location; +use crate::sst::parquet::metadata::{MetadataLoader, extract_primary_key_range}; + +/// Verifies every output, then notifies the region hook once if one is registered. +/// +/// Without a hook no footer is needed after its own verification, so each one is +/// dropped before the next output is read. With a hook every footer is kept until +/// all outputs verify, because the hook observes them in one call; footers are then +/// released when this returns, before the caller takes the manifest lock. A failed +/// verification notifies nothing. +pub(crate) async fn verify_and_notify( + region: &MitoRegionRef, + metadata: &RegionMetadataRef, + outputs: &[FileMeta], +) -> Result<()> { + let Some(hook) = region.manifest_ctx.hook() else { + for file in outputs { + verify_file(®ion.access_layer, metadata, file).await?; + } + return Ok(()); + }; + + let mut footers = Vec::with_capacity(outputs.len()); + for file in outputs { + footers.push(verify_file(®ion.access_layer, metadata, file).await?); + } + let verified = outputs + .iter() + .zip(&footers) + .map(|(file_meta, footer)| VerifiedSstFileInfo { file_meta, footer }) + .collect::>(); + hook.on_sst_files_verified(region.region_id, metadata, &verified) + .await; + Ok(()) +} + +async fn verify_file( + access: &AccessLayer, + metadata: &RegionMetadataRef, + file: &FileMeta, +) -> Result> { + let region_id = metadata.region_id; + let path = location::sst_file_path(access.table_dir(), file.file_id(), access.path_type()); + let size = access + .object_store() + .stat(&path) + .await + .context(OpenDalSnafu)? + .content_length(); + ensure!( + size == file.file_size, + InvalidRequestSnafu { + region_id, + reason: "Output SST size mismatch" + } + ); + let footer = MetadataLoader::new(access.object_store().clone(), &path, size) + .load(&mut Default::default()) + .await?; + let cached = CachedSstMeta::try_new_with_region_metadata(&path, footer, None)?; + ensure!( + cached.region_metadata() == *metadata, + InvalidRequestSnafu { + region_id, + reason: "Output footer schema or partition mismatch" + } + ); + let footer = cached.parquet_metadata(); + verify_statistics(&footer, metadata, file)?; + if file.index_file_size > 0 { + let path = + location::index_file_path(access.table_dir(), file.index_id(), access.path_type()); + let index_size = access + .object_store() + .stat(&path) + .await + .context(OpenDalSnafu)? + .content_length(); + ensure!( + index_size == file.index_file_size, + InvalidRequestSnafu { + region_id, + reason: "Output index size mismatch" + } + ); + } else { + ensure!( + file.available_indexes.is_empty() && file.indexes.is_empty(), + InvalidRequestSnafu { + region_id, + reason: "Output advertises indexes without an index object", + } + ); + } + Ok(footer) +} + +fn verify_statistics( + footer: &ParquetMetaData, + metadata: &RegionMetadataRef, + file: &FileMeta, +) -> Result<()> { + let max_uncompressed = footer + .row_groups() + .iter() + .map(|group| { + group + .columns() + .iter() + .map(|column| column.uncompressed_size() as u64) + .sum::() + }) + .max() + .unwrap_or(0); + let pk_range = extract_primary_key_range(footer, metadata); + ensure!( + footer.file_metadata().num_rows() > 0 + && footer.file_metadata().num_rows() as u64 == file.num_rows + && footer.num_row_groups() as u64 == file.num_row_groups + && max_uncompressed == file.max_row_group_uncompressed_size + && pk_range.as_ref().map(|range| &range.0) == file.primary_key_min.as_ref() + && pk_range.as_ref().map(|range| &range.1) == file.primary_key_max.as_ref() + // `num_series` is the writer's caller-reported estimate derived from + // timestamp boundaries, not a footer statistic, so it only bounds the + // count against rows instead of being verified against the footer. + && file.num_series <= file.num_rows, + InvalidRequestSnafu { + region_id: metadata.region_id, + reason: "Output statistics differ from its footer" + } + ); + verify_time_range(footer, metadata, file) +} + +fn verify_time_range( + footer: &ParquetMetaData, + metadata: &RegionMetadataRef, + file: &FileMeta, +) -> Result<()> { + let region_id = metadata.region_id; + let column_name = &metadata.time_index_column().column_schema.name; + let index = footer + .file_metadata() + .schema_descr() + .columns() + .iter() + .position(|column| column.path().parts().len() == 1 && column.name() == column_name) + .context(InvalidRequestSnafu { + region_id, + reason: "Output has no timestamp column", + })?; + let mut minimum = i64::MAX; + let mut maximum = i64::MIN; + for group in footer.row_groups() { + let Some(Statistics::Int64(stats)) = group.column(index).statistics() else { + return InvalidRequestSnafu { + region_id, + reason: "Output timestamp statistics missing", + } + .fail(); + }; + minimum = minimum.min(*stats.min_opt().context(InvalidRequestSnafu { + region_id, + reason: "Output timestamp minimum missing", + })?); + maximum = maximum.max(*stats.max_opt().context(InvalidRequestSnafu { + region_id, + reason: "Output timestamp maximum missing", + })?); + } + let unit = metadata.time_index_type().unit(); + ensure!( + file.time_range.0.value() == minimum + && file.time_range.1.value() == maximum + && file.time_range.0.unit() == unit + && file.time_range.1.unit() == unit, + InvalidRequestSnafu { + region_id, + reason: "Output timestamp range mismatch" + } + ); + Ok(()) +} diff --git a/src/mito2/src/engine/table_compaction/tests.rs b/src/mito2/src/engine/table_compaction/tests.rs new file mode 100644 index 00000000000..144c586a33e --- /dev/null +++ b/src/mito2/src/engine/table_compaction/tests.rs @@ -0,0 +1,1086 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Weak}; + +use api::v1::Rows; +use bytes::Bytes; +use common_base::Plugins; +use common_recordbatch::RecordBatches; +use common_time::Timestamp; +use common_time::timestamp::TimeUnit; +use parquet::file::metadata::ParquetMetaData; +use store_api::region_engine::RegionEngine; +use store_api::region_request::{RegionRequest, RegionTruncateRequest}; +use store_api::storage::{FileId, RegionId}; + +use super::{PublicationOutcome, TableCompactionSnapshot}; +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::engine::region_hook::{RegionHook, RegionHookRef, VerifiedSstFileInfo}; +use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; +use crate::region::MitoRegionRef; +use crate::sst::file::{ColumnIndexMetadata, FileMeta, IndexType, LEVEL_SHARED, SHARED_LEVEL_SLOT}; +use crate::sst::location; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, +}; + +fn config() -> MitoConfig { + MitoConfig { + schedule_compaction_after_edit: false, + manifest_checkpoint_distance: 0, + ..Default::default() + } +} + +async fn copy_sst(region: &MitoRegionRef, source: &FileMeta, level: u8) -> FileMeta { + let mut output = source.clone(); + output.file_id = FileId::random(); + output.level = level; + let access = ®ion.access_layer; + let path = |file: &FileMeta| { + location::sst_file_path(access.table_dir(), file.file_id(), access.path_type()) + }; + access + .object_store() + .copy(&path(source), &path(&output)) + .await + .unwrap(); + if source.index_file_size > 0 { + let path = |file: &FileMeta| { + location::index_file_path(access.table_dir(), file.index_id(), access.path_type()) + }; + access + .object_store() + .copy(&path(source), &path(&output)) + .await + .unwrap(); + } + output +} + +/// Input admission is deliberately separate from the compaction API. +async fn seed(engine: &MitoEngine, region_id: RegionId) -> (TableCompactionSnapshot, FileMeta) { + let (snapshot, outputs) = seed_shared_inputs(engine, region_id, &[(0, 3)]).await; + let output = outputs.into_iter().next().unwrap(); + (snapshot, output) +} + +/// Writes each row range as a flushed SST, promotes the flushed files to shared +/// level, and returns one region-local output per shared input under the snapshot. +async fn seed_shared_inputs( + engine: &MitoEngine, + region_id: RegionId, + row_ranges: &[(usize, usize)], +) -> (TableCompactionSnapshot, Vec) { + let create = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(create.clone())) + .await + .unwrap(); + for &(start, end) in row_ranges { + put_rows( + engine, + region_id, + Rows { + schema: rows_schema(&create), + rows: build_rows(start, end), + }, + ) + .await; + flush_region(engine, region_id, None).await; + } + let region = engine.get_region(region_id).unwrap(); + let originals = region.version().ssts.levels()[0] + .files() + .map(|file| file.meta_ref().clone()) + .collect::>(); + assert_eq!(row_ranges.len(), originals.len()); + let mut shared = Vec::with_capacity(originals.len()); + for original in &originals { + shared.push(copy_sst(®ion, original, LEVEL_SHARED).await); + } + engine + .edit_region( + region_id, + RegionEdit { + files_to_add: shared, + files_to_remove: originals, + ..Default::default() + }, + ) + .await + .unwrap(); + let snapshot = engine.table_compaction_snapshot(region_id).unwrap(); + assert_eq!(row_ranges.len(), snapshot.inputs.len()); + let mut outputs = Vec::with_capacity(snapshot.inputs.len()); + for input in &snapshot.inputs { + outputs.push(copy_sst(®ion, input, 0).await); + } + (snapshot, outputs) +} + +async fn row_count(engine: &MitoEngine, region_id: RegionId) -> usize { + RecordBatches::try_collect( + engine + .scan_to_stream(region_id, Default::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .iter() + .map(|batch| batch.num_rows()) + .sum() +} + +#[tokio::test] +async fn test_duplicate_publication_and_later_replacement_do_not_readmit_inputs() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(config()).await; + let region_id = RegionId::new(1, 1); + let (snapshot, output) = seed(&engine, region_id).await; + let region = engine.get_region(region_id).unwrap(); + let sequence = region.version_control.committed_sequence(); + let (first, second) = tokio::join!( + engine.publish_table_compaction(region_id, snapshot.clone(), vec![output.clone()]), + engine.publish_table_compaction(region_id, snapshot.clone(), vec![output.clone()]) + ); + let outcomes = [first.unwrap(), second.unwrap()]; + assert_eq!( + 1, + outcomes + .iter() + .filter(|outcome| **outcome == PublicationOutcome::Applied) + .count() + ); + assert_eq!( + 1, + outcomes + .iter() + .filter(|outcome| **outcome == PublicationOutcome::StaleInputs) + .count() + ); + assert_eq!(sequence, region.version_control.committed_sequence()); + assert_eq!(sequence, region.version_control.current().last_sequence); + assert_eq!(3, row_count(&engine, region_id).await); + let manifest = region.manifest_ctx.manifest().await; + let encoded = serde_json::to_value(manifest.as_ref()).unwrap(); + assert!(encoded.get("table_compaction").is_none() && encoded.get("extensions").is_none()); + + // Publish a later ordinary compaction, then remove the old output object. + // A late retry must reject its obsolete inputs before trying to verify C. + let replacement = copy_sst(®ion, &output, 0).await; + let edit = RegionEdit { + files_to_remove: vec![output.clone()], + files_to_add: vec![replacement.clone()], + ..Default::default() + }; + region + .manifest_ctx + .update_manifest_for_compaction(RegionMetaActionList::with_action(RegionMetaAction::Edit( + edit.clone(), + ))) + .await + .unwrap(); + region + .version_control + .apply_edit(Some(edit), &[], region.file_purger.clone()); + let path = location::sst_file_path( + region.access_layer.table_dir(), + output.file_id(), + region.path_type(), + ); + region + .access_layer + .object_store() + .delete(&path) + .await + .unwrap(); + assert_eq!( + PublicationOutcome::StaleInputs, + engine + .publish_table_compaction(region_id, snapshot, vec![output.clone()]) + .await + .unwrap() + ); + let manifest = region.manifest_ctx.manifest().await; + assert_eq!(1, manifest.files.len()); + assert!(manifest.files.contains_key(&replacement.file_id)); + assert!(!manifest.files.contains_key(&output.file_id)); + reopen_region(&engine, region_id, "test".into(), true, Default::default()).await; + assert_eq!(3, row_count(&engine, region_id).await); +} + +#[tokio::test] +async fn test_partial_success_can_retry_both_regions_without_duplicate_files() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(config()).await; + let a = RegionId::new(1, 1); + let b = RegionId::new(1, 2); + let (snapshot_a, output_a) = seed(&engine, a).await; + let (snapshot_b, output_b) = seed(&engine, b).await; + let invalid = FileMeta { + file_size: output_b.file_size + 1, + ..output_b.clone() + }; + let (first_a, first_b) = tokio::join!( + engine.publish_table_compaction(a, snapshot_a.clone(), vec![output_a.clone()]), + engine.publish_table_compaction(b, snapshot_b.clone(), vec![invalid]) + ); + assert_eq!(PublicationOutcome::Applied, first_a.unwrap()); + assert!(first_b.is_err()); + assert_eq!( + snapshot_b.inputs, + engine.table_compaction_snapshot(b).unwrap().inputs + ); + let (retry_a, retry_b) = tokio::join!( + engine.publish_table_compaction(a, snapshot_a, vec![output_a.clone()]), + engine.publish_table_compaction(b, snapshot_b, vec![output_b.clone()]) + ); + assert_eq!(PublicationOutcome::StaleInputs, retry_a.unwrap()); + assert_eq!(PublicationOutcome::Applied, retry_b.unwrap()); + for (id, output) in [(a, output_a), (b, output_b)] { + let region = engine.get_region(id).unwrap(); + let manifest = region.manifest_ctx.manifest().await; + assert_eq!(1, manifest.files.len()); + assert_eq!(&output, &manifest.files[&output.file_id]); + assert_eq!(3, row_count(&engine, id).await); + } +} + +#[tokio::test] +async fn test_invalid_results_preserve_inputs_and_reopen_can_replan() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(config()).await; + let region_id = RegionId::new(1, 1); + let (snapshot, output) = seed(&engine, region_id).await; + let region = engine.get_region(region_id).unwrap(); + let version = region.manifest_ctx.manifest_version().await; + // Snapshot-level and footer-derived mutations must all fail output admission + // with the specific reason, leaving the manifest, inputs and rows untouched. + #[derive(Clone, Copy)] + enum ExpectedError { + InvalidRequest(&'static str), + OpenDal, + } + let statistics_mismatch = + ExpectedError::InvalidRequest("Output statistics differ from its footer"); + let invalid_output = ExpectedError::InvalidRequest("Invalid region-local compaction output"); + let timestamp_mismatch = ExpectedError::InvalidRequest("Output timestamp range mismatch"); + let cases = [ + ( + FileMeta { + file_size: output.file_size + 1, + ..output.clone() + }, + ExpectedError::InvalidRequest("Output SST size mismatch"), + ), + ( + FileMeta { + num_rows: output.num_rows + 1, + ..output.clone() + }, + statistics_mismatch, + ), + ( + FileMeta { + num_row_groups: output.num_row_groups + 1, + ..output.clone() + }, + statistics_mismatch, + ), + ( + FileMeta { + max_row_group_uncompressed_size: output.max_row_group_uncompressed_size + 1, + ..output.clone() + }, + statistics_mismatch, + ), + ( + FileMeta { + primary_key_min: Some(Bytes::from_static(b"not-the-footer-minimum")), + ..output.clone() + }, + statistics_mismatch, + ), + ( + FileMeta { + primary_key_max: Some(Bytes::from_static(b"not-the-footer-maximum")), + ..output.clone() + }, + statistics_mismatch, + ), + ( + FileMeta { + time_range: ( + Timestamp::new(output.time_range.0.value() + 1, output.time_range.0.unit()), + output.time_range.1, + ), + ..output.clone() + }, + timestamp_mismatch, + ), + ( + FileMeta { + time_range: ( + output.time_range.0, + Timestamp::new(output.time_range.1.value(), TimeUnit::Second), + ), + ..output.clone() + }, + timestamp_mismatch, + ), + ( + FileMeta { + file_id: FileId::random(), + ..output.clone() + }, + ExpectedError::OpenDal, + ), + ( + FileMeta { + sequence: None, + ..output.clone() + }, + invalid_output, + ), + ( + FileMeta { + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..output.clone() + }, + ExpectedError::InvalidRequest("Output advertises indexes without an index object"), + ), + ( + FileMeta { + indexes: vec![ColumnIndexMetadata { + column_id: 0, + created_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..Default::default() + }], + ..output.clone() + }, + ExpectedError::InvalidRequest("Output advertises indexes without an index object"), + ), + ( + FileMeta { + file_id: snapshot.inputs[0].file_id, + ..output.clone() + }, + invalid_output, + ), + ]; + for (invalid, expected) in cases { + let error = engine + .publish_table_compaction(region_id, snapshot.clone(), vec![invalid]) + .await + .unwrap_err(); + let source = match error { + crate::error::Error::EditRegion { source, .. } => source, + other => panic!("unexpected error: {other}"), + }; + match expected { + ExpectedError::InvalidRequest(reason) => assert!( + matches!(source.as_ref(), crate::error::Error::InvalidRequest { reason: actual, .. } if actual == reason), + "unexpected error: {source}" + ), + ExpectedError::OpenDal => assert!( + matches!(source.as_ref(), crate::error::Error::OpenDal { .. }), + "unexpected error: {source}" + ), + } + assert_eq!(version, region.manifest_ctx.manifest_version().await); + assert_eq!( + snapshot.inputs, + engine.table_compaction_snapshot(region_id).unwrap().inputs + ); + assert_eq!(3, row_count(&engine, region_id).await); + } + // A real footer written for another region's metadata must reach the footer + // check; the stale-input shortcut must not mask it. + let foreign_region_id = RegionId::new(1, 2); + let (_, foreign_output) = seed(&engine, foreign_region_id).await; + let foreign_region = engine.get_region(foreign_region_id).unwrap(); + let mut foreign = foreign_output; + assert_eq!(0, foreign.index_file_size); + let foreign_path = location::sst_file_path( + foreign_region.access_layer.table_dir(), + foreign.file_id(), + foreign_region.access_layer.path_type(), + ); + foreign.file_id = FileId::random(); + foreign.region_id = region_id; + foreign.sequence = output.sequence; + let target_path = location::sst_file_path( + region.access_layer.table_dir(), + foreign.file_id(), + region.access_layer.path_type(), + ); + foreign_region + .access_layer + .object_store() + .copy(&foreign_path, &target_path) + .await + .unwrap(); + let error = engine + .publish_table_compaction(region_id, snapshot.clone(), vec![foreign]) + .await + .unwrap_err(); + let source = match error { + crate::error::Error::EditRegion { source, .. } => source, + other => panic!("unexpected error: {other}"), + }; + assert!( + matches!(source.as_ref(), crate::error::Error::InvalidRequest { reason, .. } + if reason == "Output footer schema or partition mismatch"), + "unexpected error: {source}" + ); + assert_eq!(version, region.manifest_ctx.manifest_version().await); + assert_eq!( + snapshot.inputs, + engine.table_compaction_snapshot(region_id).unwrap().inputs + ); + assert_eq!(3, row_count(&engine, region_id).await); + let mut stale_input = snapshot.clone(); + stale_input.inputs[0].index_version += 1; + assert_eq!( + PublicationOutcome::StaleInputs, + engine + .publish_table_compaction(region_id, stale_input, vec![output.clone()]) + .await + .unwrap() + ); + let mut duplicate = snapshot.clone(); + duplicate.inputs.push(snapshot.inputs[0].clone()); + assert!( + engine + .publish_table_compaction(region_id, duplicate, vec![output.clone()]) + .await + .is_err() + ); + let mut malformed = snapshot.clone(); + malformed.options = serde_json::json!({ "append_mode": "not-a-bool" }); + let error = engine + .publish_table_compaction(region_id, malformed, vec![output.clone()]) + .await + .unwrap_err(); + assert!(matches!(error, crate::error::Error::SerdeJson { .. })); + assert_eq!(version, region.manifest_ctx.manifest_version().await); + assert_eq!( + snapshot.inputs, + engine.table_compaction_snapshot(region_id).unwrap().inputs + ); + let mut stale = snapshot; + stale.options["append_mode"] = serde_json::json!(true); + assert_eq!( + PublicationOutcome::StaleInputs, + engine + .publish_table_compaction(region_id, stale, vec![output.clone()]) + .await + .unwrap() + ); + assert_eq!(version, region.manifest_ctx.manifest_version().await); + reopen_region(&engine, region_id, "test".into(), true, Default::default()).await; + // V1 refreshes the snapshot after restart instead of resuming a Prepared task. + // The public snapshot contract is JSON, so a serde round trip must still publish. + let fresh = engine.table_compaction_snapshot(region_id).unwrap(); + let fresh: TableCompactionSnapshot = + serde_json::from_value(serde_json::to_value(&fresh).unwrap()).unwrap(); + assert_eq!( + PublicationOutcome::Applied, + engine + .publish_table_compaction(region_id, fresh, vec![output]) + .await + .unwrap() + ); + assert_eq!(3, row_count(&engine, region_id).await); +} + +/// An output carrying a real index object written by flush is accepted. +#[tokio::test] +async fn test_indexed_output_with_real_index_object_publishes() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(config()).await; + let region_id = RegionId::new(1, 1); + let create = CreateRequestBuilder::new().build_with_index(); + engine + .handle_request(region_id, RegionRequest::Create(create.clone())) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&create), + rows: build_rows(0, 3), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let region = engine.get_region(region_id).unwrap(); + let original = region.version().ssts.levels()[0] + .files() + .next() + .unwrap() + .meta_ref() + .clone(); + assert!(original.index_file_size > 0); + assert!(!original.available_indexes.is_empty()); + assert!(!original.indexes.is_empty()); + let shared = copy_sst(®ion, &original, LEVEL_SHARED).await; + engine + .edit_region( + region_id, + RegionEdit { + files_to_add: vec![shared], + files_to_remove: vec![original], + ..Default::default() + }, + ) + .await + .unwrap(); + let snapshot = engine.table_compaction_snapshot(region_id).unwrap(); + let output = copy_sst(®ion, &snapshot.inputs[0], 0).await; + assert!(output.index_file_size > 0); + assert!(!output.available_indexes.is_empty()); + assert!(!output.indexes.is_empty()); + assert_eq!( + PublicationOutcome::Applied, + engine + .publish_table_compaction(region_id, snapshot, vec![output.clone()]) + .await + .unwrap() + ); + let manifest = region.manifest_ctx.manifest().await; + assert_eq!(Some(&output), manifest.files.get(&output.file_id)); + assert_eq!(3, row_count(&engine, region_id).await); +} + +/// A preserve-row-sequence region trusts output markers only when the typed +/// snapshot options and the admitted inputs both carry the marker. +#[tokio::test] +async fn test_preserve_row_sequence_option_drives_output_admission() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(config()).await; + let region_id = RegionId::new(1, 1); + let create = CreateRequestBuilder::new() + .insert_option("append_mode", "true") + .insert_option("preserve_row_sequence", "true") + .build(); + engine + .handle_request(region_id, RegionRequest::Create(create.clone())) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&create), + rows: build_rows(0, 3), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let region = engine.get_region(region_id).unwrap(); + let original = region.version().ssts.levels()[0] + .files() + .next() + .unwrap() + .meta_ref() + .clone(); + assert!(original.preserve_row_sequence); + // A generic import clears the trust marker, so install the admitted shared + // input directly: this is the shape table compaction receives after admission. + let shared = copy_sst(®ion, &original, LEVEL_SHARED).await; + let edit = RegionEdit { + files_to_add: vec![shared], + files_to_remove: vec![original], + ..Default::default() + }; + region + .manifest_ctx + .update_manifest_for_compaction(RegionMetaActionList::with_action(RegionMetaAction::Edit( + edit.clone(), + ))) + .await + .unwrap(); + region + .version_control + .apply_edit(Some(edit), &[], region.file_purger.clone()); + let snapshot = engine.table_compaction_snapshot(region_id).unwrap(); + assert_eq!( + Some(true), + snapshot.options["preserve_row_sequence"].as_bool() + ); + let output = copy_sst(®ion, &snapshot.inputs[0], 0).await; + assert!(output.preserve_row_sequence); + let untrusted = FileMeta { + preserve_row_sequence: false, + ..output.clone() + }; + let version = region.manifest_ctx.manifest_version().await; + let error = engine + .publish_table_compaction(region_id, snapshot.clone(), vec![untrusted]) + .await + .unwrap_err(); + assert!(matches!( + error, + crate::error::Error::EditRegion { source, .. } + if matches!(source.as_ref(), crate::error::Error::InvalidRequest { reason, .. } + if reason == "Invalid region-local compaction output") + )); + assert_eq!(version, region.manifest_ctx.manifest_version().await); + assert_eq!( + snapshot.inputs, + engine.table_compaction_snapshot(region_id).unwrap().inputs + ); + assert_eq!( + PublicationOutcome::Applied, + engine + .publish_table_compaction(region_id, snapshot, vec![output.clone()]) + .await + .unwrap() + ); + let manifest = region.manifest_ctx.manifest().await; + assert_eq!(Some(&output), manifest.files.get(&output.file_id)); + assert_eq!(3, row_count(&engine, region_id).await); +} + +#[tokio::test] +async fn test_empty_target_replacement_preserves_other_shared_reference() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(config()).await; + let source = RegionId::new(1, 1); + let target = RegionId::new(1, 2); + let (source_snapshot, _) = seed(&engine, source).await; + let expr = + partition::expr::col("tag_0").eq(datatypes::value::Value::String("not-in-source".into())); + let create = CreateRequestBuilder::new() + .partition_expr_json(Some(expr.as_json_str().unwrap())) + .build(); + engine + .handle_request(target, RegionRequest::Create(create)) + .await + .unwrap(); + let copy = store_api::region_engine::MitoCopyRegionFromRequest { + source_region_id: source, + parallelism: 1, + }; + engine.copy_region_from(target, copy.clone()).await.unwrap(); + engine.copy_region_from(target, copy).await.unwrap(); + let snapshot = engine.table_compaction_snapshot(target).unwrap(); + assert_eq!(1, snapshot.inputs.len()); + assert_eq!( + source_snapshot.inputs[0].file_id(), + snapshot.inputs[0].file_id() + ); + let region = engine.get_region(target).unwrap(); + assert_eq!(1, region.version_control.current().last_sequence); + assert_eq!(0, row_count(&engine, target).await); + assert_eq!( + PublicationOutcome::Applied, + engine + .publish_table_compaction(target, snapshot.clone(), vec![]) + .await + .unwrap() + ); + assert_eq!( + PublicationOutcome::StaleInputs, + engine + .publish_table_compaction(target, snapshot, vec![]) + .await + .unwrap() + ); + assert_eq!(1, region.version_control.current().last_sequence); + assert_eq!(0, row_count(&engine, target).await); + assert_eq!(3, row_count(&engine, source).await); +} + +#[tokio::test] +async fn test_staging_rejects_publication_without_changing_files() { + use store_api::region_request::{EnterStagingRequest, StagingPartitionDirective}; + let mut env = TestEnv::new().await; + let engine = env.create_engine(config()).await; + let region_id = RegionId::new(1, 1); + let (snapshot, output) = seed(&engine, region_id).await; + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let version = region.manifest_ctx.manifest_version().await; + assert!( + engine + .publish_table_compaction(region_id, snapshot, vec![output]) + .await + .is_err() + ); + assert!(region.is_staging()); + assert_eq!(version, region.manifest_ctx.manifest_version().await); + assert_eq!( + 1, + region.version().ssts.levels()[SHARED_LEVEL_SLOT] + .files + .len() + ); +} + +#[rstest::rstest] +#[case::before_persist(false)] +#[case::lost_ack(true)] +#[tokio::test] +async fn test_uncertain_write_recovers_files_before_retry(#[case] persisted: bool) { + use object_store::layers::mock::{ + Buffer, Error, ErrorKind, Metadata, MockLayerBuilder, Result, oio, + }; + struct FailManifestWrite { + inner: oio::Writer, + pending: Vec, + armed: Arc, + persisted: bool, + } + impl oio::Write for FailManifestWrite { + async fn write(&mut self, bytes: Buffer) -> Result<()> { + // Buffer at the I/O boundary so the pre-persist case never depends + // on whether the backing filesystem supports aborting a writer. + self.pending.push(bytes); + Ok(()) + } + async fn close(&mut self) -> Result { + let fail = self.armed.swap(false, Ordering::SeqCst); + if fail && !self.persisted { + return Err(Error::new( + ErrorKind::Unexpected, + "injected failure before persistence", + )); + } + for bytes in self.pending.drain(..) { + self.inner.write(bytes).await?; + } + let metadata = self.inner.close().await?; + if fail { + Err(Error::new( + ErrorKind::Unexpected, + "lost manifest acknowledgement", + )) + } else { + Ok(metadata) + } + } + async fn abort(&mut self) -> Result<()> { + self.inner.abort().await + } + } + let armed = Arc::new(AtomicBool::new(false)); + let fault = armed.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + if path.contains("/manifest/") { + Box::new(FailManifestWrite { + inner, + pending: Vec::new(), + armed: fault.clone(), + persisted, + }) as oio::Writer + } else { + inner + } + })) + .build() + .unwrap(); + let mut env = TestEnv::new().await.with_mock_layer(layer); + let engine = env.create_engine(config()).await; + let region_id = RegionId::new(1, 1); + let (snapshot, output) = seed(&engine, region_id).await; + let access = engine.get_region(region_id).unwrap().access_layer.clone(); + armed.store(true, Ordering::SeqCst); + assert!( + engine + .publish_table_compaction(region_id, snapshot.clone(), vec![output.clone()]) + .await + .is_err() + ); + assert!(!armed.load(Ordering::SeqCst)); + assert!(engine.get_region(region_id).is_none()); + let path = location::sst_file_path(access.table_dir(), output.file_id(), access.path_type()); + assert!(access.object_store().exists(&path).await.unwrap()); + reopen_region(&engine, region_id, "test".into(), true, Default::default()).await; + let fresh = engine.table_compaction_snapshot(region_id).unwrap(); + assert_eq!(!persisted, !fresh.inputs.is_empty()); + if persisted { + assert_eq!( + PublicationOutcome::StaleInputs, + engine + .publish_table_compaction(region_id, snapshot, vec![output]) + .await + .unwrap() + ); + } else { + assert_eq!( + PublicationOutcome::Applied, + engine + .publish_table_compaction(region_id, fresh, vec![output]) + .await + .unwrap() + ); + } + assert_eq!(3, row_count(&engine, region_id).await); +} + +#[derive(Debug)] +struct PublicationBarrier { + after_commit: bool, + armed: AtomicBool, + entered: tokio::sync::Notify, + resume: tokio::sync::Notify, +} + +impl PublicationBarrier { + async fn pause(&self) { + if self.armed.swap(false, Ordering::SeqCst) { + self.entered.notify_one(); + self.resume.notified().await; + } + } +} + +#[async_trait::async_trait] +impl RegionHook for PublicationBarrier { + async fn on_sst_files_verified( + &self, + _: RegionId, + _: &store_api::metadata::RegionMetadataRef, + _: &[VerifiedSstFileInfo<'_>], + ) { + if !self.after_commit { + self.pause().await; + } + } + async fn on_manifest_updated( + &self, + _: RegionId, + actions: &RegionMetaActionList, + _: store_api::ManifestVersion, + ) { + if self.after_commit && actions.actions.iter().any(|action| matches!(action, + RegionMetaAction::Edit(edit) if edit.files_to_remove.iter().any(|file| file.level == LEVEL_SHARED))) { + self.pause().await; + } + } +} + +#[rstest::rstest] +#[case::verification(false)] +#[case::completion(true)] +#[tokio::test] +async fn test_close_reopen_fences_old_publication(#[case] after_commit: bool) { + let hook = Arc::new(PublicationBarrier { + after_commit, + armed: AtomicBool::new(false), + entered: Default::default(), + resume: Default::default(), + }); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + let mut env = TestEnv::new().await; + let engine = env.create_engine_with_plugins(config(), plugins).await; + let region_id = RegionId::new(1, 1); + let (snapshot, output) = seed(&engine, region_id).await; + hook.armed.store(true, Ordering::SeqCst); + let (late, ()) = tokio::time::timeout(std::time::Duration::from_secs(10), async { + tokio::join!( + engine.publish_table_compaction(region_id, snapshot.clone(), vec![output.clone()]), + async { + hook.entered.notified().await; + reopen_region(&engine, region_id, "test".into(), true, Default::default()).await; + engine + .handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::All), + ) + .await + .unwrap(); + hook.resume.notify_one(); + } + ) + }) + .await + .unwrap(); + assert!(late.is_err()); + let region = engine.get_region(region_id).unwrap(); + assert!(region.manifest_ctx.manifest().await.files.is_empty()); + assert!( + region + .version() + .ssts + .levels() + .iter() + .all(|level| level.files.is_empty()) + ); + assert_eq!( + PublicationOutcome::StaleInputs, + engine + .publish_table_compaction(region_id, snapshot, vec![output]) + .await + .unwrap() + ); + assert_eq!(0, row_count(&engine, region_id).await); +} + +/// Records one notification per verification, keeping verified footers only as +/// weak references, and checks those references again when the publication's +/// manifest update is announced. `on_manifest_updated` runs after the commit, so +/// this catches footers retained across it; it does not claim the hook fires +/// before the commit. Nothing else in the test keeps the footers alive. +#[derive(Debug, Default)] +struct VerificationRecorder { + verified_files: AtomicUsize, + notifications: AtomicUsize, + manifest_updates: AtomicUsize, + released_by_commit: AtomicBool, + footers: std::sync::Mutex>>, +} + +#[async_trait::async_trait] +impl RegionHook for VerificationRecorder { + async fn on_sst_files_verified( + &self, + _: RegionId, + _: &store_api::metadata::RegionMetadataRef, + files: &[VerifiedSstFileInfo<'_>], + ) { + let mut footers = self.footers.lock().unwrap(); + footers.extend(files.iter().map(|file| Arc::downgrade(file.footer))); + self.verified_files.fetch_add(files.len(), Ordering::SeqCst); + self.notifications.fetch_add(1, Ordering::SeqCst); + } + + async fn on_manifest_updated( + &self, + _: RegionId, + actions: &RegionMetaActionList, + _: store_api::ManifestVersion, + ) { + let is_compaction = actions.actions.iter().any(|action| matches!(action, + RegionMetaAction::Edit(edit) if edit.files_to_remove.iter().any(|file| file.level == LEVEL_SHARED))); + if !is_compaction { + return; + } + let footers = self.footers.lock().unwrap(); + self.released_by_commit.store( + !footers.is_empty() && footers.iter().all(|footer| footer.upgrade().is_none()), + Ordering::SeqCst, + ); + self.manifest_updates.fetch_add(1, Ordering::SeqCst); + } +} + +#[rstest::rstest] +#[case::without_hook(false)] +#[case::with_hook(true)] +#[tokio::test] +async fn test_multiple_outputs_verify_notify_once_and_release_footers(#[case] with_hook: bool) { + let recorder = Arc::new(VerificationRecorder::default()); + let mut env = TestEnv::new().await; + let engine = if with_hook { + let plugins = Plugins::new(); + plugins.insert(recorder.clone() as RegionHookRef); + env.create_engine_with_plugins(config(), plugins).await + } else { + env.create_engine(config()).await + }; + let region_id = RegionId::new(1, 1); + let (snapshot, outputs) = seed_shared_inputs(&engine, region_id, &[(0, 3), (3, 6)]).await; + assert_eq!( + PublicationOutcome::Applied, + engine + .publish_table_compaction(region_id, snapshot, outputs.clone()) + .await + .unwrap() + ); + let expected_notifications = usize::from(with_hook); + assert_eq!( + expected_notifications * outputs.len(), + recorder.verified_files.load(Ordering::SeqCst) + ); + assert_eq!( + expected_notifications, + recorder.notifications.load(Ordering::SeqCst) + ); + if with_hook { + assert_eq!(1, recorder.manifest_updates.load(Ordering::SeqCst)); + assert!(recorder.released_by_commit.load(Ordering::SeqCst)); + assert!( + recorder + .footers + .lock() + .unwrap() + .iter() + .all(|footer| footer.upgrade().is_none()) + ); + } else { + assert_eq!(0, recorder.manifest_updates.load(Ordering::SeqCst)); + } + let region = engine.get_region(region_id).unwrap(); + let manifest = region.manifest_ctx.manifest().await; + assert_eq!(outputs.len(), manifest.files.len()); + for output in &outputs { + assert_eq!(Some(output), manifest.files.get(&output.file_id)); + } + assert_eq!(6, row_count(&engine, region_id).await); +} + +#[tokio::test] +async fn test_failed_verification_does_not_notify_hook() { + let recorder = Arc::new(VerificationRecorder::default()); + let plugins = Plugins::new(); + plugins.insert(recorder.clone() as RegionHookRef); + let mut env = TestEnv::new().await; + let engine = env.create_engine_with_plugins(config(), plugins).await; + let region_id = RegionId::new(1, 1); + let (snapshot, outputs) = seed_shared_inputs(&engine, region_id, &[(0, 3), (3, 6)]).await; + let region = engine.get_region(region_id).unwrap(); + let version = region.manifest_ctx.manifest_version().await; + let invalid = FileMeta { + file_size: outputs[1].file_size + 1, + ..outputs[1].clone() + }; + let error = engine + .publish_table_compaction( + region_id, + snapshot.clone(), + vec![outputs[0].clone(), invalid], + ) + .await + .unwrap_err(); + assert!(matches!( + error, + crate::error::Error::EditRegion { source, .. } + if matches!(source.as_ref(), crate::error::Error::InvalidRequest { reason, .. } + if reason == "Output SST size mismatch") + )); + assert_eq!(0, recorder.verified_files.load(Ordering::SeqCst)); + assert_eq!(0, recorder.notifications.load(Ordering::SeqCst)); + assert_eq!(version, region.manifest_ctx.manifest_version().await); + assert_eq!( + snapshot.inputs, + engine.table_compaction_snapshot(region_id).unwrap().inputs + ); + assert_eq!(6, row_count(&engine, region_id).await); +} diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 123133d4be0..d6777131b4f 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -1373,10 +1373,19 @@ impl Waiters { } } +/// Transient execution semantics; this is never serialized into the manifest. +#[derive(Debug)] +pub(crate) enum RegionEditKind { + Import, + #[cfg(feature = "enterprise")] + Compaction(Box), +} + /// Request to edit a region directly. #[derive(Debug)] pub(crate) struct RegionEditRequest { pub(crate) region_id: RegionId, + pub(crate) kind: RegionEditKind, pub(crate) edit: RegionEdit, /// Whether to preload SST files into the write cache. pub(crate) preload_sst_cache: bool, @@ -1398,6 +1407,7 @@ impl RegionEditRequest { region_id, edit, preload_sst_cache, + kind: RegionEditKind::Import, waiters: Waiters::one(waiter), staging_partition_expr: None, } diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 79398a6d219..d8f5b0b3620 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -29,6 +29,8 @@ mod handle_manifest; mod handle_open; mod handle_rebuild_index; mod handle_remap; +#[cfg(feature = "enterprise")] +mod handle_table_compaction; mod handle_truncate; mod handle_write; diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index f1a575d4433..2a37f898f3c 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -45,8 +45,8 @@ use crate::region::version::VersionControlRef; use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; use crate::request::{ BackgroundNotify, BuildIndexRequest, DiscardUnflushedResult, OptionOutputTx, - RegionChangeResult, RegionEditRequest, RegionEditResult, RegionSyncRequest, TruncateResult, - WorkerRequest, WorkerRequestWithTime, + RegionChangeResult, RegionEditKind, RegionEditRequest, RegionEditResult, RegionSyncRequest, + TruncateResult, WorkerRequest, WorkerRequestWithTime, }; use crate::sst::index::IndexBuildType; use crate::sst::{file, location}; @@ -99,7 +99,11 @@ impl RegionEditQueue { } fn dequeue(&mut self) -> Option { - fn can_merge(edit: &RegionEdit) -> bool { + fn can_merge(request: &RegionEditRequest) -> bool { + if !matches!(request.kind, RegionEditKind::Import) { + return false; + } + let edit = &request.edit; // Only the `RegionEdit`: // 1. contains the "raw" (file without a sequence) files to add, // 2. and no `committed_sequence`, @@ -120,14 +124,11 @@ impl RegionEditQueue { } let mut merged = self.requests.pop_front()?; - if !can_merge(&merged.edit) { + if !can_merge(&merged) { return Some(merged); } - while let Some(request) = self - .requests - .pop_front_if(|request| can_merge(&request.edit)) - { + while let Some(request) = self.requests.pop_front_if(|request| can_merge(request)) { merged.edit.files_to_add.extend(request.edit.files_to_add); merged.waiters.merge(request.waiters); } @@ -365,12 +366,20 @@ impl RegionWorkerLoop { let RegionEditRequest { region_id: _, + kind, mut edit, waiters, preload_sst_cache, staging_partition_expr: _, } = request; - if let Err(error) = prepare_region_edit(®ion, &mut edit) { + #[cfg(feature = "enterprise")] + if matches!(kind, RegionEditKind::Compaction(_)) && region.is_staging() { + waiters.reply_with(|| RegionBusySnafu { region_id }.fail()); + return RegionEditDispatch::Rejected; + } + if matches!(kind, RegionEditKind::Import) + && let Err(error) = prepare_region_edit(®ion, &mut edit) + { let error = Arc::new(error); waiters.reply_with(|| Err(error.clone()).context(EditRegionSnafu { region_id })); return RegionEditDispatch::Rejected; @@ -399,15 +408,24 @@ impl RegionWorkerLoop { // Now the region is in editing state. // Updates manifest in background. common_runtime::spawn_global(async move { - let result = edit_region( - ®ion, - edit.clone(), - cache_manager, - listener, - is_staging, - preload_sst_cache, - ) - .await + let result = match kind { + RegionEditKind::Import => { + edit_region( + ®ion, + edit.clone(), + cache_manager, + listener, + is_staging, + preload_sst_cache, + ) + .await + } + #[cfg(feature = "enterprise")] + RegionEditKind::Compaction(context) => { + crate::worker::handle_table_compaction::publish(®ion, &context, edit.clone()) + .await + } + } .map_err(Arc::new); let notify = WorkerRequest::Background { region_id, diff --git a/src/mito2/src/worker/handle_table_compaction.rs b/src/mito2/src/worker/handle_table_compaction.rs new file mode 100644 index 00000000000..eddc81636e4 --- /dev/null +++ b/src/mito2/src/worker/handle_table_compaction.rs @@ -0,0 +1,151 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +//! At-least-once requests, conditional file publication, no durable task state. + +use std::collections::HashSet; + +use snafu::ensure; + +use crate::engine::table_compaction::{PublicationContext, output}; +use crate::error::{InvalidRequestSnafu, RegionBusySnafu, Result, StaleCompactionExecutionSnafu}; +use crate::manifest::action::{RegionEdit, RegionManifest, RegionMetaAction, RegionMetaActionList}; +use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; +use crate::sst::file::LEVEL_SHARED; + +pub(crate) async fn publish( + region: &MitoRegionRef, + context: &PublicationContext, + edit: RegionEdit, +) -> Result<()> { + let snapshot = region.manifest_ctx.manifest().await; + // An old result can name outputs already compacted away. Reject stale inputs + // before inspecting those objects; absence of an output is not a task receipt. + ensure!( + is_current(region, context, &edit, &snapshot), + StaleCompactionExecutionSnafu { + region_id: region.region_id + } + ); + validate_files(region, context, &edit, &snapshot)?; + output::verify_and_notify(region, &context.metadata, &edit.files_to_add).await?; + + let mut manager = region.manifest_ctx.manifest_manager.write().await; + ensure!( + matches!( + region.state(), + RegionRoleState::Leader(RegionLeaderState::Editing | RegionLeaderState::Downgrading) + ) && !manager.is_stopped(), + RegionBusySnafu { + region_id: region.region_id + } + ); + let manifest = manager.manifest(); + ensure!( + is_current(region, context, &edit, &manifest), + StaleCompactionExecutionSnafu { + region_id: region.region_id + } + ); + validate_files(region, context, &edit, &manifest)?; + let pending = match region + .manifest_ctx + .update_locked( + &mut manager, + RegionMetaActionList::with_action(RegionMetaAction::Edit(edit)), + false, + ) + .await + { + Ok(pending) => pending, + Err(error) => { + if error.may_have_persisted_manifest_update() { + // Do not accept another publication against a potentially stale + // in-memory file set. Completion removes this incarnation. + manager.stop().await; + } + return Err(error); + } + }; + drop(manager); + pending.fire().await; + Ok(()) +} + +fn is_current( + region: &MitoRegionRef, + context: &PublicationContext, + edit: &RegionEdit, + manifest: &RegionManifest, +) -> bool { + let current = region.version_control.current(); + context.metadata == manifest.metadata + && context.options == current.version.options + && !edit.files_to_remove.is_empty() + && edit + .files_to_remove + .iter() + .all(|file| manifest.files.get(&file.file_id) == Some(file)) +} + +fn validate_files( + region: &MitoRegionRef, + context: &PublicationContext, + edit: &RegionEdit, + manifest: &RegionManifest, +) -> Result<()> { + let region_id = region.region_id; + let mut inputs = HashSet::new(); + ensure!( + edit.files_to_remove + .iter() + .all(|file| file.level == LEVEL_SHARED + && file.region_id.table_id() == region_id.table_id() + && inputs.insert(file.file_id())), + InvalidRequestSnafu { + region_id, + reason: "Expected distinct, already-admitted shared inputs" + } + ); + let trusted = context.options.preserve_row_sequence + && edit.files_to_remove.iter().all(|file| { + if file.region_id == region_id { + file.preserve_row_sequence + } else { + file.sequence.is_some() + } + }); + let partition = + crate::region::parse_partition_expr(context.metadata.partition_expr.as_deref())?; + let sequence = edit + .files_to_remove + .iter() + .map(|file| file.sequence) + .collect::>>() + .and_then(|sequences| sequences.into_iter().max()); + let mut outputs = HashSet::new(); + ensure!( + edit.files_to_add + .iter() + .all(|file| file.region_id == region_id + && file.level == 0 + && outputs.insert(file.file_id) + && !manifest.files.contains_key(&file.file_id) + && file.partition_expr == partition + && file.preserve_row_sequence == trusted + && file.sequence == sequence), + InvalidRequestSnafu { + region_id, + reason: "Invalid region-local compaction output" + } + ); + Ok(()) +} diff --git a/tests/compatibility/manifest/table_compaction/README.md b/tests/compatibility/manifest/table_compaction/README.md index 7244199462e..4eb09a451d3 100644 --- a/tests/compatibility/manifest/table_compaction/README.md +++ b/tests/compatibility/manifest/table_compaction/README.md @@ -7,6 +7,18 @@ released binary. They preserve L0/L1 metadata, sequences and flush boundaries. V1 publishes compaction results using ordinary file edits. It adds no task IDs, phases, receipts, request epochs or extension-state fields to the manifest. Task retries may be submitted more than once; publication checks current inputs. +The Enterprise `publish_table_compaction` API consumes a query-visible input +snapshot and returns `Applied` or `StaleInputs`. A stale attempt must re-snapshot +and replan, never re-import its old inputs. Even stale/error outcomes must not +trigger blind output deletion: an earlier attempt may already have published +those objects. Callers retain output protection until recovery or GC resolves +liveness. + +Publication tests (including partial success, lost ACKs and reopened regions): + +```sh +cargo nextest run -p mito2 --features enterprise -E 'test(engine::table_compaction)' +``` ```sh cargo nextest run -p mito2 -E 'test(manifest::tests::compatibility)'