1use std::fmt;
16use std::future::Future;
17use std::sync::Arc;
18use std::time::Instant;
19
20use api::v1::region::compact_request;
21use common_base::Plugins;
22use common_meta::key::SchemaMetadataManagerRef;
23use common_telemetry::{debug, error, info, warn};
24use common_time::TimeToLive;
25use common_time::range::TimestampRange;
26use futures::FutureExt;
27use snafu::ResultExt;
28use store_api::storage::RegionId;
29use tokio::sync::mpsc::{self, Sender};
30
31use crate::access_layer::AccessLayerRef;
32use crate::cache::CacheManagerRef;
33use crate::compaction::compactor::{CompactionRegion, CompactionVersion};
34use crate::compaction::last_non_null;
35use crate::compaction::picker::{PickerOutput, new_picker, new_twcs_picker};
36use crate::compaction::scheduler::history::{PickHistory, PickKey};
37use crate::compaction::scheduler::state::{
38 CompactingFiles, CompactionExecution, CompactionPhase, CompactionTaskState, CompactionTrigger,
39};
40use crate::compaction::scheduler::{CompactionScheduler, CompactionTransition};
41use crate::compaction::unit::CompactionUnit;
42use crate::compaction::{CompactionOutput, find_dynamic_options, format_panic_reason};
43use crate::config::MitoConfig;
44use crate::error::{CompactRegionSnafu, Error, RemoteCompactionSnafu, Result, UnexpectedSnafu};
45use crate::metrics::{
46 COMPACTION_MEMORY_REJECTED, COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT,
47};
48use crate::region::ManifestContextRef;
49use crate::region::options::{MergeMode, RegionOptions};
50use crate::request::{BackgroundNotify, OutputTx, WorkerRequest, WorkerRequestWithTime};
51use crate::schedule::remote_job_scheduler::{
52 CompactionJob, DefaultNotifier, RemoteJob, RemoteJobSchedulerRef,
53};
54use crate::sst::file::FileHandle;
55use crate::sst::version::SstVersion;
56use crate::worker::WorkerListener;
57
58pub struct CompactionRequest {
60 pub(crate) pick_history: Option<PickHistory>,
62 pub(crate) engine_config: Arc<MitoConfig>,
63 pub(crate) current_version: CompactionVersion,
64 pub(crate) access_layer: AccessLayerRef,
65 pub(crate) request_sender: mpsc::Sender<WorkerRequestWithTime>,
67 pub(crate) start_time: Instant,
69 pub(crate) cache_manager: CacheManagerRef,
70 pub(crate) manifest_ctx: ManifestContextRef,
71 pub(crate) listener: WorkerListener,
72 pub(crate) schema_metadata_manager: SchemaMetadataManagerRef,
73 pub(crate) max_parallelism: usize,
74}
75
76impl CompactionRequest {
77 pub(crate) fn region_id(&self) -> RegionId {
78 self.current_version.metadata.region_id
79 }
80}
81
82pub(crate) enum CompactionPlanningResult {
84 Prepared(PreparedCompaction),
85 NoPlan(Option<PickHistory>),
87 Error(Arc<Error>),
88}
89
90impl fmt::Debug for CompactionPlanningResult {
91 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92 match self {
93 Self::Prepared(prepared) => f
94 .debug_tuple("Prepared")
95 .field(&prepared.compaction_region.region_id)
96 .finish(),
97 Self::NoPlan(_) => f.write_str("NoPlan"),
98 Self::Error(err) => f.debug_tuple("Error").field(err).finish(),
99 }
100 }
101}
102
103#[derive(Debug)]
105pub(crate) struct CompactionPickFinished {
106 pub(crate) region_id: RegionId,
107 pub(crate) plan_id: u64,
108 pub(crate) result: CompactionPlanningResult,
109}
110
111pub(crate) struct PreparedCompaction {
112 pick_history: Option<PickHistory>,
115 pub(super) compaction_region: CompactionRegion,
116 pub(super) picker_output: PickerOutput,
117 start_time: Instant,
118 ttl: TimeToLive,
119}
120
121async fn pick_automatic(
125 region: &CompactionRegion,
126 history: &mut PickHistory,
127 max_outputs: Option<usize>,
128 time_range: Option<TimestampRange>,
129) -> Result<Option<PickerOutput>> {
130 let picker = new_twcs_picker(®ion.region_options, max_outputs, time_range, false);
131 let mut candidates = picker.discover(region).await?;
132 let last_non_null = region.region_options.merge_mode() == MergeMode::LastNonNull;
133 let now = Instant::now();
134 candidates.candidates = history.rank_candidates(candidates.candidates, now);
135 if !last_non_null {
136 candidates
137 .candidates
138 .truncate(max_outputs.unwrap_or(usize::MAX));
139 }
140 let keys: Vec<_> = candidates
143 .candidates
144 .iter()
145 .rev()
146 .map(PickKey::from)
147 .collect();
148 let output = candidates.finish().await?;
149 if last_non_null {
150 let (output, attempted_seeds) =
151 last_non_null::close_outputs(region, output, max_outputs).await?;
152 history.acknowledge(attempted_seeds.into_iter().map(|index| keys[index]), now);
153 Ok(output)
154 } else {
155 history.acknowledge(keys.into_iter(), now);
156 Ok((!output.outputs.is_empty() || !output.expired_ssts.is_empty()).then_some(output))
157 }
158}
159
160#[cfg(test)]
162pub(crate) async fn pick_automatic_for_test(
163 region: &CompactionRegion,
164 history: &mut PickHistory,
165 max_outputs: Option<usize>,
166) -> Result<Option<PickerOutput>> {
167 pick_automatic(region, history, max_outputs, None).await
168}
169
170impl CompactionScheduler {
171 pub(super) fn dispatch_compaction_planning(
172 &mut self,
173 plan_id: u64,
174 request: CompactionRequest,
175 options: compact_request::Options,
176 time_range: Option<TimestampRange>,
177 trigger: CompactionTrigger,
178 ) {
179 if let Some(status) = self.region_status.get_mut(&request.region_id()) {
180 status.active.task_state.get_or_insert_with(|| {
181 CompactionTaskState::new(
182 request.region_id(),
183 &options,
184 request.max_parallelism,
185 time_range,
186 trigger,
187 )
188 });
189 status.publish_state();
190 }
191 let plugins = self.plugins.clone();
192 let max_background_compactions = self.engine_config.max_background_compactions;
193 common_runtime::spawn_compact(async move {
194 let region_id = request.region_id();
195 let request_sender = request.request_sender.clone();
196 let planning = Self::prepare_compaction(
197 request,
198 options,
199 plugins,
200 max_background_compactions,
201 time_range,
202 trigger,
203 );
204 Self::notify_planning_result(region_id, plan_id, request_sender, planning).await;
205 });
206 }
207
208 pub(super) async fn notify_planning_result(
217 region_id: RegionId,
218 plan_id: u64,
219 request_sender: Sender<WorkerRequestWithTime>,
220 planning: impl Future<Output = CompactionPlanningResult> + Send,
221 ) {
222 let result = std::panic::AssertUnwindSafe(planning).catch_unwind().await.unwrap_or_else(|payload| {
224 let reason = format_panic_reason(payload.as_ref());
225 CompactionPlanningResult::Error(Arc::new(
226 UnexpectedSnafu {
227 reason: format!(
228 "Compaction planning panicked for region {region_id}, plan_id {plan_id}: {reason}"
229 ),
230 }
231 .build(),
232 ))
233 });
234 if let CompactionPlanningResult::Error(err) = &result {
235 error!(err; "Compaction planning failed for region {}, plan_id: {}", region_id, plan_id);
236 }
237 let request = WorkerRequestWithTime::new(WorkerRequest::Background {
238 region_id,
239 notify: BackgroundNotify::CompactionPickFinished(CompactionPickFinished {
240 region_id,
241 plan_id,
242 result,
243 }),
244 });
245 if request_sender.send(request).await.is_err() {
246 warn!("Failed to send compaction planning result for region {region_id}");
247 }
248 }
249
250 async fn prepare_compaction(
251 request: CompactionRequest,
252 options: compact_request::Options,
253 plugins: Plugins,
254 max_background_compactions: usize,
255 time_range: Option<TimestampRange>,
256 trigger: CompactionTrigger,
257 ) -> CompactionPlanningResult {
258 let region_id = request.region_id();
259 let (dynamic_compaction_opts, ttl) = find_dynamic_options(
260 region_id,
261 &request.current_version.options,
262 &request.schema_metadata_manager,
263 )
264 .await
265 .unwrap_or_else(|e| {
266 warn!(e; "Failed to find dynamic options for region: {}", region_id);
267 (
268 request.current_version.options.compaction.clone(),
269 request.current_version.options.ttl.unwrap_or_default(),
270 )
271 });
272
273 let max_picker_outputs = max_background_compactions.min(request.max_parallelism.max(1));
276 let region_options = RegionOptions {
277 compaction: dynamic_compaction_opts,
278 ..request.current_version.options.clone()
279 };
280 let region_id = request.region_id();
281 let CompactionRequest {
282 pick_history,
283 engine_config,
284 current_version,
285 access_layer,
286 request_sender: _,
287 start_time,
288 cache_manager,
289 manifest_ctx,
290 listener,
291 schema_metadata_manager: _,
292 max_parallelism,
293 } = request;
294
295 let compaction_region = CompactionRegion {
296 region_id,
297 current_version: current_version.clone(),
298 region_options,
299 engine_config: engine_config.clone(),
300 region_metadata: current_version.metadata.clone(),
301 cache_manager: cache_manager.clone(),
302 access_layer: access_layer.clone(),
303 manifest_ctx: manifest_ctx.clone(),
304 file_purger: None,
305 ttl: Some(ttl),
306 max_parallelism,
307 plugins,
308 };
309
310 listener.on_compaction_pick_begin(region_id).await;
311 let _pick_timer = COMPACTION_STAGE_ELAPSED
312 .with_label_values(&["pick"])
313 .start_timer();
314 let mut history = pick_history;
315 let picked = if let Some(history) = &mut history {
316 pick_automatic(
317 &compaction_region,
318 history,
319 Some(max_picker_outputs),
320 time_range,
321 )
322 .await
323 } else {
324 new_picker(
325 &options,
326 &compaction_region.region_options,
327 Some(max_picker_outputs),
328 time_range,
329 trigger == CompactionTrigger::Manual,
330 )
331 .pick(&compaction_region)
332 .await
333 };
334 let picker_output = match picked {
335 Ok(Some(output)) => output,
336 Ok(None) => return CompactionPlanningResult::NoPlan(history),
337 Err(err) => return CompactionPlanningResult::Error(Arc::new(err)),
338 };
339
340 CompactionPlanningResult::Prepared(PreparedCompaction {
341 pick_history: history,
342 compaction_region,
343 picker_output,
344 start_time,
345 ttl,
346 })
347 }
348
349 pub(super) async fn handle_compaction_pick_finished_inner(
358 &mut self,
359 finished: CompactionPickFinished,
360 manifest_ctx: &ManifestContextRef,
361 schema_metadata_manager: SchemaMetadataManagerRef,
362 ) -> CompactionTransition {
363 let region_id = finished.region_id;
364 let plan_id = finished.plan_id;
365 let Some(status) = self.region_status.get_mut(®ion_id) else {
366 return CompactionTransition::NoAction;
367 };
368
369 if !status.is_picking(finished.plan_id) {
370 return CompactionTransition::NoAction;
371 }
372 if !status.accept_plan(finished.plan_id) {
375 return CompactionTransition::from_pending_ddls(
376 self.finish_compaction_on_cancel(region_id),
377 );
378 }
379
380 match finished.result {
383 CompactionPlanningResult::Prepared(mut prepared) => {
384 if let Some(history) = prepared.pick_history.take() {
385 status.pick_history = history;
386 }
387 let current = status.version_control.current().version;
388 let Some(picker_output) =
389 refresh_picker_output(prepared.picker_output, ¤t.ssts)
390 else {
391 return self
392 .finish_compaction_planning(
393 region_id,
394 None,
395 manifest_ctx,
396 schema_metadata_manager,
397 )
398 .await;
399 };
400 let Some(files) = CompactingFiles::try_new(&picker_output) else {
401 return self
402 .finish_compaction_planning(
403 region_id,
404 None,
405 manifest_ctx,
406 schema_metadata_manager,
407 )
408 .await;
409 };
410 prepared.picker_output = picker_output;
411 let Some(status) = self.region_status.get_mut(®ion_id) else {
412 return CompactionTransition::NoAction;
413 };
414 if let Some(info) = &mut status.active.task_state {
415 info.accept_plan(&prepared.compaction_region, &prepared.picker_output);
416 }
417 status.publish_state();
418 let waiters = status.take_waiters();
419 match self
420 .prepare_or_submit_compaction(prepared, files, waiters, plan_id)
421 .await
422 {
423 Ok(Some(phase)) => {
424 if let Some(status) = self.region_status.get_mut(®ion_id) {
425 status.set_phase(phase);
426 }
427 self.dispatch_local_units(region_id);
428 self.finish_local_units(region_id, manifest_ctx, schema_metadata_manager)
431 .await
432 }
433 Ok(None) => {
434 self.finish_compaction_planning(
435 region_id,
436 None,
437 manifest_ctx,
438 schema_metadata_manager,
439 )
440 .await
441 }
442 Err(err) => {
443 self.fail_compaction(region_id, Arc::new(err));
444 CompactionTransition::NoAction
445 }
446 }
447 }
448 CompactionPlanningResult::NoPlan(history) => {
451 if let Some(history) = history {
452 status.pick_history = history;
453 }
454 self.finish_compaction_planning(
455 region_id,
456 None,
457 manifest_ctx,
458 schema_metadata_manager,
459 )
460 .await
461 }
462 CompactionPlanningResult::Error(err) => {
463 self.finish_compaction_planning(
464 region_id,
465 Some(err),
466 manifest_ctx,
467 schema_metadata_manager,
468 )
469 .await
470 }
471 }
472 }
473
474 async fn finish_compaction_planning(
475 &mut self,
476 region_id: RegionId,
477 err: Option<Arc<Error>>,
478 manifest_ctx: &ManifestContextRef,
479 schema_metadata_manager: SchemaMetadataManagerRef,
480 ) -> CompactionTransition {
481 let Some(status) = self.region_status.get_mut(®ion_id) else {
482 return CompactionTransition::NoAction;
483 };
484 for waiter in status.take_waiters() {
485 if let Some(err) = &err {
486 waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id }));
487 } else {
488 waiter.send(Ok(0));
489 }
490 }
491
492 if let Some(ddls) = self.take_ready_ddls(region_id) {
493 return CompactionTransition::DdlReady(ddls);
494 }
495 if self.handle_pending_compaction_request(
496 region_id,
497 manifest_ctx,
498 schema_metadata_manager.clone(),
499 ) {
500 return CompactionTransition::NoAction;
501 }
502
503 let Some(status) = self.region_status.get_mut(®ion_id) else {
504 return CompactionTransition::NoAction;
505 };
506
507 if status.active.reset_automatic_followup()
508 && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager)
509 {
510 return CompactionTransition::AutomaticFollowupScheduled;
511 }
512
513 if let Some(status) = self.region_status.get_mut(®ion_id) {
514 status.become_idle();
515 }
516 CompactionTransition::NoAction
517 }
518
519 async fn prepare_or_submit_compaction(
522 &mut self,
523 prepared: PreparedCompaction,
524 files: CompactingFiles,
525 waiters: Vec<OutputTx>,
526 plan_id: u64,
527 ) -> Result<Option<CompactionPhase>> {
528 let PreparedCompaction {
529 pick_history: _,
530 compaction_region,
531 picker_output,
532 start_time,
533 ttl,
534 } = prepared;
535 let region_id = compaction_region.region_id;
536 let dynamic_compaction_opts = &compaction_region.region_options.compaction;
537
538 let waiters = if dynamic_compaction_opts.remote_compaction() {
541 if let Some(remote_job_scheduler) = &self.plugins.get::<RemoteJobSchedulerRef>() {
542 let execution = CompactionExecution::new(plan_id, files.clone());
543 let remote_compaction_job = CompactionJob {
544 compaction_region: compaction_region.clone(),
545 picker_output: picker_output.clone(),
546 start_time,
547 waiters,
548 ttl,
549 };
550
551 let result = remote_job_scheduler
552 .schedule(
553 RemoteJob::CompactionJob(remote_compaction_job),
554 Box::new(DefaultNotifier::new(
555 self.request_sender.clone(),
556 execution.clone(),
557 )),
558 )
559 .await;
560
561 match result {
562 Ok(job_id) => {
563 info!(
564 "Scheduled remote compaction job {} for region {}",
565 job_id, region_id
566 );
567 INFLIGHT_COMPACTION_COUNT.inc();
568 return Ok(Some(CompactionPhase::Remote { execution }));
569 }
570 Err(e) => {
571 if !dynamic_compaction_opts.fallback_to_local() {
572 error!(e; "Failed to schedule remote compaction job for region {}", region_id);
573 if let Some(status) = self.region_status.get_mut(®ion_id) {
574 status.extend_waiters(e.waiters);
575 }
576 return RemoteCompactionSnafu {
577 region_id,
578 job_id: None,
579 reason: e.reason,
580 }
581 .fail();
582 }
583
584 error!(e; "Failed to schedule remote compaction job for region {}, fallback to local compaction", region_id);
585 e.waiters
588 }
589 }
590 } else {
591 debug!(
592 "Remote compaction is not enabled, fallback to local compaction for region {}",
593 region_id
594 );
595 waiters
596 }
597 } else {
598 waiters
599 };
600
601 let units = CompactionUnit::from_picker(picker_output);
602 if let Some(status) = self.region_status.get_mut(®ion_id) {
603 status.extend_waiters(waiters);
604 }
605 let Some(mut units) = units else {
606 return UnexpectedSnafu {
607 reason: "Compaction plan expires a merge input",
608 }
609 .fail();
610 };
611 units.retain(|unit| {
612 let estimated_bytes = unit.estimated_memory_bytes();
613 let Some(limit_bytes) = self.exceeds_compaction_memory_limit(estimated_bytes) else { return true };
614 COMPACTION_MEMORY_REJECTED
615 .with_label_values(&["oversized"])
616 .inc();
617 warn!(
618 "Skip compaction for region {} because estimated memory {} bytes exceeds compaction memory limit {} bytes",
619 region_id, estimated_bytes, limit_bytes,
620 );
621 false
622 });
623 if units.is_empty() {
624 return Ok(None);
625 }
626 Ok(Some(self.prepare_local_units(
627 compaction_region,
628 units,
629 files,
630 )))
631 }
632
633 fn exceeds_compaction_memory_limit(&self, estimated_bytes: u64) -> Option<u64> {
634 let limit_bytes = self.memory_manager.limit_bytes();
635 if limit_bytes > 0 && estimated_bytes > limit_bytes {
636 Some(limit_bytes)
637 } else {
638 None
639 }
640 }
641}
642
643fn refresh_picker_output(output: PickerOutput, current: &SstVersion) -> Option<PickerOutput> {
653 let refresh = |file: FileHandle| {
654 current
655 .file_for_compaction(&file)
656 .filter(|current| !current.is_deleted() && !current.compacting())
657 .cloned()
658 };
659 let outputs = output
660 .outputs
661 .into_iter()
662 .map(|output| {
663 let inputs = output
664 .inputs
665 .into_iter()
666 .map(&refresh)
667 .collect::<Option<Vec<_>>>()?;
668 Some(CompactionOutput { inputs, ..output })
669 })
670 .collect::<Option<Vec<_>>>()?;
671 let expired_ssts = output
672 .expired_ssts
673 .into_iter()
674 .map(refresh)
675 .collect::<Option<Vec<_>>>()?;
676
677 Some(PickerOutput {
678 outputs,
679 expired_ssts,
680 time_window_size: output.time_window_size,
681 max_file_size: output.max_file_size,
682 })
683}