1use std::collections::BTreeMap;
16use std::fmt::Debug;
17
18use common_telemetry::info;
19use common_time::range::TimestampRange;
20use common_time::timestamp::TimeUnit;
21use common_time::timestamp_millis::BucketAligned;
22use common_time::Timestamp;
23use store_api::storage::RegionId;
24
25use crate::compaction::buckets::infer_time_bucket;
26use crate::compaction::compactor::{CompactionRegion, CompactionVersion};
27use crate::compaction::picker::{Picker, PickerOutput};
28use crate::compaction::{get_expired_ssts, CompactionOutput};
29use crate::sst::file::FileHandle;
30
31#[derive(Debug)]
35pub struct WindowedCompactionPicker {
36 compaction_time_window_seconds: Option<i64>,
37}
38
39impl WindowedCompactionPicker {
40 pub fn new(window_seconds: Option<i64>) -> Self {
41 Self {
42 compaction_time_window_seconds: window_seconds,
43 }
44 }
45
46 fn calculate_time_window(
51 &self,
52 region_id: RegionId,
53 current_version: &CompactionVersion,
54 ) -> i64 {
55 self.compaction_time_window_seconds
56 .or(current_version
57 .compaction_time_window
58 .map(|t| t.as_secs() as i64))
59 .unwrap_or_else(|| {
60 let levels = current_version.ssts.levels();
61 let inferred = infer_time_bucket(levels[0].files());
62 info!(
63 "Compaction window for region {} is not present, inferring from files: {:?}",
64 region_id, inferred
65 );
66 inferred
67 })
68 }
69
70 fn pick_inner(
71 &self,
72 region_id: RegionId,
73 current_version: &CompactionVersion,
74 current_time: Timestamp,
75 ) -> (Vec<CompactionOutput>, Vec<FileHandle>, i64) {
76 let time_window = self.calculate_time_window(region_id, current_version);
77 info!(
78 "Compaction window for region: {} is {} seconds",
79 region_id, time_window
80 );
81
82 let expired_ssts = get_expired_ssts(
83 current_version.ssts.levels(),
84 current_version.options.ttl,
85 current_time,
86 );
87 if !expired_ssts.is_empty() {
88 info!("Expired SSTs in region {}: {:?}", region_id, expired_ssts);
89 expired_ssts.iter().for_each(|f| f.set_compacting(true));
91 }
92
93 let windows = assign_files_to_time_windows(
94 time_window,
95 current_version
96 .ssts
97 .levels()
98 .iter()
99 .flat_map(|level| level.files.values()),
100 );
101
102 (build_output(windows), expired_ssts, time_window)
103 }
104}
105
106impl Picker for WindowedCompactionPicker {
107 fn pick(&self, compaction_region: &CompactionRegion) -> Option<PickerOutput> {
108 let (outputs, expired_ssts, time_window) = self.pick_inner(
109 compaction_region.current_version.metadata.region_id,
110 &compaction_region.current_version,
111 Timestamp::current_millis(),
112 );
113
114 Some(PickerOutput {
115 outputs,
116 expired_ssts,
117 time_window_size: time_window,
118 })
119 }
120}
121
122fn build_output(windows: BTreeMap<i64, (i64, Vec<FileHandle>)>) -> Vec<CompactionOutput> {
123 let mut outputs = Vec::with_capacity(windows.len());
124 for (lower_bound, (upper_bound, files)) in windows {
125 let output_time_range = Some(
127 TimestampRange::new(
128 Timestamp::new_second(lower_bound),
129 Timestamp::new_second(upper_bound),
130 )
131 .unwrap(),
132 );
133
134 let output = CompactionOutput {
135 output_level: 1,
136 inputs: files,
137 filter_deleted: false,
138 output_time_range,
139 };
140 outputs.push(output);
141 }
142
143 outputs
144}
145
146fn assign_files_to_time_windows<'a>(
150 bucket_sec: i64,
151 files: impl Iterator<Item = &'a FileHandle>,
152) -> BTreeMap<i64, (i64, Vec<FileHandle>)> {
153 let mut buckets = BTreeMap::new();
154
155 for file in files {
156 if file.compacting() {
157 continue;
158 }
159 let (start, end) = file.time_range();
160 let bounds = file_time_bucket_span(
161 start.convert_to(TimeUnit::Second).unwrap().value(),
163 end.convert_to(TimeUnit::Second).unwrap().value(),
164 bucket_sec,
165 );
166 for (lower_bound, upper_bound) in bounds {
167 let (_, files) = buckets
168 .entry(lower_bound)
169 .or_insert_with(|| (upper_bound, Vec::new()));
170 files.push(file.clone());
171 }
172 }
173 buckets
174}
175
176fn file_time_bucket_span(start_sec: i64, end_sec: i64, bucket_sec: i64) -> Vec<(i64, i64)> {
178 assert!(start_sec <= end_sec);
179
180 let mut start_aligned = start_sec.align_by_bucket(bucket_sec).unwrap_or(i64::MIN);
183 let end_aligned = end_sec
184 .align_by_bucket(bucket_sec)
185 .unwrap_or(start_aligned + (end_sec - start_sec));
186
187 let mut res = Vec::with_capacity(((end_aligned - start_aligned) / bucket_sec + 1) as usize);
188 while start_aligned <= end_aligned {
189 let window_size = if start_aligned % bucket_sec == 0 {
190 bucket_sec
191 } else {
192 (start_aligned % bucket_sec).abs()
193 };
194 let upper_bound = start_aligned.checked_add(window_size).unwrap_or(i64::MAX);
195 res.push((start_aligned, upper_bound));
196 start_aligned = upper_bound;
197 }
198 res
199}
200
201#[cfg(test)]
202mod tests {
203 use std::sync::Arc;
204 use std::time::Duration;
205
206 use common_time::range::TimestampRange;
207 use common_time::Timestamp;
208 use store_api::storage::RegionId;
209
210 use crate::compaction::compactor::CompactionVersion;
211 use crate::compaction::window::{file_time_bucket_span, WindowedCompactionPicker};
212 use crate::region::options::RegionOptions;
213 use crate::sst::file::{FileId, FileMeta, Level};
214 use crate::sst::version::SstVersion;
215 use crate::test_util::memtable_util::metadata_for_test;
216 use crate::test_util::NoopFilePurger;
217
218 fn build_version(
219 files: &[(FileId, i64, i64, Level)],
220 ttl: Option<Duration>,
221 ) -> CompactionVersion {
222 let metadata = metadata_for_test();
223 let file_purger_ref = Arc::new(NoopFilePurger);
224
225 let mut ssts = SstVersion::new();
226
227 ssts.add_files(
228 file_purger_ref,
229 files.iter().map(|(file_id, start, end, level)| FileMeta {
230 file_id: *file_id,
231 time_range: (
232 Timestamp::new_millisecond(*start),
233 Timestamp::new_millisecond(*end),
234 ),
235 level: *level,
236 ..Default::default()
237 }),
238 );
239
240 CompactionVersion {
241 metadata,
242 ssts: Arc::new(ssts),
243 options: RegionOptions {
244 ttl: ttl.map(|t| t.into()),
245 compaction: Default::default(),
246 storage: None,
247 append_mode: false,
248 wal_options: Default::default(),
249 index_options: Default::default(),
250 memtable: None,
251 merge_mode: None,
252 },
253 compaction_time_window: None,
254 }
255 }
256
257 #[test]
258 fn test_pick_expired() {
259 let picker = WindowedCompactionPicker::new(None);
260 let files = vec![(FileId::random(), 0, 10, 0)];
261
262 let version = build_version(&files, Some(Duration::from_millis(1)));
263 let (outputs, expired_ssts, _window) = picker.pick_inner(
264 RegionId::new(0, 0),
265 &version,
266 Timestamp::new_millisecond(12),
267 );
268 assert!(outputs.is_empty());
269 assert_eq!(1, expired_ssts.len());
270 }
271
272 const HOUR: i64 = 60 * 60 * 1000;
273
274 #[test]
275 fn test_infer_window() {
276 let picker = WindowedCompactionPicker::new(None);
277
278 let files = vec![
279 (FileId::random(), 0, HOUR, 0),
280 (FileId::random(), HOUR, HOUR * 2 - 1, 0),
281 ];
282
283 let version = build_version(&files, Some(Duration::from_millis(3 * HOUR as u64)));
284
285 let (outputs, expired_ssts, window_seconds) = picker.pick_inner(
286 RegionId::new(0, 0),
287 &version,
288 Timestamp::new_millisecond(HOUR * 2),
289 );
290 assert!(expired_ssts.is_empty());
291 assert_eq!(2 * HOUR / 1000, window_seconds);
292 assert_eq!(1, outputs.len());
293 assert_eq!(2, outputs[0].inputs.len());
294 }
295
296 #[test]
297 fn test_assign_files_to_windows() {
298 let picker = WindowedCompactionPicker::new(Some(HOUR / 1000));
299 let files = vec![
300 (FileId::random(), 0, 2 * HOUR - 1, 0),
301 (FileId::random(), HOUR, HOUR * 3 - 1, 0),
302 ];
303 let version = build_version(&files, Some(Duration::from_millis(3 * HOUR as u64)));
304 let (outputs, expired_ssts, window_seconds) = picker.pick_inner(
305 RegionId::new(0, 0),
306 &version,
307 Timestamp::new_millisecond(HOUR * 3),
308 );
309
310 assert!(expired_ssts.is_empty());
311 assert_eq!(HOUR / 1000, window_seconds);
312 assert_eq!(3, outputs.len());
313
314 assert_eq!(1, outputs[0].inputs.len());
315 assert_eq!(files[0].0, outputs[0].inputs[0].file_id());
316 assert_eq!(
317 TimestampRange::new(
318 Timestamp::new_millisecond(0),
319 Timestamp::new_millisecond(HOUR)
320 ),
321 outputs[0].output_time_range
322 );
323
324 assert_eq!(2, outputs[1].inputs.len());
325 assert_eq!(
326 TimestampRange::new(
327 Timestamp::new_millisecond(HOUR),
328 Timestamp::new_millisecond(2 * HOUR)
329 ),
330 outputs[1].output_time_range
331 );
332
333 assert_eq!(1, outputs[2].inputs.len());
334 assert_eq!(files[1].0, outputs[2].inputs[0].file_id());
335 assert_eq!(
336 TimestampRange::new(
337 Timestamp::new_millisecond(2 * HOUR),
338 Timestamp::new_millisecond(3 * HOUR)
339 ),
340 outputs[2].output_time_range
341 );
342 }
343
344 #[test]
345 fn test_assign_compacting_files_to_windows() {
346 let picker = WindowedCompactionPicker::new(Some(HOUR / 1000));
347 let files = vec![
348 (FileId::random(), 0, 2 * HOUR - 1, 0),
349 (FileId::random(), HOUR, HOUR * 3 - 1, 0),
350 ];
351 let version = build_version(&files, Some(Duration::from_millis(3 * HOUR as u64)));
352 version.ssts.levels()[0]
353 .files()
354 .for_each(|f| f.set_compacting(true));
355 let (outputs, expired_ssts, window_seconds) = picker.pick_inner(
356 RegionId::new(0, 0),
357 &version,
358 Timestamp::new_millisecond(HOUR * 3),
359 );
360
361 assert!(expired_ssts.is_empty());
362 assert_eq!(HOUR / 1000, window_seconds);
363 assert!(outputs.is_empty());
364 }
365
366 #[test]
367 fn test_file_time_bucket_span() {
368 assert_eq!(
369 vec![(i64::MIN, i64::MIN + 8),],
370 file_time_bucket_span(i64::MIN, i64::MIN + 1, 10)
371 );
372
373 assert_eq!(
374 vec![(i64::MIN, i64::MIN + 8), (i64::MIN + 8, i64::MIN + 18)],
375 file_time_bucket_span(i64::MIN, i64::MIN + 8, 10)
376 );
377
378 assert_eq!(
379 vec![
380 (i64::MIN, i64::MIN + 8),
381 (i64::MIN + 8, i64::MIN + 18),
382 (i64::MIN + 18, i64::MIN + 28)
383 ],
384 file_time_bucket_span(i64::MIN, i64::MIN + 20, 10)
385 );
386
387 assert_eq!(
388 vec![(-10, 0), (0, 10), (10, 20)],
389 file_time_bucket_span(-1, 11, 10)
390 );
391
392 assert_eq!(
393 vec![(-3, 0), (0, 3), (3, 6)],
394 file_time_bucket_span(-1, 3, 3)
395 );
396
397 assert_eq!(vec![(0, 10)], file_time_bucket_span(0, 9, 10));
398
399 assert_eq!(
400 vec![(i64::MAX - (i64::MAX % 10), i64::MAX)],
401 file_time_bucket_span(i64::MAX - 1, i64::MAX, 10)
402 );
403 }
404}