feat: more region lifecycle hooks (#8467)

* feat: hook extension for region/file close/drop

* feat: cover region lifecycle with open/close/drop hooks

* docs(mito2): clarify on_region_opened runs in spawned task on open path

Address review feedback: the module/trait docs claimed on_region_opened
runs inline in the region worker loop. That holds for the create path
but not for the open path, where it fires inside the spawned open task
(common_runtime::spawn_global) after WAL replay and before registration.
Correct both the Notes block and the trait method doc so hook authors
don't assume worker-loop-thread affinity or strict ordering on open.

Signed-off-by: Ning Sun <sunning@greptime.com>

* docs(mito2): fix region_hook inventory and lifecycle wording

Address shuiyisong's review feedback:
- Overview no longer hardcodes 'two methods'; the lifecycle bullet list
  now includes on_region_opened.
- on_region_opened is described as firing after open/create succeeds but
  before registration (it runs before insert_region on both paths).
- on_region_closed drops the inaccurate 'follower/catchup regions'
  exclusion: remove_region fires it for any role, which is consistent
  with on_region_opened firing for followers too.

Signed-off-by: Ning Sun <sunning@greptime.com>

---------

Signed-off-by: Ning Sun <sunning@greptime.com>
This commit is contained in:
Ning Sun
2026-07-13 13:21:21 +08:00
committed by GitHub
parent 29b345fe2d
commit c70cbaa132
9 changed files with 399 additions and 8 deletions
+43
View File
@@ -12,11 +12,17 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::sync::Arc;
use std::sync::atomic::Ordering;
use common_base::Plugins;
use store_api::region_engine::RegionEngine;
use store_api::region_request::{RegionCloseRequest, RegionRequest};
use store_api::storage::RegionId;
use crate::config::MitoConfig;
use crate::engine::flush_test::MockRegionHook;
use crate::engine::region_hook::RegionHookRef;
use crate::test_util::{CreateRequestBuilder, TestEnv};
#[tokio::test]
@@ -60,3 +66,40 @@ async fn test_engine_close_region_with_format(flat_format: bool) {
.await
.unwrap();
}
#[tokio::test]
async fn test_region_hook_on_close() {
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();
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
// Sanity: no lifecycle events before closing besides the create open.
assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1);
assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0);
assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0);
engine
.handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
.await
.unwrap();
// Closing fires on_region_closed exactly once.
assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1);
// Close must not be confused with drop.
assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0);
assert_eq!(hook.files_removed_count.load(Ordering::Relaxed), 0);
}
+76
View File
@@ -13,9 +13,11 @@
// limitations under the License.
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use api::v1::Rows;
use common_base::Plugins;
use common_meta::key::SchemaMetadataManager;
use common_meta::kv_backend::KvBackendRef;
use object_store::util::join_path;
@@ -25,7 +27,9 @@ use store_api::storage::RegionId;
use crate::config::MitoConfig;
use crate::engine::MitoEngine;
use crate::engine::flush_test::MockRegionHook;
use crate::engine::listener::DropListener;
use crate::engine::region_hook::RegionHookRef;
use crate::test_util::{
CreateRequestBuilder, TestEnv, build_rows_for_key, flush_region, put_rows, rows_schema,
};
@@ -302,3 +306,75 @@ async fn test_engine_drop_region_for_custom_store_with_format(flat_format: bool)
.unwrap()
);
}
#[tokio::test]
async fn test_region_hook_on_drop() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::with_prefix("drop_hook").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);
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"test_table",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
let request = CreateRequestBuilder::new().build();
let column_schemas = rows_schema(&request);
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
// Put some rows so the region has active data (kept in memtable; we
// intentionally do not flush so the drop GC worker finds no parquet files
// and removes the directory immediately, avoiding the multi-minute GC wait).
let rows = Rows {
schema: column_schemas,
rows: build_rows_for_key("a", 0, 2, 0),
};
put_rows(&engine, region_id, rows).await;
// Sanity: nothing fired yet besides the open from create.
assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1);
assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0);
assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0);
// full drop (partial_drop = false) so the GC worker physically deletes the
// directory and fires on_region_files_removed.
engine
.handle_request(
region_id,
RegionRequest::Drop(RegionDropRequest {
fast_path: false,
force: false,
partial_drop: false,
}),
)
.await
.unwrap();
// on_region_dropped fires inline while handling the drop request.
hook.wait_for_dropped().await;
assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 1);
// Drop must not be confused with close.
assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0);
// on_region_files_removed fires from the background GC worker once the
// directory is physically deleted.
hook.wait_for_files_removed().await;
assert_eq!(hook.files_removed_count.load(Ordering::Relaxed), 1);
}
+55 -4
View File
@@ -895,24 +895,46 @@ async fn test_update_topic_latest_entry_id(factory: Option<LogStoreFactory>) {
}
#[derive(Debug)]
struct MockRegionHook {
sst_written_count: AtomicUsize,
manifest_updated_count: AtomicUsize,
pub(super) struct MockRegionHook {
pub(super) sst_written_count: AtomicUsize,
pub(super) manifest_updated_count: AtomicUsize,
pub(super) opened_count: AtomicUsize,
pub(super) closed_count: AtomicUsize,
pub(super) dropped_count: AtomicUsize,
pub(super) files_removed_count: AtomicUsize,
notify: Notify,
dropped_notify: Notify,
files_removed_notify: Notify,
}
impl MockRegionHook {
fn new() -> Self {
pub(super) fn new() -> Self {
Self {
sst_written_count: AtomicUsize::new(0),
manifest_updated_count: AtomicUsize::new(0),
opened_count: AtomicUsize::new(0),
closed_count: AtomicUsize::new(0),
dropped_count: AtomicUsize::new(0),
files_removed_count: AtomicUsize::new(0),
notify: Notify::new(),
dropped_notify: Notify::new(),
files_removed_notify: Notify::new(),
}
}
async fn wait_for_manifest_update(&self) {
self.notify.notified().await;
}
/// Waits until `on_region_dropped` has fired at least once.
pub(super) async fn wait_for_dropped(&self) {
self.dropped_notify.notified().await;
}
/// Waits until `on_region_files_removed` has fired at least once.
pub(super) async fn wait_for_files_removed(&self) {
self.files_removed_notify.notified().await;
}
}
#[async_trait]
@@ -966,6 +988,35 @@ impl RegionHook for MockRegionHook {
);
self.notify.notify_one();
}
async fn on_region_opened(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) {
self.opened_count.fetch_add(1, Ordering::Relaxed);
common_telemetry::info!("MockRegionHook::on_region_opened: region={}", region_id);
}
async fn on_region_closed(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) {
self.closed_count.fetch_add(1, Ordering::Relaxed);
common_telemetry::info!("MockRegionHook::on_region_closed: region={}", region_id);
}
async fn on_region_dropped(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) {
self.dropped_count.fetch_add(1, Ordering::Relaxed);
common_telemetry::info!("MockRegionHook::on_region_dropped: region={}", region_id);
self.dropped_notify.notify_one();
}
async fn on_region_files_removed(
&self,
region_id: RegionId,
_region_metadata: &RegionMetadataRef,
) {
self.files_removed_count.fetch_add(1, Ordering::Relaxed);
common_telemetry::info!(
"MockRegionHook::on_region_files_removed: region={}",
region_id
);
self.files_removed_notify.notify_one();
}
}
#[tokio::test]
+35
View File
@@ -32,6 +32,8 @@ use tokio::sync::oneshot;
use crate::compaction::compactor::{OpenCompactionRegionRequest, open_compaction_region};
use crate::config::MitoConfig;
use crate::engine::flush_test::MockRegionHook;
use crate::engine::region_hook::RegionHookRef;
use crate::error;
use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef};
use crate::region::options::RegionOptions;
@@ -805,3 +807,36 @@ async fn test_open_keeps_none_without_fetcher() {
let meta = engine.get_region(region_id).unwrap().metadata();
assert!(meta.partition_expr.is_none());
}
#[tokio::test]
async fn test_region_hook_on_open() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::with_prefix("open_hook").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 table_dir = request.table_dir.clone();
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
// Creating a region fires on_region_opened exactly once.
assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1);
assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0);
// `reopen_region` closes then opens; the open must fire the hook again.
reopen_region(&engine, region_id, table_dir, false, Default::default()).await;
assert_eq!(hook.opened_count.load(Ordering::Relaxed), 2);
assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1);
assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0);
}
+122 -4
View File
@@ -16,7 +16,8 @@
//!
//! ## Design
//!
//! The [`RegionHook`] trait provides two methods with clear separation of concerns:
//! The [`RegionHook`] trait observes region activity through two categories of
//! callbacks — manifest/file observation and region lifecycle:
//!
//! - [`on_sst_files_written`]: Fires when mito2 physically writes SST **data files**.
//! Provides per-file [`SstInfo`] + [`FileMeta`]; metadata richness varies by path
@@ -26,6 +27,10 @@
//! Receives the full [`RegionMetaActionList`] so consumers can inspect what changed
//! (file additions/removals, schema changes, truncation, partition expression changes, etc.).
//!
//! - [`on_region_opened`] / [`on_region_closed`] / [`on_region_dropped`] / [`on_region_files_removed`]:
//! Region **lifecycle** callbacks for open, close, logical drop, and physical file removal.
//! See [Region lifecycle](#region-lifecycle) below.
//!
//! Hook implementations are registered via the [`Plugins`](common_base::Plugins) system:
//! ```ignore
//! plugins.insert(Arc::new(MyHook) as RegionHookRef);
@@ -56,7 +61,39 @@
//!
//! 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)
//! - GC / checkpoint / remap (internal bookkeeping, not logical state changes)
//!
//! An explicit region **drop** does fire lifecycle hooks — see
//! [Region lifecycle](#region-lifecycle).
//!
//! ## Region lifecycle
//!
//! Beyond manifest/SST observation, the hook observes the high-level lifecycle of an
//! active region:
//!
//! | Event | Method | When |
//! |-------|--------|------|
//! | Open | [`on_region_opened`] | A create or open request registers the region as active (the counterpart to close/drop). Does not fire for the compactor's transient regions or the catch-up reopen. |
//! | Close | [`on_region_closed`] | A close request (or a close-after-flush) removes the region from the active set. Data files, manifest and WAL state are **preserved**; the region may be reopened. |
//! | Logical drop | [`on_region_dropped`] | A drop request has been handled: the region leaves the active set and its WAL entries are marked obsolete. Data files are **not yet deleted**. |
//! | Physical file removal | [`on_region_files_removed`] | The drop GC worker has deleted the region directory. Terminal file-lifecycle event. |
//!
//! Notes:
//! - `on_region_closed` / `on_region_dropped` run **inline in the region worker loop**,
//! so implementations must be fast (same contract as `on_manifest_updated`).
//! - `on_region_opened` runs inline in the worker loop on the **create** path, but on the
//! **open** path it fires inside the spawned open task (`common_runtime::spawn_global`),
//! i.e. concurrently with the worker loop — after WAL replay, before the region is
//! registered and its open request is acknowledged. Implementations must still be fast
//! and must not assume worker-loop-thread affinity or strict ordering against concurrent
//! requests to other regions.
//! - `on_region_files_removed` runs on the background drop GC task, outside the worker loop.
//! - When global GC is enabled and a normal table region is dropped with `partial_drop`, its
//! directory is left for global reclamation and `on_region_files_removed` is **not** fired
//! by the drop worker (observe it via the global GC path instead).
//! - Logical file removal (compaction, region edit, truncate) is already observable via
//! [`on_manifest_updated`] (`Edit.files_to_remove` / `Truncate` action); only the drop
//! worker's physical directory deletion needs a dedicated file hook.
//!
//! ## Invocation points
//!
@@ -77,11 +114,18 @@
//!
//! ## 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_region_files_removed` currently covers only the **drop** GC worker's physical
//! directory removal. A broader per-file `on_files_removed` hook covering compaction
//! removal, truncate, and the global GC reclamation path is not yet implemented
//! (though logical file removal is already observable via `on_manifest_updated`).
//! Role/leadership transitions (`on_region_role_changed`) are also not hooked.
//!
//! [`on_sst_files_written`]: RegionHook::on_sst_files_written
//! [`on_manifest_updated`]: RegionHook::on_manifest_updated
//! [`on_region_opened`]: RegionHook::on_region_opened
//! [`on_region_closed`]: RegionHook::on_region_closed
//! [`on_region_dropped`]: RegionHook::on_region_dropped
//! [`on_region_files_removed`]: RegionHook::on_region_files_removed
//! [`RegionManifestManager::update`]: crate::manifest::manager::RegionManifestManager::update
//! [`ManifestContext::update_locked`]: crate::region::ManifestContext::update_locked
//! [`ManifestContext::update_manifest`]: crate::region::ManifestContext::update_manifest
@@ -228,6 +272,80 @@ pub trait RegionHook: Send + Sync + Debug {
) {
let _ = (region_id, action_list, manifest_version);
}
/// Called once a region **open** or **create** succeeds, but **before** the
/// region is registered in the engine's active set (`insert_region` runs
/// immediately afterwards).
///
/// Fires once when a region becomes active via a create or open request —
/// the natural counterpart to [`on_region_closed`] / [`on_region_dropped`].
/// It does **not** fire for the compactor's transient compaction regions
/// (`open_compaction_region`), nor for the internal reopen performed during
/// follower catch-up / leadership promotion.
///
/// On the **create** path it runs inline in the region worker loop; on the
/// **open** path it runs inside the spawned open task
/// (`common_runtime::spawn_global`), concurrently with the worker loop
/// (after WAL replay, before the region is registered/acknowledged).
/// Implementations must be fast and must **not** assume worker-loop-thread
/// affinity or strict ordering against concurrent requests to other regions.
///
/// [`on_region_closed`]: RegionHook::on_region_closed
/// [`on_region_dropped`]: RegionHook::on_region_dropped
async fn on_region_opened(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) {
let _ = (region_id, region_metadata);
}
/// Called after a region is **closed** via a close request.
///
/// The region is removed from the engine's active set, but its data files,
/// manifest, and WAL state are **preserved**; the region may be reopened
/// later. Fires once per successful close, after the region's background
/// tasks (flush/compaction) have been stopped.
///
/// Fires for a region of **any** role (leader or follower) that is closed.
/// Does **not** fire when a region is dropped (see [`on_region_dropped`]).
///
/// Runs inline in the region worker loop; implementations should be fast.
///
/// [`on_region_dropped`]: RegionHook::on_region_dropped
async fn on_region_closed(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) {
let _ = (region_id, region_metadata);
}
/// Called after a region is **logically dropped** (a drop request has been
/// handled).
///
/// The region is removed from the active set and its WAL entries are marked
/// obsolete. Its data files are **not yet deleted** — they are scheduled for
/// asynchronous removal by the GC worker. Observe physical deletion via
/// [`on_region_files_removed`].
///
/// Runs inline in the region worker loop; implementations should be fast.
///
/// [`on_region_files_removed`]: RegionHook::on_region_files_removed
async fn on_region_dropped(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) {
let _ = (region_id, region_metadata);
}
/// Called after a dropped region's data files are **physically removed** by
/// the drop GC worker (the region directory has been deleted).
///
/// This is the terminal event in a region's file lifecycle; no further
/// callbacks fire for this region id afterwards. Fires only when the drop
/// worker itself deletes the directory. When global GC is enabled and the
/// region is a normal table region dropped with `partial_drop`, the
/// directory is left for global reclamation and this hook is **not** fired
/// by the drop worker.
///
/// Runs on a background task, outside the region worker loop.
async fn on_region_files_removed(
&self,
region_id: RegionId,
region_metadata: &RegionMetadataRef,
) {
let _ = (region_id, region_metadata);
}
}
pub type RegionHookRef = Arc<dyn RegionHook>;
+8
View File
@@ -89,5 +89,13 @@ impl<S: LogStore> RegionWorkerLoop<S> {
// clean index build status.
self.index_build_scheduler.on_region_closed(region_id).await;
self.region_count.dec();
// Notify the region hook that the region has been closed. The region is
// fully stopped and unregistered, but its files/manifest are preserved.
// Runs inline; the hook contract requires it to be fast.
if let Some(hook) = region.manifest_ctx.hook() {
let metadata = region.metadata();
hook.on_region_closed(region_id, &metadata).await;
}
}
}
+6
View File
@@ -86,6 +86,12 @@ impl<S: LogStore> RegionWorkerLoop<S> {
self.region_count.inc();
// Notify the region hook that the region has been opened (created).
// Fires before registration; allocates nothing when no hook is registered.
if let Some(hook) = region.manifest_ctx.hook() {
hook.on_region_opened(region_id, &region.metadata()).await;
}
// Insert the MitoRegion into the RegionMap.
self.regions.insert_region(region);
+47
View File
@@ -23,10 +23,12 @@ use object_store::util::join_path;
use object_store::{EntryMode, ObjectStore};
use snafu::ResultExt;
use store_api::logstore::LogStore;
use store_api::metadata::RegionMetadataRef;
use store_api::region_request::{AffectedRows, PathType};
use store_api::storage::RegionId;
use tokio::time::sleep;
use crate::engine::region_hook::RegionHookRef;
use crate::error::{OpenDalSnafu, Result};
use crate::region::{RegionLeaderState, RegionMapRef};
use crate::worker::{DROPPING_MARKER_FILE, RegionWorkerLoop};
@@ -106,6 +108,22 @@ where
self.region_count.dec();
// Notify registered hooks that the region has been logically dropped, and
// prepare a payload for the background GC task to fire the terminal
// `on_region_files_removed`. When no hook is registered this allocates
// nothing — no metadata snapshot is taken and no payload is built.
let hook_payload = match region.manifest_ctx.hook() {
Some(hook) => {
let region_metadata = region.metadata();
hook.on_region_dropped(region_id, &region_metadata).await;
Some(DropHookPayload {
hook,
metadata: region_metadata,
})
}
None => None,
};
let object_store = region.access_layer.object_store().clone();
let dropping_regions = self.dropping_regions.clone();
let listener = self.listener.clone();
@@ -122,6 +140,7 @@ where
object_store,
dropping_regions,
partial_drop,
hook_payload.as_ref(),
)
.await
} else {
@@ -135,6 +154,7 @@ where
object_store,
dropping_regions,
gc_duration,
hook_payload.as_ref(),
)
.await
};
@@ -156,6 +176,19 @@ where
}
}
/// Carries the region hook and region metadata into the background GC task so
/// it can fire [`RegionHook::on_region_files_removed`] once the dropped region's
/// directory is physically deleted.
///
/// Only constructed when a hook is registered; the task receives it as
/// `Option<&DropHookPayload>` so the no-hook path allocates nothing.
///
/// [`RegionHook::on_region_files_removed`]: crate::engine::region_hook::RegionHook::on_region_files_removed
struct DropHookPayload {
hook: RegionHookRef,
metadata: RegionMetadataRef,
}
/// Background GC task to remove the entire region path once one of the following
/// conditions is true:
/// - It finds there is no parquet file left.
@@ -172,6 +205,7 @@ async fn later_drop_task_without_global_gc(
object_store: ObjectStore,
dropping_regions: RegionMapRef,
gc_duration: Duration,
hook_payload: Option<&DropHookPayload>,
) -> bool {
remove_region_with_retry(
region_id,
@@ -180,6 +214,7 @@ async fn later_drop_task_without_global_gc(
dropping_regions,
Some(gc_duration),
false,
hook_payload,
)
.await
}
@@ -191,6 +226,7 @@ async fn remove_region_with_retry(
dropping_regions: std::sync::Arc<crate::region::RegionMap>,
gc_duration: Option<Duration>,
mut force: bool,
hook_payload: Option<&DropHookPayload>,
) -> bool {
for _ in 0..MAX_RETRY_TIMES {
let result = remove_region_dir_once(&region_path, &object_store, force).await;
@@ -204,6 +240,15 @@ async fn remove_region_with_retry(
Ok(true) => {
dropping_regions.remove_region(region_id);
info!("Region {} is dropped, force: {}", region_path, force);
// The region directory has been physically deleted; fire the
// terminal file-lifecycle event. The partial-drop/global-GC path
// never reaches here (it does not delete the directory itself).
if let Some(hook_payload) = hook_payload {
hook_payload
.hook
.on_region_files_removed(region_id, &hook_payload.metadata)
.await;
}
return true;
}
Ok(false) => (),
@@ -230,6 +275,7 @@ async fn later_drop_task_with_global_gc(
object_store: ObjectStore,
dropping_regions: RegionMapRef,
partial_drop: bool,
hook_payload: Option<&DropHookPayload>,
) -> bool {
// For metadata regions or regions marked for full deletion (such as when dropping a table)
// the region directory is forcefully removed immediately.
@@ -243,6 +289,7 @@ async fn later_drop_task_with_global_gc(
dropping_regions,
None,
true,
hook_payload,
)
.await
} else {
+7
View File
@@ -149,6 +149,13 @@ impl<S: LogStore> RegionWorkerLoop<S> {
);
region_count.inc();
// Notify the region hook that the region has been opened.
// Fires before registration; allocates nothing when no hook
// is registered.
if let Some(hook) = region.manifest_ctx.hook() {
hook.on_region_opened(region_id, &region.metadata()).await;
}
// Insert the Region into the RegionMap.
regions.insert_region(region);