rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-08 14:39:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(mito2): run compaction picking in background with plan tracking

Browse Source 
Move the compaction picker out of the region worker's critical path by
dispatching planning to a background task and reporting the result back
via CompactionPickFinished. CompactionStatus now tracks an explicit
picking phase keyed by a monotonic plan id, so stale planning results
are rejected and duplicate regular triggers coalesce while picking.

Before submitting a prepared compaction, the picker output is refreshed
against the current SST version (file handles are re-resolved and
conflicts roll back reservations), ensuring the plan still matches live
state. CompactionExecution identifies the running task by
(plan id, kind, version control) so finish/cancel/fail notifications
from outdated executions are ignored.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
This commit is contained in:
Lei, HUANG
2026-07-23 14:23:36 +08:00
parent 598a412fbb
commit b82aa10f09
13 changed files with 4063 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/src/compaction.rs
+2702 -232
View File
File diff suppressed because it is too large Load Diff
src/mito2/src/compaction/task.rs
+12 -16
View File
@@ -24,10 +24,10 @@ 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::compaction::{CompactionExecution, LocalCompactionState};
use crate::error::{CompactRegionSnafu, CompactionMemoryExhaustedSnafu};
use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList};
use crate::metrics::{COMPACTION_FAILURE_COUNT, COMPACTION_MEMORY_WAIT, COMPACTION_STAGE_ELAPSED};
@@ -46,6 +46,8 @@ pub const MAX_PARALLEL_COMPACTION: usize = 1;
pub(crate) struct CompactionTaskImpl {
/// Shared local-compaction state for cooperative cancellation.
pub(crate) state: LocalCompactionState,
/// Identity and reservation lease of this accepted execution.
pub(crate) execution: CompactionExecution,
pub compaction_region: CompactionRegion,
/// Request sender to notify the worker.
pub(crate) request_sender: mpsc::Sender<WorkerRequestWithTime>,
@@ -80,20 +82,7 @@ impl Debug for CompactionTaskImpl {
}
}
impl Drop for CompactionTaskImpl {
fn drop(&mut self) {
self.mark_files_compacting(false)
}
}
impl CompactionTaskImpl {
fn mark_files_compacting(&self, compacting: bool) {
self.picker_output
.outputs
.iter()
.for_each(|o| o.inputs.iter().for_each(|f| f.set_compacting(compacting)));
}
/// Acquires memory budget based on the configured policy.
///
/// Returns an error if memory cannot be acquired according to the policy.
@@ -301,6 +290,7 @@ impl CompactionTask for CompactionTaskImpl {
self.on_failure(err.clone());
let notify = BackgroundNotify::CompactionFailed(CompactionFailed {
region_id: self.compaction_region.region_id,
execution: self.execution.clone(),
err,
});
self.send_to_worker(WorkerRequest::Background {
@@ -312,8 +302,6 @@ impl CompactionTask for CompactionTaskImpl {
}
};
// Marks files compacting before compaction and unmark after compaction (even if compaction is cancelled or failed), so that they won't be picked by other compaction tasks.
self.mark_files_compacting(true);
self.handle_expiration().await;
let cancel_handle = self.state.cancel_handle();
@@ -331,14 +319,19 @@ impl CompactionTask for CompactionTaskImpl {
let senders = std::mem::take(&mut self.waiters);
BackgroundNotify::CompactionCancelled(CompactionCancelled {
region_id: self.compaction_region.region_id,
execution: self.execution.clone(),
senders,
})
} else {
self.listener
.on_compaction_commit_begin(self.compaction_region.region_id)
.await;
match self.update_manifest(merge_output).await {
Ok((edit, _manifest_version)) => {
let senders = std::mem::take(&mut self.waiters);
BackgroundNotify::CompactionFinished(CompactionFinished {
region_id: self.compaction_region.region_id,
execution: self.execution.clone(),
senders,
start_time: self.start_time,
edit,
@@ -350,6 +343,7 @@ impl CompactionTask for CompactionTaskImpl {
self.on_failure(err.clone());
BackgroundNotify::CompactionFailed(CompactionFailed {
region_id: self.compaction_region.region_id,
execution: self.execution.clone(),
err,
})
}
@@ -364,6 +358,7 @@ impl CompactionTask for CompactionTaskImpl {
let senders = std::mem::take(&mut self.waiters);
BackgroundNotify::CompactionCancelled(CompactionCancelled {
region_id: self.compaction_region.region_id,
execution: self.execution.clone(),
senders,
})
}
@@ -374,6 +369,7 @@ impl CompactionTask for CompactionTaskImpl {
self.on_failure(err.clone());
BackgroundNotify::CompactionFailed(CompactionFailed {
region_id: self.compaction_region.region_id,
execution: self.execution.clone(),
err,
})
}
src/mito2/src/compaction/twcs.rs
+82 -6
View File
@@ -13,7 +13,7 @@
// limitations under the License.
use std::collections::hash_map::Entry;
use std::collections::{BTreeMap, HashMap};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fmt::Debug;
use std::num::NonZeroU64;
@@ -246,9 +246,11 @@ impl Picker for TwcsPicker {
get_expired_ssts(levels, compaction_region.ttl, Timestamp::current_millis());
if !expired_ssts.is_empty() {
info!("Expired SSTs in region {}: {:?}", region_id, expired_ssts);
// here we mark expired SSTs as compacting to avoid them being picked.
expired_ssts.iter().for_each(|f| f.set_compacting(true));
}
let expired_file_ids = expired_ssts
.iter()
.map(|file| file.file_id())
.collect::<HashSet<_>>();
let compaction_time_window = compaction_region
.current_version
@@ -268,8 +270,13 @@ impl Picker for TwcsPicker {
// Find active window from files in level 0.
let active_window = find_latest_window_in_seconds(levels[0].files(), time_window_size);
// Assign files to windows
let mut windows =
assign_to_windows(levels.iter().flat_map(LevelMeta::files), time_window_size);
let mut windows = assign_to_windows(
levels
.iter()
.flat_map(LevelMeta::files)
.filter(|file| !expired_file_ids.contains(&file.file_id())),
time_window_size,
);
let outputs = self.build_output(region_id, &mut windows, active_window);
if outputs.is_empty() && expired_ssts.is_empty() {
@@ -411,16 +418,85 @@ fn find_latest_window_in_seconds<'a>(
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use std::num::NonZeroU64;
use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use common_base::Plugins;
use store_api::storage::FileId;
use super::*;
use crate::cache::CacheManager;
use crate::compaction::compactor::CompactionVersion;
use crate::compaction::test_util::{
new_file_handle, new_file_handle_with_sequence, new_file_handle_with_size_and_sequence,
new_file_handle_with_size_sequence_and_primary_key_range,
};
use crate::sst::file::Level;
use crate::config::MitoConfig;
use crate::region::options::RegionOptions;
use crate::sst::file::{FileMeta, Level};
use crate::sst::version::SstVersion;
use crate::test_util::memtable_util::metadata_for_test;
use crate::test_util::scheduler_util::SchedulerEnv;
async fn compaction_region_with_expired_sst() -> CompactionRegion {
let env = SchedulerEnv::new().await;
let metadata = metadata_for_test();
let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
let mut ssts = SstVersion::new();
ssts.add_files(
Arc::new(crate::sst::file_purger::NoopFilePurger),
(1..=4).map(|sequence| FileMeta {
file_id: FileId::random(),
time_range: (
Timestamp::new_millisecond(0),
Timestamp::new_millisecond(10),
),
level: 0,
sequence: NonZeroU64::new(sequence),
..Default::default()
}),
);
CompactionRegion {
region_id: metadata.region_id,
region_options: RegionOptions::default(),
engine_config: Arc::new(MitoConfig::default()),
region_metadata: metadata.clone(),
cache_manager: Arc::new(CacheManager::default()),
access_layer: env.access_layer,
manifest_ctx,
current_version: CompactionVersion {
metadata,
options: RegionOptions::default(),
ssts: Arc::new(ssts),
compaction_time_window: None,
},
file_purger: None,
ttl: Some(Duration::from_millis(1).into()),
max_parallelism: 1,
plugins: Plugins::new(),
}
}
#[tokio::test]
async fn test_pick_expired_ssts_without_marking_compacting() {
let picker = TwcsPicker {
trigger_file_num: 4,
time_window_seconds: Some(3),
max_output_file_size: None,
append_mode: false,
max_background_tasks: None,
};
let compaction_region = compaction_region_with_expired_sst().await;
let output = picker.pick(&compaction_region).unwrap();
assert!(output.outputs.is_empty());
assert!(!output.expired_ssts.is_empty());
assert!(output.expired_ssts.iter().all(|file| !file.compacting()));
}
#[test]
fn test_get_latest_window_in_seconds() {
src/mito2/src/compaction/window.rs
+37 -15
View File
@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::BTreeMap;
use std::collections::{BTreeMap, HashSet};
use std::fmt::Debug;
use common_telemetry::info;
@@ -86,9 +86,11 @@ impl WindowedCompactionPicker {
);
if !expired_ssts.is_empty() {
info!("Expired SSTs in region {}: {:?}", region_id, expired_ssts);
// here we mark expired SSTs as compacting to avoid them being picked.
expired_ssts.iter().for_each(|f| f.set_compacting(true));
}
let expired_file_ids = expired_ssts
.iter()
.map(|file| file.file_id())
.collect::<HashSet<_>>();
let windows = assign_files_to_time_windows(
time_window,
@@ -96,7 +98,8 @@ impl WindowedCompactionPicker {
.ssts
.levels()
.iter()
.flat_map(|level| level.files.values()),
.flat_map(|level| level.files.values())
.filter(|file| !expired_file_ids.contains(&file.file_id())),
);
(build_output(windows), expired_ssts, time_window)
@@ -204,17 +207,22 @@ mod tests {
use std::sync::Arc;
use std::time::Duration;
use common_base::Plugins;
use common_time::Timestamp;
use common_time::range::TimestampRange;
use store_api::storage::{FileId, RegionId};
use crate::compaction::compactor::CompactionVersion;
use crate::cache::CacheManager;
use crate::compaction::compactor::{CompactionRegion, CompactionVersion};
use crate::compaction::picker::Picker;
use crate::compaction::window::{WindowedCompactionPicker, file_time_bucket_span};
use crate::config::MitoConfig;
use crate::region::options::RegionOptions;
use crate::sst::file::{FileMeta, Level};
use crate::sst::file_purger::NoopFilePurger;
use crate::sst::version::SstVersion;
use crate::test_util::memtable_util::metadata_for_test;
use crate::test_util::scheduler_util::SchedulerEnv;
fn build_version(
files: &[(FileId, i64, i64, Level)],
@@ -260,19 +268,33 @@ mod tests {
}
}
#[test]
fn test_pick_expired() {
#[tokio::test]
async fn test_pick_expired_ssts_without_marking_compacting() {
let picker = WindowedCompactionPicker::new(None);
let files = vec![(FileId::random(), 0, 10, 0)];
let version = build_version(&files, Some(Duration::from_millis(1)));
let (outputs, expired_ssts, _window) = picker.pick_inner(
RegionId::new(0, 0),
&version,
Timestamp::new_millisecond(12),
);
assert!(outputs.is_empty());
assert_eq!(1, expired_ssts.len());
let env = SchedulerEnv::new().await;
let manifest_ctx = env.mock_manifest_context(version.metadata.clone()).await;
let compaction_region = CompactionRegion {
region_id: version.metadata.region_id,
region_options: RegionOptions::default(),
engine_config: Arc::new(MitoConfig::default()),
region_metadata: version.metadata.clone(),
cache_manager: Arc::new(CacheManager::default()),
access_layer: env.access_layer,
manifest_ctx,
current_version: version,
file_purger: None,
ttl: None,
max_parallelism: 1,
plugins: Plugins::new(),
};
let output = picker.pick(&compaction_region).unwrap();
assert!(output.outputs.is_empty());
assert!(!output.expired_ssts.is_empty());
assert!(output.expired_ssts.iter().all(|file| !file.compacting()));
}
const HOUR: i64 = 60 * 60 * 1000;
src/mito2/src/engine/compaction_test.rs
+694 -40
View File
@@ -17,6 +17,7 @@ use std::ops::Range;
use std::sync::Arc;
use std::time::Duration;
use api::v1::region::{StrictWindow, compact_request};
use api::v1::{ColumnSchema, Rows};
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
@@ -26,8 +27,8 @@ use datatypes::arrow::datatypes::TimestampMillisecondType;
use store_api::region_engine::{RegionEngine, RegionRole};
use store_api::region_request::AlterKind::SetRegionOptions;
use store_api::region_request::{
EnterStagingRequest, PathType, RegionAlterRequest, RegionCompactRequest, RegionDeleteRequest,
RegionFlushRequest, RegionOpenRequest, RegionRequest, SetRegionOption,
EnterStagingRequest, PathType, RegionAlterRequest, RegionCloseRequest, RegionCompactRequest,
RegionDeleteRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest, SetRegionOption,
StagingPartitionDirective,
};
use store_api::storage::{RegionId, ScanRequest};
@@ -35,7 +36,7 @@ use tokio::sync::Notify;
use crate::config::MitoConfig;
use crate::engine::MitoEngine;
use crate::engine::listener::CompactionListener;
use crate::engine::listener::{CompactionListener, CompactionPlanningGate};
use crate::test_util::{
CreateRequestBuilder, TestEnv, build_rows_for_key, column_metadata_to_column_schema, put_rows,
};
@@ -132,6 +133,656 @@ async fn collect_stream_ts(stream: SendableRecordBatchStream) -> Vec<i64> {
res
}
struct CompactionListenerGuard(Option<Arc<CompactionListener>>);
impl CompactionListenerGuard {
fn new(listener: Arc<CompactionListener>) -> Self {
Self(Some(listener))
}
fn release(mut self) {
self.0.take().unwrap().wake();
}
}
impl Drop for CompactionListenerGuard {
fn drop(&mut self) {
if let Some(listener) = self.0.take() {
listener.wake();
}
}
}
#[tokio::test]
async fn test_region_b_progresses_while_same_worker_region_a_is_picking() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::new().await;
let region_a = RegionId::new(1, 1);
let region_b = RegionId::new(2, 1);
let gate = Arc::new(CompactionPlanningGate::new(region_a));
let engine = env
.create_engine_with(
MitoConfig {
num_workers: 1,
min_compaction_interval: Duration::ZERO,
..Default::default()
},
None,
Some(gate.clone()),
None,
)
.await;
for (region_id, table_name) in [(region_a, "region_a"), (region_b, "region_b")] {
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
table_name,
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
engine
.handle_request(
region_id,
RegionRequest::Create(
CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.build(),
),
)
.await
.unwrap();
}
let request = CreateRequestBuilder::new().build();
let column_schemas = request
.column_metadatas
.iter()
.map(column_metadata_to_column_schema)
.collect::<Vec<_>>();
let gate_guard = gate.arm();
let engine_for_compaction = engine.clone();
let region_a_compaction = tokio::spawn(async move {
engine_for_compaction
.handle_request(
region_a,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("region A planning did not reach the gate");
tokio::time::timeout(
Duration::from_secs(5),
put_and_flush(&engine, region_a, &column_schemas, 10..20),
)
.await
.expect("region A automatic compaction trigger did not finish");
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_schedule_attempts(2))
.await
.expect("region A automatic trigger did not reach the scheduler");
assert_eq!(2, gate.schedule_attempt_count());
assert_eq!(1, gate.invocation_count());
let engine_for_region_b = engine.clone();
let mut region_b_work = tokio::spawn(async move {
put_and_flush(&engine_for_region_b, region_b, &column_schemas, 0..10).await;
});
tokio::time::timeout(Duration::from_secs(5), &mut region_b_work)
.await
.expect("region B was blocked by region A compaction planning")
.expect("region B work task panicked");
let followup_guard = gate.arm();
gate_guard.release();
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("coalesced region A trigger did not start its follow-up plan");
tokio::time::timeout(Duration::from_secs(5), region_a_compaction)
.await
.expect("region A compaction task did not finish after gate release")
.expect("region A compaction task panicked")
.expect("region A compaction failed");
assert_eq!(2, gate.invocation_count());
followup_guard.release();
}
#[tokio::test]
async fn test_regular_trigger_while_picking_replans_after_no_plan() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::new().await;
let region_id = RegionId::new(7, 1);
let gate = Arc::new(CompactionPlanningGate::new(region_id));
let engine = env
.create_engine_with(
MitoConfig {
min_compaction_interval: Duration::ZERO,
..Default::default()
},
None,
Some(gate.clone()),
None,
)
.await;
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"replan_after_no_plan",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
let create = CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.build();
let column_schemas = create
.column_metadatas
.iter()
.map(column_metadata_to_column_schema)
.collect::<Vec<_>>();
engine
.handle_request(region_id, RegionRequest::Create(create))
.await
.unwrap();
let first_plan_guard = gate.arm();
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("first automatic compaction did not reach the planning gate");
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_schedule_attempts(2))
.await
.expect("second flush did not trigger automatic compaction");
assert_eq!(1, gate.invocation_count());
let second_plan_guard = gate.arm();
first_plan_guard.release();
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("coalesced regular trigger was lost after the first plan returned no plan");
assert_eq!(2, gate.invocation_count());
second_plan_guard.release();
}
#[tokio::test]
async fn test_regular_trigger_while_picking_replans_after_prepared_execution() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::new().await;
let region_id = RegionId::new(8, 1);
let gate = Arc::new(CompactionPlanningGate::new(region_id));
let engine = env
.create_engine_with(
MitoConfig {
min_compaction_interval: Duration::from_secs(60 * 60),
..Default::default()
},
None,
Some(gate.clone()),
None,
)
.await;
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"replan_after_prepared",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
let create = CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.build();
let column_schemas = create
.column_metadatas
.iter()
.map(column_metadata_to_column_schema)
.collect::<Vec<_>>();
engine
.handle_request(region_id, RegionRequest::Create(create))
.await
.unwrap();
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
let first_plan_guard = gate.arm();
let first_engine = engine.clone();
let first = tokio::spawn(async move {
first_engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("first regular compaction did not reach the planning gate");
let expected_attempts = gate.schedule_attempt_count() + 1;
let second_engine = engine.clone();
let second = tokio::spawn(async move {
second_engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
});
tokio::time::timeout(
Duration::from_secs(5),
gate.wait_until_schedule_attempts(expected_attempts),
)
.await
.expect("second regular compaction did not reach the scheduler");
let commit_guard = gate.arm_commit();
first_plan_guard.release();
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_commit_entered())
.await
.expect("first regular compaction did not produce a prepared execution");
let followup_plan_guard = gate.arm();
commit_guard.release();
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("prepared execution did not immediately admit the retained regular follow-up");
tokio::time::timeout(Duration::from_secs(5), first)
.await
.expect("first regular compaction waiter was not notified")
.expect("first regular compaction task panicked")
.expect("first regular compaction failed");
assert!(!second.is_finished());
followup_plan_guard.release();
tokio::time::timeout(Duration::from_secs(5), second)
.await
.expect("retained regular compaction waiter was not notified")
.expect("retained regular compaction task panicked")
.expect("retained regular compaction failed");
}
#[tokio::test]
async fn test_pending_manual_compaction_finishes_before_queued_ddl() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::new().await;
let region_id = RegionId::new(6, 1);
let gate = Arc::new(CompactionPlanningGate::new(region_id));
let engine = env
.create_engine_with(
MitoConfig {
num_workers: 1,
min_compaction_interval: Duration::from_secs(60 * 60),
..Default::default()
},
None,
Some(gate.clone()),
None,
)
.await;
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"pending_manual_before_ddl",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
let create = CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.build();
let column_schemas = create
.column_metadatas
.iter()
.map(column_metadata_to_column_schema)
.collect::<Vec<_>>();
engine
.handle_request(region_id, RegionRequest::Create(create))
.await
.unwrap();
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
let commit_guard = gate.arm_commit();
let regular_engine = engine.clone();
let regular_task = tokio::spawn(async move {
regular_engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_commit_entered())
.await
.expect("regular compaction did not reach its non-cancellable commit gate");
let manual_plan_guard = gate.arm();
let manual_engine = engine.clone();
let manual_task = tokio::spawn(async move {
manual_engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest {
options: compact_request::Options::StrictWindow(StrictWindow {
window_seconds: 60,
}),
..Default::default()
}),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_schedule_attempts(2))
.await
.expect("manual compaction did not reach the scheduler");
assert!(!manual_task.is_finished());
let ddl_engine = engine.clone();
let ddl_task = tokio::spawn(async move {
ddl_engine
.handle_request(
region_id,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::RejectAllWrites,
}),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_cancel_requested())
.await
.expect("enter-staging DDL was not queued behind regular compaction");
assert!(!ddl_task.is_finished());
commit_guard.release();
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("pending manual compaction was not planned after regular completion");
assert_eq!(2, gate.invocation_count());
assert!(!manual_task.is_finished());
assert!(!ddl_task.is_finished());
let pending_ddl_guard = gate.arm_pending_ddl_dispatch();
manual_plan_guard.release();
tokio::time::timeout(
Duration::from_secs(5),
gate.wait_until_pending_ddl_dispatch(),
)
.await
.expect("manual result was not notified before pending DDL dispatch");
tokio::time::timeout(Duration::from_secs(5), manual_task)
.await
.expect("manual compaction did not notify its waiter")
.expect("manual compaction task panicked")
.expect("manual compaction failed");
assert!(!ddl_task.is_finished());
pending_ddl_guard.release();
tokio::time::timeout(Duration::from_secs(5), regular_task)
.await
.expect("regular compaction did not finish after commit release")
.expect("regular compaction task panicked")
.expect("regular compaction failed");
tokio::time::timeout(Duration::from_secs(5), ddl_task)
.await
.expect("queued DDL did not finish after manual compaction")
.expect("queued DDL task panicked")
.expect("queued DDL failed");
assert!(engine.get_region(region_id).unwrap().is_staging());
}
#[tokio::test]
async fn test_picking_close_reopen_ignores_old_plan() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::new().await;
let region_id = RegionId::new(3, 1);
let gate = Arc::new(CompactionPlanningGate::new(region_id));
let engine = env
.create_engine_with(
MitoConfig {
num_workers: 1,
..Default::default()
},
None,
Some(gate.clone()),
None,
)
.await;
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"close_reopen",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
let create = CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.build();
let table_dir = create.table_dir.clone();
let options = create.options.clone();
engine
.handle_request(region_id, RegionRequest::Create(create))
.await
.unwrap();
let gate_guard = gate.arm();
let compact_engine = engine.clone();
let compact_task = tokio::spawn(async move {
compact_engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("planning did not reach the gate");
engine
.handle_request(
region_id,
RegionRequest::Close(RegionCloseRequest::default()),
)
.await
.unwrap();
let compact_err = tokio::time::timeout(Duration::from_secs(5), compact_task)
.await
.expect("closed region compaction waiter was not released")
.expect("closed region compaction task panicked")
.unwrap_err();
assert_eq!(compact_err.status_code(), StatusCode::Cancelled);
engine
.handle_request(
region_id,
RegionRequest::Open(RegionOpenRequest {
engine: String::new(),
table_dir,
path_type: PathType::Bare,
options,
skip_wal_replay: false,
checkpoint: None,
requirements: Default::default(),
}),
)
.await
.unwrap();
engine
.set_region_role(region_id, RegionRole::Leader)
.unwrap();
gate_guard.release();
tokio::time::timeout(Duration::from_secs(5), compact(&engine, region_id))
.await
.expect("replacement compaction was blocked by the stale plan");
assert!(engine.is_region_exists(region_id));
assert_eq!(2, gate.invocation_count());
}
#[tokio::test]
async fn test_enter_staging_waits_for_picking_logical_cancellation_ack() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::new().await;
let region_id = RegionId::new(4, 1);
let gate = Arc::new(CompactionPlanningGate::new(region_id));
let engine = env
.create_engine_with(
MitoConfig {
num_workers: 1,
..Default::default()
},
None,
Some(gate.clone()),
None,
)
.await;
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"enter_staging",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
engine
.handle_request(
region_id,
RegionRequest::Create(
CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.build(),
),
)
.await
.unwrap();
let gate_guard = gate.arm();
let compact_engine = engine.clone();
let compact_task = tokio::spawn(async move {
compact_engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("planning did not reach the gate");
let staging_engine = engine.clone();
let staging_task = tokio::spawn(async move {
staging_engine
.handle_request(
region_id,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::RejectAllWrites,
}),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_cancel_requested())
.await
.expect("enter-staging did not request picking cancellation");
assert!(!compact_task.is_finished());
assert!(!staging_task.is_finished());
gate_guard.release();
let compact_err = tokio::time::timeout(Duration::from_secs(5), compact_task)
.await
.expect("cancelled compaction waiter was not released")
.expect("cancelled compaction task panicked")
.unwrap_err();
assert_eq!(compact_err.status_code(), StatusCode::Cancelled);
tokio::time::timeout(Duration::from_secs(5), staging_task)
.await
.expect("enter-staging did not finish after cancellation acknowledgment")
.expect("enter-staging task panicked")
.expect("enter-staging request failed");
assert!(engine.get_region(region_id).unwrap().is_staging());
}
#[tokio::test]
async fn test_worker_shutdown_fails_picking_waiter() {
common_telemetry::init_default_ut_logging();
let mut env = TestEnv::new().await;
let region_id = RegionId::new(5, 1);
let gate = Arc::new(CompactionPlanningGate::new(region_id));
let engine = env
.create_engine_with(
MitoConfig {
num_workers: 1,
..Default::default()
},
None,
Some(gate.clone()),
None,
)
.await;
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"worker_shutdown",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
engine
.handle_request(
region_id,
RegionRequest::Create(
CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.build(),
),
)
.await
.unwrap();
let gate_guard = gate.arm();
let compact_engine = engine.clone();
let compact_task = tokio::spawn(async move {
compact_engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
});
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_entered())
.await
.expect("planning did not reach the gate");
tokio::time::timeout(Duration::from_secs(5), engine.stop())
.await
.expect("worker shutdown blocked on picking")
.unwrap();
let compact_err = tokio::time::timeout(Duration::from_secs(5), compact_task)
.await
.expect("worker shutdown did not release the compaction waiter")
.expect("compaction task panicked during worker shutdown")
.unwrap_err();
assert_eq!(compact_err.status_code(), StatusCode::Cancelled);
gate_guard.release();
}
#[tokio::test]
async fn test_compaction_region() {
test_compaction_region_with_format(false).await;
@@ -603,6 +1254,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) {
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
let listener_guard = CompactionListenerGuard::new(listener.clone());
// Flush 2 SSTs for compaction.
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
@@ -615,7 +1267,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) {
.set_region_role(region_id, RegionRole::Follower)
.unwrap();
// Wakes up the listener.
listener.wake();
listener_guard.release();
let notify = Arc::new(Notify::new());
// We already sets max background purges to 1, so we can submit a task to the
@@ -686,28 +1338,28 @@ async fn test_enter_staging_cancels_inflight_local_compaction_before_commit() {
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
let _listener_guard = CompactionListenerGuard::new(listener.clone());
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
listener.wait_handle_finished().await;
tokio::time::timeout(Duration::from_secs(5), listener.wait_handle_finished())
.await
.expect("local compaction did not reach its pre-commit gate");
let engine_cloned = engine.clone();
let enter_staging = tokio::spawn(async move {
engine_cloned
.handle_request(
region_id,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::RejectAllWrites,
}),
)
.await
});
tokio::time::sleep(Duration::from_millis(100)).await;
// The enter staging should finished, and the compaction should be cancelled.
assert!(enter_staging.is_finished());
let _ = enter_staging.await.unwrap().unwrap();
tokio::time::timeout(
Duration::from_secs(5),
engine.handle_request(
region_id,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::RejectAllWrites,
}),
),
)
.await
.expect("enter-staging waited for the blocked local compaction")
.expect("enter-staging request failed");
assert!(engine.get_region(region_id).unwrap().is_staging());
}
#[tokio::test]
@@ -751,6 +1403,7 @@ async fn test_manual_compaction_returns_cancelled_when_enter_staging_cancels_it(
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
let _listener_guard = CompactionListenerGuard::new(listener.clone());
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
@@ -765,28 +1418,29 @@ async fn test_manual_compaction_returns_cancelled_when_enter_staging_cancels_it(
.await
});
listener.wait_handle_finished().await;
tokio::time::timeout(Duration::from_secs(5), listener.wait_handle_finished())
.await
.expect("manual compaction did not reach its pre-commit gate");
let engine_cloned = engine.clone();
let enter_staging = tokio::spawn(async move {
engine_cloned
.handle_request(
region_id,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::RejectAllWrites,
}),
)
.await
});
tokio::time::timeout(
Duration::from_secs(5),
engine.handle_request(
region_id,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::RejectAllWrites,
}),
),
)
.await
.expect("enter-staging waited for the blocked manual compaction")
.expect("enter-staging request failed");
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(compact.is_finished());
assert!(enter_staging.is_finished());
let err = compact.await.unwrap().unwrap_err();
let err = tokio::time::timeout(Duration::from_secs(5), compact)
.await
.expect("cancelled manual compaction waiter was not released")
.expect("manual compaction task panicked")
.unwrap_err();
assert_eq!(err.status_code(), StatusCode::Cancelled);
let _ = enter_staging.await.unwrap();
}
#[tokio::test]
src/mito2/src/engine/listener.rs
+240 -1
View File
@@ -15,7 +15,7 @@
//! Engine event listener for tests.
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;
use async_trait::async_trait;
@@ -71,6 +71,21 @@ pub trait EventListener: Send + Sync {
/// Notifies the listener that the compaction is scheduled.
fn on_compaction_scheduled(&self, _region_id: RegionId) {}
/// Notifies the listener immediately before a worker asks the scheduler for compaction.
fn on_compaction_schedule_attempt(&self, _region_id: RegionId) {}
/// Notifies the listener immediately before compaction planning invokes the picker.
async fn on_compaction_pick_begin(&self, _region_id: RegionId) {}
/// Notifies the listener after local compaction becomes non-cancellable and before commit.
async fn on_compaction_commit_begin(&self, _region_id: RegionId) {}
/// Notifies the listener after compaction results are sent and before pending DDL dispatch.
async fn on_compaction_result_notified(&self, _region_id: RegionId) {}
/// Notifies the listener after compaction cancellation is requested and its DDL is queued.
fn on_compaction_cancel_requested(&self, _region_id: RegionId) {}
/// Notifies the listener that region starts to send a region change result to worker.
async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {}
@@ -89,6 +104,213 @@ pub trait EventListener: Send + Sync {
pub type EventListenerRef = Arc<dyn EventListener>;
/// Test gate that blocks compaction planning for one region before picker execution.
pub struct CompactionPlanningGate {
region_id: RegionId,
armed: AtomicBool,
entered: Notify,
cancel_requested: Notify,
permits: Semaphore,
invocation_count: AtomicUsize,
schedule_attempted: Notify,
schedule_attempt_count: AtomicUsize,
commit_armed: AtomicBool,
commit_entered: Notify,
commit_permits: Semaphore,
pending_ddl_armed: AtomicBool,
pending_ddl_entered: Notify,
pending_ddl_permits: Semaphore,
}
/// Releases an armed [`CompactionPlanningGate`] when a test exits unexpectedly.
pub struct CompactionPlanningGateGuard {
gate: Option<Arc<CompactionPlanningGate>>,
}
impl CompactionPlanningGateGuard {
pub fn release(mut self) {
self.gate.take().unwrap().release();
}
}
impl Drop for CompactionPlanningGateGuard {
fn drop(&mut self) {
if let Some(gate) = self.gate.take() {
gate.release();
}
}
}
/// Releases an armed commit gate when a test exits unexpectedly.
pub struct CompactionCommitGateGuard {
gate: Option<Arc<CompactionPlanningGate>>,
}
impl CompactionCommitGateGuard {
pub fn release(mut self) {
self.gate.take().unwrap().release_commit();
}
}
impl Drop for CompactionCommitGateGuard {
fn drop(&mut self) {
if let Some(gate) = self.gate.take() {
gate.release_commit();
}
}
}
/// Releases an armed pending-DDL dispatch gate when a test exits unexpectedly.
pub struct CompactionPendingDdlGateGuard {
gate: Option<Arc<CompactionPlanningGate>>,
}
impl CompactionPendingDdlGateGuard {
pub fn release(mut self) {
self.gate.take().unwrap().release_pending_ddl_dispatch();
}
}
impl Drop for CompactionPendingDdlGateGuard {
fn drop(&mut self) {
if let Some(gate) = self.gate.take() {
gate.release_pending_ddl_dispatch();
}
}
}
impl CompactionPlanningGate {
pub fn new(region_id: RegionId) -> Self {
Self {
region_id,
armed: AtomicBool::new(false),
entered: Notify::new(),
cancel_requested: Notify::new(),
permits: Semaphore::new(0),
invocation_count: AtomicUsize::new(0),
schedule_attempted: Notify::new(),
schedule_attempt_count: AtomicUsize::new(0),
commit_armed: AtomicBool::new(false),
commit_entered: Notify::new(),
commit_permits: Semaphore::new(0),
pending_ddl_armed: AtomicBool::new(false),
pending_ddl_entered: Notify::new(),
pending_ddl_permits: Semaphore::new(0),
}
}
pub fn arm(self: &Arc<Self>) -> CompactionPlanningGateGuard {
self.armed.store(true, Ordering::Relaxed);
CompactionPlanningGateGuard {
gate: Some(self.clone()),
}
}
pub async fn wait_until_entered(&self) {
self.entered.notified().await;
}
pub async fn wait_until_cancel_requested(&self) {
self.cancel_requested.notified().await;
}
pub async fn wait_until_schedule_attempts(&self, expected: usize) {
while self.schedule_attempt_count() < expected {
self.schedule_attempted.notified().await;
}
}
pub fn schedule_attempt_count(&self) -> usize {
self.schedule_attempt_count.load(Ordering::Relaxed)
}
pub fn arm_commit(self: &Arc<Self>) -> CompactionCommitGateGuard {
self.commit_armed.store(true, Ordering::Relaxed);
CompactionCommitGateGuard {
gate: Some(self.clone()),
}
}
pub async fn wait_until_commit_entered(&self) {
self.commit_entered.notified().await;
}
pub fn arm_pending_ddl_dispatch(self: &Arc<Self>) -> CompactionPendingDdlGateGuard {
self.pending_ddl_armed.store(true, Ordering::Relaxed);
CompactionPendingDdlGateGuard {
gate: Some(self.clone()),
}
}
pub async fn wait_until_pending_ddl_dispatch(&self) {
self.pending_ddl_entered.notified().await;
}
pub fn release(&self) {
self.permits.add_permits(1);
}
pub fn invocation_count(&self) -> usize {
self.invocation_count.load(Ordering::Relaxed)
}
fn release_commit(&self) {
self.commit_permits.add_permits(1);
}
fn release_pending_ddl_dispatch(&self) {
self.pending_ddl_permits.add_permits(1);
}
}
#[async_trait]
impl EventListener for CompactionPlanningGate {
fn on_compaction_schedule_attempt(&self, region_id: RegionId) {
if region_id == self.region_id {
self.schedule_attempt_count.fetch_add(1, Ordering::Relaxed);
self.schedule_attempted.notify_one();
}
}
async fn on_compaction_pick_begin(&self, region_id: RegionId) {
if region_id != self.region_id {
return;
}
self.invocation_count.fetch_add(1, Ordering::Relaxed);
if !self.armed.swap(false, Ordering::Relaxed) {
return;
}
self.entered.notify_one();
self.permits.acquire().await.unwrap().forget();
}
async fn on_compaction_commit_begin(&self, region_id: RegionId) {
if region_id != self.region_id || !self.commit_armed.swap(false, Ordering::Relaxed) {
return;
}
self.commit_entered.notify_one();
self.commit_permits.acquire().await.unwrap().forget();
}
async fn on_compaction_result_notified(&self, region_id: RegionId) {
if region_id != self.region_id || !self.pending_ddl_armed.swap(false, Ordering::Relaxed) {
return;
}
self.pending_ddl_entered.notify_one();
self.pending_ddl_permits.acquire().await.unwrap().forget();
}
fn on_compaction_cancel_requested(&self, region_id: RegionId) {
if region_id == self.region_id {
self.cancel_requested.notify_one();
}
}
}
/// Listener to watch flush events.
#[derive(Default)]
pub struct FlushListener {
@@ -566,3 +788,20 @@ impl EventListener for GateIndexBuildListener {
self.stop_notify.notify_one();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_compaction_planning_gate_counts_every_matching_callback() {
let region_id = RegionId::new(1, 1);
let gate = CompactionPlanningGate::new(region_id);
gate.on_compaction_pick_begin(region_id).await;
gate.on_compaction_pick_begin(RegionId::new(2, 1)).await;
gate.on_compaction_pick_begin(region_id).await;
assert_eq!(2, gate.invocation_count());
}
}
src/mito2/src/request.rs
+15
View File
@@ -46,6 +46,7 @@ use store_api::region_request::{
use store_api::storage::{FileId, RegionId};
use tokio::sync::oneshot::{self, Receiver, Sender};
use crate::compaction::{CompactionExecution, CompactionPickFinished};
use crate::error::{
CompactRegionSnafu, CompactionCancelledSnafu, ConvertColumnDataTypeSnafu, CreateDefaultSnafu,
Error, FillDefaultSnafu, FlushRegionSnafu, InvalidPartitionExprSnafu, InvalidRequestSnafu,
@@ -895,6 +896,8 @@ pub(crate) struct SenderDdlRequest {
/// Notification from a background job.
#[derive(Debug)]
pub(crate) enum BackgroundNotify {
/// Compaction planning has finished.
CompactionPickFinished(CompactionPickFinished),
/// Flush has finished.
FlushFinished(FlushFinished),
/// Flush has failed.
@@ -997,6 +1000,8 @@ pub(crate) struct IndexBuildFailed {
pub(crate) struct CompactionFinished {
/// Region id.
pub(crate) region_id: RegionId,
/// Identity and reservation lease of the accepted execution.
pub(crate) execution: CompactionExecution,
/// Compaction result senders.
pub(crate) senders: Vec<OutputTx>,
/// Start time of compaction task.
@@ -1010,6 +1015,8 @@ pub(crate) struct CompactionFinished {
pub(crate) struct CompactionCancelled {
/// Region id.
pub(crate) region_id: RegionId,
/// Identity and reservation lease of the accepted execution.
pub(crate) execution: CompactionExecution,
/// Waiters to wake once the cancellation has been observed by the worker.
pub(crate) senders: Vec<OutputTx>,
}
@@ -1051,6 +1058,8 @@ impl OnFailure for CompactionFinished {
#[derive(Debug)]
pub(crate) struct CompactionFailed {
pub(crate) region_id: RegionId,
/// Identity and reservation lease of the accepted execution.
pub(crate) execution: CompactionExecution,
/// The error source of the failure.
pub(crate) err: Arc<Error>,
}
@@ -1341,9 +1350,15 @@ mod tests {
#[test]
fn test_compaction_cancelled_sends_cancelled_error() {
let version_control =
Arc::new(crate::test_util::version_util::VersionControlBuilder::new().build());
let (tx, rx) = oneshot::channel();
let request = CompactionCancelled {
region_id: RegionId::new(1, 1),
execution: crate::compaction::CompactionExecution::for_test(
version_control,
crate::compaction::CompactionExecutionKind::Local,
),
senders: vec![OutputTx::new(tx)],
};
src/mito2/src/schedule/remote_job_scheduler.rs
+103 -1
View File
@@ -13,7 +13,7 @@
// limitations under the License.
use std::fmt;
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use common_telemetry::error;
@@ -24,6 +24,7 @@ use store_api::storage::RegionId;
use tokio::sync::mpsc::Sender;
use uuid::Uuid;
use crate::compaction::CompactionExecution;
use crate::compaction::compactor::CompactionRegion;
use crate::compaction::picker::PickerOutput;
use crate::error::{CompactRegionSnafu, Error, ParseJobIdSnafu, Result};
@@ -134,9 +135,20 @@ pub struct CompactionJobResult {
pub(crate) struct DefaultNotifier {
/// The sender to send WorkerRequest to the mito engine. This is used to notify the mito engine when a remote job is completed.
pub(crate) request_sender: Sender<WorkerRequestWithTime>,
execution: Mutex<Option<CompactionExecution>>,
}
impl DefaultNotifier {
pub(crate) fn new(
request_sender: Sender<WorkerRequestWithTime>,
execution: CompactionExecution,
) -> Self {
Self {
request_sender,
execution: Mutex::new(Some(execution)),
}
}
fn on_failure(&self, err: Arc<Error>, region_id: RegionId, mut waiters: Vec<OutputTx>) {
COMPACTION_FAILURE_COUNT.inc();
for waiter in waiters.drain(..) {
@@ -149,12 +161,22 @@ impl DefaultNotifier {
impl Notifier for DefaultNotifier {
async fn notify(&self, result: RemoteJobResult, waiters: Vec<OutputTx>) {
INFLIGHT_COMPACTION_COUNT.dec();
let Some(execution) = self
.execution
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.take()
else {
error!("Remote compaction notifier invoked more than once");
return;
};
match result {
RemoteJobResult::CompactionJobResult(result) => {
let notify = {
match result.region_edit {
Ok(edit) => BackgroundNotify::CompactionFinished(CompactionFinished {
region_id: result.region_id,
execution,
senders: waiters,
start_time: result.start_time,
edit,
@@ -168,6 +190,7 @@ impl Notifier for DefaultNotifier {
self.on_failure(err.clone(), result.region_id, waiters);
BackgroundNotify::CompactionFailed(CompactionFailed {
region_id: result.region_id,
execution,
err,
})
}
@@ -195,6 +218,9 @@ impl Notifier for DefaultNotifier {
#[cfg(test)]
mod tests {
use super::*;
use crate::compaction::{CompactionExecution, CompactionExecutionKind};
use crate::error::InvalidSchedulerStateSnafu;
use crate::test_util::version_util::VersionControlBuilder;
#[test]
fn test_job_id() {
@@ -202,4 +228,80 @@ mod tests {
let job_id = JobId::parse_str(&id).unwrap();
assert_eq!(job_id.to_string(), id);
}
#[tokio::test]
async fn test_default_notifier_carries_remote_execution_on_success() {
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
let execution = CompactionExecution::for_test(
Arc::new(VersionControlBuilder::new().build()),
CompactionExecutionKind::Remote,
);
let expected = execution.clone();
let notifier = DefaultNotifier::new(tx, execution);
let region_id = RegionId::new(1, 1);
notifier
.notify(
RemoteJobResult::CompactionJobResult(CompactionJobResult {
job_id: JobId::parse_str("00000000-0000-0000-0000-000000000003").unwrap(),
region_id,
start_time: Instant::now(),
region_edit: Ok(RegionEdit {
files_to_add: Vec::new(),
files_to_remove: Vec::new(),
timestamp_ms: None,
compaction_time_window: None,
flushed_entry_id: None,
flushed_sequence: None,
committed_sequence: None,
}),
}),
Vec::new(),
)
.await;
let request = rx.recv().await.unwrap();
let WorkerRequest::Background {
notify: BackgroundNotify::CompactionFinished(finished),
..
} = request.request
else {
panic!("expected remote compaction success notification");
};
assert!(finished.execution.matches(&expected));
}
#[tokio::test]
async fn test_default_notifier_carries_remote_execution_on_failure() {
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
let execution = CompactionExecution::for_test(
Arc::new(VersionControlBuilder::new().build()),
CompactionExecutionKind::Remote,
);
let expected = execution.clone();
let notifier = DefaultNotifier::new(tx, execution);
let region_id = RegionId::new(1, 1);
notifier
.notify(
RemoteJobResult::CompactionJobResult(CompactionJobResult {
job_id: JobId::parse_str("00000000-0000-0000-0000-000000000004").unwrap(),
region_id,
start_time: Instant::now(),
region_edit: Err(InvalidSchedulerStateSnafu.build()),
}),
Vec::new(),
)
.await;
let request = rx.recv().await.unwrap();
let WorkerRequest::Background {
notify: BackgroundNotify::CompactionFailed(failed),
..
} = request.request
else {
panic!("expected remote compaction failure notification");
};
assert!(failed.execution.matches(&expected));
}
}
src/mito2/src/sst/file.rs
+21
View File
@@ -538,6 +538,14 @@ impl FileHandle {
self.inner.compacting.store(compacting, Ordering::Relaxed);
}
/// Atomically marks this file as compacting if it is currently available.
pub fn try_set_compacting(&self) -> bool {
self.inner
.compacting
.compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
}
pub fn index_outdated(&self) -> bool {
self.inner.index_outdated.load(Ordering::Relaxed)
}
@@ -830,6 +838,19 @@ mod tests {
}
}
#[test]
fn test_try_set_compacting() {
let file = FileHandle::new(
create_file_meta(FileId::random(), 0),
crate::test_util::new_noop_file_purger(),
);
assert!(file.try_set_compacting());
assert!(!file.try_set_compacting());
file.set_compacting(false);
assert!(file.try_set_compacting());
}
#[test]
fn test_deserialize_file_meta() {
let file_meta = create_file_meta(FileId::random(), 0);
src/mito2/src/sst/version.rs
+48
View File
@@ -45,6 +45,22 @@ impl SstVersion {
&self.levels
}
/// Returns the unique current handle matching the selected file's identity and level.
pub(crate) fn file_for_compaction(&self, selected: &FileHandle) -> Option<&FileHandle> {
let mut files = self
.levels
.iter()
.filter_map(|level| level.files.get(&selected.file_id().file_id()));
let current = files.next()?;
if files.next().is_some()
|| current.file_id() != selected.file_id()
|| current.level() != selected.level()
{
return None;
}
Some(current)
}
/// Add files to the version. If a file with the same `file_id` already exists,
/// it will be overwritten with the new file.
///
@@ -270,6 +286,38 @@ mod tests {
});
}
#[test]
fn test_file_for_compaction_returns_unambiguous_matching_level() {
let purger = new_noop_file_purger();
let file_id = FileId::random();
let file = FileMeta {
file_id,
level: 1,
..Default::default()
};
let selected = FileHandle::new(file.clone(), purger.clone());
let mut version = SstVersion::new();
version.add_files(purger, std::iter::once(file));
assert_eq!(
version
.file_for_compaction(&selected)
.unwrap()
.file_id()
.file_id(),
file_id
);
let missing = FileHandle::new(
FileMeta {
file_id: FileId::random(),
level: 1,
..Default::default()
},
new_noop_file_purger(),
);
assert!(version.file_for_compaction(&missing).is_none());
}
#[test]
fn test_usage_only_counts_owned_files() {
let purger = new_noop_file_purger();
src/mito2/src/worker.rs
+38
View File
@@ -1235,6 +1235,9 @@ impl<S: LogStore> RegionWorkerLoop<S> {
/// Handles region background request
async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) {
match notify {
BackgroundNotify::CompactionPickFinished(req) => {
self.handle_compaction_pick_finished(region_id, req).await
}
BackgroundNotify::FlushFinished(req) => {
self.handle_flush_finished(region_id, req).await
}
@@ -1415,6 +1418,41 @@ impl WorkerListener {
}
}
pub(crate) fn on_compaction_schedule_attempt(&self, _region_id: RegionId) {
#[cfg(any(test, feature = "test"))]
if let Some(listener) = &self.listener {
listener.on_compaction_schedule_attempt(_region_id);
}
}
pub(crate) async fn on_compaction_pick_begin(&self, _region_id: RegionId) {
#[cfg(any(test, feature = "test"))]
if let Some(listener) = &self.listener {
listener.on_compaction_pick_begin(_region_id).await;
}
}
pub(crate) async fn on_compaction_commit_begin(&self, _region_id: RegionId) {
#[cfg(any(test, feature = "test"))]
if let Some(listener) = &self.listener {
listener.on_compaction_commit_begin(_region_id).await;
}
}
pub(crate) async fn on_compaction_result_notified(&self, _region_id: RegionId) {
#[cfg(any(test, feature = "test"))]
if let Some(listener) = &self.listener {
listener.on_compaction_result_notified(_region_id).await;
}
}
pub(crate) fn on_compaction_cancel_requested(&self, _region_id: RegionId) {
#[cfg(any(test, feature = "test"))]
if let Some(listener) = &self.listener {
listener.on_compaction_cancel_requested(_region_id);
}
}
pub(crate) async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {
#[cfg(any(test, feature = "test"))]
if let Some(listener) = &self.listener {
src/mito2/src/worker/handle_compaction.rs
+70 -11
View File
@@ -18,8 +18,9 @@ use store_api::logstore::LogStore;
use store_api::region_request::RegionCompactRequest;
use store_api::storage::RegionId;
use crate::compaction::{CompactionExecution, CompactionPickFinished};
use crate::config::IndexBuildMode;
use crate::error::RegionNotFoundSnafu;
use crate::error::{RegionClosedSnafu, RegionNotFoundSnafu};
use crate::metrics::COMPACTION_REQUEST_COUNT;
use crate::region::MitoRegionRef;
use crate::request::{
@@ -30,6 +31,43 @@ use crate::sst::index::IndexBuildType;
use crate::worker::RegionWorkerLoop;
impl<S> RegionWorkerLoop<S> {
fn is_current_compaction_execution(
&self,
region: &MitoRegionRef,
execution: &CompactionExecution,
) -> bool {
self.compaction_scheduler.is_current_region_execution(
region.region_id,
&region.version_control,
execution,
)
}
pub(crate) async fn handle_compaction_pick_finished(
&mut self,
region_id: RegionId,
request: CompactionPickFinished,
) where
S: LogStore,
{
let Some(region) = self.regions.get_region(region_id) else {
return;
};
let mut pending_ddls = self
.compaction_scheduler
.accept_compaction_pick_finished(
request,
&region.version_control,
&region.manifest_ctx,
self.schema_metadata_manager.clone(),
)
.await;
if !pending_ddls.is_empty() {
self.listener.on_compaction_result_notified(region_id).await;
}
self.handle_ddl_requests(&mut pending_ddls).await;
}
/// Handles compaction request submitted to region worker.
pub(crate) async fn handle_compaction_request(
&mut self,
@@ -42,6 +80,7 @@ impl<S> RegionWorkerLoop<S> {
};
COMPACTION_REQUEST_COUNT.inc();
let parallelism = req.parallelism.unwrap_or(1) as usize;
self.listener.on_compaction_schedule_attempt(region_id);
if let Err(e) = self
.compaction_scheduler
.schedule_compaction(
@@ -80,6 +119,11 @@ impl<S> RegionWorkerLoop<S> {
return;
}
};
if !self.is_current_compaction_execution(&region, &request.execution) {
request.on_failure(RegionClosedSnafu { region_id }.build());
return;
}
let execution = request.execution.clone();
region.version_control.apply_edit(
Some(request.edit.clone()),
@@ -91,6 +135,7 @@ impl<S> RegionWorkerLoop<S> {
// compaction finished.
request.on_success();
self.listener.on_compaction_result_notified(region_id).await;
// In async mode, create indexes after compact if new files are created.
if self.config.index.build_mode == IndexBuildMode::Async
@@ -110,8 +155,9 @@ impl<S> RegionWorkerLoop<S> {
// Schedule next compaction if necessary.
let mut pending_ddls = self
.compaction_scheduler
.on_compaction_finished(
.on_execution_finished(
region_id,
&execution,
&region.manifest_ctx,
self.schema_metadata_manager.clone(),
)
@@ -151,17 +197,22 @@ impl<S> RegionWorkerLoop<S> {
) where
S: LogStore,
{
let execution = request.execution.clone();
let is_current = self
.regions
.get_region(region_id)
.is_some_and(|region| self.is_current_compaction_execution(&region, &execution));
request.on_success();
if !is_current {
return;
}
// Reuse the scheduler's finish path to wake pending DDLs after a cooperative stop.
let mut pending_ddls = match self.regions.get_region(region_id) {
Some(_) => {
self.compaction_scheduler
.on_compaction_cancelled(region_id)
.await
}
None => Vec::new(),
};
let mut pending_ddls = self
.compaction_scheduler
.on_execution_cancelled(region_id, &execution)
.await;
self.handle_ddl_requests(&mut pending_ddls).await;
}
@@ -170,8 +221,14 @@ impl<S> RegionWorkerLoop<S> {
pub(crate) async fn handle_compaction_failure(&mut self, req: CompactionFailed) {
error!(req.err; "Failed to compact region: {}", req.region_id);
let Some(region) = self.regions.get_region(req.region_id) else {
return;
};
if !self.is_current_compaction_execution(&region, &req.execution) {
return;
}
self.compaction_scheduler
.on_compaction_failed(req.region_id, req.err);
.on_execution_failed(req.region_id, &req.execution, req.err);
}
/// Schedule compaction for the region if necessary.
@@ -191,6 +248,8 @@ impl<S> RegionWorkerLoop<S> {
"minimal compaction interval time {:?} has passed, scheduling next compaction",
self.config.min_compaction_interval
);
self.listener
.on_compaction_schedule_attempt(region.region_id);
match self
.compaction_scheduler
.schedule_compaction(
src/mito2/src/worker/handle_enter_staging.rs
+1
View File
@@ -113,6 +113,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
partition_directive,
}),
});
self.listener.on_compaction_cancel_requested(region_id);
return;
}