1use std::{
5 collections::{BTreeMap, BTreeSet, HashMap},
6 fmt::Debug,
7 hash::Hash,
8 marker::PhantomData,
9};
10
11use reifydb_codec::key::{
12 encode_u64,
13 encoded::{EncodedKey, IntoEncodedKey},
14};
15use reifydb_value::{Result, reifydb_assertions, value::row_number::RowNumber};
16use serde::{Deserialize, Serialize, de::DeserializeOwned};
17
18use crate::window::{
19 accumulator::WindowAccumulator,
20 engine::{
21 AccumulatorEvent, EmitKind, GroupMeta, MetaKey, WindowResult, WindowStateKey,
22 config::WindowEngineConfig, expiry_due_range, expiry_key, meta_key_for, sweep_stale_meta,
23 },
24 span::{Slot, WindowSpan},
25 state::StateCache,
26 store::WindowStore,
27};
28
29pub type TumblingBuckets<G, C, Contribution> = BTreeMap<(G, WindowSpan<C>), Vec<AccumulatorEvent<Contribution>>>;
30
31type MetaLoaded<G, C> = HashMap<G, GroupMeta<C>>;
32type SlotResolved = Vec<Option<(RowNumber, bool)>>;
33
34pub struct ExpiredWindow<G, C, Output> {
35 pub row_number: RowNumber,
36 pub group: G,
37 pub window_start: C,
38 pub value: Option<Output>,
39}
40
41#[derive(Serialize, Deserialize)]
42#[serde(bound(serialize = "G: Serialize, C: Serialize", deserialize = "G: DeserializeOwned, C: DeserializeOwned"))]
43struct TumblingIndexEntry<G, C> {
44 group: G,
45 window_start: C,
46 row_number: u64,
47}
48
49pub fn reindex_window<S, G, C>(
50 store: &mut S,
51 group: &G,
52 window_start: C,
53 row_number: RowNumber,
54 prior: Option<u64>,
55 new: Option<u64>,
56) -> Result<()>
57where
58 S: WindowStore,
59 G: Clone + Serialize,
60 C: Slot + Serialize,
61 for<'a> &'a G: IntoEncodedKey,
62{
63 if prior == new {
64 return Ok(());
65 }
66 let suffix = encode_u64(window_start.order_key());
67 if let Some(old) = prior {
68 store.internal_drop(&expiry_key(old, group, &suffix))?;
69 }
70 if let Some(new) = new {
71 store.internal_set(
72 &expiry_key(new, group, &suffix),
73 &TumblingIndexEntry {
74 group: group.clone(),
75 window_start,
76 row_number: row_number.0,
77 },
78 )?;
79 }
80 Ok(())
81}
82
83pub struct TumblingEngine<G, C, Accumulator> {
84 accumulators: StateCache<WindowStateKey, Accumulator>,
85 meta: StateCache<MetaKey, GroupMeta<C>>,
86 meta_low_water: Option<u64>,
87 expire_batch: usize,
88 _pd: PhantomData<G>,
89}
90
91impl<G, C, Accumulator> TumblingEngine<G, C, Accumulator>
92where
93 G: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
94 C: Slot + Hash + Serialize + DeserializeOwned,
95 Accumulator: WindowAccumulator,
96 for<'a> &'a G: IntoEncodedKey,
97{
98 pub fn new(config: WindowEngineConfig) -> Self {
99 Self {
100 accumulators: StateCache::<WindowStateKey, Accumulator>::new_internal(
101 config.state_cache_capacity(),
102 ),
103 meta: StateCache::<MetaKey, GroupMeta<C>>::new_internal(config.internal_state_cache_capacity()),
104 meta_low_water: None,
105 expire_batch: config.expire_batch(),
106 _pd: PhantomData,
107 }
108 }
109
110 pub fn apply<S, K, NA>(
111 &mut self,
112 store: &mut S,
113 buckets: TumblingBuckets<G, C, Accumulator::Contribution>,
114 row_key: K,
115 new_accumulator: NA,
116 ) -> Result<Vec<WindowResult<G, C, Accumulator::Output>>>
117 where
118 S: WindowStore,
119 K: Fn(&G, C) -> EncodedKey,
120 NA: Fn() -> Accumulator,
121 {
122 if buckets.is_empty() {
123 return Ok(Vec::new());
124 }
125 let mut meta_loaded = self.warm_and_load_meta(store, &buckets)?;
126 let slot_resolved = self.resolve_survivor_rows(store, &buckets, &meta_loaded, &row_key)?;
127 let results =
128 self.apply_events(store, buckets, slot_resolved, &mut meta_loaded, &row_key, &new_accumulator)?;
129 self.persist_meta(store, meta_loaded)?;
130 Ok(results)
131 }
132
133 pub fn flush<S: WindowStore>(&mut self, store: &mut S) -> Result<()> {
134 self.accumulators.flush(store)?;
135 self.meta.flush(store)?;
136 Ok(())
137 }
138
139 fn warm_and_load_meta<S: WindowStore>(
140 &mut self,
141 store: &mut S,
142 buckets: &TumblingBuckets<G, C, Accumulator::Contribution>,
143 ) -> Result<MetaLoaded<G, C>> {
144 let meta_keys: Vec<MetaKey> = buckets
145 .keys()
146 .map(|(group, _)| group)
147 .collect::<BTreeSet<_>>()
148 .into_iter()
149 .map(meta_key_for)
150 .collect();
151 self.meta.warm(store, &meta_keys)?;
152
153 let mut meta_loaded: MetaLoaded<G, C> = HashMap::new();
154 for (group, _) in buckets.keys() {
155 if !meta_loaded.contains_key(group) {
156 let m = self.meta.get(store, &meta_key_for(group))?.unwrap_or_default();
157 meta_loaded.insert(group.clone(), m);
158 }
159 }
160 Ok(meta_loaded)
161 }
162
163 fn resolve_survivor_rows<S, K>(
164 &mut self,
165 store: &mut S,
166 buckets: &TumblingBuckets<G, C, Accumulator::Contribution>,
167 meta_loaded: &MetaLoaded<G, C>,
168 row_key: &K,
169 ) -> Result<SlotResolved>
170 where
171 S: WindowStore,
172 K: Fn(&G, C) -> EncodedKey,
173 {
174 let mut survivor_keys: Vec<EncodedKey> = Vec::new();
175 let mut slot_survives: Vec<bool> = Vec::with_capacity(buckets.len());
176 for (group, span) in buckets.keys() {
177 let initial_high_water = meta_loaded.get(group).and_then(|m| m.high_water);
178 let survives = initial_high_water.is_none_or(|hw| span.start >= hw);
179 slot_survives.push(survives);
180 if survives {
181 survivor_keys.push(row_key(group, span.start));
182 }
183 }
184 let resolved_rows = store.get_or_create_row_numbers(&survivor_keys)?;
185 reifydb_assertions! {
186 let survivors = survivor_keys.len();
187 let resolved = resolved_rows.len();
188 assert!(
189 resolved == survivors,
190 "get_or_create_row_numbers must return exactly one row per survivor key; a short batch would \
191 leave a surviving slot with no resolved row, so the slot_resolved zip below pairs it with None \
192 and apply_events silently re-creates a fresh row instead of reusing the existing window \
193 state, double-counting it (survivor_keys={survivors}, resolved_rows={resolved})"
194 );
195 }
196 let accumulator_keys: Vec<WindowStateKey> =
197 resolved_rows.iter().map(|(rn, _)| WindowStateKey(*rn)).collect();
198 self.accumulators.warm(store, &accumulator_keys)?;
199 let mut resolved_rows = resolved_rows.into_iter();
200 let slot_resolved: SlotResolved = slot_survives
201 .into_iter()
202 .map(|survives| {
203 if survives {
204 resolved_rows.next()
205 } else {
206 None
207 }
208 })
209 .collect();
210 Ok(slot_resolved)
211 }
212
213 fn apply_events<S, K, NA>(
214 &mut self,
215 store: &mut S,
216 buckets: TumblingBuckets<G, C, Accumulator::Contribution>,
217 slot_resolved: SlotResolved,
218 meta_loaded: &mut MetaLoaded<G, C>,
219 row_key: &K,
220 new_accumulator: &NA,
221 ) -> Result<Vec<WindowResult<G, C, Accumulator::Output>>>
222 where
223 S: WindowStore,
224 K: Fn(&G, C) -> EncodedKey,
225 NA: Fn() -> Accumulator,
226 {
227 let mut results: Vec<WindowResult<G, C, Accumulator::Output>> = Vec::new();
228
229 for (((group, span), events), slot_pre) in buckets.into_iter().zip(slot_resolved) {
230 let entry = meta_loaded.entry(group.clone()).or_default();
231 match entry.high_water {
232 Some(hw) if span.start > hw => entry.high_water = Some(span.start),
233 None => entry.high_water = Some(span.start),
234 _ => {}
235 }
236
237 let (row_number, is_new) = match slot_pre {
238 Some(resolved) => resolved,
239 None => {
240 let key = row_key(&group, span.start);
241 store.get_or_create_row_number(&key)?
242 }
243 };
244
245 let mut accumulator: Accumulator = self
246 .accumulators
247 .get(store, &WindowStateKey(row_number))?
248 .unwrap_or_else(new_accumulator);
249 let was_empty_before = accumulator.is_empty();
250 let prior = if was_empty_before {
251 None
252 } else {
253 accumulator.finalize()
254 };
255
256 for event in events {
257 match event {
258 AccumulatorEvent::Add(c) => {
259 accumulator.add(&c);
260 }
261 AccumulatorEvent::Remove(c) => {
262 if accumulator.is_empty() {
263 continue;
264 }
265 accumulator.remove(&c);
266 }
267 }
268 }
269
270 let value = accumulator.finalize();
271 self.accumulators.put(store, &WindowStateKey(row_number), accumulator)?;
272
273 match value {
274 Some(value) => {
275 let kind = if is_new || was_empty_before {
276 EmitKind::Insert
277 } else {
278 EmitKind::Update
279 };
280 results.push(WindowResult {
281 row_number,
282 group,
283 span,
284 value,
285 prior,
286 kind,
287 });
288 }
289 None => {
290 if let Some(p) = prior.clone() {
291 results.push(WindowResult {
292 row_number,
293 group,
294 span,
295 value: p,
296 prior,
297 kind: EmitKind::Remove,
298 });
299 }
300 }
301 }
302 }
303 Ok(results)
304 }
305
306 pub fn expire<S: WindowStore>(
307 &mut self,
308 store: &mut S,
309 threshold: u64,
310 ) -> Result<Vec<ExpiredWindow<G, C, Accumulator::Output>>> {
311 let mut due: Vec<(EncodedKey, TumblingIndexEntry<G, C>)> = Vec::new();
312 store.internal_range_visit::<TumblingIndexEntry<G, C>>(
313 expiry_due_range(threshold),
314 Some(self.expire_batch),
315 &mut |key, entry| {
316 due.push((key, entry));
317 Ok(())
318 },
319 )?;
320
321 let mut out: Vec<ExpiredWindow<G, C, Accumulator::Output>> = Vec::new();
322 for (index_key, entry) in due {
323 let row_number = RowNumber(entry.row_number);
324 store.internal_drop(&index_key)?;
325 let value = self
326 .accumulators
327 .get(store, &WindowStateKey(row_number))?
328 .and_then(|accumulator| accumulator.finalize());
329 self.accumulators.remove(store, &WindowStateKey(row_number))?;
330 out.push(ExpiredWindow {
331 row_number,
332 group: entry.group,
333 window_start: entry.window_start,
334 value,
335 });
336 }
337 Ok(out)
338 }
339
340 fn persist_meta<S: WindowStore>(&mut self, store: &mut S, meta_loaded: MetaLoaded<G, C>) -> Result<()> {
341 for (group, meta) in meta_loaded {
342 self.meta.set(store, &meta_key_for(&group), &meta)?;
343 }
344 Ok(())
345 }
346
347 pub fn expire_meta<S: WindowStore>(&mut self, store: &mut S, threshold: u64) -> Result<usize> {
348 sweep_stale_meta(store, &mut self.meta, threshold, &mut self.meta_low_water)
349 }
350}
351
352#[cfg(test)]
353mod tests {
354 use std::collections::BTreeMap;
355
356 use reifydb_codec::key::encoded::EncodedKey;
357
358 use crate::window::{
359 engine::{
360 AccumulatorEvent, EmitKind, WindowResult,
361 config::WindowEngineConfig,
362 test_support::{MockStore, SumAccumulator},
363 tumbling::{TumblingBuckets, TumblingEngine, reindex_window},
364 },
365 span::WindowSpan,
366 };
367
368 fn test_config() -> WindowEngineConfig {
369 WindowEngineConfig::builder().state_cache_capacity(8).internal_state_cache_capacity(64).build()
370 }
371
372 fn row_key(group: &u32, window_start: u64) -> EncodedKey {
373 EncodedKey::builder().u32(*group).u64(window_start).build()
374 }
375
376 fn seed_window(store: &mut MockStore, window_start: u64, contribution: i64) -> WindowResult<u32, u64, i64> {
377 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
378 let mut buckets: TumblingBuckets<u32, u64, i64> = BTreeMap::new();
379 buckets.insert(
380 (1u32, WindowSpan::new(window_start, window_start + 1)),
381 vec![AccumulatorEvent::Add(contribution)],
382 );
383 let mut results = engine.apply(store, buckets, row_key, SumAccumulator::default).expect("apply");
384 engine.flush(store).expect("flush");
385 results.pop().expect("one window")
386 }
387
388 #[test]
389 fn expire_returns_only_due_windows_and_clears_their_state() {
390 let mut store = MockStore::default();
391 let w0 = seed_window(&mut store, 0, 5);
393 reindex_window(&mut store, &w0.group, w0.span.start, w0.row_number, None, Some(10)).unwrap();
394 let w100 = seed_window(&mut store, 100, 7);
395 reindex_window(&mut store, &w100.group, w100.span.start, w100.row_number, None, Some(90)).unwrap();
396 assert_eq!(store.index_entry_count(), 2, "both live windows are indexed");
397
398 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
400 let expired = engine.expire(&mut store, 10).unwrap();
401 engine.flush(&mut store).unwrap();
402 assert_eq!(expired.len(), 1, "exactly one window is due, not the whole population");
403 assert_eq!(expired[0].window_start, 0);
404 assert_eq!(expired[0].value, Some(5));
405 assert_eq!(store.index_entry_count(), 1, "the due window's index entry is gone, the other remains");
406
407 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
409 let later = engine.expire(&mut store, 1000).unwrap();
410 assert_eq!(later.len(), 1);
411 assert_eq!(later[0].window_start, 100);
412 assert_eq!(later[0].value, Some(7));
413 assert_eq!(store.index_entry_count(), 0);
414 }
415
416 #[test]
417 fn meta_reclaimed_when_group_stale_past_threshold() {
418 let mut store = MockStore::default();
423 seed_window(&mut store, 0, 5);
424 assert_eq!(store.meta_entry_count(), 1, "applying a window persisted the group's meta");
425
426 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
427 let dropped = engine.expire_meta(&mut store, 100).unwrap();
428 assert_eq!(dropped, 1, "the group's high water (0) is below the threshold (100)");
429 assert_eq!(store.meta_entry_count(), 0, "a stale group must not leak its GroupMeta");
430 }
431
432 #[test]
433 fn meta_survives_while_group_high_water_at_or_after_threshold() {
434 let mut store = MockStore::default();
437 seed_window(&mut store, 100, 7);
438 assert_eq!(store.meta_entry_count(), 1);
439
440 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
441 let dropped = engine.expire_meta(&mut store, 50).unwrap();
442 assert_eq!(dropped, 0, "high water (100) is not below the threshold (50)");
443 assert_eq!(store.meta_entry_count(), 1, "a group within the staleness horizon keeps its meta");
444 }
445
446 #[test]
447 fn meta_sweep_leaves_row_number_mappings_intact() {
448 let mut store = MockStore::default();
452 seed_window(&mut store, 0, 5);
453 store.seed_mapping_key(0x01);
454 assert_eq!(store.mapping_entry_count(), 1);
455
456 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
457 engine.expire_meta(&mut store, 100).unwrap();
458 assert_eq!(store.meta_entry_count(), 0, "the stale group's meta is swept");
459 assert_eq!(store.mapping_entry_count(), 1, "the sweep must not touch row-number mapping keys");
460 }
461
462 #[test]
463 fn meta_sweep_skips_then_reclaims_as_threshold_advances() {
464 let mut store = MockStore::default();
468 seed_window(&mut store, 100, 7);
469
470 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
471 assert_eq!(engine.expire_meta(&mut store, 50).unwrap(), 0);
473 assert_eq!(store.meta_entry_count(), 1);
474 assert_eq!(engine.expire_meta(&mut store, 100).unwrap(), 0);
476 assert_eq!(store.meta_entry_count(), 1);
477 assert_eq!(engine.expire_meta(&mut store, 101).unwrap(), 1);
479 assert_eq!(store.meta_entry_count(), 0, "the guard must not permanently skip a group that goes stale");
480 }
481
482 #[test]
483 fn expire_threshold_is_inclusive() {
484 let mut store = MockStore::default();
485 let w = seed_window(&mut store, 0, 4);
486 reindex_window(&mut store, &w.group, w.span.start, w.row_number, None, Some(50)).unwrap();
487
488 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
490 assert!(engine.expire(&mut store, 49).unwrap().is_empty());
491 engine.flush(&mut store).unwrap();
492 assert_eq!(store.index_entry_count(), 1);
493
494 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
496 assert_eq!(engine.expire(&mut store, 50).unwrap().len(), 1);
497 }
498
499 #[test]
500 fn expire_processes_at_most_expire_batch_then_resumes_next_tick() {
501 let mut store = MockStore::default();
508 for (start, due) in [(0u64, 10u64), (100, 20), (200, 30)] {
509 let w = seed_window(&mut store, start, 1);
510 reindex_window(&mut store, &w.group, w.span.start, w.row_number, None, Some(due)).unwrap();
511 }
512 assert_eq!(store.index_entry_count(), 3);
513
514 let capped = WindowEngineConfig::builder()
515 .state_cache_capacity(8)
516 .internal_state_cache_capacity(64)
517 .expire_batch(2)
518 .build();
519
520 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(capped);
521 let first = engine.expire(&mut store, 1000).unwrap();
522 engine.flush(&mut store).unwrap();
523 assert_eq!(first.len(), 2, "one tick drains at most expire_batch windows");
524 assert_eq!(first[0].window_start, 200, "inverted key order: newest due drains first");
525 assert_eq!(first[1].window_start, 100);
526 assert_eq!(store.index_entry_count(), 1, "the deferred window keeps its index entry");
527
528 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(capped);
529 let second = engine.expire(&mut store, 1000).unwrap();
530 engine.flush(&mut store).unwrap();
531 assert_eq!(second.len(), 1, "the next tick picks up the deferred backlog");
532 assert_eq!(second[0].window_start, 0);
533 assert_eq!(second[0].value, Some(1), "a deferred window still finalizes with its state intact");
534 assert_eq!(store.index_entry_count(), 0);
535 }
536
537 #[test]
538 fn reindex_rekeys_without_leaving_a_stale_entry() {
539 let mut store = MockStore::default();
540 let w = seed_window(&mut store, 0, 9);
541 reindex_window(&mut store, &w.group, w.span.start, w.row_number, None, Some(10)).unwrap();
543 reindex_window(&mut store, &w.group, w.span.start, w.row_number, Some(10), Some(80)).unwrap();
544 assert_eq!(store.index_entry_count(), 1, "re-keying must not leave the old entry behind");
545
546 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
547 assert!(engine.expire(&mut store, 10).unwrap().is_empty(), "no longer due at the old expiry");
548 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
549 assert_eq!(engine.expire(&mut store, 80).unwrap().len(), 1, "due at the new expiry");
550 }
551
552 #[test]
553 fn accumulator_survives_restart() {
554 let mut store = MockStore::default();
561
562 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
563 let mut buckets: TumblingBuckets<u32, u64, i64> = BTreeMap::new();
564 buckets.insert((1u32, WindowSpan::new(0, 1)), vec![AccumulatorEvent::Add(5)]);
565 let published: Vec<WindowResult<u32, u64, i64>> =
566 engine.apply(&mut store, buckets, row_key, SumAccumulator::default).unwrap();
567 engine.flush(&mut store).unwrap();
568 assert_eq!(published.len(), 1);
569 assert!(matches!(published[0].kind, EmitKind::Insert));
570 assert_eq!(published[0].value, 5);
571
572 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
574 let mut buckets: TumblingBuckets<u32, u64, i64> = BTreeMap::new();
575 buckets.insert((1u32, WindowSpan::new(0, 1)), vec![AccumulatorEvent::Remove(5)]);
576 let withdrawn: Vec<WindowResult<u32, u64, i64>> =
577 engine.apply(&mut store, buckets, row_key, SumAccumulator::default).unwrap();
578 engine.flush(&mut store).unwrap();
579
580 assert_eq!(withdrawn.len(), 1, "emptying the window emits exactly one terminal diff");
581 assert!(
582 matches!(withdrawn[0].kind, EmitKind::Remove),
583 "the window emptied under retraction, so the last published row must be withdrawn"
584 );
585 assert_eq!(withdrawn[0].value, 5, "the withdrawn value is the reloaded pre-batch accumulator output");
586 assert_eq!(
587 withdrawn[0].row_number, published[0].row_number,
588 "the withdrawal targets the same row that was published"
589 );
590 }
591
592 #[test]
593 fn accumulator_survives_lru_eviction() {
594 let mut store = MockStore::default();
599 let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new(test_config());
600
601 let mut published_group_1: Vec<WindowResult<u32, u64, i64>> = Vec::new();
602 for group in 1u32..=11u32 {
603 let mut buckets: TumblingBuckets<u32, u64, i64> = BTreeMap::new();
604 buckets.insert((group, WindowSpan::new(0, 1)), vec![AccumulatorEvent::Add(i64::from(group))]);
605 let out: Vec<WindowResult<u32, u64, i64>> =
606 engine.apply(&mut store, buckets, row_key, SumAccumulator::default).unwrap();
607 if group == 1 {
608 published_group_1 = out;
609 }
610 }
611 engine.flush(&mut store).unwrap();
612 assert_eq!(published_group_1.len(), 1);
613 assert!(matches!(published_group_1[0].kind, EmitKind::Insert));
614 assert_eq!(published_group_1[0].value, 1);
615
616 let mut buckets: TumblingBuckets<u32, u64, i64> = BTreeMap::new();
619 buckets.insert((1u32, WindowSpan::new(0, 1)), vec![AccumulatorEvent::Remove(1)]);
620 let withdrawn: Vec<WindowResult<u32, u64, i64>> =
621 engine.apply(&mut store, buckets, row_key, SumAccumulator::default).unwrap();
622 engine.flush(&mut store).unwrap();
623
624 assert_eq!(withdrawn.len(), 1, "emptying the evicted window emits exactly one terminal diff");
625 assert!(
626 matches!(withdrawn[0].kind, EmitKind::Remove),
627 "the evicted window emptied under retraction, so the last published row must be withdrawn"
628 );
629 assert_eq!(withdrawn[0].value, 1, "the withdrawn value is the reloaded accumulator output for group 1");
630 assert_eq!(
631 withdrawn[0].row_number, published_group_1[0].row_number,
632 "the withdrawal targets the same row that was published for group 1"
633 );
634 }
635}