krabka_client_streams/runtime/iqv2/
request.rs

1//! `IQv2` request envelope: `Position`/`PositionBound` (KIP-796 bounded
2//! staleness), partition selection, and the finalized `StateQuery<Q>` that
3//! `KafkaStreams::query` consumes.
4
5use std::collections::{BTreeMap, BTreeSet};
6
7use super::query::Query;
8
9/// Source-topic consumed offsets folded into a store, as topic to partition to
10/// offset. The offset is the next offset to read, one past the last consumed
11/// record.
12#[derive(Debug, Clone, Default, PartialEq, Eq)]
13pub struct Position(pub BTreeMap<String, BTreeMap<i32, i64>>);
14
15impl Position {
16    /// Offset recorded for one topic-partition, if any.
17    #[must_use]
18    pub fn offset(&self, topic: &str, partition: i32) -> Option<i64> {
19        self.0.get(topic).and_then(|m| m.get(&partition)).copied()
20    }
21
22    /// True when `self` meets or exceeds every `(topic, partition)` offset in
23    /// `bound`. A bound that names a partition `self` has never advanced fails.
24    #[must_use]
25    pub fn dominates(&self, bound: &Position) -> bool {
26        bound.0.iter().all(|(topic, parts)| {
27            parts
28                .iter()
29                .all(|(p, off)| self.offset(topic, *p).is_some_and(|cur| cur >= *off))
30        })
31    }
32}
33
34/// Freshness bound for a query (KIP-796). `At` requires each partition's
35/// `Position` to dominate the given one. Otherwise that partition fails fast with
36/// `NotUpToBound`, and the query never blocks.
37#[derive(Debug, Clone, Default)]
38pub enum PositionBound {
39    #[default]
40    Unbounded,
41    At(Position),
42}
43
44/// Which local partitions to query.
45#[derive(Debug, Clone, Default)]
46pub(crate) enum PartitionSel {
47    #[default]
48    All,
49    Set(BTreeSet<i32>),
50}
51
52/// A finalized `IQv2` request, built with
53/// `StateQueryRequest::in_store(name).with_query(q)`.
54pub struct StateQuery<Q: Query> {
55    pub(crate) store: String,
56    pub(crate) query: Q,
57    pub(crate) partitions: PartitionSel,
58    pub(crate) bound: PositionBound,
59    pub(crate) require_active: bool,
60}
61
62impl<Q: Query> StateQuery<Q> {
63    /// Restrict the query to a set of local partitions. Default: all of them.
64    #[must_use]
65    pub fn with_partitions(mut self, set: BTreeSet<i32>) -> Self {
66        self.partitions = PartitionSel::Set(set);
67        self
68    }
69
70    /// The partitions the query is restricted to, or `None` when it reads
71    /// every locally assigned partition.
72    #[must_use]
73    pub fn partitions(&self) -> Option<&BTreeSet<i32>> {
74        match &self.partitions {
75            PartitionSel::All => None,
76            PartitionSel::Set(set) => Some(set),
77        }
78    }
79
80    /// Query all locally assigned partitions (the default).
81    #[must_use]
82    pub fn with_all_partitions(mut self) -> Self {
83        self.partitions = PartitionSel::All;
84        self
85    }
86
87    /// Require each queried partition to meet a freshness bound.
88    #[must_use]
89    pub fn with_position_bound(mut self, bound: PositionBound) -> Self {
90        self.bound = bound;
91        self
92    }
93
94    /// Serve only from active tasks, and not from standby or restoring tasks.
95    #[must_use]
96    pub fn require_active(mut self) -> Self {
97        self.require_active = true;
98        self
99    }
100}
101
102/// Entry point namespace: `StateQueryRequest::in_store("s").with_query(q)`.
103pub struct StateQueryRequest;
104
105impl StateQueryRequest {
106    /// Begin a request against state store `name`.
107    #[must_use]
108    pub fn in_store(name: impl Into<String>) -> StateQueryRequestBuilder {
109        StateQueryRequestBuilder { store: name.into() }
110    }
111}
112
113/// Half-built request that waits for `.with_query(q)`.
114pub struct StateQueryRequestBuilder {
115    store: String,
116}
117
118impl StateQueryRequestBuilder {
119    /// Attach the query and finalize the request.
120    #[must_use]
121    pub fn with_query<Q: Query>(self, query: Q) -> StateQuery<Q> {
122        StateQuery {
123            store: self.store,
124            query,
125            partitions: PartitionSel::All,
126            bound: PositionBound::Unbounded,
127            require_active: false,
128        }
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    fn pos(entries: &[(&str, i32, i64)]) -> Position {
137        let mut m: BTreeMap<String, BTreeMap<i32, i64>> = BTreeMap::new();
138        for (t, p, o) in entries {
139            m.entry((*t).to_string()).or_default().insert(*p, *o);
140        }
141        Position(m)
142    }
143
144    #[test]
145    fn dominates_requires_all_bound_partitions_met() {
146        let cur = pos(&[("in", 0, 10), ("in", 1, 5)]);
147        for (bound, want) in [
148            (pos(&[("in", 0, 10)]), true),
149            (pos(&[("in", 0, 9), ("in", 1, 5)]), true),
150            (pos(&[("in", 0, 11)]), false),   // behind
151            (pos(&[("other", 0, 1)]), false), // unknown tp
152        ] {
153            assert_eq!(cur.dominates(&bound), want);
154        }
155    }
156
157    #[test]
158    fn position_offset_reads_present_and_absent_tps() {
159        let p = pos(&[("in", 0, 10), ("in", 1, 5)]);
160        assert_eq!(p.offset("in", 0), Some(10));
161        assert_eq!(p.offset("in", 1), Some(5));
162        assert_eq!(p.offset("in", 2), None); // present topic, absent partition
163        assert_eq!(p.offset("other", 0), None); // absent topic
164    }
165
166    #[test]
167    fn partitions_reports_the_restriction() {
168        use assert2::check;
169
170        use crate::runtime::iqv2::query::KeyQuery;
171
172        let q = StateQueryRequest::in_store("s")
173            .with_query(KeyQuery::<String, i64>::with_key("k".into()));
174        check!(q.partitions().is_none());
175        let set: BTreeSet<i32> = [1, 3].into_iter().collect();
176        let q = q.with_partitions(set.clone());
177        check!(q.partitions() == Some(&set));
178        check!(q.with_all_partitions().partitions().is_none());
179    }
180
181    #[test]
182    fn builder_defaults_and_store_name() {
183        use crate::runtime::iqv2::query::KeyQuery;
184
185        let q = StateQueryRequest::in_store("s")
186            .with_query(KeyQuery::<String, i64>::with_key("k".into()));
187        assert_eq!(q.store, "s");
188        assert!(matches!(q.partitions, PartitionSel::All));
189        assert!(matches!(q.bound, PositionBound::Unbounded));
190        assert!(!q.require_active);
191    }
192
193    #[test]
194    fn builder_chain_sets_each_field() {
195        use crate::runtime::iqv2::query::KeyQuery;
196
197        let set: BTreeSet<i32> = [0, 2].into_iter().collect();
198        let q = StateQueryRequest::in_store("s")
199            .with_query(KeyQuery::<String, i64>::with_key("k".into()))
200            .with_partitions(set.clone());
201        match &q.partitions {
202            PartitionSel::Set(s) => assert_eq!(s, &set),
203            PartitionSel::All => panic!("expected explicit partition set"),
204        }
205
206        // with_all_partitions resets back to All.
207        let q = q.with_all_partitions();
208        assert!(matches!(q.partitions, PartitionSel::All));
209
210        // with_position_bound stores the bound.
211        let bound_pos = pos(&[("in", 0, 7)]);
212        let q = q.with_position_bound(PositionBound::At(bound_pos.clone()));
213        match &q.bound {
214            PositionBound::At(p) => assert_eq!(p, &bound_pos),
215            PositionBound::Unbounded => panic!("expected At bound"),
216        }
217
218        // require_active flips the flag.
219        let q = q.require_active();
220        assert!(q.require_active);
221    }
222}
krabka_client_streams/runtime/iqv2/request.rs

krabka_client_streams/runtime/iqv2/
request.rs
