1use std::any::Any;
14
15use async_trait::async_trait;
16
17use crate::processor::{
18 api::{Processor, ProcessorContext},
19 record::{Record, RecordContext},
20};
21
22#[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 #[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
49pub 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 pub(crate) fn new(inner: &'a mut ProcessorContext<'ctx, 'd, K, VOut>) -> Self {
65 Self { inner }
66 }
67
68 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 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 #[must_use]
91 pub fn record_context(&self) -> &RecordContext {
92 self.inner.record_context()
93 }
94}
95
96#[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#[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
134pub 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
156pub(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 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 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 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 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 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}