1use 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#[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#[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
65pub trait ProcessorSupplier<KIn, VIn, KOut, VOut>: Send + Sync + 'static {
68 fn get(&self) -> Box<dyn Processor<KIn, VIn, KOut, VOut>>;
69}
70
71impl<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
91pub 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 pub fn forward(&mut self, record: Record<KOut, VOut>) {
123 let children = self.dispatch.children;
125 let Some((&last, rest)) = children.split_last() else {
126 return; };
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 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 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 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 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 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 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 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 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 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 #[must_use]
255 pub fn record_context(&self) -> &RecordContext {
256 self.dispatch.record_ctx
257 }
258
259 #[must_use]
264 pub fn store_is_cached(&self, name: &str) -> bool {
265 self.dispatch.stores.kv_is_cached(name)
266 }
267
268 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 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; boxed
419 .process(&mut ctx, Record::new(None, "hi".into(), 5))
420 .await; boxed.close().await; 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; check!(ctx.record_context().timestamp == 9);
456 ctx.forward(Record::new(None, "x".to_string(), 0)); check!(buffer.is_empty());
458 p.close().await; }
460}