tinyflow_framework/runtime/chain_segment/
sink.rs

1use std::any::Any;
2use std::marker::PhantomData;
3
4use async_trait::async_trait;
5
6use tinyflow_api::context::RuntimeContext;
7use tinyflow_api::error::StreamResult;
8use tinyflow_api::functions::SinkFunction;
9
10use super::{
11    ChainSegmentFactory, ChainStepKindInner, ChainTailSegment, NamedStep, downcast_input,
12    impl_chain_segment_lifecycle,
13};
14
15struct SinkChainSegment<T, W> {
16    function: W,
17    task_id: u32,
18    input_cnt: u64,
19    output_cnt: u64,
20    last_input_cnt: u64,
21    last_output_cnt: u64,
22    _marker: PhantomData<T>,
23}
24
25impl_chain_segment_lifecycle! {
26    impl<T, W>
27    ChainSegment for SinkChainSegment<T, W>,
28    with_metrics
29    where
30        T: Send + 'static,
31        W: SinkFunction<T> + Clone + Send + 'static,
32}
33
34#[async_trait]
35impl<T, W> ChainTailSegment for SinkChainSegment<T, W>
36where
37    T: Send + 'static,
38    W: SinkFunction<T> + Clone + Send + 'static,
39{
40    async fn process_record(
41        &mut self,
42        input: Box<dyn Any + Send>,
43        runtime: &mut RuntimeContext,
44    ) -> StreamResult<Vec<Box<dyn Any + Send>>> {
45        let record = downcast_input::<T>(input, self.task_id)?;
46        self.input_cnt += 1;
47        self.function.accept(record, runtime).await?;
48        Ok(vec![])
49    }
50}
51
52pub fn sink_chain_factory<T, F>(f: F) -> ChainSegmentFactory
53where
54    T: Send + 'static,
55    F: SinkFunction<T> + Clone + Send + 'static,
56{
57    Box::new(move || NamedStep {
58        name: "Sink".into(),
59        kind: ChainStepKindInner::Tail(Box::new(SinkChainSegment {
60            function: f.clone(),
61            task_id: 0,
62            input_cnt: 0,
63            output_cnt: 0,
64            last_input_cnt: 0,
65            last_output_cnt: 0,
66            _marker: PhantomData,
67        })),
68    })
69}
tinyflow_framework/runtime/chain_segment/sink.rs

tinyflow_framework/runtime/chain_segment/
sink.rs
