rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-30 15:36:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: flush hook extension point (#8145)

Browse Source 
* feat: introduce flush and compaction hook

* feat: update hook definitions

* chore: fmt

* fix: address review comments

* feat: remove primary key collection

* fix: revert changes on flatsource

* feat: load plugins for metasrv http builder

* chore: resolve compilation issue

* feat: load manifest version from compaction task

* feat: use region_hook to cover all possible cases

* chore: fmt

* refactor: improve lock

* fix: provide plugins for open region of remote compactor; catch manifest update in exit staging

* fix: avoid lock when calling hook

* fix: add backfill hook change calling
This commit is contained in:
Ning Sun
2026-06-12 01:53:54 +00:00
committed by GitHub
parent ab552e24b5
commit 827b904189
19 changed files with 536 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/src/compaction.rs
+2
View File
@@ -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,
))
};
src/mito2/src/compaction/compactor.rs
+35 -15
View File
@@ -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<RegionHookRef> = 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<FileMeta>,
pub files_to_remove: Vec<FileMeta>,
pub compaction_time_window: Option<i64>,
#[serde(skip)]
pub sst_infos: Vec<SstInfo>,
}
impl MergeOutput {
@@ -286,7 +298,7 @@ pub trait Compactor: Send + Sync + 'static {
&self,
compaction_region: &CompactionRegion,
merge_output: MergeOutput,
) -> Result<RegionEdit>;
) -> 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<Vec<FileMeta>>;
) -> Result<(Vec<FileMeta>, Vec<SstInfo>)>;
}
/// 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<Vec<FileMeta>> {
) -> Result<(Vec<FileMeta>, Vec<SstInfo>)> {
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<RegionHookRef> = compaction_region.plugins.get();
let mut output_files = Vec::with_capacity(tasks.len());
let mut all_sst_infos: Vec<SstInfo> = Vec::new();
let mut compacted_inputs = Vec::with_capacity(
tasks.iter().map(|(inputs, _)| inputs.len()).sum::<usize>()
+ 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<RegionEdit> {
) -> 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<Vec<FileMeta>> {
) -> Result<(Vec<FileMeta>, Vec<SstInfo>)> {
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<Vec<FileMeta>> {
) -> Result<(Vec<FileMeta>, Vec<SstInfo>)> {
self.call_idx.fetch_add(1, Ordering::SeqCst);
std::future::pending().await
}
src/mito2/src/compaction/task.rs
+26 -2
View File
@@ -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<RegionEdit> {
) -> 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<RegionHookRef> = self.compaction_region.plugins.get();
if let Some(hook) = hook {
let files: Vec<SstFileInfo<'_>> = 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,
src/mito2/src/engine.rs
+1
View File
@@ -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)]
src/mito2/src/engine/flush_test.rs
+138 -1
View File
@@ -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<LogStoreFactory>) {
.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,
);
}
src/mito2/src/engine/open_test.rs
+2
View File
@@ -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(
src/mito2/src/engine/region_hook.rs
+141
View File
@@ -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<dyn RegionHook>;
src/mito2/src/flush.rs
+27 -4
View File
@@ -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<SstFileInfo<'_>> = 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<PartitionExpr>,
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<SstInfo>,
}
struct FlatSources {
src/mito2/src/manifest/manager.rs
+1 -1
View File
@@ -693,7 +693,7 @@ impl RegionManifestManager {
version
}
fn last_version(&self) -> ManifestVersion {
pub fn last_version(&self) -> ManifestVersion {
self.last_version.load(Ordering::Relaxed)
}
src/mito2/src/region.rs
+95 -19
View File
@@ -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<Option<(RegionMetaActionList, ManifestVersion)>> {
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<Option<StagingPartitionInfo>>,
/// Optional region hook for observing manifest mutations.
hook: Option<RegionHookRef>,
}
impl ManifestContext {
pub(crate) fn new(manager: RegionManifestManager, state: RegionRoleState) -> Self {
pub(crate) fn new(
manager: RegionManifestManager,
state: RegionRoleState,
hook: Option<RegionHookRef>,
) -> 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<RegionHookRef> {
self.hook.clone()
}
pub(crate) fn staging_partition_info(&self) -> Option<StagingPartitionInfo> {
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 {
src/mito2/src/region/opener.rs
+11
View File
@@ -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<u64>,
file_ref_manager: FileReferenceManagerRef,
partition_expr_fetcher: PartitionExprFetcherRef,
hook: Option<RegionHookRef>,
}
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<RegionHookRef>) -> 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(),
src/mito2/src/test_util.rs
+28 -3
View File
@@ -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<S: LogStore>(
zelf: &TestEnv,
config: MitoConfig,
log_store: Arc<S>,
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
src/mito2/src/test_util/scheduler_util.rs
+1
View File
@@ -144,6 +144,7 @@ impl SchedulerEnv {
.await
.unwrap(),
RegionRoleState::Leader(RegionLeaderState::Writable),
None,
))
}
src/mito2/src/worker.rs
+3
View File
@@ -621,6 +621,7 @@ impl<S: LogStore> WorkerStarter<S> {
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<S> {
partition_expr_fetcher: PartitionExprFetcherRef,
/// Semaphore to control flush concurrency.
flush_semaphore: Arc<Semaphore>,
/// Plugins for flush hooks.
plugins: Plugins,
}
impl<S: LogStore> RegionWorkerLoop<S> {
src/mito2/src/worker/handle_apply_staging.rs
+14 -4
View File
@@ -152,12 +152,22 @@ impl<S: LogStore> RegionWorkerLoop<S> {
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 {
src/mito2/src/worker/handle_catchup.rs
+1
View File
@@ -125,6 +125,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
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)
src/mito2/src/worker/handle_copy_region.rs
+8 -6
View File
@@ -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<S> RegionWorkerLoop<S> {
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"
src/mito2/src/worker/handle_create.rs
+1
View File
@@ -70,6 +70,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
.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?;
src/mito2/src/worker/handle_open.rs
+1
View File
@@ -107,6 +107,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
)
.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())