1use std::collections::BTreeMap;
22use std::path::PathBuf;
23use std::sync::atomic::{AtomicU64, Ordering};
24
25use serde::{Deserialize, Serialize};
26
27use crate::index::SegmentCatalogEntry;
28use chronix_core::{SegmentFile, ShardId};
29
30pub use chronix_core::SegmentState;
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
35pub enum CompactionLevel {
36 L0 = 0,
38 L1 = 1,
40 L2 = 2,
42}
43
44impl CompactionLevel {
45 #[must_use]
47 pub fn next(self) -> Self {
48 match self {
49 Self::L0 => Self::L1,
50 Self::L1 | Self::L2 => Self::L2,
51 }
52 }
53}
54
55impl std::fmt::Display for CompactionLevel {
56 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57 match self {
58 Self::L0 => write!(f, "L0"),
59 Self::L1 => write!(f, "L1"),
60 Self::L2 => write!(f, "L2"),
61 }
62 }
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
75pub enum SizeTier {
76 Tiny,
78 Small,
80 Medium,
82 Large,
84}
85
86impl SizeTier {
87 #[must_use]
89 pub fn classify(byte_size: u64) -> Self {
90 match byte_size {
91 0..1_048_576 => Self::Tiny, 1_048_576..10_485_760 => Self::Small, 10_485_760..104_857_600 => Self::Medium, _ => Self::Large, }
96 }
97}
98
99impl std::fmt::Display for SizeTier {
100 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101 match self {
102 Self::Tiny => write!(f, "Tiny(<1MB)"),
103 Self::Small => write!(f, "Small(1-10MB)"),
104 Self::Medium => write!(f, "Medium(10-100MB)"),
105 Self::Large => write!(f, "Large(≥100MB)"),
106 }
107 }
108}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
116pub enum CompactionStrategy {
117 Twcs,
119 SizeTiered,
121 #[default]
123 Hybrid,
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct CompactionPolicy {
133 pub strategy: CompactionStrategy,
135 pub count_trigger_threshold: usize,
137 pub size_tier_threshold: usize,
139 pub max_write_amplification: f64,
143}
144
145impl Default for CompactionPolicy {
146 fn default() -> Self {
147 Self {
148 strategy: CompactionStrategy::default(),
149 count_trigger_threshold: 4,
150 size_tier_threshold: 4,
151 max_write_amplification: 10.0,
152 }
153 }
154}
155
156impl CompactionPolicy {
157 #[must_use]
162 pub fn validated(mut self) -> Self {
163 self.count_trigger_threshold = self.count_trigger_threshold.clamp(2, 1000);
164 self.size_tier_threshold = self.size_tier_threshold.clamp(2, 1000);
165 self.max_write_amplification = self.max_write_amplification.clamp(1.0, 1000.0);
166 self
167 }
168}
169
170#[derive(Debug)]
180pub struct WriteAmpTracker {
181 total_input_bytes: AtomicU64,
183 total_output_bytes: AtomicU64,
185}
186
187impl Default for WriteAmpTracker {
188 fn default() -> Self {
189 Self::new()
190 }
191}
192
193impl WriteAmpTracker {
194 #[must_use]
196 pub fn new() -> Self {
197 Self {
198 total_input_bytes: AtomicU64::new(0),
199 total_output_bytes: AtomicU64::new(0),
200 }
201 }
202
203 pub fn record(&self, input_bytes: u64, output_bytes: u64) {
205 self.total_input_bytes
206 .fetch_add(input_bytes, Ordering::Relaxed);
207 self.total_output_bytes
208 .fetch_add(output_bytes, Ordering::Relaxed);
209 }
210
211 #[must_use]
215 pub fn write_amplification(&self) -> f64 {
216 let input = self.total_input_bytes.load(Ordering::Relaxed);
217 let output = self.total_output_bytes.load(Ordering::Relaxed);
218 if input == 0 {
219 return 0.0;
220 }
221 output as f64 / input as f64
222 }
223
224 #[must_use]
226 pub fn budget_exhausted(&self, max_ratio: f64) -> bool {
227 self.write_amplification() > max_ratio
228 }
229
230 #[must_use]
232 pub fn input_bytes(&self) -> u64 {
233 self.total_input_bytes.load(Ordering::Relaxed)
234 }
235
236 #[must_use]
238 pub fn output_bytes(&self) -> u64 {
239 self.total_output_bytes.load(Ordering::Relaxed)
240 }
241
242 pub fn reset(&self) {
244 self.total_input_bytes.store(0, Ordering::Relaxed);
245 self.total_output_bytes.store(0, Ordering::Relaxed);
246 }
247}
248
249#[derive(Debug, Clone)]
251pub struct CompactionTask {
252 pub shard_id: ShardId,
254 pub input_segments: Vec<SegmentCatalogEntry>,
256 pub target_level: CompactionLevel,
258 pub source_level: CompactionLevel,
260 pub segments_dir: PathBuf,
265 pub output: SegmentFile,
267}
268
269impl CompactionTask {
270 #[must_use]
272 pub fn total_input_bytes(&self) -> u64 {
273 self.input_segments.iter().map(|s| s.byte_size).sum()
274 }
275
276 #[must_use]
278 pub fn output_path(&self) -> PathBuf {
279 self.output.resolve(&self.segments_dir)
280 }
281}
282
283#[derive(Debug, Clone, Default)]
299pub struct CompactionPicker {
300 policy: CompactionPolicy,
302}
303
304impl CompactionPicker {
305 #[must_use]
309 pub fn new(trigger_threshold: usize) -> Self {
310 Self {
311 policy: CompactionPolicy {
312 count_trigger_threshold: trigger_threshold,
313 size_tier_threshold: trigger_threshold,
314 ..CompactionPolicy::default()
315 },
316 }
317 }
318
319 #[must_use]
321 pub fn with_policy(policy: CompactionPolicy) -> Self {
322 Self { policy }
323 }
324
325 #[must_use]
327 pub fn policy(&self) -> &CompactionPolicy {
328 &self.policy
329 }
330
331 #[must_use]
333 pub fn trigger_threshold(&self) -> usize {
334 self.policy.count_trigger_threshold
335 }
336
337 #[must_use]
345 pub fn pick(
346 &self,
347 segments: &[SegmentCatalogEntry],
348 output_dir: &std::path::Path,
349 ) -> Vec<CompactionTask> {
350 self.pick_level(segments, output_dir, CompactionLevel::L0)
351 }
352
353 #[must_use]
357 pub fn pick_level(
358 &self,
359 segments: &[SegmentCatalogEntry],
360 output_dir: &std::path::Path,
361 source_level: CompactionLevel,
362 ) -> Vec<CompactionTask> {
363 let target_level = source_level.next();
364
365 let mut groups: BTreeMap<(ShardId, &str), Vec<&SegmentCatalogEntry>> = BTreeMap::new();
367 for seg in segments {
368 groups
369 .entry((seg.shard_id, seg.measurement.as_str()))
370 .or_default()
371 .push(seg);
372 }
373
374 let mut tasks = Vec::new();
375
376 match self.policy.strategy {
377 CompactionStrategy::Twcs => {
378 self.pick_twcs(&groups, output_dir, source_level, target_level, &mut tasks);
379 }
380 CompactionStrategy::SizeTiered => {
381 self.pick_size_tiered(&groups, output_dir, source_level, target_level, &mut tasks);
382 }
383 CompactionStrategy::Hybrid => {
384 let twcs_shards =
386 self.pick_twcs(&groups, output_dir, source_level, target_level, &mut tasks);
387 let remaining: BTreeMap<(ShardId, &str), Vec<&SegmentCatalogEntry>> = groups
389 .into_iter()
390 .filter(|(key, _)| !twcs_shards.contains(key))
391 .collect();
392 self.pick_size_tiered(
393 &remaining,
394 output_dir,
395 source_level,
396 target_level,
397 &mut tasks,
398 );
399 }
400 }
401
402 tasks
403 }
404
405 fn pick_twcs<'a>(
409 &self,
410 groups: &BTreeMap<(ShardId, &'a str), Vec<&SegmentCatalogEntry>>,
411 output_dir: &std::path::Path,
412 source_level: CompactionLevel,
413 target_level: CompactionLevel,
414 tasks: &mut Vec<CompactionTask>,
415 ) -> std::collections::BTreeSet<(ShardId, &'a str)> {
416 let mut selected = std::collections::BTreeSet::new();
417
418 for ((shard_id, measurement), group) in groups {
419 if group.len() >= self.policy.count_trigger_threshold {
420 let output = Self::output_file(*shard_id, measurement, group);
421 tasks.push(CompactionTask {
422 shard_id: *shard_id,
423 input_segments: group.iter().map(|s| (*s).clone()).collect(),
424 source_level,
425 target_level,
426 segments_dir: output_dir.to_path_buf(),
427 output,
428 });
429 selected.insert((*shard_id, *measurement));
430 }
431 }
432
433 selected
434 }
435
436 fn pick_size_tiered(
438 &self,
439 groups: &BTreeMap<(ShardId, &str), Vec<&SegmentCatalogEntry>>,
440 output_dir: &std::path::Path,
441 source_level: CompactionLevel,
442 target_level: CompactionLevel,
443 tasks: &mut Vec<CompactionTask>,
444 ) {
445 for ((shard_id, measurement), group) in groups {
446 let mut tiers: BTreeMap<SizeTier, Vec<&SegmentCatalogEntry>> = BTreeMap::new();
448 for seg in group {
449 let tier = SizeTier::classify(seg.byte_size);
450 tiers.entry(tier).or_default().push(seg);
451 }
452
453 for tier_segments in tiers.values() {
454 if tier_segments.len() >= self.policy.size_tier_threshold {
455 let output = Self::output_file(*shard_id, measurement, tier_segments);
456 tasks.push(CompactionTask {
457 shard_id: *shard_id,
458 input_segments: tier_segments.iter().map(|s| (*s).clone()).collect(),
459 source_level,
460 target_level,
461 segments_dir: output_dir.to_path_buf(),
462 output,
463 });
464 }
465 }
466 }
467 }
468
469 fn output_file(
477 shard_id: ShardId,
478 measurement: &str,
479 segments: &[&SegmentCatalogEntry],
480 ) -> SegmentFile {
481 let max_id = segments.iter().map(|s| s.segment_id.0).max().unwrap_or(0);
482 let name = format!("{measurement}_compacted_{max_id}.csx");
483 SegmentFile::new(shard_id, &name).unwrap_or_else(|_| {
484 SegmentFile::new(shard_id, &format!("compacted_{max_id}.csx"))
485 .expect("a name built from digits is always a valid segment file")
486 })
487 }
488
489 #[must_use]
495 pub fn backpressure_delay_ms(&self, l0_count: usize) -> u64 {
496 let heavy_threshold = self.policy.count_trigger_threshold * 4;
497 if l0_count <= heavy_threshold {
498 return 0;
499 }
500 let excess = l0_count - heavy_threshold;
501 (excess as u64 * 10).min(500)
503 }
504}
505
506#[cfg(test)]
507mod tests {
508 use super::*;
509 use chronix_core::SegmentId;
510 use std::path::PathBuf;
511
512 fn test_entry(id: u64, shard: i64, measurement: &str) -> SegmentCatalogEntry {
513 test_entry_sized(id, shard, measurement, 4096)
514 }
515
516 fn test_entry_sized(id: u64, shard: i64, measurement: &str, size: u64) -> SegmentCatalogEntry {
517 SegmentCatalogEntry {
518 segment_id: SegmentId(id),
519 shard_id: ShardId(shard),
520 measurement: measurement.to_string(),
521 file: SegmentFile::new(ShardId(shard), &format!("seg_{id}.csx")).unwrap(),
522 min_timestamp: id as i64 * 1000,
523 max_timestamp: id as i64 * 1000 + 999,
524 row_count: 1000,
525 series_count: 10,
526 byte_size: size,
527 row_group_count: 1,
528 column_count: 5,
529 column_stats: Vec::new(),
530 state: SegmentState::default(),
531 }
532 }
533
534 #[test]
537 fn picker_no_compaction_below_threshold() {
538 let picker = CompactionPicker::new(4);
539 let segments = vec![
540 test_entry(1, 0, "cpu"),
541 test_entry(2, 0, "cpu"),
542 test_entry(3, 0, "cpu"),
543 ];
544 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
545 assert!(tasks.is_empty());
546 }
547
548 #[test]
549 fn picker_triggers_at_threshold() {
550 let picker = CompactionPicker::new(4);
551 let segments = vec![
552 test_entry(1, 0, "cpu"),
553 test_entry(2, 0, "cpu"),
554 test_entry(3, 0, "cpu"),
555 test_entry(4, 0, "cpu"),
556 ];
557 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
558 assert_eq!(tasks.len(), 1);
559 assert_eq!(tasks[0].input_segments.len(), 4);
560 assert_eq!(tasks[0].shard_id, ShardId(0));
561 assert_eq!(tasks[0].target_level, CompactionLevel::L1);
562 }
563
564 #[test]
565 fn picker_groups_by_shard_and_measurement() {
566 let picker = CompactionPicker::new(2);
567 let segments = vec![
568 test_entry(1, 0, "cpu"),
569 test_entry(2, 0, "cpu"),
570 test_entry(3, 0, "mem"),
571 test_entry(4, 1, "cpu"),
572 test_entry(5, 1, "cpu"),
573 ];
574 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
575 assert_eq!(tasks.len(), 2);
577 }
578
579 #[test]
580 fn picker_never_crosses_shard_boundary() {
581 let picker = CompactionPicker::new(3);
582 let segments = vec![
583 test_entry(1, 0, "cpu"),
584 test_entry(2, 0, "cpu"),
585 test_entry(3, 1, "cpu"),
586 test_entry(4, 1, "cpu"),
587 ];
588 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
590 assert!(tasks.is_empty());
591 }
592
593 #[test]
594 fn backpressure_none_below_heavy_threshold() {
595 let picker = CompactionPicker::new(4);
596 assert_eq!(picker.backpressure_delay_ms(0), 0);
597 assert_eq!(picker.backpressure_delay_ms(15), 0);
598 assert_eq!(picker.backpressure_delay_ms(16), 0);
599 }
600
601 #[test]
602 fn backpressure_escalates_linearly() {
603 let picker = CompactionPicker::new(4);
604 assert_eq!(picker.backpressure_delay_ms(17), 10);
606 assert_eq!(picker.backpressure_delay_ms(18), 20);
607 assert_eq!(picker.backpressure_delay_ms(26), 100);
608 }
609
610 #[test]
611 fn backpressure_caps_at_500ms() {
612 let picker = CompactionPicker::new(4);
613 assert_eq!(picker.backpressure_delay_ms(100), 500);
614 }
615
616 #[test]
617 fn compaction_level_next() {
618 assert_eq!(CompactionLevel::L0.next(), CompactionLevel::L1);
619 assert_eq!(CompactionLevel::L1.next(), CompactionLevel::L2);
620 assert_eq!(CompactionLevel::L2.next(), CompactionLevel::L2);
621 }
622
623 #[test]
624 fn segment_state_display() {
625 assert_eq!(SegmentState::Active.to_string(), "Active");
626 assert_eq!(SegmentState::Compacting.to_string(), "Compacting");
627 assert_eq!(
628 SegmentState::SoftDeleted { deleted_at_ms: 123 }.to_string(),
629 "SoftDeleted(at=123)"
630 );
631 }
632
633 #[test]
636 fn size_tier_classification() {
637 assert_eq!(SizeTier::classify(0), SizeTier::Tiny);
638 assert_eq!(SizeTier::classify(500_000), SizeTier::Tiny);
639 assert_eq!(SizeTier::classify(1_048_575), SizeTier::Tiny);
640 assert_eq!(SizeTier::classify(1_048_576), SizeTier::Small);
641 assert_eq!(SizeTier::classify(5_000_000), SizeTier::Small);
642 assert_eq!(SizeTier::classify(10_485_760), SizeTier::Medium);
643 assert_eq!(SizeTier::classify(50_000_000), SizeTier::Medium);
644 assert_eq!(SizeTier::classify(104_857_600), SizeTier::Large);
645 assert_eq!(SizeTier::classify(1_000_000_000), SizeTier::Large);
646 }
647
648 #[test]
649 fn size_tiered_picker_groups_by_tier() {
650 let policy = CompactionPolicy {
651 strategy: CompactionStrategy::SizeTiered,
652 size_tier_threshold: 2,
653 ..CompactionPolicy::default()
654 };
655 let picker = CompactionPicker::with_policy(policy);
656
657 let segments = vec![
659 test_entry_sized(1, 0, "cpu", 100), test_entry_sized(2, 0, "cpu", 200), test_entry_sized(3, 0, "cpu", 5_000_000), ];
663 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
664 assert_eq!(tasks.len(), 1);
665 assert_eq!(tasks[0].input_segments.len(), 2); }
667
668 #[test]
669 fn size_tiered_never_crosses_shard_boundary() {
670 let policy = CompactionPolicy {
671 strategy: CompactionStrategy::SizeTiered,
672 size_tier_threshold: 2,
673 ..CompactionPolicy::default()
674 };
675 let picker = CompactionPicker::with_policy(policy);
676
677 let segments = vec![
679 test_entry_sized(1, 0, "cpu", 100),
680 test_entry_sized(2, 1, "cpu", 200),
681 ];
682 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
683 assert!(tasks.is_empty());
684 }
685
686 #[test]
687 fn hybrid_deduplicates_twcs_and_size_tiered() {
688 let policy = CompactionPolicy {
689 strategy: CompactionStrategy::Hybrid,
690 count_trigger_threshold: 3,
691 size_tier_threshold: 3,
692 ..CompactionPolicy::default()
693 };
694 let picker = CompactionPicker::with_policy(policy);
695
696 let segments = vec![
699 test_entry_sized(1, 0, "cpu", 100),
700 test_entry_sized(2, 0, "cpu", 200),
701 test_entry_sized(3, 0, "cpu", 300),
702 ];
703 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
704 assert_eq!(tasks.len(), 1); }
706
707 #[test]
708 fn hybrid_size_tiered_catches_remaining_groups() {
709 let policy = CompactionPolicy {
710 strategy: CompactionStrategy::Hybrid,
711 count_trigger_threshold: 10, size_tier_threshold: 3,
713 ..CompactionPolicy::default()
714 };
715 let picker = CompactionPicker::with_policy(policy);
716
717 let segments = vec![
719 test_entry_sized(1, 0, "cpu", 100),
720 test_entry_sized(2, 0, "cpu", 200),
721 test_entry_sized(3, 0, "cpu", 300),
722 test_entry_sized(4, 0, "cpu", 400),
723 ];
724 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
725 assert_eq!(tasks.len(), 1); }
727
728 #[test]
729 fn twcs_only_strategy_ignores_size_tiers() {
730 let policy = CompactionPolicy {
731 strategy: CompactionStrategy::Twcs,
732 count_trigger_threshold: 10,
733 size_tier_threshold: 2,
734 ..CompactionPolicy::default()
735 };
736 let picker = CompactionPicker::with_policy(policy);
737
738 let segments = vec![
740 test_entry_sized(1, 0, "cpu", 100),
741 test_entry_sized(2, 0, "cpu", 200),
742 test_entry_sized(3, 0, "cpu", 300),
743 ];
744 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
745 assert!(tasks.is_empty()); }
747
748 #[test]
751 fn write_amp_tracker_initially_zero() {
752 let tracker = WriteAmpTracker::new();
753 assert_eq!(tracker.write_amplification(), 0.0);
754 assert!(!tracker.budget_exhausted(10.0));
755 }
756
757 #[test]
758 fn write_amp_tracker_records_correctly() {
759 let tracker = WriteAmpTracker::new();
760 tracker.record(1000, 1000); assert!((tracker.write_amplification() - 1.0).abs() < f64::EPSILON);
762
763 tracker.record(1000, 2000); assert!((tracker.write_amplification() - 1.5).abs() < f64::EPSILON);
765 }
766
767 #[test]
768 fn write_amp_tracker_budget_exhausted() {
769 let tracker = WriteAmpTracker::new();
770 tracker.record(1000, 15_000); assert!(tracker.budget_exhausted(10.0));
772 assert!(!tracker.budget_exhausted(20.0));
773 }
774
775 #[test]
776 fn write_amp_tracker_reset() {
777 let tracker = WriteAmpTracker::new();
778 tracker.record(1000, 5000);
779 assert!(tracker.input_bytes() > 0);
780 tracker.reset();
781 assert_eq!(tracker.input_bytes(), 0);
782 assert_eq!(tracker.output_bytes(), 0);
783 }
784
785 #[test]
788 fn default_policy_is_hybrid() {
789 let policy = CompactionPolicy::default();
790 assert_eq!(policy.strategy, CompactionStrategy::Hybrid);
791 assert_eq!(policy.count_trigger_threshold, 4);
792 assert_eq!(policy.size_tier_threshold, 4);
793 assert!((policy.max_write_amplification - 10.0).abs() < f64::EPSILON);
794 }
795
796 #[test]
797 fn with_policy_constructor() {
798 let policy = CompactionPolicy {
799 strategy: CompactionStrategy::SizeTiered,
800 count_trigger_threshold: 8,
801 size_tier_threshold: 6,
802 max_write_amplification: 5.0,
803 };
804 let picker = CompactionPicker::with_policy(policy.clone());
805 assert_eq!(picker.policy().strategy, CompactionStrategy::SizeTiered);
806 assert_eq!(picker.policy().size_tier_threshold, 6);
807 }
808
809 #[test]
810 fn compaction_task_total_input_bytes() {
811 let segments = vec![
812 test_entry_sized(1, 0, "cpu", 1000),
813 test_entry_sized(2, 0, "cpu", 2000),
814 test_entry_sized(3, 0, "cpu", 3000),
815 ];
816 let task = CompactionTask {
817 shard_id: ShardId(0),
818 input_segments: segments,
819 source_level: CompactionLevel::L0,
820 target_level: CompactionLevel::L1,
821 segments_dir: PathBuf::from("/out"),
822 output: SegmentFile::new(ShardId(0), "out.csx").unwrap(),
823 };
824 assert_eq!(task.total_input_bytes(), 6000);
825 }
826
827 #[test]
828 fn policy_validated_clamps_thresholds() {
829 let policy = CompactionPolicy {
830 strategy: CompactionStrategy::Hybrid,
831 count_trigger_threshold: 0,
832 size_tier_threshold: 100_000,
833 max_write_amplification: 0.1,
834 }
835 .validated();
836
837 assert_eq!(policy.count_trigger_threshold, 2);
838 assert_eq!(policy.size_tier_threshold, 1000);
839 assert!((policy.max_write_amplification - 1.0).abs() < f64::EPSILON);
840 }
841
842 #[test]
843 fn policy_validated_keeps_valid_values() {
844 let policy = CompactionPolicy::default().validated();
845 assert_eq!(policy.count_trigger_threshold, 4);
846 assert_eq!(policy.size_tier_threshold, 4);
847 assert!((policy.max_write_amplification - 10.0).abs() < f64::EPSILON);
848 }
849}