krabka_client_streams/runtime/
iq_view.rs

1//! Typed, read-only composite views over the local state stores, for
2//! Interactive Queries.
3//!
4//! Each view owns its serdes and round-trips byte-level `IqRequest`s to the
5//! supervisor. The results are eagerly materialized `Vec`s. This is one
6//! deliberate difference from the JVM's lazy `KeyValueIterator`.
7
8use bytes::Bytes;
9use tokio::sync::{mpsc, oneshot};
10
11use crate::{
12    dsl::windows::{Window, Windowed},
13    error::StreamsClientError,
14    processor::serde::Serde,
15    runtime::iq::{IqError, IqOp, IqPayload, IqRequest},
16    store::iq::StoreKind,
17};
18
19/// Round-trip one operation to the supervisor. All three views share this
20/// function.
21async fn query(
22    tx: &mpsc::Sender<IqRequest>,
23    store: &str,
24    kind: StoreKind,
25    op: IqOp,
26) -> Result<IqPayload, StreamsClientError> {
27    let (reply, rx) = oneshot::channel();
28    tx.send(IqRequest {
29        store: store.to_string(),
30        kind,
31        op,
32        reply,
33    })
34    .await
35    .map_err(|_| StreamsClientError::InteractiveQuery(IqError::RebalanceInProgress))?;
36    rx.await
37        .map_err(|_| StreamsClientError::InteractiveQuery(IqError::RebalanceInProgress))?
38        .map_err(StreamsClientError::InteractiveQuery)
39}
40
41fn deser<T: 'static>(
42    topic: &str,
43    serde: &dyn Serde<T>,
44    bytes: &[u8],
45) -> Result<T, StreamsClientError> {
46    serde
47        .deserialize(topic, bytes)
48        .map_err(|e| StreamsClientError::Runtime(format!("iq deserialize: {e}")))
49}
50
51pub(crate) fn unexpected(p: &IqPayload) -> StreamsClientError {
52    StreamsClientError::Runtime(format!("iq: unexpected payload {p:?}"))
53}
54
55/// Check that a store exists locally and has the requested kind. The accessors
56/// run this check eagerly.
57#[tracing::instrument(
58    name = "streams.iq.validate",
59    level = "debug",
60    skip_all,
61    fields(store = %store, kind = ?kind),
62    err,
63)]
64pub(crate) async fn validate(
65    tx: &mpsc::Sender<IqRequest>,
66    store: &str,
67    kind: StoreKind,
68) -> Result<(), StreamsClientError> {
69    match query(tx, store, kind, IqOp::Validate).await? {
70        IqPayload::Validated => Ok(()),
71        other => Err(unexpected(&other)),
72    }
73}
74
75/// Read-only composite KV store view, for Interactive Queries.
76pub struct ReadOnlyKeyValueStore<K, V> {
77    pub(crate) tx: mpsc::Sender<IqRequest>,
78    pub(crate) store: String,
79    pub(crate) key_serde: Box<dyn Serde<K>>,
80    pub(crate) value_serde: Box<dyn Serde<V>>,
81}
82
83impl<K: 'static, V: 'static> ReadOnlyKeyValueStore<K, V> {
84    /// The value for `key`, or `None` when the key is absent.
85    /// # Errors
86    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
87    pub async fn get(&self, key: &K) -> Result<Option<V>, StreamsClientError> {
88        let kb = self.key_serde.serialize(&self.store, key);
89        match query(
90            &self.tx,
91            &self.store,
92            StoreKind::KeyValue,
93            IqOp::KvGet { key: kb },
94        )
95        .await?
96        {
97            IqPayload::Value(Some(vb)) => Ok(Some(deser(&self.store, &*self.value_serde, &vb)?)),
98            IqPayload::Value(None) => Ok(None),
99            other => Err(unexpected(&other)),
100        }
101    }
102
103    /// Inclusive `[lo, hi]` range, ascending memcmp key order.
104    /// # Errors
105    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
106    pub async fn range(&self, lo: &K, hi: &K) -> Result<Vec<(K, V)>, StreamsClientError> {
107        let lo_b = self.key_serde.serialize(&self.store, lo);
108        let hi_b = self.key_serde.serialize(&self.store, hi);
109        match query(
110            &self.tx,
111            &self.store,
112            StoreKind::KeyValue,
113            IqOp::KvRange { lo: lo_b, hi: hi_b },
114        )
115        .await?
116        {
117            IqPayload::Entries(pairs) => self.decode_pairs(pairs),
118            other => Err(unexpected(&other)),
119        }
120    }
121
122    /// Every entry.
123    /// # Errors
124    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
125    pub async fn all(&self) -> Result<Vec<(K, V)>, StreamsClientError> {
126        match query(&self.tx, &self.store, StoreKind::KeyValue, IqOp::KvAll).await? {
127            IqPayload::Entries(pairs) => self.decode_pairs(pairs),
128            other => Err(unexpected(&other)),
129        }
130    }
131
132    /// The approximate entry count. It is exact for an in-memory store, and it
133    /// is summed across the partitions.
134    /// # Errors
135    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
136    pub async fn approximate_num_entries(&self) -> Result<u64, StreamsClientError> {
137        match query(
138            &self.tx,
139            &self.store,
140            StoreKind::KeyValue,
141            IqOp::KvApproxCount,
142        )
143        .await?
144        {
145            IqPayload::Count(n) => Ok(n),
146            other => Err(unexpected(&other)),
147        }
148    }
149
150    fn decode_pairs(&self, pairs: Vec<(Bytes, Bytes)>) -> Result<Vec<(K, V)>, StreamsClientError> {
151        pairs
152            .into_iter()
153            .map(|(kb, vb)| {
154                Ok((
155                    deser(&self.store, &*self.key_serde, &kb)?,
156                    deser(&self.store, &*self.value_serde, &vb)?,
157                ))
158            })
159            .collect()
160    }
161}
162
163/// Read-only composite window store view. `fetch` gives `(windowStart, V)`
164/// pairs.
165pub struct ReadOnlyWindowStore<K, V> {
166    pub(crate) tx: mpsc::Sender<IqRequest>,
167    pub(crate) store: String,
168    pub(crate) key_serde: Box<dyn Serde<K>>,
169    pub(crate) value_serde: Box<dyn Serde<V>>,
170}
171
172impl<K: 'static, V: 'static> ReadOnlyWindowStore<K, V> {
173    /// The value of the window for `key` that starts exactly at `window_start`,
174    /// or `None`.
175    /// # Errors
176    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
177    pub async fn fetch_single(
178        &self,
179        key: &K,
180        window_start: i64,
181    ) -> Result<Option<V>, StreamsClientError> {
182        let kb = self.key_serde.serialize(&self.store, key);
183        match query(
184            &self.tx,
185            &self.store,
186            StoreKind::Window,
187            IqOp::WindowFetchSingle {
188                key: kb,
189                window_start,
190            },
191        )
192        .await?
193        {
194            IqPayload::Value(Some(vb)) => Ok(Some(deser(&self.store, &*self.value_serde, &vb)?)),
195            IqPayload::Value(None) => Ok(None),
196            other => Err(unexpected(&other)),
197        }
198    }
199
200    /// The windows for `key` whose start is in the inclusive range
201    /// `[time_from, time_to]`, in ascending order by start. Each item is a
202    /// `(windowStart, value)` pair.
203    /// # Errors
204    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
205    pub async fn fetch(
206        &self,
207        key: &K,
208        time_from: i64,
209        time_to: i64,
210    ) -> Result<Vec<(i64, V)>, StreamsClientError> {
211        let kb = self.key_serde.serialize(&self.store, key);
212        match query(
213            &self.tx,
214            &self.store,
215            StoreKind::Window,
216            IqOp::WindowFetch {
217                key: kb,
218                time_from,
219                time_to,
220            },
221        )
222        .await?
223        {
224            IqPayload::WindowEntries(rows) => rows
225                .into_iter()
226                .map(|(t, vb)| Ok((t, deser(&self.store, &*self.value_serde, &vb)?)))
227                .collect(),
228            other => Err(unexpected(&other)),
229        }
230    }
231}
232
233/// Read-only composite session store view. `fetch` gives each session as a
234/// `Windowed<K>`, which holds the key and `[start, end]`, together with its
235/// value.
236pub struct ReadOnlySessionStore<K, V> {
237    pub(crate) tx: mpsc::Sender<IqRequest>,
238    pub(crate) store: String,
239    pub(crate) key_serde: Box<dyn Serde<K>>,
240    pub(crate) value_serde: Box<dyn Serde<V>>,
241}
242
243impl<K: 'static, V: 'static> ReadOnlySessionStore<K, V> {
244    /// All sessions for `key`, in store order.
245    /// # Errors
246    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
247    pub async fn fetch(&self, key: &K) -> Result<Vec<(Windowed<K>, V)>, StreamsClientError> {
248        let kb = self.key_serde.serialize(&self.store, key);
249        match query(
250            &self.tx,
251            &self.store,
252            StoreKind::Session,
253            IqOp::SessionFetchKey { key: kb },
254        )
255        .await?
256        {
257            IqPayload::SessionEntries(rows) => rows
258                .into_iter()
259                .map(|((start, end), vb)| {
260                    // Re-deserialize the key per row (avoids a `K: Clone` bound).
261                    let k = deser(
262                        &self.store,
263                        &*self.key_serde,
264                        &self.key_serde.serialize(&self.store, key),
265                    )?;
266                    Ok((
267                        Windowed {
268                            key: k,
269                            window: Window { start, end },
270                        },
271                        deser(&self.store, &*self.value_serde, &vb)?,
272                    ))
273                })
274                .collect(),
275            other => Err(unexpected(&other)),
276        }
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use krabka_units::prelude::*;
283
284    use super::*;
285    use crate::{
286        processor::serde::{I64Serde, StringSerde},
287        runtime::iq::answer_iq,
288        store::{api::KeyValueStore, kv::KeyValueBytesStore, registry::StoreRegistry},
289    };
290
291    /// Spawn a small servicer over one registry. It returns the sender that the
292    /// views use.
293    pub(super) fn servicer(reg: StoreRegistry) -> mpsc::Sender<IqRequest> {
294        let (tx, mut rx) = mpsc::channel::<IqRequest>(16);
295        tokio::spawn(async move {
296            while let Some(req) = rx.recv().await {
297                let matching = reg.iq_get(&req.store).into_iter().collect::<Vec<_>>();
298                let res = answer_iq(matching, req.kind, &req.op, &req.store, true).await;
299                let _ = req.reply.send(res);
300            }
301        });
302        tx
303    }
304
305    async fn kv_registry() -> StoreRegistry {
306        let mut s = KeyValueBytesStore::<String, i64>::in_memory(
307            "counts".into(),
308            Box::new(StringSerde),
309            Box::new(I64Serde),
310            "counts-changelog".into(),
311        );
312        for (k, v) in [("a", 1), ("b", 2), ("c", 3)] {
313            s.put(k.into(), v).await;
314        }
315        let mut reg = StoreRegistry::default();
316        reg.insert(Box::new(s));
317        reg
318    }
319
320    #[tokio::test]
321    async fn kv_view_get_range_all_count() {
322        let tx = servicer(kv_registry().await);
323        let view = ReadOnlyKeyValueStore::<String, i64> {
324            tx,
325            store: "counts".into(),
326            key_serde: Box::new(StringSerde),
327            value_serde: Box::new(I64Serde),
328        };
329        assert_eq!(view.get(&"b".to_string()).await.unwrap(), Some(2));
330        assert_eq!(view.get(&"z".to_string()).await.unwrap(), None);
331        let r = view
332            .range(&"a".to_string(), &"b".to_string())
333            .await
334            .unwrap();
335        assert_eq!(r, vec![("a".to_string(), 1), ("b".to_string(), 2)]);
336        assert_eq!(view.all().await.unwrap().len(), 3);
337        assert_eq!(view.approximate_num_entries().await.unwrap(), 3);
338    }
339
340    async fn window_registry() -> StoreRegistry {
341        use crate::store::window::{WindowBytesStore, WindowStore};
342        let mut s = WindowBytesStore::<String, i64>::in_memory(
343            "wc".into(),
344            Box::new(StringSerde),
345            Box::new(I64Serde),
346            "wc-changelog".into(),
347            secs(1),
348        );
349        s.put("k".into(), 0, 10, 5).await;
350        s.put("k".into(), 1000, 20, 1005).await;
351        let mut reg = StoreRegistry::default();
352        reg.insert(Box::new(s));
353        reg
354    }
355
356    #[tokio::test]
357    async fn window_view_fetch() {
358        let tx = servicer(window_registry().await);
359        let view = ReadOnlyWindowStore::<String, i64> {
360            tx,
361            store: "wc".into(),
362            key_serde: Box::new(StringSerde),
363            value_serde: Box::new(I64Serde),
364        };
365        assert_eq!(
366            view.fetch_single(&"k".to_string(), 0).await.unwrap(),
367            Some(10)
368        );
369        assert_eq!(view.fetch_single(&"k".to_string(), 5).await.unwrap(), None);
370        let r = view.fetch(&"k".to_string(), 0, 1000).await.unwrap();
371        assert_eq!(r, vec![(0, 10), (1000, 20)]);
372    }
373
374    async fn session_registry() -> StoreRegistry {
375        use crate::store::session::{SessionBytesStore, SessionStore};
376        let mut s = SessionBytesStore::<String, i64>::in_memory(
377            "sc".into(),
378            Box::new(StringSerde),
379            Box::new(I64Serde),
380            "sc-changelog".into(),
381        );
382        s.put("k".into(), 0, 10, 1).await;
383        s.put("k".into(), 20, 30, 2).await;
384        let mut reg = StoreRegistry::default();
385        reg.insert(Box::new(s));
386        reg
387    }
388
389    #[tokio::test]
390    async fn session_view_fetch() {
391        use crate::dsl::windows::Window;
392        let tx = servicer(session_registry().await);
393        let view = ReadOnlySessionStore::<String, i64> {
394            tx,
395            store: "sc".into(),
396            key_serde: Box::new(StringSerde),
397            value_serde: Box::new(I64Serde),
398        };
399        let rows = view.fetch(&"k".to_string()).await.unwrap();
400        let got: Vec<(Window, i64)> = rows.into_iter().map(|(w, v)| (w.window, v)).collect();
401        assert!(got.contains(&(Window { start: 0, end: 10 }, 1)));
402        assert!(got.contains(&(Window { start: 20, end: 30 }, 2)));
403        assert_eq!(got.len(), 2);
404    }
405}
krabka_client_streams/runtime/iq_view.rs

krabka_client_streams/runtime/
iq_view.rs
