Skip to main content

ruda_runtime/runtime/memory_management/drop_queue/
policy.rs

1use ruda_core::bytes::Bytes;
2
3/// Defines the thresholds that determine when a [`PendingDropQueue`] should be
4/// flushed.
5///
6/// A flush is triggered when **either** limit is exceeded — whichever comes
7/// first. Set a field to `u32::MAX` / `usize::MAX` to effectively disable it.
8#[derive(Debug)]
9pub struct FlushingPolicy {
10    /// Flush when this many allocations have been staged.
11    pub max_bytes_count: u32,
12    /// Flush when the total staged size reaches this many bytes.
13    pub max_bytes_size: u32,
14}
15
16impl Default for FlushingPolicy {
17    fn default() -> Self {
18        Self {
19            max_bytes_count: 64,
20            max_bytes_size: 64 * 1024 * 1024, // 64 MiB
21        }
22    }
23}
24
25/// Tracks staged allocations and evaluates them against a [`FlushingPolicy`].
26#[derive(Default, Debug)]
27pub(crate) struct FlushingPolicyState {
28    bytes_count: u32,
29    bytes_size: usize,
30}
31
32impl FlushingPolicyState {
33    /// Record a newly staged [`Bytes`] allocation.
34    pub(crate) fn register(&mut self, bytes: &Bytes) {
35        self.bytes_count = self.bytes_count.saturating_add(1);
36        self.bytes_size = self.bytes_size.saturating_add(bytes.len());
37    }
38
39    /// Reset all counters, typically called after a flush.
40    pub(crate) fn reset(&mut self) {
41        self.bytes_count = 0;
42        self.bytes_size = 0;
43    }
44
45    /// Returns `true` if either threshold in `policy` has been reached.
46    pub(crate) fn should_flush(&self, policy: &FlushingPolicy) -> bool {
47        self.bytes_count >= policy.max_bytes_count
48            || self.bytes_size >= policy.max_bytes_size as usize
49    }
50}
51
52#[cfg(test)]
53mod policy_tests {
54    use std::vec;
55
56    use super::*;
57
58    fn policy() -> FlushingPolicy {
59        FlushingPolicy {
60            max_bytes_count: 4,
61            max_bytes_size: 100,
62        }
63    }
64
65    fn state() -> FlushingPolicyState {
66        FlushingPolicyState {
67            bytes_count: 0,
68            bytes_size: 0,
69        }
70    }
71
72    #[test]
73    fn no_flush_when_below_both_thresholds() {
74        let s = state();
75        assert!(!s.should_flush(&policy()));
76    }
77
78    #[test]
79    fn flush_when_count_threshold_reached() {
80        let mut s = state();
81        for _ in 0..4 {
82            s.register(&Bytes::from_elems(vec![0u8]));
83        }
84        assert!(s.should_flush(&policy()));
85    }
86
87    #[test]
88    fn flush_when_size_threshold_reached() {
89        let mut s = state();
90        s.register(&Bytes::from_elems(vec![0u8; 101]));
91        assert!(s.should_flush(&policy()));
92    }
93
94    #[test]
95    fn flush_triggered_by_whichever_limit_comes_first() {
96        let mut s = state();
97        // Only 2 allocations but already over the size limit.
98        s.register(&Bytes::from_elems(vec![0u8; 60]));
99        s.register(&Bytes::from_elems(vec![0u8; 60]));
100        assert!(s.should_flush(&policy()));
101    }
102
103    #[test]
104    fn reset_clears_state() {
105        let mut s = state();
106        for _ in 0..4 {
107            s.register(&Bytes::from_elems(vec![0u8]));
108        }
109        assert!(s.should_flush(&policy()));
110        s.reset();
111        assert!(!s.should_flush(&policy()));
112    }
113
114    #[test]
115    fn counters_do_not_wrap_before_flush() {
116        let mut s = FlushingPolicyState {
117            bytes_count: u32::MAX,
118            bytes_size: u32::MAX as usize - 1,
119        };
120        s.register(&Bytes::from_elems(vec![0u8; 8]));
121        assert_eq!(s.bytes_count, u32::MAX);
122        assert_eq!(s.bytes_size, (u32::MAX as usize - 1).saturating_add(8));
123        assert!(s.should_flush(&FlushingPolicy {
124            max_bytes_count: u32::MAX,
125            max_bytes_size: u32::MAX,
126        }));
127        s.bytes_size = usize::MAX;
128        s.register(&Bytes::from_elems(vec![0u8; 8]));
129        assert_eq!(s.bytes_size, usize::MAX);
130    }
131}