diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index c91809e2531..e743ef82baf 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -536,6 +536,7 @@ impl CompactionScheduler { file_purger: None, ttl: Some(ttl), max_parallelism, + plugins: self.plugins.clone(), }; let picker_output = { @@ -1884,6 +1885,7 @@ mod tests { Arc::new(ManifestContext::new( manager, RegionRoleState::Leader(RegionLeaderState::Staging), + None, )) }; diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index d78b2cc8fe6..b712b229aee 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -16,6 +16,7 @@ use std::num::NonZero; use std::sync::Arc; use std::time::Duration; +use common_base::Plugins; use common_base::cancellation::{CancellableFuture, CancellationHandle}; use common_meta::key::SchemaMetadataManagerRef; use common_telemetry::{debug, info, warn}; @@ -26,6 +27,7 @@ use object_store::manager::ObjectStoreManagerRef; use partition::expr::PartitionExpr; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt}; +use store_api::ManifestVersion; use store_api::metadata::RegionMetadataRef; use store_api::region_request::PathType; use store_api::storage::RegionId; @@ -37,6 +39,7 @@ use crate::cache::{CacheManager, CacheManagerRef}; use crate::compaction::picker::PickerOutput; use crate::compaction::{CompactionOutput, CompactionSstReaderBuilder, find_dynamic_options}; use crate::config::MitoConfig; +use crate::engine::region_hook::RegionHookRef; use crate::error; use crate::error::{ EmptyRegionDirSnafu, InvalidPartitionExprSnafu, ObjectStoreNotFoundSnafu, Result, @@ -53,8 +56,8 @@ use crate::sst::file_purger::LocalFilePurger; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::PuffinManagerFactory; use crate::sst::location::region_dir_from_table_dir; -use crate::sst::parquet::WriteOptions; use crate::sst::parquet::metadata::extract_primary_key_range; +use crate::sst::parquet::{SstInfo, WriteOptions}; use crate::sst::version::{SstVersion, SstVersionRef}; /// Region version for compaction that does not hold memtables. @@ -106,16 +109,20 @@ pub struct CompactionRegion { /// The parallel is inside this compaction task, not across different compaction tasks. /// It can be different windows of the same compaction task or something like this. pub max_parallelism: usize, + + pub(crate) plugins: Plugins, } /// OpenCompactionRegionRequest represents the request to open a compaction region. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct OpenCompactionRegionRequest { pub region_id: RegionId, pub table_dir: String, pub path_type: PathType, pub region_options: RegionOptions, pub max_parallelism: usize, + /// Plugins for the compaction region, used to look up the [`RegionHook`](crate::engine::region_hook::RegionHook). + pub plugins: Plugins, } /// Open a compaction region from a compaction request. @@ -174,9 +181,11 @@ pub async fn open_compaction_region( let manifest = manifest_manager.manifest(); let region_metadata = manifest.metadata.clone(); + let hook: Option = req.plugins.get(); let manifest_ctx = Arc::new(ManifestContext::new( manifest_manager, RegionRoleState::Leader(RegionLeaderState::Writable), + hook, )); let file_purger = { @@ -230,6 +239,7 @@ pub async fn open_compaction_region( file_purger: Some(file_purger), ttl: Some(ttl), max_parallelism: req.max_parallelism, + plugins: req.plugins.clone(), }) } @@ -250,11 +260,13 @@ impl CompactionRegion { } /// `[MergeOutput]` represents the output of merging SST files. -#[derive(Default, Clone, Debug, Serialize, Deserialize)] +#[derive(Default, Debug, Serialize, Deserialize)] pub struct MergeOutput { pub files_to_add: Vec, pub files_to_remove: Vec, pub compaction_time_window: Option, + #[serde(skip)] + pub sst_infos: Vec, } impl MergeOutput { @@ -286,7 +298,7 @@ pub trait Compactor: Send + Sync + 'static { &self, compaction_region: &CompactionRegion, merge_output: MergeOutput, - ) -> Result; + ) -> Result<(RegionEdit, ManifestVersion)>; } /// Trait for merging a single compaction output into SST files. @@ -300,7 +312,7 @@ pub trait SstMerger: Send + Sync + 'static { compaction_region: CompactionRegion, output: CompactionOutput, write_opts: WriteOptions, - ) -> Result>; + ) -> Result<(Vec, Vec)>; } /// The production [`SstMerger`] that reads, merges, and writes SST files. @@ -314,7 +326,7 @@ impl SstMerger for DefaultSstMerger { compaction_region: CompactionRegion, output: CompactionOutput, write_opts: WriteOptions, - ) -> Result> { + ) -> Result<(Vec, Vec)> { let region_id = compaction_region.region_id; let storage = compaction_region.region_options.storage.clone(); let index_options = compaction_region @@ -359,6 +371,7 @@ impl SstMerger for DefaultSstMerger { merge_mode, }; let source = builder.build_flat_sst_reader().await?; + let mut metrics = Metrics::new(WriteType::Compaction); let region_metadata = compaction_region.region_metadata.clone(); let sst_infos = compaction_region @@ -400,7 +413,7 @@ impl SstMerger for DefaultSstMerger { }; let output_files = sst_infos - .into_iter() + .iter() .map(|sst_info| { let pk_range = sst_info .file_metadata @@ -438,7 +451,7 @@ impl SstMerger for DefaultSstMerger { region_id, input_file_names, output_file_names, flat_format, metrics ); metrics.observe(); - Ok(output_files) + Ok((output_files, sst_infos.into_iter().collect())) } } @@ -506,7 +519,9 @@ where tasks.push((inputs_to_remove, fut)); } + let hook: Option = compaction_region.plugins.get(); let mut output_files = Vec::with_capacity(tasks.len()); + let mut all_sst_infos: Vec = Vec::new(); let mut compacted_inputs = Vec::with_capacity( tasks.iter().map(|(inputs, _)| inputs.len()).sum::() + picker_output.expired_ssts.len(), @@ -530,8 +545,11 @@ where while let Some((inputs, handle)) = spawned.pop() { let abort_handle = handle.abort_handle(); match CancellableFuture::new(handle, self.cancel_handle.clone()).await { - Ok(Ok(Ok(files))) => { + Ok(Ok(Ok((files, infos)))) => { output_files.extend(files); + if hook.is_some() { + all_sst_infos.extend(infos); + } compacted_inputs.extend(inputs); } Ok(Ok(Err(e))) => { @@ -591,6 +609,7 @@ where files_to_add: output_files, files_to_remove: compacted_inputs, compaction_time_window: Some(compaction_time_window), + sst_infos: all_sst_infos, }) } @@ -598,7 +617,7 @@ where &self, compaction_region: &CompactionRegion, merge_output: MergeOutput, - ) -> Result { + ) -> Result<(RegionEdit, ManifestVersion)> { // Write region edit to manifest. let edit = RegionEdit { files_to_add: merge_output.files_to_add, @@ -615,12 +634,12 @@ where let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())); // TODO: We might leak files if we fail to update manifest. We can add a cleanup task to remove them later. - compaction_region + let manifest_version = compaction_region .manifest_ctx .update_manifest_for_compaction(action_list) .await?; - Ok(edit) + Ok((edit, manifest_version)) } } @@ -679,6 +698,7 @@ mod tests { file_purger: None, ttl: None, max_parallelism: 1, + plugins: Plugins::new(), } } @@ -707,10 +727,10 @@ mod tests { _compaction_region: CompactionRegion, _output: CompactionOutput, _write_opts: WriteOptions, - ) -> Result> { + ) -> Result<(Vec, Vec)> { let idx = self.call_idx.fetch_add(1, Ordering::SeqCst); match self.results.lock().unwrap().get(idx) { - Some(Ok(files)) => Ok(files.clone()), + Some(Ok(files)) => Ok((files.clone(), Vec::new())), Some(Err(_)) => error::InvalidMetaSnafu { reason: format!("simulated failure at index {idx}"), } @@ -879,7 +899,7 @@ mod tests { _compaction_region: CompactionRegion, _output: CompactionOutput, _write_opts: WriteOptions, - ) -> Result> { + ) -> Result<(Vec, Vec)> { self.call_idx.fetch_add(1, Ordering::SeqCst); std::future::pending().await } diff --git a/src/mito2/src/compaction/task.rs b/src/mito2/src/compaction/task.rs index e48866b275c..bc024dd6dbd 100644 --- a/src/mito2/src/compaction/task.rs +++ b/src/mito2/src/compaction/task.rs @@ -21,12 +21,14 @@ use common_memory_manager::OnExhaustedPolicy; use common_telemetry::{error, info, warn}; use itertools::Itertools; use snafu::ResultExt; +use store_api::ManifestVersion; use tokio::sync::mpsc; use crate::compaction::LocalCompactionState; use crate::compaction::compactor::{CompactionRegion, Compactor, MergeOutput}; use crate::compaction::memory_manager::{CompactionMemoryGuard, CompactionMemoryManager}; use crate::compaction::picker::{CompactionTask, PickerOutput}; +use crate::engine::region_hook::{RegionHookRef, SstFileInfo}; use crate::error::{CompactRegionSnafu, CompactionMemoryExhaustedSnafu}; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::metrics::{COMPACTION_FAILURE_COUNT, COMPACTION_MEMORY_WAIT, COMPACTION_STAGE_ELAPSED}; @@ -249,7 +251,7 @@ impl CompactionTaskImpl { async fn update_manifest( &self, compaction_result: crate::compaction::compactor::MergeOutput, - ) -> error::Result { + ) -> error::Result<(RegionEdit, ManifestVersion)> { let _manifest_timer = COMPACTION_STAGE_ELAPSED .with_label_values(&["write_manifest"]) .start_timer(); @@ -282,6 +284,27 @@ impl CompactionTaskImpl { ); } } + + async fn invoke_sst_hook(&self, merge_output: &MergeOutput) { + let hook: Option = self.compaction_region.plugins.get(); + if let Some(hook) = hook { + let files: Vec> = merge_output + .sst_infos + .iter() + .zip(merge_output.files_to_add.iter()) + .map(|(info, meta)| SstFileInfo { + sst_info_ref: info, + file_meta: meta, + }) + .collect(); + hook.on_sst_files_written( + self.compaction_region.region_id, + &self.compaction_region.region_metadata, + &files, + ) + .await; + } + } } #[async_trait::async_trait] @@ -320,6 +343,7 @@ impl CompactionTask for CompactionTaskImpl { .await { Ok(Ok(merge_output)) => { + self.invoke_sst_hook(&merge_output).await; // Stop accepting cancellation once we are about to publish the compaction edit. if !self.state.mark_commit_started() { let senders = std::mem::take(&mut self.waiters); @@ -329,7 +353,7 @@ impl CompactionTask for CompactionTaskImpl { }) } else { match self.update_manifest(merge_output).await { - Ok(edit) => { + Ok((edit, _manifest_version)) => { let senders = std::mem::take(&mut self.waiters); BackgroundNotify::CompactionFinished(CompactionFinished { region_id: self.compaction_region.region_id, diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index b70b8087533..ee679cb2250 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -56,6 +56,7 @@ mod parallel_test; mod projection_test; #[cfg(test)] mod prune_test; +pub mod region_hook; #[cfg(test)] mod row_selector_test; #[cfg(test)] diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index f08628bcd08..214060b1bb6 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -16,24 +16,31 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering}; use std::time::Duration; use api::v1::Rows; +use async_trait::async_trait; +use common_base::Plugins; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; use rstest::rstest; use rstest_reuse::{self, apply}; +use store_api::ManifestVersion; +use store_api::metadata::RegionMetadataRef; use store_api::region_engine::RegionEngine; use store_api::region_request::{ PathType, RegionCloseRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest, ReplayCheckpoint, }; use store_api::storage::{RegionId, ScanRequest}; +use tokio::sync::Notify; use crate::config::MitoConfig; use crate::engine::listener::{FlushListener, StallListener}; +use crate::engine::region_hook::{RegionHook, RegionHookRef, SstFileInfo}; +use crate::manifest::action::RegionMetaActionList; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, build_rows, build_rows_for_key, flush_region, kafka_log_store_factory, multiple_log_store_factories, @@ -891,3 +898,133 @@ async fn test_update_topic_latest_entry_id(factory: Option) { .unwrap(); assert_eq!(region.topic_latest_entry_id.load(Ordering::Relaxed), 1); } + +#[derive(Debug)] +struct MockRegionHook { + sst_written_count: AtomicUsize, + manifest_updated_count: AtomicUsize, + notify: Notify, +} + +impl MockRegionHook { + fn new() -> Self { + Self { + sst_written_count: AtomicUsize::new(0), + manifest_updated_count: AtomicUsize::new(0), + notify: Notify::new(), + } + } + + async fn wait_for_manifest_update(&self) { + self.notify.notified().await; + } +} + +#[async_trait] +impl RegionHook for MockRegionHook { + async fn on_sst_files_written( + &self, + region_id: RegionId, + _region_metadata: &RegionMetadataRef, + files: &[SstFileInfo<'_>], + ) { + self.sst_written_count + .fetch_add(files.len(), Ordering::Relaxed); + common_telemetry::info!( + "MockRegionHook::on_sst_files_written: region={}, files={}", + region_id, + files.len(), + ); + for (i, file) in files.iter().enumerate() { + common_telemetry::info!( + " file[{}]: file_id={}, num_rows={}, num_series={}, file_size={}", + i, + file.sst_info_ref.file_id, + file.sst_info_ref.num_rows, + file.sst_info_ref.num_series, + file.sst_info_ref.file_size, + ); + } + } + + async fn on_manifest_updated( + &self, + region_id: RegionId, + action_list: &RegionMetaActionList, + manifest_version: ManifestVersion, + ) { + self.manifest_updated_count.fetch_add(1, Ordering::Relaxed); + // Count files added across all Edit actions. + let files_added: usize = action_list + .actions + .iter() + .map(|action| match action { + crate::manifest::action::RegionMetaAction::Edit(edit) => edit.files_to_add.len(), + _ => 0, + }) + .sum(); + common_telemetry::info!( + "MockRegionHook::on_manifest_updated: region={}, manifest_version={}, files_added={}", + region_id, + manifest_version, + files_added, + ); + self.notify.notify_one(); + } +} + +#[tokio::test] +async fn test_region_hook() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + + let hook = Arc::new(MockRegionHook::new()); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + + let engine = env + .create_engine_with_plugins(MitoConfig::default(), plugins) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("a", 0, 10, 0), + }; + put_rows(&engine, region_id, rows).await; + + let rows2 = Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("b", 0, 10, 0), + }; + put_rows(&engine, region_id, rows2).await; + + flush_region(&engine, region_id, None).await; + + hook.wait_for_manifest_update().await; + + let sst_count = hook.sst_written_count.load(Ordering::Relaxed); + let manifest_count = hook.manifest_updated_count.load(Ordering::Relaxed); + + assert!( + sst_count > 0, + "Expected at least 1 SST file, got {sst_count}" + ); + assert_eq!( + manifest_count, 1, + "Expected exactly 1 manifest update, got {manifest_count}" + ); + + common_telemetry::info!( + "test_region_hook passed: sst_count={}, manifest_count={}", + sst_count, + manifest_count, + ); +} diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index 11279954a9b..b22e17dd015 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -18,6 +18,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use api::v1::Rows; +use common_base::Plugins; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; @@ -619,6 +620,7 @@ async fn test_open_compaction_region_with_format(flat_format: bool) { path_type: PathType::Bare, region_options: RegionOptions::default(), max_parallelism: 1, + plugins: Plugins::new(), }; let compaction_region = open_compaction_region( diff --git a/src/mito2/src/engine/region_hook.rs b/src/mito2/src/engine/region_hook.rs new file mode 100644 index 00000000000..0abf3fc1104 --- /dev/null +++ b/src/mito2/src/engine/region_hook.rs @@ -0,0 +1,141 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Region hook extension point for observing SST writes and manifest mutations. +//! +//! ## Design +//! +//! The [`RegionHook`] trait provides two methods with clear separation of concerns: +//! +//! - [`on_sst_files_written`]: Fires when mito2 physically writes SST **data files**. +//! Provides per-file [`SstInfo`] + [`FileMeta`] — the richest available metadata +//! (row counts, index metadata, Parquet metadata, etc.). +//! +//! - [`on_manifest_updated`]: Fires after **any** manifest write is successfully committed. +//! Receives the full [`RegionMetaActionList`] so consumers can inspect what changed +//! (file additions/removals, schema changes, truncation, partition expression changes, etc.). +//! +//! Hook implementations are registered via the [`Plugins`](common_base::Plugins) system: +//! ```ignore +//! plugins.insert(Arc::new(MyHook) as RegionHookRef); +//! ``` +//! +//! ## Coverage +//! +//! | Scenario | `on_sst_files_written` | `on_manifest_updated` | +//! |------------------------------|:----------------------:|:---------------------:| +//! | Flush (memtable → SST) | ✅ Yes | ✅ Yes | +//! | Local compaction | ✅ Yes | ✅ Yes | +//! | Remote compaction | ✅ (compactor node) ¹ | ✅ (compactor node) ¹ | +//! | RegionEdit / bulk ingestion | ❌ (files pre-written) | ✅ Yes | +//! | Copy region | ❌ (object-store copy) | ✅ Yes | +//! | Apply staging | ❌ (delegates to edit) | ✅ Yes ² | +//! | Alter (schema change) | ❌ (no SST files) | ✅ Yes | +//! | Truncate | ❌ (removes files) | ✅ Yes | +//! | Enter staging | ❌ (no SST files) | ✅ Yes | +//! | Async index build | ❌ (index files only) | ✅ Yes | +//! +//! ¹ Remote compaction runs on a dedicated compactor node via `open_compaction_region()`. +//! The caller must pass plugins via `OpenCompactionRegionRequest` to enable hooks on the +//! compactor node. +//! ² Apply staging fires `on_manifest_updated` twice: once when the staging SST files are +//! committed via `RegionEdit`, and once when `exit_staging_on_success` merges all staged +//! manifest actions into the live manifest. +//! +//! The following paths do **not** trigger any hook: +//! - Follower region sync / catchup (manifest read-only; followers don't author changes) +//! - GC / checkpoint / drop / remap (internal bookkeeping, not logical state changes) +//! +//! ## Invocation points +//! +//! `on_sst_files_written` is invoked at the SST write site (flush task or compaction task), +//! immediately after SST files are written but **before** the manifest is committed. +//! +//! `on_manifest_updated` is centralized in [`ManifestContext::update_manifest_with_state_check`], +//! so it automatically covers all manifest write paths that go through `ManifestContext`. +//! The sole exception is [`MitoRegion::exit_staging_on_success`], which returns the hook +//! payload to the caller so the hook can be invoked after dropping the manifest write lock. +//! +//! ## Future work +//! +//! A future `on_files_removed` hook may be added to observe file lifecycle end +//! (GC, drop, truncate, compaction removal). This is not yet implemented. +//! +//! [`on_sst_files_written`]: RegionHook::on_sst_files_written +//! [`on_manifest_updated`]: RegionHook::on_manifest_updated + +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use store_api::ManifestVersion; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::RegionId; + +use crate::manifest::action::RegionMetaActionList; +use crate::sst::file::FileMeta; +use crate::sst::parquet::SstInfo; + +/// Information about a single SST data file written during flush or compaction. +pub struct SstFileInfo<'a> { + pub sst_info_ref: &'a SstInfo, + pub file_meta: &'a FileMeta, +} + +/// Hook for observing region mutations in mito2. +/// +/// Implementations can be registered via the `Plugins` system: +/// ```ignore +/// use std::sync::Arc; +/// use common_base::Plugins; +/// use mito2::engine::region_hook::{RegionHook, RegionHookRef}; +/// +/// plugins.insert(Arc::new(MyHook) as RegionHookRef); +/// ``` +#[async_trait] +pub trait RegionHook: Send + Sync + Debug { + /// Called after SST **data files** are physically written, before manifest commit. + /// + /// This fires only when mito2 itself writes SST files (flush and compaction). + /// It does **not** fire when SST files are pre-written externally (bulk ingestion, + /// copy region) or when only index files are written (async index build). + async fn on_sst_files_written( + &self, + region_id: RegionId, + region_metadata: &RegionMetadataRef, + files: &[SstFileInfo<'_>], + ) { + let _ = (region_id, region_metadata, files); + } + + /// Called after the region manifest is successfully committed. + /// + /// Fires for **all** manifest write paths: flush, compaction, region edit, + /// copy region, alter, truncate, enter staging, index build, etc. + /// + /// Does **not** fire for: + /// - Manifest reads / follower sync (no write) + /// - GC / checkpoint (internal bookkeeping) + /// - Failed manifest updates + async fn on_manifest_updated( + &self, + region_id: RegionId, + action_list: &RegionMetaActionList, + manifest_version: ManifestVersion, + ) { + let _ = (region_id, action_list, manifest_version); + } +} + +pub type RegionHookRef = Arc; diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 16260327f5a..3105816393d 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -36,6 +36,7 @@ use crate::access_layer::{ }; use crate::cache::CacheManagerRef; use crate::config::MitoConfig; +use crate::engine::region_hook::SstFileInfo; use crate::error::{ Error, FlushRegionSnafu, JoinSnafu, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, Result, @@ -386,6 +387,7 @@ impl RegionFlushTask { series_count, encoded_part_count, flush_metrics, + sst_infos, } = self.do_flush_memtables(version, write_opts).await?; if !file_metas.is_empty() { @@ -414,6 +416,20 @@ impl RegionFlushTask { ); flush_metrics.observe(); + let hook = self.manifest_ctx.hook(); + if let Some(hook) = &hook { + let files: Vec> = sst_infos + .iter() + .zip(file_metas.iter()) + .map(|(sst_info, file_meta)| SstFileInfo { + sst_info_ref: sst_info, + file_meta, + }) + .collect(); + hook.on_sst_files_written(self.region_id, &version.metadata, &files) + .await; + } + let edit = RegionEdit { files_to_add: file_metas, files_to_remove: Vec::new(), @@ -444,12 +460,12 @@ impl RegionFlushTask { }; // We will leak files if the manifest update fails, but we ignore them for simplicity. We can // add a cleanup job to remove them later. - let version = self + let manifest_version = self .manifest_ctx .update_manifest(expected_state, action_list, self.is_staging) .await?; info!( - "Successfully update manifest version to {version}, region: {}, is_staging: {}, reason: {}", + "Successfully update manifest version to {manifest_version}, region: {}, is_staging: {}, reason: {}", self.region_id, self.is_staging, self.reason.as_str() @@ -470,6 +486,8 @@ impl RegionFlushTask { let mut encoded_part_count = 0; let mut flush_metrics = Metrics::new(WriteType::Flush); let partition_expr = parse_partition_expr(self.partition_expr.as_deref())?; + let hook = self.manifest_ctx.hook(); + let mut all_sst_infos = Vec::new(); for mem in memtables { if mem.is_empty() { // Skip empty memtables. @@ -523,7 +541,7 @@ impl RegionFlushTask { flush_metrics = flush_metrics.merge(metrics); - for sst_info in ssts_written { + for sst_info in &ssts_written { flushed_bytes += sst_info.file_size; let pk_range = sst_info .file_metadata @@ -537,6 +555,9 @@ impl RegionFlushTask { pk_range, )); } + if hook.is_some() { + all_sst_infos.extend(ssts_written); + } } common_telemetry::debug!( @@ -558,6 +579,7 @@ impl RegionFlushTask { series_count, encoded_part_count, flush_metrics, + sst_infos: all_sst_infos, }) } @@ -626,7 +648,7 @@ impl RegionFlushTask { fn new_file_meta( region_id: RegionId, max_sequence: u64, - sst_info: SstInfo, + sst_info: &SstInfo, partition_expr: Option, primary_key_range: Option<(Bytes, Bytes)>, ) -> FileMeta { @@ -722,6 +744,7 @@ struct DoFlushMemtablesResult { series_count: usize, encoded_part_count: usize, flush_metrics: Metrics, + sst_infos: Vec, } struct FlatSources { diff --git a/src/mito2/src/manifest/manager.rs b/src/mito2/src/manifest/manager.rs index a69f08693c9..c471efb4307 100644 --- a/src/mito2/src/manifest/manager.rs +++ b/src/mito2/src/manifest/manager.rs @@ -693,7 +693,7 @@ impl RegionManifestManager { version } - fn last_version(&self) -> ManifestVersion { + pub fn last_version(&self) -> ManifestVersion { self.last_version.load(Ordering::Relaxed) } diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 4acec5a8931..f58d09d1cca 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -45,6 +45,7 @@ use tokio::sync::RwLockWriteGuard; pub use utils::*; use crate::access_layer::AccessLayerRef; +use crate::engine::region_hook::RegionHookRef; use crate::error::{ FlushableRegionStateSnafu, InvalidPartitionExprSnafu, RegionNotFoundSnafu, RegionStateSnafu, RegionTruncatedSnafu, Result, UnexpectedSnafu, UpdateManifestSnafu, @@ -441,7 +442,7 @@ impl MitoRegion { self.manifest_ctx.manifest_manager.write().await; let current_state = self.state(); - match state { + let hook_payload: Option<(RegionMetaActionList, ManifestVersion)> = match state { SettableRegionRoleState::Leader => { // Exit staging mode and return to normal writable leader // Only allowed from staging state @@ -449,11 +450,12 @@ impl MitoRegion { RegionRoleState::Leader(RegionLeaderState::Staging) => { info!("Exiting staging mode for region {}", self.region_id); // Use the success exit path that merges all staged manifests - self.exit_staging_on_success(&mut manager).await?; + self.exit_staging_on_success(&mut manager).await? } RegionRoleState::Leader(RegionLeaderState::Writable) => { // Already in desired state - no-op info!("Region {} already in normal leader mode", self.region_id); + None } _ => { // Only staging -> leader transition is allowed @@ -488,6 +490,7 @@ impl MitoRegion { .build()); } } + None } SettableRegionRoleState::Follower => { @@ -510,6 +513,7 @@ impl MitoRegion { info!("Region {} already in follower mode", self.region_id); } } + None } SettableRegionRoleState::DowngradingLeader => { @@ -538,10 +542,12 @@ impl MitoRegion { ); } } + None } - } + }; // Hack(zhongzc): If we have just become leader (writable), persist any backfilled metadata. + let mut backfill_hook_payload: Option<(RegionMetaActionList, ManifestVersion)> = None; if self.state() == RegionRoleState::Leader(RegionLeaderState::Writable) { // Persist backfilled metadata if manifest is missing fields (e.g., partition_expr) let manifest_meta = &manager.manifest().metadata; @@ -553,9 +559,10 @@ impl MitoRegion { sst_format: current_version.options.sst_format.unwrap_or_default(), append_mode: None, }); - let result = manager - .update(RegionMetaActionList::with_action(action), false) - .await; + let action_list = RegionMetaActionList::with_action(action); + let hook = self.manifest_ctx.hook(); + let al_for_hook = hook.as_ref().map(|_| action_list.clone()); + let result = manager.update(action_list, false).await; match result { Ok(version) => { @@ -563,6 +570,7 @@ impl MitoRegion { "Successfully persisted backfilled metadata for region {}, version: {}", self.region_id, version ); + backfill_hook_payload = al_for_hook.map(|al| (al, version)); } Err(e) => { warn!(e; "Failed to persist backfilled metadata for region {}", self.region_id); @@ -573,6 +581,27 @@ impl MitoRegion { drop(manager); + // Merge both payloads into a single hook invocation so consumers see + // the complete set of actions and the latest manifest version. + let merged = match (hook_payload, backfill_hook_payload) { + (Some((al, _v)), Some((backfill_al, backfill_v))) => { + // Combine action lists; use the later version (backfill happens after staging merge). + let mut combined = al; + combined.actions.extend(backfill_al.actions); + Some((combined, backfill_v)) + } + (Some(payload), None) => Some(payload), + (None, Some(payload)) => Some(payload), + (None, None) => None, + }; + + if let Some((action_list, version)) = merged + && let Some(hook) = self.manifest_ctx.hook() + { + hook.on_manifest_updated(self.region_id, &action_list, version) + .await; + } + Ok(()) } @@ -784,10 +813,18 @@ impl MitoRegion { } /// Exit staging mode successfully by merging all staged manifests and making them visible. + /// Merges staged manifest actions into the live manifest and exits staging mode. + /// + /// The caller must hold the manifest write lock and pass it via `manager`. + /// Returns `Ok(Some(action_list, version))` when staging manifests were merged, + /// or `Ok(None)` when there were no staged manifests to merge. + /// + /// **Important:** the returned hook payload must be invoked **after** the caller + /// drops the manifest write lock, to avoid deadlocking if the hook reads the manifest. pub(crate) async fn exit_staging_on_success( &self, manager: &mut RwLockWriteGuard<'_, RegionManifestManager>, - ) -> Result<()> { + ) -> Result> { let current_state = self.manifest_ctx.current_state(); ensure!( current_state == RegionRoleState::Leader(RegionLeaderState::Staging), @@ -808,7 +845,7 @@ impl MitoRegion { ); // Even if no manifests to merge, we still need to exit staging mode self.exit_staging()?; - return Ok(()); + return Ok(None); } }; let expect_change = merged_actions.actions.iter().any(|a| a.is_change()); @@ -851,8 +888,13 @@ impl MitoRegion { ); } - // Submit merged actions using the manifest manager's update method - // Pass the `false` so it saves to normal directory, not staging + // Submit merged actions using the manifest manager's update method. + // Pass `false` so it saves to normal directory, not staging. + // + // Clone the action list before it is consumed so we can return it to + // the caller for hook invocation after the lock is dropped. + let hook = self.manifest_ctx.hook(); + let action_list_for_hook = hook.as_ref().map(|_| merged_actions.clone()); let new_version = manager.update(merged_actions, false).await?; info!( "Successfully submitted merged staged manifests for region {}, new version: {}", @@ -877,7 +919,9 @@ impl MitoRegion { } self.exit_staging()?; - Ok(()) + // Return the hook payload; the caller invokes the hook after dropping the lock. + let hook_payload = action_list_for_hook.map(|action_list| (action_list, new_version)); + Ok(hook_payload) } /// Returns the partition expression string for this region. @@ -947,17 +991,29 @@ pub(crate) struct ManifestContext { /// During the staging mode, the region metadata in [`VersionControlRef`] is not updated, /// so we need to store the partition info separately. staging_partition_info: Mutex>, + /// Optional region hook for observing manifest mutations. + hook: Option, } impl ManifestContext { - pub(crate) fn new(manager: RegionManifestManager, state: RegionRoleState) -> Self { + pub(crate) fn new( + manager: RegionManifestManager, + state: RegionRoleState, + hook: Option, + ) -> Self { ManifestContext { manifest_manager: tokio::sync::RwLock::new(manager), state: AtomicCell::new(state), staging_partition_info: Mutex::new(None), + hook, } } + /// Returns the region hook if one is registered. + pub(crate) fn hook(&self) -> Option { + self.hook.clone() + } + pub(crate) fn staging_partition_info(&self) -> Option { self.staging_partition_info.lock().unwrap().clone() } @@ -1173,18 +1229,35 @@ impl ManifestContext { } } - // Now we can update the manifest. - let version = manager.update(action_list, is_staging).await.inspect_err( - |e| error!(e; "Failed to update manifest, region_id: {}", manifest.metadata.region_id), - )?; + // If a region hook is registered, clone the action list before it is consumed + // so we can pass a reference to the hook after a successful update. + let hook = self.hook.clone(); + let action_list_for_hook = hook.as_ref().map(|_| action_list.clone()); + let region_id = manifest.metadata.region_id; + let version = manager + .update(action_list, is_staging) + .await + .inspect_err(|e| error!(e; "Failed to update manifest, region_id: {}", region_id))?; + + // Drop the write lock before invoking the hook. Hook implementations may + // read the manifest or send region requests that acquire this lock; + // holding it would deadlock. Slow hooks also must not block manifest updates. + drop(manager); if self.state.load() == RegionRoleState::Follower { warn!( "Region {} becomes follower while updating manifest which may cause inconsistency, manifest version: {version}", - manifest.metadata.region_id + region_id ); } + if let Some(hook) = hook + && let Some(action_list) = action_list_for_hook + { + hook.on_manifest_updated(region_id, &action_list, version) + .await; + } + Ok(version) } @@ -1718,6 +1791,7 @@ mod tests { let manifest_ctx = Arc::new(ManifestContext::new( manager, RegionRoleState::Leader(RegionLeaderState::Writable), + None, )); MitoRegion { @@ -1806,7 +1880,7 @@ mod tests { .await .unwrap(); - region.exit_staging_on_success(&mut manager).await.unwrap(); + let _hook_payload = region.exit_staging_on_success(&mut manager).await.unwrap(); drop(manager); assert_eq!( @@ -1845,7 +1919,7 @@ mod tests { .await .unwrap(); - region.exit_staging_on_success(&mut manager).await.unwrap(); + let _hook_payload = region.exit_staging_on_success(&mut manager).await.unwrap(); drop(manager); assert_eq!( @@ -2007,6 +2081,7 @@ mod tests { Arc::new(ManifestContext::new( manager, RegionRoleState::Leader(RegionLeaderState::Staging), + None, )) }; @@ -2075,6 +2150,7 @@ mod tests { let manifest_ctx = Arc::new(ManifestContext::new( manager, RegionRoleState::Leader(RegionLeaderState::Writable), + None, )); let region = MitoRegion { diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 1aedc5ae258..84364d6baf7 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -46,6 +46,7 @@ use crate::access_layer::AccessLayer; use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileCache, FileType, IndexKey}; use crate::config::MitoConfig; +use crate::engine::region_hook::RegionHookRef; use crate::error; use crate::error::{ EmptyRegionDirSnafu, InvalidMetadataSnafu, InvalidRegionOptionsSnafu, ObjectStoreNotFoundSnafu, @@ -115,6 +116,7 @@ pub(crate) struct RegionOpener { replay_checkpoint: Option, file_ref_manager: FileReferenceManagerRef, partition_expr_fetcher: PartitionExprFetcherRef, + hook: Option, } impl RegionOpener { @@ -153,9 +155,16 @@ impl RegionOpener { replay_checkpoint: None, file_ref_manager, partition_expr_fetcher, + hook: None, } } + /// Sets the region hook for observing manifest mutations. + pub(crate) fn hook(mut self, hook: Option) -> Self { + self.hook = hook; + self + } + /// Sets metadata builder of the region to create. pub(crate) fn metadata_builder(mut self, builder: RegionMetadataBuilder) -> Self { self.metadata_builder = Some(builder); @@ -360,6 +369,7 @@ impl RegionOpener { manifest_ctx: Arc::new(ManifestContext::new( manifest_manager, RegionRoleState::Leader(RegionLeaderState::Writable), + self.hook.clone(), )), file_purger: create_file_purger( config.gc.enable, @@ -611,6 +621,7 @@ impl RegionOpener { manifest_ctx: Arc::new(ManifestContext::new( manifest_manager, RegionRoleState::Follower, + self.hook.clone(), )), file_purger, provider: provider.clone(), diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index 45f2e9c9059..9d7600ba579 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -300,10 +300,20 @@ impl TestEnv { } pub(crate) async fn new_mito_engine(&self, config: MitoConfig) -> MitoEngine { + self.new_mito_engine_with_plugins(config, Plugins::new()) + .await + } + + pub(crate) async fn new_mito_engine_with_plugins( + &self, + config: MitoConfig, + plugins: Plugins, + ) -> MitoEngine { async fn create( zelf: &TestEnv, config: MitoConfig, log_store: Arc, + plugins: Plugins, ) -> MitoEngine { let data_home = zelf.data_home().display().to_string(); MitoEngine::new( @@ -314,15 +324,15 @@ impl TestEnv { zelf.schema_metadata_manager.clone(), zelf.file_ref_manager.clone(), zelf.partition_expr_fetcher.clone(), - Plugins::new(), + plugins, ) .await .unwrap() } match self.log_store.as_ref().unwrap().clone() { - LogStoreImpl::RaftEngine(log_store) => create(self, config, log_store).await, - LogStoreImpl::Kafka(log_store) => create(self, config, log_store).await, + LogStoreImpl::RaftEngine(log_store) => create(self, config, log_store, plugins).await, + LogStoreImpl::Kafka(log_store) => create(self, config, log_store, plugins).await, } } @@ -337,6 +347,21 @@ impl TestEnv { self.new_mito_engine(config).await } + /// Creates a new engine with specific config and plugins. + pub async fn create_engine_with_plugins( + &mut self, + config: MitoConfig, + plugins: Plugins, + ) -> MitoEngine { + let (log_store, object_store_manager) = self.create_log_and_object_store_manager().await; + + let object_store_manager = Arc::new(object_store_manager); + self.log_store = Some(log_store.clone()); + self.object_store_manager = Some(object_store_manager.clone()); + + self.new_mito_engine_with_plugins(config, plugins).await + } + /// Creates a new engine with specific config and existing logstore and object store manager. pub async fn create_follower_engine(&mut self, config: MitoConfig) -> MitoEngine { self.new_mito_engine(config).await diff --git a/src/mito2/src/test_util/scheduler_util.rs b/src/mito2/src/test_util/scheduler_util.rs index 94a16f0b1d2..10d4bffec53 100644 --- a/src/mito2/src/test_util/scheduler_util.rs +++ b/src/mito2/src/test_util/scheduler_util.rs @@ -144,6 +144,7 @@ impl SchedulerEnv { .await .unwrap(), RegionRoleState::Leader(RegionLeaderState::Writable), + None, )) } diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 75279acf8bc..3e1d42caf2c 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -621,6 +621,7 @@ impl WorkerStarter { file_ref_manager: self.file_ref_manager.clone(), partition_expr_fetcher: self.partition_expr_fetcher, flush_semaphore: self.flush_semaphore, + plugins: self.plugins, }; let handle = common_runtime::spawn_global(async move { worker_thread.run().await; @@ -894,6 +895,8 @@ struct RegionWorkerLoop { partition_expr_fetcher: PartitionExprFetcherRef, /// Semaphore to control flush concurrency. flush_semaphore: Arc, + /// Plugins for flush hooks. + plugins: Plugins, } impl RegionWorkerLoop { diff --git a/src/mito2/src/worker/handle_apply_staging.rs b/src/mito2/src/worker/handle_apply_staging.rs index dee66752cf8..cefdda1fea1 100644 --- a/src/mito2/src/worker/handle_apply_staging.rs +++ b/src/mito2/src/worker/handle_apply_staging.rs @@ -152,12 +152,22 @@ impl RegionWorkerLoop { return; }; let mut manager = region.manifest_ctx.manifest_manager.write().await; - match region.exit_staging_on_success(&mut manager).await { - Ok(()) => { - sender.send(Ok(0)); + let hook_payload = match region.exit_staging_on_success(&mut manager).await { + Ok(payload) => payload, + Err(e) => { + sender.send(Err(e)); + return; } - Err(e) => sender.send(Err(e)), + }; + // Drop the write lock before invoking the hook to avoid deadlock. + drop(manager); + if let Some((action_list, version)) = hook_payload + && let Some(hook) = region.manifest_ctx.hook() + { + hook.on_manifest_updated(region.region_id, &action_list, version) + .await; } + sender.send(Ok(0)); } else { sender.send( UnexpectedSnafu { diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index ad517c9ef72..71d3f05c690 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -125,6 +125,7 @@ impl RegionWorkerLoop { self.partition_expr_fetcher.clone(), ) .cache(Some(self.cache_manager.clone())) + .hook(self.plugins.get()) .options(region.version().options.clone())? .skip_wal_replay(true) .open(&self.config, &self.wal) diff --git a/src/mito2/src/worker/handle_copy_region.rs b/src/mito2/src/worker/handle_copy_region.rs index aa5b26448e8..739bef79046 100644 --- a/src/mito2/src/worker/handle_copy_region.rs +++ b/src/mito2/src/worker/handle_copy_region.rs @@ -19,7 +19,9 @@ use store_api::storage::{FileId, RegionId}; use crate::error::{InvalidRequestSnafu, MissingManifestSnafu, Result}; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; -use crate::region::{FileDescriptor, MitoRegionRef, RegionFileCopier, RegionMetadataLoader}; +use crate::region::{ + FileDescriptor, MitoRegionRef, RegionFileCopier, RegionLeaderState, RegionMetadataLoader, +}; use crate::request::{ BackgroundNotify, CopyRegionFromFinished, CopyRegionFromRequest, WorkerRequest, }; @@ -228,13 +230,13 @@ impl RegionWorkerLoop { committed_sequence: None, }; let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())); - info!("Applying {edit:?} to region {target_region_id}, reason: CopyRegionFrom"); + info!( + "Applying {:?} to region {target_region_id}, reason: CopyRegionFrom", + edit + ); let version = region .manifest_ctx - .manifest_manager - .write() - .await - .update(action_list, false) + .update_manifest(RegionLeaderState::Writable, action_list, false) .await?; info!( "Successfully update manifest version to {version}, region: {target_region_id}, reason: CopyRegionFrom" diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 9e812ba88fb..c1ab18d59d1 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -70,6 +70,7 @@ impl RegionWorkerLoop { .metadata_builder(builder) .parse_options(request.options)? .cache(Some(self.cache_manager.clone())) + .hook(self.plugins.get()) .create_or_open(&self.config, &self.wal) .await?; diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index a154140d98a..15e800c270e 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -107,6 +107,7 @@ impl RegionWorkerLoop { ) .skip_wal_replay(request.skip_wal_replay) .cache(Some(self.cache_manager.clone())) + .hook(self.plugins.get()) .wal_entry_reader(wal_entry_receiver.map(|receiver| Box::new(receiver) as _)) .replay_checkpoint(request.checkpoint.map(|checkpoint| checkpoint.entry_id)) .parse_options(request.options.clone())