Skip to main content

krabka_client_streams/processor/
api.rs

1//! The typed Processor API: `Processor`, `ProcessorSupplier`, and the
2//! `ProcessorContext` users call `forward` on.
3
4use std::{any::Any, marker::PhantomData};
5
6use async_trait::async_trait;
7use krabka_units::prelude::*;
8
9use super::{
10    erased::{Dispatch, ErasedRecord},
11    record::{Record, RecordContext},
12};
13
14/// A stateless record processor. [`ProcessorSupplier::get`] creates one instance
15/// per task. This mirrors
16/// `org.apache.kafka.streams.processor.api.Processor`.
17///
18/// ## Lifecycle
19///
20/// The runtime calls `init` once before the first record and `close` once at
21/// task shutdown. [`TopologyTestDriver`](crate::TopologyTestDriver) calls `init`
22/// when it instantiates a topology for tests.
23#[async_trait]
24pub trait Processor<KIn: Send, VIn: Send, KOut: Send, VOut: Send>: Send + 'static {
25    async fn init(&mut self, _ctx: &mut ProcessorContext<'_, '_, KOut, VOut>) {}
26    async fn process(
27        &mut self,
28        ctx: &mut ProcessorContext<'_, '_, KOut, VOut>,
29        record: Record<KIn, VIn>,
30    );
31    async fn close(&mut self) {}
32}
33
34/// A boxed processor is itself a [`Processor`] and delegates to the inner value.
35///
36/// This lets a [`ProcessorSupplier`] closure return `Box<dyn Processor<…>>` when
37/// the code chooses the concrete type at runtime. The boxed value still
38/// satisfies the supplier blanket impl, which needs only *some* `Processor` as
39/// the closure's return type. In the common case, return the concrete processor
40/// directly with `|| MyProc` and skip the box.
41#[async_trait]
42impl<KIn, VIn, KOut, VOut> Processor<KIn, VIn, KOut, VOut>
43    for Box<dyn Processor<KIn, VIn, KOut, VOut>>
44where
45    KIn: Send + 'static,
46    VIn: Send + 'static,
47    KOut: Send + 'static,
48    VOut: Send + 'static,
49{
50    async fn init(&mut self, ctx: &mut ProcessorContext<'_, '_, KOut, VOut>) {
51        (**self).init(ctx).await;
52    }
53    async fn process(
54        &mut self,
55        ctx: &mut ProcessorContext<'_, '_, KOut, VOut>,
56        record: Record<KIn, VIn>,
57    ) {
58        (**self).process(ctx, record).await;
59    }
60    async fn close(&mut self) {
61        (**self).close().await;
62    }
63}
64
65/// Factory for [`Processor`] instances. It makes one instance per task, which
66/// gives per-task isolation.
67pub trait ProcessorSupplier<KIn, VIn, KOut, VOut>: Send + Sync + 'static {
68    fn get(&self) -> Box<dyn Processor<KIn, VIn, KOut, VOut>>;
69}
70
71// Blanket impl so a closure `|| MyProc` is a supplier. The closure returns a
72// *concrete* `P: Processor`, which we box. Because `P` is concrete, the four KV
73// type parameters are inferred from `P`'s single `Processor` impl — callers
74// never annotate them. A closure returning `Box<dyn Processor<…>>` also works
75// (the boxed value is itself a `Processor`, see the impl above), covering the
76// rarer case of picking the concrete processor type at runtime.
77impl<F, P, KIn, VIn, KOut, VOut> ProcessorSupplier<KIn, VIn, KOut, VOut> for F
78where
79    F: Fn() -> P + Send + Sync + 'static,
80    KIn: Send,
81    VIn: Send,
82    KOut: Send,
83    VOut: Send,
84    P: Processor<KIn, VIn, KOut, VOut>,
85{
86    fn get(&self) -> Box<dyn Processor<KIn, VIn, KOut, VOut>> {
87        Box::new(self())
88    }
89}
90
91/// The context handed to [`Processor::process`]. `forward` boxes the record and
92/// queues it for each child node, and the driver drains the queue.
93///
94/// There are two lifetimes. `'ctx` is the borrow of the `Dispatch` reference
95/// itself. `'d` is the lifetime of the data inside `Dispatch`, such as the
96/// buffers and the slices. Separate lifetimes avoid lifetime-invariance problems
97/// when the code builds a `ProcessorContext` from a `&mut Dispatch<'d>` with an
98/// independently-scoped outer borrow `'ctx`.
99pub struct ProcessorContext<'ctx, 'd, KOut, VOut> {
100    dispatch: &'ctx mut Dispatch<'d>,
101    _pd: PhantomData<fn(KOut, VOut)>,
102}
103
104impl<'ctx, 'd, KOut, VOut> ProcessorContext<'ctx, 'd, KOut, VOut>
105where
106    KOut: Any + Send + Clone,
107    VOut: Any + Send + Clone,
108{
109    pub(crate) fn new(dispatch: &'ctx mut Dispatch<'d>) -> Self {
110        Self {
111            dispatch,
112            _pd: PhantomData,
113        }
114    }
115
116    /// Forward a record to all child nodes.
117    ///
118    /// The fan-out clones the record once per child. The last child receives the
119    /// original by move, so the common single-child case makes no clone. This
120    /// mirrors the JVM `ProcessorContext.forward(Record)`, which takes the
121    /// record by value.
122    pub fn forward(&mut self, record: Record<KOut, VOut>) {
123        // Copy the child-slice reference out so we can mutably borrow `buffer`.
124        let children = self.dispatch.children;
125        let Some((&last, rest)) = children.split_last() else {
126            return; // no children — drop the record
127        };
128        for &child in rest {
129            let key: Option<Box<dyn Any + Send>> = record
130                .key
131                .clone()
132                .map(|k| Box::new(k) as Box<dyn Any + Send>);
133            let value: Box<dyn Any + Send> = Box::new(record.value.clone());
134            self.dispatch
135                .buffer
136                .push_back((child, ErasedRecord::new(key, value, record.timestamp)));
137        }
138        let ts = record.timestamp;
139        let key: Option<Box<dyn Any + Send>> =
140            record.key.map(|k| Box::new(k) as Box<dyn Any + Send>);
141        let value: Box<dyn Any + Send> = Box::new(record.value);
142        self.dispatch
143            .buffer
144            .push_back((last, ErasedRecord::new(key, value, ts)));
145    }
146
147    /// Access a connected state store, typed. It returns `None` when the store
148    /// is absent or the K and V types do not match. Fetch it once per record and
149    /// do not hold it across `process` calls.
150    pub fn get_state_store<K2: Send + Sync + 'static, V2: Send + 'static>(
151        &mut self,
152        name: &str,
153    ) -> Option<&mut dyn crate::store::api::KeyValueStore<K2, V2>> {
154        self.dispatch.stores.get_kv::<K2, V2>(name)
155    }
156
157    /// Look up a value in a connected GLOBAL store, which is fully replicated
158    /// and shared across tasks.
159    ///
160    /// The method returns an owned value. No borrow escapes the shared manager's
161    /// lock, so the caller need not hold the lookup future across `forward`. It
162    /// returns `None` on a miss or a type mismatch. Fetch it once per record and
163    /// do not hold it across `process` calls.
164    pub async fn global_get<GK: Send + Sync + 'static, VG: Send + 'static>(
165        &mut self,
166        store: &str,
167        key: &GK,
168    ) -> Option<VG> {
169        self.dispatch.globals.get::<GK, VG>(store, key).await
170    }
171
172    /// Access a connected window store, typed. It returns `None` when the store
173    /// is absent or the K and V types do not match. Fetch it once per record and
174    /// do not hold it across `process` calls.
175    pub fn get_window_store<K2: Send + Sync + 'static, V2: Send + 'static>(
176        &mut self,
177        name: &str,
178    ) -> Option<&mut dyn crate::store::window::WindowStore<K2, V2>> {
179        self.dispatch.stores.get_window::<K2, V2>(name)
180    }
181
182    /// Access a connected join-window store (retainDuplicates), typed. It
183    /// returns `None` when the store is absent or the K and V types do not
184    /// match. Fetch it once per record and do not hold it across `process`
185    /// calls.
186    pub fn get_join_window_store<K2: Send + Sync + 'static, V2: Send + 'static>(
187        &mut self,
188        name: &str,
189    ) -> Option<&mut dyn crate::store::join_window::JoinWindowStore<K2, V2>> {
190        self.dispatch.stores.get_join_window::<K2, V2>(name)
191    }
192
193    /// Access a connected session store, typed. It returns `None` when the store
194    /// is absent or the K and V types do not match. Fetch it once per record and
195    /// do not hold it across `process` calls.
196    pub fn get_session_store<K2: Send + Sync + 'static, V2: Send + 'static>(
197        &mut self,
198        name: &str,
199    ) -> Option<&mut dyn crate::store::session::SessionStore<K2, V2>> {
200        self.dispatch.stores.get_session::<K2, V2>(name)
201    }
202
203    /// Access a connected versioned store (KIP-889), typed. It returns `None`
204    /// when the store is absent or the K and V types do not match. Fetch it once
205    /// per record and do not hold it across `process` calls.
206    pub fn get_versioned_store<K2: Send + Sync + 'static, V2: Send + 'static>(
207        &mut self,
208        name: &str,
209    ) -> Option<&mut dyn crate::store::versioned::VersionedKeyValueStore<K2, V2>> {
210        self.dispatch.stores.get_versioned::<K2, V2>(name)
211    }
212
213    /// Access a connected suppress store, typed. It returns `None` when the
214    /// store is absent or the K and V types do not match. Fetch it once per
215    /// record and do not hold it across `process` calls.
216    ///
217    /// This method is `pub(crate)` for two reasons. The returned trait surfaces
218    /// the crate-internal `Change<V>`, and the suppress store is a built-in DSL
219    /// mechanism and not a user-facing store.
220    pub(crate) fn get_suppress_store<K2: Send + Sync + 'static, V2: Send + 'static>(
221        &mut self,
222        name: &str,
223    ) -> Option<&mut dyn crate::store::suppress_store::SuppressStore<K2, V2>> {
224        self.dispatch.stores.get_suppress::<K2, V2>(name)
225    }
226
227    /// Access a connected join-grace buffer store (KIP-923), typed. It returns
228    /// `None` when the store is absent or the K and V types do not match. Fetch
229    /// it once per record and do not hold it across `process` calls.
230    ///
231    /// This method is `pub(crate)` because the grace buffer is a built-in DSL
232    /// mechanism, not a user-facing store. The stream-table join's grace-flush
233    /// processor reaches it through the context.
234    pub(crate) fn get_join_grace_store<K2: Send + Sync + 'static, V2: Send + 'static>(
235        &mut self,
236        name: &str,
237    ) -> Option<&mut crate::store::join_grace_buffer::JoinGraceBufferStore<K2, V2>> {
238        self.dispatch.stores.get_join_grace::<K2, V2>(name)
239    }
240
241    /// Access the connected FK subscription store. It returns `None` when the
242    /// store is absent.
243    ///
244    /// This method is `pub(crate)` because the subscription store is an internal
245    /// KIP-213 FK-join mechanism, not a user-facing store.
246    pub(crate) fn get_fk_subscription_store(
247        &mut self,
248        name: &str,
249    ) -> Option<&mut crate::store::fk_subscription::SubscriptionBytesStore> {
250        self.dispatch.stores.get_fk_subscription(name)
251    }
252
253    /// The metadata of the source record that this context processes now.
254    #[must_use]
255    pub fn record_context(&self) -> &RecordContext {
256        self.dispatch.record_ctx
257    }
258
259    /// Whether the named KV state store is record-cached. When it is, this
260    /// processor should suppress its immediate forward and let the cache flush
261    /// forward the deduped change. The method returns false for an absent store,
262    /// a non-KV store, and an uncached store.
263    #[must_use]
264    pub fn store_is_cached(&self, name: &str) -> bool {
265        self.dispatch.stores.kv_is_cached(name)
266    }
267
268    /// Schedule a periodic [`Punctuator`]. You can call this from `init` or from
269    /// `process`. `interval` must be positive. The method returns a
270    /// [`Cancellable`] that stops the punctuator.
271    ///
272    /// [`Punctuator`]: crate::processor::punctuation::Punctuator
273    /// [`Cancellable`]: crate::processor::punctuation::Cancellable
274    /// # Panics
275    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
276    pub fn schedule<P>(
277        &mut self,
278        interval: std::time::Duration,
279        ty: crate::processor::punctuation::PunctuationType,
280        punctuator: P,
281    ) -> crate::processor::punctuation::Cancellable
282    where
283        P: crate::processor::punctuation::Punctuator<KOut, VOut>,
284    {
285        use crate::processor::punctuation::PunctuationType;
286        let interval = interval.as_time();
287        assert!(
288            interval >= millis(1),
289            "schedule interval must be positive (>= 1ms)"
290        );
291        let base = match ty {
292            PunctuationType::StreamTime => self.dispatch.sched_stream_time,
293            PunctuationType::WallClockTime => self.dispatch.sched_wall_clock,
294        };
295        let next_time = base.saturating_add(interval.millis_i64());
296        let cancel = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
297        let erased: Box<dyn crate::processor::punctuation::ErasedPunctuator> =
298            Box::new(crate::processor::punctuation::TypedPunctuator::<
299                KOut,
300                VOut,
301                P,
302            >::new(punctuator));
303        self.dispatch
304            .schedules
305            .push(crate::processor::punctuation::ScheduleEntry {
306                node_idx: self.dispatch.node_idx,
307                interval,
308                ty,
309                next_time,
310                punctuator: erased,
311                cancel: cancel.clone(),
312            });
313        crate::processor::punctuation::Cancellable::new(cancel)
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use std::collections::VecDeque;
320
321    use assert2::check;
322
323    use super::*;
324    use crate::processor::{
325        erased::{Dispatch, ErasedRecord},
326        record::{Record, RecordContext},
327    };
328
329    struct Upper;
330    crate::impl_processor! {
331        impl Upper: (String, String) -> (String, String) {
332            async fn process(&mut self, ctx, r) {
333                ctx.forward(Record::new(r.key, r.value.to_uppercase(), r.timestamp));
334            }
335        }
336    }
337
338    struct Noop;
339    #[async_trait]
340    impl Processor<String, String, String, String> for Noop {
341        async fn process(
342            &mut self,
343            _ctx: &mut ProcessorContext<'_, '_, String, String>,
344            _r: Record<String, String>,
345        ) {
346        }
347    }
348
349    #[tokio::test]
350    async fn forward_pushes_erased_record_to_each_child() {
351        let mut buffer: VecDeque<(usize, ErasedRecord)> = VecDeque::new();
352        let mut output = Vec::new();
353        let rc = RecordContext {
354            topic: "t".into(),
355            partition: 0,
356            offset: 0,
357            timestamp: 5,
358        };
359        let children = [3usize, 4usize];
360        let mut stores = crate::store::registry::StoreRegistry::default();
361        let globals = crate::runtime::global::GlobalStateManager::default();
362        let mut scheds = Vec::new();
363        let mut dispatch = Dispatch {
364            buffer: &mut buffer,
365            children: &children,
366            output: &mut output,
367            record_ctx: &rc,
368            stores: &mut stores,
369            globals: &globals,
370            node_idx: 0,
371            schedules: &mut scheds,
372            sched_stream_time: i64::MIN,
373            sched_wall_clock: 0,
374        };
375        let mut ctx = ProcessorContext::<'_, '_, String, String>::new(&mut dispatch);
376        Upper
377            .process(&mut ctx, Record::new(Some("k".into()), "hi".into(), 5))
378            .await;
379        check!(buffer.len() == 2);
380        let (child, rec) = buffer.pop_front().unwrap();
381        check!(child == 3);
382        check!(*rec.value.downcast::<String>().unwrap() == "HI");
383    }
384
385    #[tokio::test]
386    async fn boxed_dyn_processor_delegates_init_process_close() {
387        // A `Box<dyn Processor>` is itself a `Processor`, forwarding every method
388        // to the inner value. This is the runtime-dispatch path a
389        // `ProcessorSupplier` closure takes when it returns `Box<dyn Processor<…>>`
390        // instead of a concrete processor.
391        let mut boxed: Box<dyn Processor<String, String, String, String>> = Box::new(Upper);
392        let mut buffer: VecDeque<(usize, ErasedRecord)> = VecDeque::new();
393        let mut output = Vec::new();
394        let rc = RecordContext {
395            topic: "t".into(),
396            partition: 0,
397            offset: 0,
398            timestamp: 5,
399        };
400        let children = [1usize];
401        let mut stores = crate::store::registry::StoreRegistry::default();
402        let globals = crate::runtime::global::GlobalStateManager::default();
403        let mut scheds = Vec::new();
404        let mut dispatch = Dispatch {
405            buffer: &mut buffer,
406            children: &children,
407            output: &mut output,
408            record_ctx: &rc,
409            stores: &mut stores,
410            globals: &globals,
411            node_idx: 0,
412            schedules: &mut scheds,
413            sched_stream_time: i64::MIN,
414            sched_wall_clock: 0,
415        };
416        let mut ctx = ProcessorContext::<'_, '_, String, String>::new(&mut dispatch);
417        boxed.init(&mut ctx).await; // forwards to Upper's default no-op
418        boxed
419            .process(&mut ctx, Record::new(None, "hi".into(), 5))
420            .await; // forwards → uppercases
421        boxed.close().await; // forwards to Upper's default no-op
422        check!(buffer.len() == 1);
423        let (_child, rec) = buffer.pop_front().unwrap();
424        check!(*rec.value.downcast::<String>().unwrap() == "HI");
425    }
426
427    #[tokio::test]
428    async fn default_init_and_close_are_noops_and_forward_with_no_children_drops() {
429        let mut p = Noop;
430        let mut buffer: VecDeque<(usize, ErasedRecord)> = VecDeque::new();
431        let mut output = Vec::new();
432        let rc = RecordContext {
433            topic: "t".into(),
434            partition: 0,
435            offset: 0,
436            timestamp: 9,
437        };
438        let mut stores = crate::store::registry::StoreRegistry::default();
439        let globals = crate::runtime::global::GlobalStateManager::default();
440        let mut scheds = Vec::new();
441        let mut dispatch = Dispatch {
442            buffer: &mut buffer,
443            children: &[],
444            output: &mut output,
445            record_ctx: &rc,
446            stores: &mut stores,
447            globals: &globals,
448            node_idx: 0,
449            schedules: &mut scheds,
450            sched_stream_time: i64::MIN,
451            sched_wall_clock: 0,
452        };
453        let mut ctx = ProcessorContext::<'_, '_, String, String>::new(&mut dispatch);
454        p.init(&mut ctx).await; // default no-op
455        check!(ctx.record_context().timestamp == 9);
456        ctx.forward(Record::new(None, "x".to_string(), 0)); // no children → dropped, no panic
457        check!(buffer.is_empty());
458        p.close().await; // default no-op
459    }
460}