Skip to main content

fraiseql_wire/stream/
memory_estimator.rs

1//! Memory estimator for buffered items
2//!
3//! Provides pluggable memory estimation strategy for buffered JSON items.
4//! Default: conservative 2KB per item estimation.
5
6/// Trait for estimating memory usage of buffered items
7///
8/// This allows customization of memory estimation if workload characteristics differ
9/// from the default conservative 2KB per item assumption.
10pub trait MemoryEstimator: Send + Sync {
11    /// Estimate total memory in bytes for given number of buffered items
12    fn estimate_bytes(&self, items_buffered: usize) -> usize;
13
14    /// Human-readable name for this estimator (for debugging/logging)
15    fn name(&self) -> &'static str;
16}
17
18/// Default conservative memory estimator: 2KB per item
19///
20/// Used by default for all streams. Suitable for typical JSON documents (1-5KB).
21/// - Underestimates small objects (< 2KB) → hits limit later (safe)
22/// - Overestimates large objects (> 2KB) → hits limit earlier (safe)
23#[derive(Debug, Clone)]
24pub struct ConservativeEstimator;
25
26impl MemoryEstimator for ConservativeEstimator {
27    fn estimate_bytes(&self, items_buffered: usize) -> usize {
28        items_buffered * 2048 // 2 KB per item
29    }
30
31    fn name(&self) -> &'static str {
32        "conservative_2kb"
33    }
34}
35
36/// Custom memory estimator using fixed bytes per item
37///
38/// Use this if your analysis shows different typical item sizes.
39/// Example: if your JSON averages 4KB, use `FixedEstimator::new(4096)`
40#[derive(Debug, Clone)]
41pub struct FixedEstimator {
42    bytes_per_item: usize,
43}
44
45impl FixedEstimator {
46    /// Create estimator with custom bytes-per-item
47    pub const fn new(bytes_per_item: usize) -> Self {
48        Self { bytes_per_item }
49    }
50}
51
52impl MemoryEstimator for FixedEstimator {
53    fn estimate_bytes(&self, items_buffered: usize) -> usize {
54        items_buffered.saturating_mul(self.bytes_per_item)
55    }
56
57    fn name(&self) -> &'static str {
58        "fixed_custom"
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn test_conservative_estimator() {
68        let est = ConservativeEstimator;
69        assert_eq!(est.estimate_bytes(0), 0);
70        assert_eq!(est.estimate_bytes(1), 2048);
71        assert_eq!(est.estimate_bytes(100), 204_800);
72        assert_eq!(est.estimate_bytes(256), 524_288);
73    }
74
75    #[test]
76    fn test_conservative_name() {
77        let est = ConservativeEstimator;
78        assert_eq!(est.name(), "conservative_2kb");
79    }
80
81    #[test]
82    fn test_fixed_estimator() {
83        let est = FixedEstimator::new(4096);
84        assert_eq!(est.estimate_bytes(0), 0);
85        assert_eq!(est.estimate_bytes(1), 4096);
86        assert_eq!(est.estimate_bytes(100), 409_600);
87        assert_eq!(est.estimate_bytes(256), 1_048_576);
88    }
89
90    #[test]
91    fn test_fixed_estimator_custom_sizes() {
92        for size in &[1024, 2048, 4096, 8192] {
93            let est = FixedEstimator::new(*size);
94            assert_eq!(est.estimate_bytes(10), 10 * size);
95        }
96    }
97
98    #[test]
99    fn test_fixed_estimator_overflow_safe() {
100        let est = FixedEstimator::new(usize::MAX / 2);
101        // Should not panic on overflow
102        let _ = est.estimate_bytes(usize::MAX);
103    }
104}