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::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 output_path: PathBuf,
262}
263
264impl CompactionTask {
265 #[must_use]
267 pub fn total_input_bytes(&self) -> u64 {
268 self.input_segments.iter().map(|s| s.byte_size).sum()
269 }
270}
271
272#[derive(Debug, Clone, Default)]
288pub struct CompactionPicker {
289 policy: CompactionPolicy,
291}
292
293impl CompactionPicker {
294 #[must_use]
298 pub fn new(trigger_threshold: usize) -> Self {
299 Self {
300 policy: CompactionPolicy {
301 count_trigger_threshold: trigger_threshold,
302 size_tier_threshold: trigger_threshold,
303 ..CompactionPolicy::default()
304 },
305 }
306 }
307
308 #[must_use]
310 pub fn with_policy(policy: CompactionPolicy) -> Self {
311 Self { policy }
312 }
313
314 #[must_use]
316 pub fn policy(&self) -> &CompactionPolicy {
317 &self.policy
318 }
319
320 #[must_use]
322 pub fn trigger_threshold(&self) -> usize {
323 self.policy.count_trigger_threshold
324 }
325
326 #[must_use]
334 pub fn pick(
335 &self,
336 segments: &[SegmentCatalogEntry],
337 output_dir: &std::path::Path,
338 ) -> Vec<CompactionTask> {
339 self.pick_level(segments, output_dir, CompactionLevel::L0)
340 }
341
342 #[must_use]
346 pub fn pick_level(
347 &self,
348 segments: &[SegmentCatalogEntry],
349 output_dir: &std::path::Path,
350 source_level: CompactionLevel,
351 ) -> Vec<CompactionTask> {
352 let target_level = source_level.next();
353
354 let mut groups: BTreeMap<(ShardId, &str), Vec<&SegmentCatalogEntry>> = BTreeMap::new();
356 for seg in segments {
357 groups
358 .entry((seg.shard_id, seg.measurement.as_str()))
359 .or_default()
360 .push(seg);
361 }
362
363 let mut tasks = Vec::new();
364
365 match self.policy.strategy {
366 CompactionStrategy::Twcs => {
367 self.pick_twcs(&groups, output_dir, source_level, target_level, &mut tasks);
368 }
369 CompactionStrategy::SizeTiered => {
370 self.pick_size_tiered(&groups, output_dir, source_level, target_level, &mut tasks);
371 }
372 CompactionStrategy::Hybrid => {
373 let twcs_shards =
375 self.pick_twcs(&groups, output_dir, source_level, target_level, &mut tasks);
376 let remaining: BTreeMap<(ShardId, &str), Vec<&SegmentCatalogEntry>> = groups
378 .into_iter()
379 .filter(|(key, _)| !twcs_shards.contains(key))
380 .collect();
381 self.pick_size_tiered(
382 &remaining,
383 output_dir,
384 source_level,
385 target_level,
386 &mut tasks,
387 );
388 }
389 }
390
391 tasks
392 }
393
394 fn pick_twcs<'a>(
398 &self,
399 groups: &BTreeMap<(ShardId, &'a str), Vec<&SegmentCatalogEntry>>,
400 output_dir: &std::path::Path,
401 source_level: CompactionLevel,
402 target_level: CompactionLevel,
403 tasks: &mut Vec<CompactionTask>,
404 ) -> std::collections::BTreeSet<(ShardId, &'a str)> {
405 let mut selected = std::collections::BTreeSet::new();
406
407 for ((shard_id, measurement), group) in groups {
408 if group.len() >= self.policy.count_trigger_threshold {
409 let output_path = Self::output_path(output_dir, *shard_id, measurement, group);
410 tasks.push(CompactionTask {
411 shard_id: *shard_id,
412 input_segments: group.iter().map(|s| (*s).clone()).collect(),
413 source_level,
414 target_level,
415 output_path,
416 });
417 selected.insert((*shard_id, *measurement));
418 }
419 }
420
421 selected
422 }
423
424 fn pick_size_tiered(
426 &self,
427 groups: &BTreeMap<(ShardId, &str), Vec<&SegmentCatalogEntry>>,
428 output_dir: &std::path::Path,
429 source_level: CompactionLevel,
430 target_level: CompactionLevel,
431 tasks: &mut Vec<CompactionTask>,
432 ) {
433 for ((shard_id, measurement), group) in groups {
434 let mut tiers: BTreeMap<SizeTier, Vec<&SegmentCatalogEntry>> = BTreeMap::new();
436 for seg in group {
437 let tier = SizeTier::classify(seg.byte_size);
438 tiers.entry(tier).or_default().push(seg);
439 }
440
441 for tier_segments in tiers.values() {
442 if tier_segments.len() >= self.policy.size_tier_threshold {
443 let output_path =
444 Self::output_path(output_dir, *shard_id, measurement, tier_segments);
445 tasks.push(CompactionTask {
446 shard_id: *shard_id,
447 input_segments: tier_segments.iter().map(|s| (*s).clone()).collect(),
448 source_level,
449 target_level,
450 output_path,
451 });
452 }
453 }
454 }
455 }
456
457 fn output_path(
459 output_dir: &std::path::Path,
460 shard_id: ShardId,
461 measurement: &str,
462 segments: &[&SegmentCatalogEntry],
463 ) -> PathBuf {
464 output_dir
465 .join(format!("shard_{}", shard_id.0))
466 .join(format!(
467 "{measurement}_compacted_{}.csx",
468 segments.iter().map(|s| s.segment_id.0).max().unwrap_or(0)
469 ))
470 }
471
472 #[must_use]
478 pub fn backpressure_delay_ms(&self, l0_count: usize) -> u64 {
479 let heavy_threshold = self.policy.count_trigger_threshold * 4;
480 if l0_count <= heavy_threshold {
481 return 0;
482 }
483 let excess = l0_count - heavy_threshold;
484 (excess as u64 * 10).min(500)
486 }
487}
488
489#[cfg(test)]
490mod tests {
491 use super::*;
492 use chronix_core::SegmentId;
493 use std::path::PathBuf;
494
495 fn test_entry(id: u64, shard: i64, measurement: &str) -> SegmentCatalogEntry {
496 test_entry_sized(id, shard, measurement, 4096)
497 }
498
499 fn test_entry_sized(id: u64, shard: i64, measurement: &str, size: u64) -> SegmentCatalogEntry {
500 SegmentCatalogEntry {
501 segment_id: SegmentId(id),
502 shard_id: ShardId(shard),
503 measurement: measurement.to_string(),
504 path: PathBuf::from(format!("shard_{shard}/seg_{id}.csx")),
505 min_timestamp: id as i64 * 1000,
506 max_timestamp: id as i64 * 1000 + 999,
507 row_count: 1000,
508 series_count: 10,
509 byte_size: size,
510 row_group_count: 1,
511 column_count: 5,
512 column_stats: Vec::new(),
513 state: SegmentState::default(),
514 }
515 }
516
517 #[test]
520 fn picker_no_compaction_below_threshold() {
521 let picker = CompactionPicker::new(4);
522 let segments = vec![
523 test_entry(1, 0, "cpu"),
524 test_entry(2, 0, "cpu"),
525 test_entry(3, 0, "cpu"),
526 ];
527 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
528 assert!(tasks.is_empty());
529 }
530
531 #[test]
532 fn picker_triggers_at_threshold() {
533 let picker = CompactionPicker::new(4);
534 let segments = vec![
535 test_entry(1, 0, "cpu"),
536 test_entry(2, 0, "cpu"),
537 test_entry(3, 0, "cpu"),
538 test_entry(4, 0, "cpu"),
539 ];
540 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
541 assert_eq!(tasks.len(), 1);
542 assert_eq!(tasks[0].input_segments.len(), 4);
543 assert_eq!(tasks[0].shard_id, ShardId(0));
544 assert_eq!(tasks[0].target_level, CompactionLevel::L1);
545 }
546
547 #[test]
548 fn picker_groups_by_shard_and_measurement() {
549 let picker = CompactionPicker::new(2);
550 let segments = vec![
551 test_entry(1, 0, "cpu"),
552 test_entry(2, 0, "cpu"),
553 test_entry(3, 0, "mem"),
554 test_entry(4, 1, "cpu"),
555 test_entry(5, 1, "cpu"),
556 ];
557 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
558 assert_eq!(tasks.len(), 2);
560 }
561
562 #[test]
563 fn picker_never_crosses_shard_boundary() {
564 let picker = CompactionPicker::new(3);
565 let segments = vec![
566 test_entry(1, 0, "cpu"),
567 test_entry(2, 0, "cpu"),
568 test_entry(3, 1, "cpu"),
569 test_entry(4, 1, "cpu"),
570 ];
571 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
573 assert!(tasks.is_empty());
574 }
575
576 #[test]
577 fn backpressure_none_below_heavy_threshold() {
578 let picker = CompactionPicker::new(4);
579 assert_eq!(picker.backpressure_delay_ms(0), 0);
580 assert_eq!(picker.backpressure_delay_ms(15), 0);
581 assert_eq!(picker.backpressure_delay_ms(16), 0);
582 }
583
584 #[test]
585 fn backpressure_escalates_linearly() {
586 let picker = CompactionPicker::new(4);
587 assert_eq!(picker.backpressure_delay_ms(17), 10);
589 assert_eq!(picker.backpressure_delay_ms(18), 20);
590 assert_eq!(picker.backpressure_delay_ms(26), 100);
591 }
592
593 #[test]
594 fn backpressure_caps_at_500ms() {
595 let picker = CompactionPicker::new(4);
596 assert_eq!(picker.backpressure_delay_ms(100), 500);
597 }
598
599 #[test]
600 fn compaction_level_next() {
601 assert_eq!(CompactionLevel::L0.next(), CompactionLevel::L1);
602 assert_eq!(CompactionLevel::L1.next(), CompactionLevel::L2);
603 assert_eq!(CompactionLevel::L2.next(), CompactionLevel::L2);
604 }
605
606 #[test]
607 fn segment_state_display() {
608 assert_eq!(SegmentState::Active.to_string(), "Active");
609 assert_eq!(SegmentState::Compacting.to_string(), "Compacting");
610 assert_eq!(
611 SegmentState::SoftDeleted { deleted_at_ms: 123 }.to_string(),
612 "SoftDeleted(at=123)"
613 );
614 }
615
616 #[test]
619 fn size_tier_classification() {
620 assert_eq!(SizeTier::classify(0), SizeTier::Tiny);
621 assert_eq!(SizeTier::classify(500_000), SizeTier::Tiny);
622 assert_eq!(SizeTier::classify(1_048_575), SizeTier::Tiny);
623 assert_eq!(SizeTier::classify(1_048_576), SizeTier::Small);
624 assert_eq!(SizeTier::classify(5_000_000), SizeTier::Small);
625 assert_eq!(SizeTier::classify(10_485_760), SizeTier::Medium);
626 assert_eq!(SizeTier::classify(50_000_000), SizeTier::Medium);
627 assert_eq!(SizeTier::classify(104_857_600), SizeTier::Large);
628 assert_eq!(SizeTier::classify(1_000_000_000), SizeTier::Large);
629 }
630
631 #[test]
632 fn size_tiered_picker_groups_by_tier() {
633 let policy = CompactionPolicy {
634 strategy: CompactionStrategy::SizeTiered,
635 size_tier_threshold: 2,
636 ..CompactionPolicy::default()
637 };
638 let picker = CompactionPicker::with_policy(policy);
639
640 let segments = vec![
642 test_entry_sized(1, 0, "cpu", 100), test_entry_sized(2, 0, "cpu", 200), test_entry_sized(3, 0, "cpu", 5_000_000), ];
646 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
647 assert_eq!(tasks.len(), 1);
648 assert_eq!(tasks[0].input_segments.len(), 2); }
650
651 #[test]
652 fn size_tiered_never_crosses_shard_boundary() {
653 let policy = CompactionPolicy {
654 strategy: CompactionStrategy::SizeTiered,
655 size_tier_threshold: 2,
656 ..CompactionPolicy::default()
657 };
658 let picker = CompactionPicker::with_policy(policy);
659
660 let segments = vec![
662 test_entry_sized(1, 0, "cpu", 100),
663 test_entry_sized(2, 1, "cpu", 200),
664 ];
665 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
666 assert!(tasks.is_empty());
667 }
668
669 #[test]
670 fn hybrid_deduplicates_twcs_and_size_tiered() {
671 let policy = CompactionPolicy {
672 strategy: CompactionStrategy::Hybrid,
673 count_trigger_threshold: 3,
674 size_tier_threshold: 3,
675 ..CompactionPolicy::default()
676 };
677 let picker = CompactionPicker::with_policy(policy);
678
679 let segments = vec![
682 test_entry_sized(1, 0, "cpu", 100),
683 test_entry_sized(2, 0, "cpu", 200),
684 test_entry_sized(3, 0, "cpu", 300),
685 ];
686 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
687 assert_eq!(tasks.len(), 1); }
689
690 #[test]
691 fn hybrid_size_tiered_catches_remaining_groups() {
692 let policy = CompactionPolicy {
693 strategy: CompactionStrategy::Hybrid,
694 count_trigger_threshold: 10, size_tier_threshold: 3,
696 ..CompactionPolicy::default()
697 };
698 let picker = CompactionPicker::with_policy(policy);
699
700 let segments = vec![
702 test_entry_sized(1, 0, "cpu", 100),
703 test_entry_sized(2, 0, "cpu", 200),
704 test_entry_sized(3, 0, "cpu", 300),
705 test_entry_sized(4, 0, "cpu", 400),
706 ];
707 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
708 assert_eq!(tasks.len(), 1); }
710
711 #[test]
712 fn twcs_only_strategy_ignores_size_tiers() {
713 let policy = CompactionPolicy {
714 strategy: CompactionStrategy::Twcs,
715 count_trigger_threshold: 10,
716 size_tier_threshold: 2,
717 ..CompactionPolicy::default()
718 };
719 let picker = CompactionPicker::with_policy(policy);
720
721 let segments = vec![
723 test_entry_sized(1, 0, "cpu", 100),
724 test_entry_sized(2, 0, "cpu", 200),
725 test_entry_sized(3, 0, "cpu", 300),
726 ];
727 let tasks = picker.pick(&segments, std::path::Path::new("/out"));
728 assert!(tasks.is_empty()); }
730
731 #[test]
734 fn write_amp_tracker_initially_zero() {
735 let tracker = WriteAmpTracker::new();
736 assert_eq!(tracker.write_amplification(), 0.0);
737 assert!(!tracker.budget_exhausted(10.0));
738 }
739
740 #[test]
741 fn write_amp_tracker_records_correctly() {
742 let tracker = WriteAmpTracker::new();
743 tracker.record(1000, 1000); assert!((tracker.write_amplification() - 1.0).abs() < f64::EPSILON);
745
746 tracker.record(1000, 2000); assert!((tracker.write_amplification() - 1.5).abs() < f64::EPSILON);
748 }
749
750 #[test]
751 fn write_amp_tracker_budget_exhausted() {
752 let tracker = WriteAmpTracker::new();
753 tracker.record(1000, 15_000); assert!(tracker.budget_exhausted(10.0));
755 assert!(!tracker.budget_exhausted(20.0));
756 }
757
758 #[test]
759 fn write_amp_tracker_reset() {
760 let tracker = WriteAmpTracker::new();
761 tracker.record(1000, 5000);
762 assert!(tracker.input_bytes() > 0);
763 tracker.reset();
764 assert_eq!(tracker.input_bytes(), 0);
765 assert_eq!(tracker.output_bytes(), 0);
766 }
767
768 #[test]
771 fn default_policy_is_hybrid() {
772 let policy = CompactionPolicy::default();
773 assert_eq!(policy.strategy, CompactionStrategy::Hybrid);
774 assert_eq!(policy.count_trigger_threshold, 4);
775 assert_eq!(policy.size_tier_threshold, 4);
776 assert!((policy.max_write_amplification - 10.0).abs() < f64::EPSILON);
777 }
778
779 #[test]
780 fn with_policy_constructor() {
781 let policy = CompactionPolicy {
782 strategy: CompactionStrategy::SizeTiered,
783 count_trigger_threshold: 8,
784 size_tier_threshold: 6,
785 max_write_amplification: 5.0,
786 };
787 let picker = CompactionPicker::with_policy(policy.clone());
788 assert_eq!(picker.policy().strategy, CompactionStrategy::SizeTiered);
789 assert_eq!(picker.policy().size_tier_threshold, 6);
790 }
791
792 #[test]
793 fn compaction_task_total_input_bytes() {
794 let segments = vec![
795 test_entry_sized(1, 0, "cpu", 1000),
796 test_entry_sized(2, 0, "cpu", 2000),
797 test_entry_sized(3, 0, "cpu", 3000),
798 ];
799 let task = CompactionTask {
800 shard_id: ShardId(0),
801 input_segments: segments,
802 source_level: CompactionLevel::L0,
803 target_level: CompactionLevel::L1,
804 output_path: PathBuf::from("/out"),
805 };
806 assert_eq!(task.total_input_bytes(), 6000);
807 }
808
809 #[test]
810 fn policy_validated_clamps_thresholds() {
811 let policy = CompactionPolicy {
812 strategy: CompactionStrategy::Hybrid,
813 count_trigger_threshold: 0,
814 size_tier_threshold: 100_000,
815 max_write_amplification: 0.1,
816 }
817 .validated();
818
819 assert_eq!(policy.count_trigger_threshold, 2);
820 assert_eq!(policy.size_tier_threshold, 1000);
821 assert!((policy.max_write_amplification - 1.0).abs() < f64::EPSILON);
822 }
823
824 #[test]
825 fn policy_validated_keeps_valid_values() {
826 let policy = CompactionPolicy::default().validated();
827 assert_eq!(policy.count_trigger_threshold, 4);
828 assert_eq!(policy.size_tier_threshold, 4);
829 assert!((policy.max_write_amplification - 10.0).abs() < f64::EPSILON);
830 }
831}