mito2/compaction/
task.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::{Debug, Formatter};
16use std::sync::Arc;
17
18use common_base::cancellation::CancellableFuture;
19use common_memory_manager::OnExhaustedPolicy;
20use common_telemetry::error;
21use futures::FutureExt;
22use snafu::ResultExt;
23use tokio::sync::{mpsc, oneshot};
24
25use crate::compaction::compactor::{CompactionContext, DefaultCompactor, MergeOutput};
26use crate::compaction::memory_manager::CompactionMemoryManager;
27use crate::compaction::unit::CompactionUnit;
28use crate::compaction::{CompactionExecution, format_panic_reason};
29use crate::error::{self, CompactionCancelledSnafu, CompactionMemoryExhaustedSnafu, Result};
30use crate::manifest::action::RegionEdit;
31use crate::metrics::{
32    COMPACTION_FAILURE_COUNT, COMPACTION_INPUT_BYTES, COMPACTION_OUTPUT_BYTES,
33    COMPACTION_STAGE_ELAPSED,
34};
35use crate::request::{
36    BackgroundNotify, CompactionUnitNotification, WorkerRequest, WorkerRequestWithTime,
37};
38use crate::schedule::CancellableTaskState;
39use crate::sst::file::UncommittedSsts;
40use crate::worker::WorkerListener;
41
42pub const MAX_PARALLEL_COMPACTION: usize = 1;
43
44/// One local atomic unit. The scheduler owns request waiters and sibling units.
45pub(crate) struct CompactionTask {
46    pub(crate) state: CancellableTaskState,
47    pub(crate) execution: CompactionExecution,
48    pub(crate) compaction_context: CompactionContext,
49    pub(crate) request_sender: mpsc::Sender<WorkerRequestWithTime>,
50    pub(crate) listener: WorkerListener,
51    pub(crate) unit: CompactionUnit,
52    pub(crate) memory_manager: Arc<CompactionMemoryManager>,
53    pub(crate) memory_policy: OnExhaustedPolicy,
54    /// Estimated memory bytes needed for this compaction.
55    pub(crate) estimated_memory_bytes: u64,
56    /// Finalized output SSTs not committed to the manifest yet.
57    pub(crate) uncommitted: UncommittedSsts,
58}
59
60impl Debug for CompactionTask {
61    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
62        f.debug_struct("CompactionTask")
63            .field("region_id", &self.compaction_context.region_id)
64            .field("unit", &self.unit)
65            .field(
66                "append_mode",
67                &self.compaction_context.region_options.append_mode,
68            )
69            .finish()
70    }
71}
72
73impl CompactionTask {
74    /// Executes one unit and retains its resources until worker apply or failure acknowledgement.
75    pub(crate) async fn run(&mut self) {
76        let start = std::time::Instant::now();
77        let result = std::panic::AssertUnwindSafe(self.run_once())
78            .catch_unwind()
79            .await
80            .unwrap_or_else(|payload| Err(self.panic_to_error(payload.as_ref())));
81        self.complete(result, start).await;
82    }
83
84    /// Acquires the unit's memory budget and merges on the compaction runtime.
85    async fn merge(&self) -> Result<MergeOutput> {
86        let region_id = self.compaction_context.region_id;
87        self.listener
88            .on_compaction_unit_merge_begin(region_id, self.execution.plan_id())
89            .await;
90        let bytes = self.estimated_memory_bytes;
91        let _memory_guard = if bytes == 0 {
92            None
93        } else {
94            let _wait_timer = crate::metrics::COMPACTION_MEMORY_WAIT.start_timer();
95            Some(
96                CancellableFuture::new(
97                    self.memory_manager
98                        .acquire_with_policy(bytes, self.memory_policy),
99                    self.state.cancel_handle(),
100                )
101                .await
102                .map_err(|_| CompactionCancelledSnafu.build())?
103                .context(CompactionMemoryExhaustedSnafu {
104                    region_id,
105                    policy: format!("{:?}", self.memory_policy),
106                })?,
107            )
108        };
109        let _timer = COMPACTION_STAGE_ELAPSED
110            .with_label_values(&["merge"])
111            .start_timer();
112        let compactor = DefaultCompactor::with_cancel_handle(
113            self.state.cancel_handle(),
114            self.uncommitted.clone(),
115        );
116        let context = self.compaction_context.clone();
117        let unit = self.unit.clone();
118        let output =
119            common_runtime::spawn_compact(
120                async move { compactor.merge_unit(&context, &unit).await },
121            )
122            .await
123            .context(error::JoinSnafu)??;
124        COMPACTION_INPUT_BYTES.inc_by(output.input_file_size() as f64);
125        COMPACTION_OUTPUT_BYTES.inc_by(output.output_file_size() as f64);
126        self.listener.on_merge_ssts_finished(region_id).await;
127        Ok(output)
128    }
129
130    /// Commits this unit through the existing manifest path, independently of sibling applies.
131    async fn publish(&self, output: MergeOutput) -> Result<RegionEdit> {
132        if !self.state.mark_commit_started() {
133            return CompactionCancelledSnafu.fail();
134        }
135        // Report this unit's SSTs before its manifest update; sibling units may progress independently.
136        self.compaction_context.invoke_sst_hook(&output).await;
137        self.listener
138            .on_compaction_commit_begin(self.compaction_context.region_id)
139            .await;
140        let _timer = COMPACTION_STAGE_ELAPSED
141            .with_label_values(&["write_manifest"])
142            .start_timer();
143        let (edit, _manifest_version) = self.compaction_context.update_manifest(output).await?;
144        self.uncommitted.disarm_cleanup();
145        self.listener
146            .on_compaction_unit_committed(
147                self.compaction_context.region_id,
148                self.execution.plan_id(),
149            )
150            .await;
151        Ok(edit)
152    }
153
154    /// Converts a caught panic into the unit's terminal error.
155    fn panic_to_error(&self, payload: &(dyn std::any::Any + Send + 'static)) -> error::Error {
156        let region_id = self.compaction_context.region_id;
157        let plan_id = self.execution.plan_id();
158        error::UnexpectedSnafu {
159            reason: format!(
160                "Compaction unit panicked for region {region_id}, plan_id {plan_id}: {}",
161                format_panic_reason(payload)
162            ),
163        }
164        .build()
165    }
166
167    /// Runs merge and publication.
168    async fn run_once(&mut self) -> Result<RegionEdit> {
169        let output = self.merge().await?;
170        self.publish(output).await
171    }
172
173    /// Cleans up and reports a terminal result before task exit.
174    async fn complete(&self, result: Result<RegionEdit>, start: std::time::Instant) {
175        if let Err(err) = &result {
176            if !matches!(err, error::Error::CompactionCancelled { .. }) {
177                COMPACTION_FAILURE_COUNT.inc();
178                error!(err; "Compaction unit failed, region: {}, plan_id: {}",
179                    self.compaction_context.region_id, self.execution.plan_id());
180            }
181            if self.state.commit_started() && err.may_have_persisted_manifest_update() {
182                self.uncommitted.disarm_cleanup();
183            } else {
184                self.uncommitted.cleanup().await;
185            }
186        }
187        let succeeded = result.is_ok();
188        let (applied, receiver) = oneshot::channel();
189        let notify = CompactionUnitNotification::Finished {
190            plan_id: self.execution.plan_id(),
191            result: result.map_err(Arc::new),
192            applied,
193        };
194        if let Err(err) = self.notify_result(notify).await {
195            error!(err; "Failed to notify compaction unit completion, region: {}", self.compaction_context.region_id);
196            return;
197        }
198        // This keeps the execution slot and inputs alive through worker apply.
199        if receiver.await.is_ok() && succeeded {
200            crate::metrics::COMPACTION_ELAPSED_TOTAL
201                .with_label_values(&["local"])
202                .observe(start.elapsed().as_secs_f64());
203        }
204    }
205
206    /// Sends a unit lifecycle event to the owning region worker.
207    async fn notify_result(&self, notify: CompactionUnitNotification) -> Result<()> {
208        self.request_sender
209            .send(WorkerRequestWithTime::new(WorkerRequest::Background {
210                region_id: self.compaction_context.region_id,
211                notify: BackgroundNotify::CompactionUnit(notify),
212            }))
213            .await
214            .map_err(|_| error::InvalidSenderSnafu.build())
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use std::time::Duration;
221
222    use store_api::storage::RegionId;
223
224    use super::*;
225    use crate::compaction::compactor::new_test_compaction_region;
226    use crate::compaction::memory_manager::new_compaction_memory_manager;
227    use crate::engine::listener::EventListener;
228    use crate::error::Error;
229
230    struct PanicOnUnitMergeBegin;
231
232    #[async_trait::async_trait]
233    impl EventListener for PanicOnUnitMergeBegin {
234        async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) {
235            panic!("unit merge listener boom");
236        }
237    }
238
239    #[tokio::test]
240    async fn test_compaction_unit_panic_reports_terminal_error() {
241        let region = new_test_compaction_region().await;
242        let (request_sender, mut requests) = mpsc::channel(1);
243        let mut task = CompactionTask {
244            state: CancellableTaskState::new(),
245            execution: CompactionExecution::for_test(7),
246            compaction_context: CompactionContext::from(&region),
247            request_sender,
248            listener: WorkerListener::new(Some(Arc::new(PanicOnUnitMergeBegin))),
249            unit: CompactionUnit {
250                inputs: Vec::new(),
251                outputs: Vec::new(),
252                time_window_size: 0,
253                max_file_size: None,
254            },
255            memory_manager: Arc::new(new_compaction_memory_manager(0)),
256            memory_policy: OnExhaustedPolicy::default(),
257            estimated_memory_bytes: 0,
258            uncommitted: UncommittedSsts::new(region.region_id, region.access_layer.clone(), None),
259        };
260        let task_handle = tokio::spawn(async move { task.run().await });
261
262        let request = tokio::time::timeout(Duration::from_secs(5), requests.recv())
263            .await
264            .unwrap()
265            .unwrap();
266        let WorkerRequest::Background {
267            notify:
268                BackgroundNotify::CompactionUnit(CompactionUnitNotification::Finished {
269                    plan_id,
270                    result,
271                    applied,
272                }),
273            ..
274        } = request.request
275        else {
276            panic!("expected unit terminal notification");
277        };
278        assert_eq!(7, plan_id);
279        let err = result.unwrap_err();
280        assert!(matches!(err.as_ref(), Error::Unexpected { .. }));
281        assert!(err.to_string().contains("unit merge listener boom"));
282        applied.send(()).unwrap();
283        task_handle.await.unwrap();
284    }
285}
mito2/compaction/task.rs

mito2/compaction/
task.rs
