Skip to main content

cu_ratelimit/
lib.rs

1use bincode::de::Decoder;
2use bincode::enc::Encoder;
3use bincode::error::{DecodeError, EncodeError};
4use bincode::{Decode, Encode};
5use cu29::prelude::*;
6use std::marker::PhantomData;
7
8fn phase_at_or_before(tov: CuTime, interval: CuDuration) -> CuTime {
9    let interval_ns = interval.as_nanos();
10    CuTime::from(tov.as_nanos() / interval_ns * interval_ns)
11}
12
13#[derive(Reflect)]
14#[reflect(no_field_bounds, from_reflect = false, type_path = false)]
15pub struct CuRateLimit<T>
16where
17    T: for<'a> CuMsgPayload + 'static,
18{
19    #[reflect(ignore)]
20    _marker: PhantomData<fn() -> T>,
21    interval: CuDuration,
22    last_tov: Option<CuTime>,
23}
24
25impl<T> TypePath for CuRateLimit<T>
26where
27    T: CuMsgPayload + 'static,
28{
29    fn type_path() -> &'static str {
30        "cu_ratelimit::CuRateLimit"
31    }
32
33    fn short_type_path() -> &'static str {
34        "CuRateLimit"
35    }
36
37    fn type_ident() -> Option<&'static str> {
38        Some("CuRateLimit")
39    }
40
41    fn crate_name() -> Option<&'static str> {
42        Some("cu_ratelimit")
43    }
44
45    fn module_path() -> Option<&'static str> {
46        Some("")
47    }
48}
49
50impl<T> Freezable for CuRateLimit<T>
51where
52    T: CuMsgPayload,
53{
54    fn freeze<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
55        Encode::encode(&self.last_tov, encoder)
56    }
57
58    fn thaw<D: Decoder>(&mut self, decoder: &mut D) -> Result<(), DecodeError> {
59        self.last_tov = Decode::decode(decoder)?;
60        Ok(())
61    }
62}
63
64impl<T> CuTask for CuRateLimit<T>
65where
66    T: CuMsgPayload,
67{
68    type Resources<'r> = ();
69    type Input<'m> = input_msg!(T);
70    type Output<'m> = output_msg!(T);
71
72    fn new(config: Option<&ComponentConfig>, _resources: Self::Resources<'_>) -> CuResult<Self> {
73        let hz = match config {
74            Some(cfg) => cfg
75                .get::<f64>("rate")?
76                .ok_or("Missing required 'rate' config for CuRateLimiter")?,
77            None => return Err("Missing required 'rate' config for CuRateLimiter".into()),
78        };
79        if !hz.is_finite() || hz <= 0.0 {
80            return Err("CuRateLimiter 'rate' must be finite and greater than zero".into());
81        }
82        let interval_ns = (1e9 / hz) as u64;
83        if interval_ns == 0 {
84            return Err("CuRateLimiter 'rate' cannot exceed 1 GHz".into());
85        }
86        Ok(Self {
87            _marker: PhantomData,
88            interval: CuDuration::from(interval_ns),
89            last_tov: None,
90        })
91    }
92
93    fn process<'m>(
94        &mut self,
95        _ctx: &CuContext,
96        input: &Self::Input<'m>,
97        output: &mut Self::Output<'m>,
98    ) -> CuResult<()> {
99        let tov = match input.tov {
100            Tov::Time(ts) => ts,
101            _ => return Err("Expected single timestamp TOV".into()),
102        };
103        output.tov = input.tov;
104
105        let allow = match self.last_tov {
106            None => {
107                self.last_tov = Some(phase_at_or_before(tov, self.interval));
108                true
109            }
110            Some(last) if tov < last => {
111                // Re-anchor after replay/simulation time is reset.
112                self.last_tov = Some(phase_at_or_before(tov, self.interval));
113                true
114            }
115            Some(last) => {
116                let elapsed = tov - last;
117                if elapsed < self.interval {
118                    false
119                } else {
120                    let elapsed_intervals = elapsed.as_nanos() / self.interval.as_nanos();
121                    self.last_tov = Some(last + elapsed_intervals * self.interval);
122                    true
123                }
124            }
125        };
126
127        if allow {
128            if let Some(payload) = input.payload() {
129                output.set_payload(payload.clone());
130            } else {
131                output.clear_payload();
132            }
133        } else {
134            output.clear_payload();
135        }
136
137        Ok(())
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    fn create_test_ratelimiter(rate: f64) -> CuRateLimit<i32> {
146        let mut cfg = ComponentConfig::new();
147        cfg.set("rate", rate);
148        CuRateLimit::new(Some(&cfg), ()).unwrap()
149    }
150
151    #[test]
152    fn test_rate_limiting() {
153        let ctx = CuContext::new_with_clock();
154        let mut limiter = create_test_ratelimiter(10.0); // 10 Hz = 100ms interval
155        let mut input = CuMsg::<i32>::new(Some(42));
156        let mut output = CuMsg::<i32>::new(None);
157
158        // First message should pass
159        input.tov = Tov::Time(CuTime::from(0));
160        limiter.process(&ctx, &input, &mut output).unwrap();
161        assert_eq!(output.payload(), Some(&42));
162
163        // Message within the interval should be blocked
164        input.tov = Tov::Time(CuTime::from(50_000_000)); // 50ms
165        limiter.process(&ctx, &input, &mut output).unwrap();
166        assert_eq!(output.payload(), None);
167
168        // Message after the interval should pass
169        input.tov = Tov::Time(CuTime::from(100_000_000)); // 100ms
170        limiter.process(&ctx, &input, &mut output).unwrap();
171        assert_eq!(output.payload(), Some(&42));
172        assert_eq!(output.tov, input.tov);
173    }
174
175    #[test]
176    fn test_payload_propagation() {
177        let ctx = CuContext::new_with_clock();
178        let mut limiter = create_test_ratelimiter(10.0);
179        let mut input = CuMsg::<i32>::new(None);
180        let mut output = CuMsg::<i32>::new(None);
181
182        // Test payload propagation
183        input.set_payload(123);
184        input.tov = Tov::Time(CuTime::from(0));
185        limiter.process(&ctx, &input, &mut output).unwrap();
186        assert_eq!(output.payload(), Some(&123));
187
188        // Test empty payload propagation
189        input.clear_payload();
190        input.tov = Tov::Time(CuTime::from(100_000_000));
191        limiter.process(&ctx, &input, &mut output).unwrap();
192        assert_eq!(output.payload(), None);
193    }
194
195    #[test]
196    fn test_phase_preserving_64_hz_to_30_hz() {
197        let ctx = CuContext::new_with_clock();
198        let mut limiter = create_test_ratelimiter(30.0);
199        let mut input = CuMsg::<i32>::new(Some(42));
200        let mut output = CuMsg::<i32>::new(None);
201        let mut emitted = 0;
202
203        // 64 source samples span [0, 984.375 ms], which contains exactly 30
204        // points of the 30 Hz output phase including the initial sample.
205        for source_tick in 0..64 {
206            input.tov = Tov::Time(CuTime::from(source_tick * 15_625_000));
207            limiter.process(&ctx, &input, &mut output).unwrap();
208            emitted += usize::from(output.payload().is_some());
209        }
210
211        assert_eq!(emitted, 30);
212    }
213
214    #[test]
215    fn test_first_sample_uses_absolute_rate_phase() {
216        let ctx = CuContext::new_with_clock();
217        let mut limiter = create_test_ratelimiter(30.0);
218        let mut input = CuMsg::<i32>::new(Some(42));
219        let mut output = CuMsg::<i32>::new(None);
220
221        // A camera that becomes ready off-phase must join the same 30 Hz
222        // schedule as other graph inputs instead of creating its own phase.
223        input.tov = Tov::Time(CuTime::from(78_125_000));
224        limiter.process(&ctx, &input, &mut output).unwrap();
225        assert!(output.payload().is_some());
226
227        input.tov = Tov::Time(CuTime::from(93_750_000));
228        limiter.process(&ctx, &input, &mut output).unwrap();
229        assert!(output.payload().is_none());
230
231        input.tov = Tov::Time(CuTime::from(109_375_000));
232        limiter.process(&ctx, &input, &mut output).unwrap();
233        assert!(output.payload().is_some());
234    }
235
236    #[test]
237    fn test_time_rewind_reanchors_limiter() {
238        let ctx = CuContext::new_with_clock();
239        let mut limiter = create_test_ratelimiter(30.0);
240        let mut input = CuMsg::<i32>::new(Some(42));
241        let mut output = CuMsg::<i32>::new(None);
242
243        input.tov = Tov::Time(CuTime::from(1_000_000_000));
244        limiter.process(&ctx, &input, &mut output).unwrap();
245        assert!(output.payload().is_some());
246
247        input.tov = Tov::Time(CuTime::from(1_000_000));
248        limiter.process(&ctx, &input, &mut output).unwrap();
249        assert!(output.payload().is_some());
250    }
251
252    #[test]
253    fn test_invalid_rates_are_rejected() {
254        for rate in [0.0, -1.0, f64::NAN, f64::INFINITY, 1_000_000_001.0] {
255            let mut cfg = ComponentConfig::new();
256            cfg.set("rate", rate);
257            assert!(CuRateLimit::<i32>::new(Some(&cfg), ()).is_err());
258        }
259    }
260}