1use std::collections::BTreeMap;
28
29use crate::datastore::StreamAccess;
30
31#[derive(Debug, Default)]
35pub struct StreamBufferTracker {
36 pub per_stream: BTreeMap<u64, Vec<(f64, StreamAccess)>>,
38
39 pub flat: BTreeMap<u64, u64>,
41
42 pub sum_series: Vec<(f64, u64)>,
44
45 running_sum: u64,
47}
48
49impl StreamBufferTracker {
50 pub fn new() -> Self {
51 Self::default()
52 }
53
54 pub fn update(
57 &mut self, stream_id: u64, access: StreamAccess, ev_time: f64,
58 ) -> Option<(u64, u64)> {
59 self.per_stream
60 .entry(stream_id)
61 .or_default()
62 .push((ev_time, access.clone()));
63
64 let new_end = access.offset + access.length;
65 let old_max = self.flat.get(&stream_id).copied().unwrap_or(0);
66 let new_max = old_max.max(new_end);
67
68 let result = if new_max > old_max {
69 self.flat.insert(stream_id, new_max);
70 self.running_sum += new_max - old_max;
71 Some((old_max, new_max))
72 } else {
73 None
74 };
75
76 self.sum_series.push((ev_time, self.running_sum));
77
78 result
79 }
80
81 pub fn get_stream_max(&self, stream_id: u64) -> Option<u64> {
82 self.flat.get(&stream_id).copied()
83 }
84
85 pub fn current_sum(&self) -> u64 {
86 self.running_sum
87 }
88}
89
90#[cfg(test)]
91mod tests {
92 use super::*;
93
94 #[test]
95 fn test_tracker_empty() {
96 let tracker = StreamBufferTracker::new();
97 assert_eq!(tracker.current_sum(), 0);
98 assert_eq!(tracker.get_stream_max(0), None);
99 assert!(tracker.sum_series.is_empty());
100 assert!(tracker.flat.is_empty());
101 }
102
103 #[test]
104 fn test_tracker_single_stream_in_order() {
105 let mut tracker = StreamBufferTracker::new();
106
107 let result = tracker.update(
108 0,
109 StreamAccess {
110 offset: 0,
111 length: 1000,
112 },
113 10.0,
114 );
115 assert_eq!(result, Some((0, 1000)));
116 assert_eq!(tracker.get_stream_max(0), Some(1000));
117 assert_eq!(tracker.current_sum(), 1000);
118 assert_eq!(tracker.sum_series, vec![(10.0, 1000)]);
119
120 let result = tracker.update(
121 0,
122 StreamAccess {
123 offset: 1000,
124 length: 1000,
125 },
126 20.0,
127 );
128 assert_eq!(result, Some((1000, 2000)));
129 assert_eq!(tracker.get_stream_max(0), Some(2000));
130 assert_eq!(tracker.current_sum(), 2000);
131 assert_eq!(tracker.sum_series, vec![(10.0, 1000), (20.0, 2000)]);
132 }
133
134 #[test]
135 fn test_tracker_out_of_order_events() {
136 let mut tracker = StreamBufferTracker::new();
137
138 let result = tracker.update(
140 0,
141 StreamAccess {
142 offset: 1000,
143 length: 1000,
144 },
145 20.0,
146 );
147 assert_eq!(result, Some((0, 2000)));
148 assert_eq!(tracker.get_stream_max(0), Some(2000));
149 assert_eq!(tracker.current_sum(), 2000);
150
151 let result = tracker.update(
153 0,
154 StreamAccess {
155 offset: 0,
156 length: 1000,
157 },
158 10.0,
159 );
160 assert_eq!(result, None);
161 assert_eq!(tracker.get_stream_max(0), Some(2000));
162 assert_eq!(tracker.current_sum(), 2000);
163
164 assert_eq!(tracker.sum_series, vec![(20.0, 2000), (10.0, 2000)]);
165
166 let result = tracker.update(
168 0,
169 StreamAccess {
170 offset: 2000,
171 length: 1000,
172 },
173 30.0,
174 );
175 assert_eq!(result, Some((2000, 3000)));
176 assert_eq!(tracker.get_stream_max(0), Some(3000));
177 assert_eq!(tracker.current_sum(), 3000);
178 }
179
180 #[test]
181 fn test_tracker_multiple_streams() {
182 let mut tracker = StreamBufferTracker::new();
183
184 tracker.update(
185 0,
186 StreamAccess {
187 offset: 0,
188 length: 1000,
189 },
190 10.0,
191 );
192 assert_eq!(tracker.current_sum(), 1000);
193
194 tracker.update(
195 1,
196 StreamAccess {
197 offset: 0,
198 length: 500,
199 },
200 15.0,
201 );
202 assert_eq!(tracker.current_sum(), 1500);
203
204 tracker.update(
205 0,
206 StreamAccess {
207 offset: 1000,
208 length: 500,
209 },
210 20.0,
211 );
212 assert_eq!(tracker.current_sum(), 2000);
213
214 tracker.update(
215 2,
216 StreamAccess {
217 offset: 0,
218 length: 300,
219 },
220 25.0,
221 );
222 assert_eq!(tracker.current_sum(), 2300);
223
224 assert_eq!(tracker.get_stream_max(0), Some(1500));
225 assert_eq!(tracker.get_stream_max(1), Some(500));
226 assert_eq!(tracker.get_stream_max(2), Some(300));
227 assert_eq!(tracker.get_stream_max(999), None);
228 }
229
230 #[test]
231 fn test_tracker_duplicate_values() {
232 let mut tracker = StreamBufferTracker::new();
233
234 let result = tracker.update(
235 0,
236 StreamAccess {
237 offset: 0,
238 length: 1000,
239 },
240 10.0,
241 );
242 assert_eq!(result, Some((0, 1000)));
243
244 let result = tracker.update(
245 0,
246 StreamAccess {
247 offset: 0,
248 length: 1000,
249 },
250 20.0,
251 );
252 assert_eq!(result, None);
253
254 assert_eq!(tracker.get_stream_max(0), Some(1000));
255 assert_eq!(tracker.current_sum(), 1000);
256
257 assert_eq!(tracker.sum_series.len(), 2);
258 assert_eq!(tracker.sum_series, vec![(10.0, 1000), (20.0, 1000)]);
259 }
260
261 #[test]
262 fn test_tracker_running_sum_correctness() {
263 let mut tracker = StreamBufferTracker::new();
264
265 tracker.update(
266 0,
267 StreamAccess {
268 offset: 0,
269 length: 1000,
270 },
271 1.0,
272 );
273 tracker.update(
274 1,
275 StreamAccess {
276 offset: 0,
277 length: 2000,
278 },
279 2.0,
280 );
281 tracker.update(
282 2,
283 StreamAccess {
284 offset: 0,
285 length: 1500,
286 },
287 3.0,
288 );
289 tracker.update(
290 0,
291 StreamAccess {
292 offset: 1000,
293 length: 200,
294 },
295 4.0,
296 );
297
298 let manual_sum: u64 = tracker.flat.values().sum();
299
300 assert_eq!(tracker.current_sum(), manual_sum);
301 assert_eq!(tracker.current_sum(), 4700);
302 }
303
304 #[test]
305 fn test_tracker_preserves_stream_access() {
306 let mut tracker = StreamBufferTracker::new();
307
308 tracker.update(
309 0,
310 StreamAccess {
311 offset: 100,
312 length: 500,
313 },
314 10.0,
315 );
316 tracker.update(
317 0,
318 StreamAccess {
319 offset: 600,
320 length: 300,
321 },
322 20.0,
323 );
324
325 let stream_data = tracker.per_stream.get(&0).unwrap();
326 assert_eq!(stream_data.len(), 2);
327 assert_eq!(stream_data[0].1.offset, 100);
328 assert_eq!(stream_data[0].1.length, 500);
329 assert_eq!(stream_data[1].1.offset, 600);
330 assert_eq!(stream_data[1].1.length, 300);
331
332 assert_eq!(tracker.get_stream_max(0), Some(900));
333 }
334}