Skip to main content

krabka_client_streams/processor/
fixed_key.rs

1//! Fixed-key Processor API (KIP-820 `processValues`).
2//!
3//! A fixed-key processor may change the value but NOT the key. This module is a
4//! thin typed facade over the regular [`Processor`] runtime. An internal adapter
5//! bridges a [`FixedKeyProcessor`] into a `Processor<KIn, VIn, KIn, VOut>`. The
6//! output key type equals the input key type, so the runtime keeps the partition
7//! assignment unchanged.
8//!
9//! The DSL surface `KStream::process_values` wraps a
10//! [`FixedKeyProcessorSupplier`] in that adapter. This module holds the typed
11//! facade and the adapter that it builds on.
12
13use std::any::Any;
14
15use async_trait::async_trait;
16
17use crate::processor::{
18    api::{Processor, ProcessorContext},
19    record::{Record, RecordContext},
20};
21
22/// A record whose KEY is immutable.
23///
24/// [`with_value`](FixedKeyRecord::with_value) changes the value, and it can also
25/// change the value type. It keeps the key and the timestamp. Unlike [`Record`],
26/// the key is **not** optional. A record reaches a fixed-key processor only when
27/// it carries a key, because a null-key record has no key to preserve.
28/// `process_values` needs a non-null key, which matches the JVM.
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct FixedKeyRecord<K, V> {
31    pub key: K,
32    pub value: V,
33    pub timestamp: i64,
34}
35
36impl<K, V> FixedKeyRecord<K, V> {
37    /// Replace the value, and possibly its type. This keeps the key and the
38    /// timestamp.
39    #[must_use]
40    pub fn with_value<V2>(self, value: V2) -> FixedKeyRecord<K, V2> {
41        FixedKeyRecord {
42            key: self.key,
43            value,
44            timestamp: self.timestamp,
45        }
46    }
47}
48
49/// The context handed to [`FixedKeyProcessor::process`].
50///
51/// The only `forward` re-attaches the unchanged key, so a fixed-key processor
52/// cannot repartition by accident. All other accessors delegate verbatim to the
53/// underlying [`ProcessorContext`].
54pub struct FixedKeyProcessorContext<'a, 'ctx, 'd, K, VOut> {
55    inner: &'a mut ProcessorContext<'ctx, 'd, K, VOut>,
56}
57
58impl<'a, 'ctx, 'd, K, VOut> FixedKeyProcessorContext<'a, 'ctx, 'd, K, VOut>
59where
60    K: Any + Send + Clone,
61    VOut: Any + Send + Clone,
62{
63    // Constructed by FixedKeyAdapter (wired up by KStream::process_values).
64    pub(crate) fn new(inner: &'a mut ProcessorContext<'ctx, 'd, K, VOut>) -> Self {
65        Self { inner }
66    }
67
68    /// Forward a record downstream and re-attach the unchanged key. This mirrors
69    /// the JVM `FixedKeyProcessorContext.forward(FixedKeyRecord)`.
70    pub fn forward(&mut self, record: FixedKeyRecord<K, VOut>) {
71        self.inner.forward(Record::new(
72            Some(record.key),
73            record.value,
74            record.timestamp,
75        ));
76    }
77
78    /// Access a connected key/value state store, typed. This method delegates to
79    /// [`ProcessorContext::get_state_store`]. This context has no window-store or
80    /// session-store accessor, because no DSL path connects those to a
81    /// `process_values` node yet.
82    pub fn get_state_store<K2: Send + Sync + 'static, V2: Send + 'static>(
83        &mut self,
84        name: &str,
85    ) -> Option<&mut dyn crate::store::api::KeyValueStore<K2, V2>> {
86        self.inner.get_state_store::<K2, V2>(name)
87    }
88
89    /// Metadata of the source record currently being processed.
90    #[must_use]
91    pub fn record_context(&self) -> &RecordContext {
92        self.inner.record_context()
93    }
94}
95
96/// A fixed-key record processor (KIP-820).
97///
98/// It may change the value but not the key. It is the typed analogue of
99/// [`Processor`]. [`FixedKeyProcessorSupplier::get`] creates one instance per
100/// task.
101#[async_trait]
102pub trait FixedKeyProcessor<KIn: Send, VIn: Send, VOut: Send>: Send + 'static {
103    async fn process(
104        &mut self,
105        ctx: &mut FixedKeyProcessorContext<'_, '_, '_, KIn, VOut>,
106        record: FixedKeyRecord<KIn, VIn>,
107    );
108}
109
110/// A boxed fixed-key processor is itself a [`FixedKeyProcessor`]. It delegates
111/// to the inner value.
112///
113/// This impl mirrors the `Box<dyn Processor>` blanket impl in
114/// [`crate::processor::api`]. It lets a [`FixedKeyProcessorSupplier`] closure
115/// return a `Box<dyn FixedKeyProcessor<…>>` chosen at runtime and still feed the
116/// internal fixed-key adapter, which needs `P: FixedKeyProcessor`.
117#[async_trait]
118impl<KIn, VIn, VOut> FixedKeyProcessor<KIn, VIn, VOut>
119    for Box<dyn FixedKeyProcessor<KIn, VIn, VOut>>
120where
121    KIn: Send + 'static,
122    VIn: Send + 'static,
123    VOut: Send + 'static,
124{
125    async fn process(
126        &mut self,
127        ctx: &mut FixedKeyProcessorContext<'_, '_, '_, KIn, VOut>,
128        record: FixedKeyRecord<KIn, VIn>,
129    ) {
130        (**self).process(ctx, record).await;
131    }
132}
133
134/// Factory for [`FixedKeyProcessor`] instances.
135///
136/// It creates one instance per task, which gives per-task isolation. It mirrors
137/// [`ProcessorSupplier`](crate::processor::ProcessorSupplier). A closure
138/// `|| MyFixedProc` is a supplier through the blanket impl below.
139pub trait FixedKeyProcessorSupplier<KIn, VIn, VOut>: Send + Sync + 'static {
140    fn get(&self) -> Box<dyn FixedKeyProcessor<KIn, VIn, VOut>>;
141}
142
143impl<F, P, KIn, VIn, VOut> FixedKeyProcessorSupplier<KIn, VIn, VOut> for F
144where
145    F: Fn() -> P + Send + Sync + 'static,
146    KIn: Send,
147    VIn: Send,
148    VOut: Send,
149    P: FixedKeyProcessor<KIn, VIn, VOut>,
150{
151    fn get(&self) -> Box<dyn FixedKeyProcessor<KIn, VIn, VOut>> {
152        Box::new(self())
153    }
154}
155
156/// Bridges a [`FixedKeyProcessor`] into the regular [`Processor`] runtime with
157/// `KOut = KIn`, so the key stays unchanged.
158///
159/// [`KStream::process_values`](crate::dsl::kstream::KStream::process_values)
160/// wraps this adapter around a [`FixedKeyProcessorSupplier`].
161pub(crate) struct FixedKeyAdapter<P> {
162    pub inner: P,
163}
164
165#[async_trait]
166impl<P, KIn, VIn, VOut> Processor<KIn, VIn, KIn, VOut> for FixedKeyAdapter<P>
167where
168    KIn: Any + Send + Clone + 'static,
169    VIn: Send + 'static,
170    VOut: Any + Send + Clone + 'static,
171    P: FixedKeyProcessor<KIn, VIn, VOut>,
172{
173    async fn process(
174        &mut self,
175        ctx: &mut ProcessorContext<'_, '_, KIn, VOut>,
176        record: Record<KIn, VIn>,
177    ) {
178        let key = record.key.expect("process_values requires a non-null key");
179        let fkr = FixedKeyRecord {
180            key,
181            value: record.value,
182            timestamp: record.timestamp,
183        };
184        let mut fk_ctx = FixedKeyProcessorContext::new(ctx);
185        self.inner.process(&mut fk_ctx, fkr).await;
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use std::collections::VecDeque;
192
193    use assert2::check;
194
195    use super::*;
196    use crate::processor::{
197        erased::{Dispatch, ErasedRecord},
198        record::{Record, RecordContext},
199    };
200
201    #[test]
202    fn with_value_keeps_key_and_timestamp() {
203        let r = FixedKeyRecord {
204            key: "k".to_string(),
205            value: 1i32,
206            timestamp: 42,
207        };
208        let r2 = r.with_value("v".to_string());
209        check!(r2.key == "k");
210        check!(r2.value == "v");
211        check!(r2.timestamp == 42);
212    }
213
214    /// A fixed-key processor that uppercases the value and keeps the key.
215    struct UpperValue;
216    crate::impl_fixed_key_processor! {
217        impl UpperValue: (String, String) -> String {
218            async fn process(&mut self, ctx, r) {
219                let v = r.value.clone();
220                ctx.forward(r.with_value(v.to_uppercase()));
221            }
222        }
223    }
224
225    #[tokio::test]
226    async fn adapter_bridges_fixed_key_processor_preserving_key() {
227        // Drive `UpperValue` through `FixedKeyAdapter` over a real
228        // `ProcessorContext`, constructed exactly as the `api.rs` tests do.
229        let mut adapter = FixedKeyAdapter { inner: UpperValue };
230        let mut buffer: VecDeque<(usize, ErasedRecord)> = VecDeque::new();
231        let mut output = Vec::new();
232        let rc = RecordContext {
233            topic: "t".into(),
234            partition: 0,
235            offset: 0,
236            timestamp: 5,
237        };
238        let children = [1usize];
239        let mut stores = crate::store::registry::StoreRegistry::default();
240        let globals = crate::runtime::global::GlobalStateManager::default();
241        let mut scheds = Vec::new();
242        let mut dispatch = Dispatch {
243            buffer: &mut buffer,
244            children: &children,
245            output: &mut output,
246            record_ctx: &rc,
247            stores: &mut stores,
248            globals: &globals,
249            node_idx: 0,
250            schedules: &mut scheds,
251            sched_stream_time: i64::MIN,
252            sched_wall_clock: 0,
253        };
254        let mut ctx = ProcessorContext::<'_, '_, String, String>::new(&mut dispatch);
255        adapter
256            .process(&mut ctx, Record::new(Some("k".into()), "hi".into(), 5))
257            .await;
258
259        check!(buffer.len() == 1);
260        let (child, rec) = buffer.pop_front().unwrap();
261        check!(child == 1);
262        // The key is preserved (Some("k")) and the value uppercased.
263        let key = rec.key.expect("forwarded record must carry the key");
264        check!(*key.downcast::<String>().unwrap() == "k");
265        check!(*rec.value.downcast::<String>().unwrap() == "HI");
266        check!(rec.timestamp == 5);
267    }
268
269    #[tokio::test]
270    async fn supplier_blanket_impl_boxes_processor() {
271        // A closure `|| UpperValue` is a `FixedKeyProcessorSupplier`; `get()`
272        // returns a boxed processor that itself impls `FixedKeyProcessor`.
273        let supplier = || UpperValue;
274        let mut boxed = FixedKeyProcessorSupplier::get(&supplier);
275
276        let mut buffer: VecDeque<(usize, ErasedRecord)> = VecDeque::new();
277        let mut output = Vec::new();
278        let rc = RecordContext {
279            topic: "t".into(),
280            partition: 0,
281            offset: 0,
282            timestamp: 7,
283        };
284        let children = [2usize];
285        let mut stores = crate::store::registry::StoreRegistry::default();
286        let globals = crate::runtime::global::GlobalStateManager::default();
287        let mut scheds = Vec::new();
288        let mut dispatch = Dispatch {
289            buffer: &mut buffer,
290            children: &children,
291            output: &mut output,
292            record_ctx: &rc,
293            stores: &mut stores,
294            globals: &globals,
295            node_idx: 0,
296            schedules: &mut scheds,
297            sched_stream_time: i64::MIN,
298            sched_wall_clock: 0,
299        };
300        let mut ctx = ProcessorContext::<'_, '_, String, String>::new(&mut dispatch);
301        let mut fk_ctx = FixedKeyProcessorContext::new(&mut ctx);
302
303        // Drive the boxed processor directly to exercise the Box<dyn …> blanket impl.
304        boxed
305            .process(
306                &mut fk_ctx,
307                FixedKeyRecord {
308                    key: "k".to_string(),
309                    value: "lo".to_string(),
310                    timestamp: 7,
311                },
312            )
313            .await;
314
315        check!(buffer.len() == 1);
316        let (_child, rec) = buffer.pop_front().unwrap();
317        check!(*rec.value.downcast::<String>().unwrap() == "LO");
318    }
319}