Skip to main content

mito2/compaction/scheduler/
planning.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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
58/// Region compaction request.
59pub struct CompactionRequest {
60    /// Private snapshot for automatic regular planning; manual requests carry none.
61    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    /// Sender to send notification to the region worker.
66    pub(crate) request_sender: mpsc::Sender<WorkerRequestWithTime>,
67    /// Start time of compaction task.
68    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
82/// Result returned to the worker after background compaction planning.
83pub(crate) enum CompactionPlanningResult {
84    Prepared(PreparedCompaction),
85    /// Returns automatic history even for an empty plan, to prune stale eligibility.
86    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/// Pure planning completion sent back to the owning region worker.
104#[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    /// Updated automatic snapshot, installed only after the worker accepts the plan.
113    /// Manual planning returns `None` and leaves resident waiting ages unchanged.
114    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
121/// Scheduler policy applied between stateless discovery and semantic closure.
122/// Runs on the planning runtime with a private history snapshot, never on the
123/// worker or through a shared mutex. Stale results cannot mutate resident state.
124async 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(&region.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    // finish() reverses candidates for the compactor's pop order. Keep seed
141    // identities in that same order without retaining them in execution state.
142    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/// Exercises automatic selection and closure from cross-module regression tests.
161#[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    /// Runs the planning future and always sends the planning result back to
209    /// the worker, even if the planning panics.
210    ///
211    /// The worker only leaves the picking phase after it receives the
212    /// `CompactionPickFinished` notification. If a panicked planning task
213    /// swallowed the notification, the region would be stuck in the picking
214    /// phase forever, blocking all future compactions and pending DDLs (e.g.
215    /// entering staging) of the region.
216    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        // The idiomatic way to handle a panic result.
223        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        // Bound accepted work independently of discovery. Waiting compensation
274        // persists across batches, including when automatic execution is serial.
275        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    /// Applies a background planning result to the current compaction lifecycle.
350    ///
351    /// # Returns
352    ///
353    /// Reports an automatic follow-up dispatched after Picking, or DDL requests
354    /// released when Picking terminates without creating an execution. The
355    /// owning worker must dispatch returned DDLs immediately because no
356    /// execution callback will follow.
357    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(&region_id) else {
366            return CompactionTransition::NoAction;
367        };
368
369        if !status.is_picking(finished.plan_id) {
370            return CompactionTransition::NoAction;
371        }
372        // Cancellation during Picking is completed by this notification. No
373        // execution callback will follow, so return DDLs released by the fence.
374        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        // Manual/error results retain the resident history. Automatic snapshots
381        // replace it below only after passing the same lifecycle fence.
382        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, &current.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(&region_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(&region_id) {
425                            status.set_phase(phase);
426                        }
427                        self.dispatch_local_units(region_id);
428                        // The execution now owns the terminal transition. Any
429                        // fenced DDLs remain queued until its callback arrives.
430                        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            // These paths never create an execution. Finish the Picking
449            // lifecycle here and release DDLs if no follow-up was scheduled.
450            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(&region_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(&region_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(&region_id) {
514            status.become_idle();
515        }
516        CompactionTransition::NoAction
517    }
518
519    /// Submits a remote job, or prepares local unit tasks without submitting them.
520    /// The caller dispatches local units after installing the returned phase.
521    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        // If specified to run compaction remotely, we schedule the compaction job remotely.
539        // It will fall back to local compaction if there is no remote job scheduler.
540        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(&region_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                        // An error may be ambiguous after the remote scheduler consumed
586                        // the notifier. Fence a delayed remote callback from the local fallback.
587                        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(&region_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
643/// Rebuilds picker output with current SST handles while preserving the picker's grouping.
644///
645/// Picking runs in background on a version snapshot that may be stale by the
646/// time the plan is accepted: a concurrent flush, compaction or index rebuild
647/// can replace a selected file with a new handle carrying updated metadata
648/// (e.g. `index_version`), or remove the file entirely. The handles in the
649/// picker output therefore cannot be used as-is; re-resolving them against the
650/// current version both detects gone files (aborting the plan) and ensures the
651/// execution reads and reserves the up-to-date handle.
652fn 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}