krabka_client_streams/store/
join_window.rs

1//! The retainDuplicates window store for stream-stream joins.
2//!
3//! It uses composite `WindowKeySchema` keys with a per-store incrementing
4//! seqnum, and it stores RAW values with no `ValueAndTimestamp` wrap. `fetch`
5//! returns every duplicate in a range.
6use std::any::Any;
7
8use async_trait::async_trait;
9use bytes::{BufMut, Bytes};
10
11use crate::{
12    processor::serde::Serde,
13    store::{
14        api::StateStore,
15        byte::{ByteKeyValueStore, InMemoryBytes},
16        snapshot::{SnapshotReader, put_byte_entries, read_byte_entries, store_payload},
17        window_schema::{key_bytes_of, store_key, window_start_of},
18    },
19};
20
21/// A typed windowed store that retains duplicates.
22///
23/// Every `put` with the same `(key, timestamp)` builds a distinct composite key
24/// from an incrementing seqnum, so the duplicates coexist in the backend. The
25/// store keeps the values RAW, with no `ValueAndTimestamp` wrap. The
26/// stream-stream join processors use it.
27#[async_trait]
28pub trait JoinWindowStore<K: Send + Sync, V: Send>: StateStore {
29    async fn put(&mut self, key: K, timestamp: i64, value: V);
30    async fn fetch(&self, key: &K, time_from: i64, time_to: i64) -> Vec<(i64, V)>;
31}
32
33pub struct JoinWindowBytesStore<K, V> {
34    name: String,
35    changelog_topic: String,
36    backend: Box<dyn ByteKeyValueStore>,
37    key_serde: Box<dyn Serde<K>>,
38    value_serde: Box<dyn Serde<V>>,
39    changelog: Vec<(Bytes, Option<Bytes>)>,
40    logging: bool,
41    seqnum: u32,
42}
43
44impl<K: 'static, V: 'static> JoinWindowBytesStore<K, V> {
45    #[must_use]
46    pub(crate) fn new(
47        name: String,
48        backend: Box<dyn ByteKeyValueStore>,
49        key_serde: Box<dyn Serde<K>>,
50        value_serde: Box<dyn Serde<V>>,
51        changelog_topic: String,
52    ) -> Self {
53        Self {
54            name,
55            changelog_topic,
56            backend,
57            key_serde,
58            value_serde,
59            changelog: Vec::new(),
60            logging: true,
61            seqnum: 0,
62        }
63    }
64
65    #[must_use]
66    pub fn in_memory(
67        name: String,
68        key_serde: Box<dyn Serde<K>>,
69        value_serde: Box<dyn Serde<V>>,
70        changelog_topic: String,
71    ) -> Self {
72        Self::new(
73            name,
74            Box::new(InMemoryBytes::default()),
75            key_serde,
76            value_serde,
77            changelog_topic,
78        )
79    }
80
81    fn next_seqnum(&mut self) -> u32 {
82        let s = self.seqnum;
83        self.seqnum = (self.seqnum + 1) & 0x7FFF_FFFF;
84        s
85    }
86}
87
88#[async_trait]
89impl<K: 'static, V: 'static> StateStore for JoinWindowBytesStore<K, V> {
90    fn name(&self) -> &str {
91        &self.name
92    }
93    async fn flush(&mut self) {}
94    fn close(&mut self) {}
95    fn as_any_mut(&mut self) -> &mut dyn Any {
96        self
97    }
98    fn changelog_topic(&self) -> &str {
99        &self.changelog_topic
100    }
101    fn take_changelog(&mut self) -> Vec<(Bytes, Option<Bytes>)> {
102        std::mem::take(&mut self.changelog)
103    }
104    async fn apply_changelog(&mut self, key: Bytes, value: Option<Bytes>) {
105        match value {
106            Some(v) => self.backend.put(key, v).await,
107            None => {
108                self.backend.delete(&key).await;
109            }
110        }
111    }
112    fn set_logging(&mut self, on: bool) {
113        self.logging = on;
114    }
115    async fn clear(&mut self) {
116        self.backend.clear().await;
117        self.changelog.clear();
118        self.seqnum = 0;
119    }
120    async fn snapshot(&mut self) -> Bytes {
121        // The seqnum rides with the entries. Without it a restored store would
122        // hand out composite keys that a retained duplicate already holds.
123        let mut buffer = store_payload();
124        buffer.put_u32(self.seqnum);
125        put_byte_entries(&mut buffer, &self.backend.scan_all().await);
126        buffer.freeze()
127    }
128    async fn restore_snapshot(
129        &mut self,
130        data: Bytes,
131    ) -> Result<(), crate::error::StreamsClientError> {
132        let mut reader = SnapshotReader::new(&data);
133        reader.store_version()?;
134        let seqnum = reader.u32()?;
135        let entries = read_byte_entries(&mut reader)?;
136        reader.finish()?;
137        self.backend.clear().await;
138        for (key, value) in entries {
139            self.backend.put(key, value).await;
140        }
141        self.seqnum = seqnum;
142        self.changelog.clear();
143        Ok(())
144    }
145}
146
147#[async_trait]
148impl<K: Send + Sync + 'static, V: Send + 'static> JoinWindowStore<K, V>
149    for JoinWindowBytesStore<K, V>
150{
151    async fn put(&mut self, key: K, timestamp: i64, value: V) {
152        let kb = self.key_serde.serialize(&self.changelog_topic, &key);
153        let seq = self.next_seqnum();
154        let sk = store_key(&kb, timestamp, seq);
155        let raw = self.value_serde.serialize(&self.changelog_topic, &value); // RAW — no ValueAndTimestamp wrap
156        self.backend.put(sk.clone(), raw.clone()).await;
157        if self.logging {
158            self.changelog.push((sk, Some(raw)));
159        }
160    }
161
162    async fn fetch(&self, key: &K, time_from: i64, time_to: i64) -> Vec<(i64, V)> {
163        let kb = self.key_serde.serialize(&self.changelog_topic, key);
164        // Clamp to 0: negative timestamps encode to 0xFF… in BE, which sorts
165        // after all positive entries in the lexicographic BTreeMap backend.
166        // Kafka's segment-scoped range fetch does the same (`Math.max(0, timeFrom)`).
167        let lo = store_key(&kb, time_from.max(0), 0);
168        let hi = store_key(&kb, time_to.saturating_add(1), 0);
169        let mut out = Vec::new();
170        for (k, raw) in self.backend.range(&lo, &hi).await {
171            // guard prefix collisions: only return entries whose inner key matches
172            if key_bytes_of(&k) != kb.as_ref() {
173                continue;
174            }
175            out.push((
176                window_start_of(&k),
177                self.value_serde
178                    .deserialize(&self.changelog_topic, &raw)
179                    .expect("join window value deserialize"),
180            ));
181        }
182        out
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use crate::processor::serde::StringSerde;
190
191    #[tokio::test]
192    async fn put_keeps_duplicates_and_fetch_returns_all() {
193        let mut s = JoinWindowBytesStore::<String, String>::in_memory(
194            "j".into(),
195            Box::new(StringSerde),
196            Box::new(StringSerde),
197            "app-j-changelog".into(),
198        );
199        s.put("k".into(), 5, "a".into()).await;
200        s.put("k".into(), 5, "b".into()).await; // SAME (key, ts) → kept (seqnum increments)
201        s.put("k".into(), 7, "c".into()).await;
202        assert_eq!(
203            s.fetch(&"k".to_string(), 5, 7).await,
204            vec![
205                (5, "a".to_string()),
206                (5, "b".to_string()),
207                (7, "c".to_string())
208            ]
209        );
210        assert_eq!(
211            s.fetch(&"k".to_string(), 5, 5).await,
212            vec![(5, "a".to_string()), (5, "b".to_string())]
213        );
214        assert_eq!(s.take_changelog().len(), 3);
215    }
216}
krabka_client_streams/store/join_window.rs

krabka_client_streams/store/
join_window.rs
