Skip to main content

mito2/compaction/
picker.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;
16use std::sync::Arc;
17
18use api::v1::region::compact_request;
19use common_time::range::TimestampRange;
20use common_time::{TimeToLive, Timestamp};
21use serde::{Deserialize, Serialize};
22
23use crate::compaction::compactor::CompactionRegion;
24use crate::compaction::last_non_null::LastNonNullPicker;
25use crate::compaction::twcs::TwcsPicker;
26use crate::compaction::window::WindowedCompactionPicker;
27use crate::compaction::{CompactionOutput, SerializedCompactionOutput};
28use crate::error::Result;
29use crate::region::options::{CompactionOptions, MergeMode, RegionOptions};
30use crate::sst::file::{FileHandle, FileMeta};
31use crate::sst::file_purger::FilePurger;
32use crate::sst::version::LevelMeta;
33
34/// Picker picks input SST files for compaction.
35/// Different compaction strategy may implement different pickers.
36#[async_trait::async_trait]
37pub trait Picker: Debug + Send + Sync + 'static {
38    /// Picks input SST files for compaction.
39    async fn pick(&self, compaction_region: &CompactionRegion) -> Result<Option<PickerOutput>>;
40}
41
42/// PickerOutput is the output of a [`Picker`].
43/// It contains the outputs of the compaction and the expired SST files.
44#[derive(Default, Clone, Debug)]
45pub struct PickerOutput {
46    pub outputs: Vec<CompactionOutput>,
47    pub expired_ssts: Vec<FileHandle>,
48    pub time_window_size: i64,
49    /// Max single output file size in bytes.
50    pub max_file_size: Option<usize>,
51}
52
53/// SerializedPickerOutput is a serialized version of PickerOutput by replacing [CompactionOutput] and [FileHandle] with [SerializedCompactionOutput] and [FileMeta].
54#[derive(Default, Clone, Debug, Serialize, Deserialize)]
55pub struct SerializedPickerOutput {
56    pub outputs: Vec<SerializedCompactionOutput>,
57    pub expired_ssts: Vec<FileMeta>,
58    pub time_window_size: i64,
59    pub max_file_size: Option<usize>,
60}
61
62impl From<&PickerOutput> for SerializedPickerOutput {
63    fn from(input: &PickerOutput) -> Self {
64        let outputs = input
65            .outputs
66            .iter()
67            .map(|output| SerializedCompactionOutput {
68                output_level: output.output_level,
69                inputs: output.inputs.iter().map(|s| s.meta_ref().clone()).collect(),
70                filter_deleted: output.filter_deleted,
71                output_time_range: output.output_time_range,
72            })
73            .collect();
74        let expired_ssts = input
75            .expired_ssts
76            .iter()
77            .map(|s| s.meta_ref().clone())
78            .collect();
79        Self {
80            outputs,
81            expired_ssts,
82            time_window_size: input.time_window_size,
83            max_file_size: input.max_file_size,
84        }
85    }
86}
87
88impl PickerOutput {
89    /// Converts a [SerializedPickerOutput] to a [PickerOutput]. File statistics
90    /// retain their original encoding; comparisons supply their own schema context.
91    pub fn from_serialized(
92        input: SerializedPickerOutput,
93        file_purger: Arc<dyn FilePurger>,
94    ) -> Self {
95        let outputs = input
96            .outputs
97            .into_iter()
98            .map(|output| CompactionOutput {
99                output_level: output.output_level,
100                inputs: output
101                    .inputs
102                    .into_iter()
103                    .map(|file_meta| FileHandle::new(file_meta, file_purger.clone()))
104                    .collect(),
105                filter_deleted: output.filter_deleted,
106                output_time_range: output.output_time_range,
107            })
108            .collect();
109
110        let expired_ssts = input
111            .expired_ssts
112            .into_iter()
113            .map(|file_meta| FileHandle::new(file_meta, file_purger.clone()))
114            .collect();
115
116        Self {
117            outputs,
118            expired_ssts,
119            time_window_size: input.time_window_size,
120            max_file_size: input.max_file_size,
121        }
122    }
123}
124
125/// Creates a picker for the request and the region's compaction and merge modes.
126pub fn new_picker(
127    compact_request_options: &compact_request::Options,
128    region_options: &RegionOptions,
129    max_background_tasks: Option<usize>,
130    time_range: Option<TimestampRange>,
131    manual: bool,
132) -> Arc<dyn Picker> {
133    if let compact_request::Options::StrictWindow(window) = compact_request_options {
134        let window = if window.window_seconds == 0 {
135            None
136        } else {
137            Some(window.window_seconds)
138        };
139        // Strict-window outputs already include every overlapping file segment
140        // within each disjoint, half-open output range, including LastNonNull.
141        Arc::new(WindowedCompactionPicker::new(window).with_time_range(time_range)) as Arc<_>
142    } else {
143        let picker = new_twcs_picker(region_options, max_background_tasks, time_range, manual);
144        if region_options.merge_mode() == MergeMode::LastNonNull {
145            // LastNonNull correctness spans windows and levels, so wrap the seed
146            // picker instead of trusting TWCS's independent windows.
147            Arc::new(LastNonNullPicker::new(picker))
148        } else {
149            Arc::new(picker)
150        }
151    }
152}
153
154/// Builds the stateless TWCS strategy used for discovery and legacy picks.
155pub(crate) fn new_twcs_picker(
156    region_options: &RegionOptions,
157    max_background_tasks: Option<usize>,
158    time_range: Option<TimestampRange>,
159    manual: bool,
160) -> TwcsPicker {
161    match &region_options.compaction {
162        CompactionOptions::Twcs(twcs_opts) => TwcsPicker {
163            trigger_file_num: twcs_opts.active_window_trigger_file_num,
164            active_window_l1_merge_trigger: twcs_opts.active_window_l1_merge_trigger,
165            inactive_window_trigger_file_num: twcs_opts.inactive_window_trigger_file_num,
166            inactive_window_l1_merge_trigger: twcs_opts.inactive_window_l1_merge_trigger,
167            time_window_seconds: twcs_opts.time_window_seconds(),
168            max_output_file_size: twcs_opts.max_output_file_size.map(|r| r.as_bytes()),
169            append_mode: region_options.append_mode,
170            merge_mode: region_options.merge_mode(),
171            max_background_tasks,
172            time_range,
173            all_windows_inactive: manual,
174            metrics: Default::default(),
175        },
176    }
177}
178
179/// Finds all expired SSTs across levels.
180pub(super) fn get_expired_ssts(
181    levels: &[LevelMeta],
182    ttl: Option<TimeToLive>,
183    now: Timestamp,
184) -> Vec<FileHandle> {
185    let Some(ttl) = ttl else {
186        return vec![];
187    };
188
189    levels
190        .iter()
191        .flat_map(|l| l.get_expired_files(&now, &ttl).into_iter())
192        .collect()
193}
194
195#[cfg(test)]
196mod tests {
197    use store_api::storage::FileId;
198
199    use super::*;
200    use crate::compaction::test_util::new_file_handle;
201    use crate::test_util::new_noop_file_purger;
202
203    #[test]
204    fn test_picker_output_serialization() {
205        let inputs_file_handle = vec![
206            new_file_handle(FileId::random(), 0, 999, 0),
207            new_file_handle(FileId::random(), 0, 999, 0),
208            new_file_handle(FileId::random(), 0, 999, 0),
209        ];
210        let expired_ssts_file_handle = vec![
211            new_file_handle(FileId::random(), 0, 999, 0),
212            new_file_handle(FileId::random(), 0, 999, 0),
213        ];
214
215        let picker_output = PickerOutput {
216            outputs: vec![
217                CompactionOutput {
218                    output_level: 0,
219                    inputs: inputs_file_handle.clone(),
220                    filter_deleted: false,
221                    output_time_range: None,
222                },
223                CompactionOutput {
224                    output_level: 0,
225                    inputs: inputs_file_handle.clone(),
226                    filter_deleted: false,
227                    output_time_range: None,
228                },
229            ],
230            expired_ssts: expired_ssts_file_handle.clone(),
231            time_window_size: 1000,
232            max_file_size: None,
233        };
234
235        let picker_output_str =
236            serde_json::to_string(&SerializedPickerOutput::from(&picker_output)).unwrap();
237        let serialized_picker_output: SerializedPickerOutput =
238            serde_json::from_str(&picker_output_str).unwrap();
239        let picker_output_from_serialized =
240            PickerOutput::from_serialized(serialized_picker_output, new_noop_file_purger());
241
242        picker_output
243            .expired_ssts
244            .iter()
245            .zip(picker_output_from_serialized.expired_ssts.iter())
246            .for_each(|(expected, actual)| {
247                assert_eq!(expected.meta_ref(), actual.meta_ref());
248            });
249
250        picker_output
251            .outputs
252            .iter()
253            .zip(picker_output_from_serialized.outputs.iter())
254            .for_each(|(expected, actual)| {
255                assert_eq!(expected.output_level, actual.output_level);
256                expected
257                    .inputs
258                    .iter()
259                    .zip(actual.inputs.iter())
260                    .for_each(|(expected, actual)| {
261                        assert_eq!(expected.meta_ref(), actual.meta_ref());
262                    });
263                assert_eq!(expected.filter_deleted, actual.filter_deleted);
264                assert_eq!(expected.output_time_range, actual.output_time_range);
265            });
266    }
267}