krabka_client_streams/runtime/
iq.rs

1//! The interactive-query channel protocol.
2//!
3//! The `KafkaStreams` handle sends byte-level `IqRequest`s to the supervisor
4//! task. The supervisor resolves them against the local stores with `answer_iq`
5//! and replies on a `oneshot`.
6
7use bytes::Bytes;
8use tokio::sync::oneshot;
9
10use crate::store::iq::{IqQueryable, StoreKind};
11
12/// A byte-level query operation. It carries no `K` and no `V`, because the typed
13/// view serializes and deserializes.
14#[derive(Debug)]
15pub(crate) enum IqOp {
16    Validate,
17    KvGet {
18        key: Bytes,
19    },
20    KvRange {
21        lo: Bytes,
22        hi: Bytes,
23    },
24    KvAll,
25    KvApproxCount,
26    WindowFetchSingle {
27        key: Bytes,
28        window_start: i64,
29    },
30    WindowFetch {
31        key: Bytes,
32        time_from: i64,
33        time_to: i64,
34    },
35    SessionFetchKey {
36        key: Bytes,
37    },
38}
39
40/// A byte-level query result.
41#[derive(Debug, PartialEq, Eq)]
42pub(crate) enum IqPayload {
43    Validated,
44    Value(Option<Bytes>),
45    Entries(Vec<(Bytes, Bytes)>),
46    WindowEntries(Vec<(i64, Bytes)>),
47    SessionEntries(Vec<((i64, i64), Bytes)>),
48    Count(u64),
49}
50
51/// Why an interactive query failed.
52#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
53pub enum IqError {
54    #[error("state store {0:?} is not assigned to this instance")]
55    StoreNotFound(String),
56    #[error("state store {name:?} is a {found:?} store, not {requested:?}")]
57    WrongStoreKind {
58        name: String,
59        found: StoreKind,
60        requested: StoreKind,
61    },
62    #[error("a rebalance is in progress; retry the query")]
63    RebalanceInProgress,
64}
65
66/// One query addressed to the supervisor.
67pub(crate) struct IqRequest {
68    pub store: String,
69    pub kind: StoreKind,
70    pub op: IqOp,
71    pub reply: oneshot::Sender<Result<IqPayload, IqError>>,
72}
73
74/// Resolve one operation against every local store named `store`, which gives a
75/// composite result across the partitions.
76///
77/// `matching` is the set of `IqQueryable` views for that name on this instance.
78/// `any_tasks` separates two cases: the instance is rebalancing and has no tasks
79/// at all, or the instance has tasks but none of them hosts this store.
80pub(crate) async fn answer_iq(
81    matching: Vec<&dyn IqQueryable>,
82    kind: StoreKind,
83    op: &IqOp,
84    store: &str,
85    any_tasks: bool,
86) -> Result<IqPayload, IqError> {
87    if matching.is_empty() {
88        return Err(if any_tasks {
89            IqError::StoreNotFound(store.to_string())
90        } else {
91            IqError::RebalanceInProgress
92        });
93    }
94    let found = matching[0].kind();
95    if found != kind {
96        return Err(IqError::WrongStoreKind {
97            name: store.to_string(),
98            found,
99            requested: kind,
100        });
101    }
102    Ok(match op {
103        IqOp::Validate => IqPayload::Validated,
104        IqOp::KvGet { key } => {
105            let mut hit = None;
106            for s in &matching {
107                if let Some(v) = s.iq_kv_get(key).await {
108                    hit = Some(v);
109                    break;
110                }
111            }
112            IqPayload::Value(hit)
113        }
114        IqOp::KvRange { lo, hi } => {
115            let mut out = Vec::new();
116            for s in &matching {
117                out.extend(s.iq_kv_range(lo, hi).await);
118            }
119            IqPayload::Entries(out)
120        }
121        IqOp::KvAll => {
122            let mut out = Vec::new();
123            for s in &matching {
124                out.extend(s.iq_kv_all().await);
125            }
126            IqPayload::Entries(out)
127        }
128        IqOp::KvApproxCount => {
129            let mut n = 0;
130            for s in &matching {
131                n += s.iq_kv_approx_count().await;
132            }
133            IqPayload::Count(n)
134        }
135        IqOp::WindowFetchSingle { key, window_start } => {
136            let mut hit = None;
137            for s in &matching {
138                if let Some(v) = s.iq_window_fetch_single(key, *window_start).await {
139                    hit = Some(v);
140                    break;
141                }
142            }
143            IqPayload::Value(hit)
144        }
145        IqOp::WindowFetch {
146            key,
147            time_from,
148            time_to,
149        } => {
150            let mut out = Vec::new();
151            for s in &matching {
152                out.extend(s.iq_window_fetch(key, *time_from, *time_to).await);
153            }
154            IqPayload::WindowEntries(out)
155        }
156        IqOp::SessionFetchKey { key } => {
157            let mut out = Vec::new();
158            for s in &matching {
159                out.extend(s.iq_session_fetch_key(key).await);
160            }
161            IqPayload::SessionEntries(out)
162        }
163    })
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::{
170        processor::serde::{I64Serde, Serde, StringSerde},
171        store::{
172            api::{KeyValueStore, StateStore},
173            kv::KeyValueBytesStore,
174        },
175    };
176
177    #[tokio::test]
178    async fn answer_kv_get_validate_wrongkind_notfound() {
179        let mut s = KeyValueBytesStore::<String, i64>::in_memory(
180            "c".into(),
181            Box::new(StringSerde),
182            Box::new(I64Serde),
183            "c-changelog".into(),
184        );
185        s.put("x".into(), 7).await;
186        let q = s.as_iq().unwrap();
187
188        assert_eq!(
189            answer_iq(vec![q], StoreKind::KeyValue, &IqOp::Validate, "c", true).await,
190            Ok(IqPayload::Validated)
191        );
192        let got = answer_iq(
193            vec![q],
194            StoreKind::KeyValue,
195            &IqOp::KvGet {
196                key: StringSerde.serialize("t", &"x".to_string()),
197            },
198            "c",
199            true,
200        )
201        .await;
202        assert_eq!(got, Ok(IqPayload::Value(Some(I64Serde.serialize("t", &7)))));
203        assert!(matches!(
204            answer_iq(vec![q], StoreKind::Window, &IqOp::Validate, "c", true).await,
205            Err(IqError::WrongStoreKind { .. })
206        ));
207        assert_eq!(
208            answer_iq(
209                vec![],
210                StoreKind::KeyValue,
211                &IqOp::Validate,
212                "missing",
213                true
214            )
215            .await,
216            Err(IqError::StoreNotFound("missing".into()))
217        );
218        assert_eq!(
219            answer_iq(
220                vec![],
221                StoreKind::KeyValue,
222                &IqOp::Validate,
223                "missing",
224                false
225            )
226            .await,
227            Err(IqError::RebalanceInProgress)
228        );
229    }
230}
krabka_client_streams/runtime/iq.rs

krabka_client_streams/runtime/
iq.rs
