tinyflow_framework/runtime/chain_segment/
source.rs

1use std::any::Any;
2use std::marker::PhantomData;
3use std::sync::{Arc, Mutex};
4
5use async_trait::async_trait;
6
7use tinyflow_api::error::StreamResult;
8use tinyflow_api::functions::SourceFunction;
9
10use super::{
11    ChainSegmentFactory, ChainSourceSegment, ChainStepKindInner, NamedStep,
12    impl_chain_segment_lifecycle,
13};
14
15struct SourceChainSegment<T, S>
16where
17    T: Send,
18    S: SourceFunction<T>,
19{
20    function: S,
21    task_id: u32,
22    input_cnt: u64,
23    output_cnt: u64,
24    last_input_cnt: u64,
25    last_output_cnt: u64,
26    _marker: PhantomData<T>,
27}
28
29impl<T, S> SourceChainSegment<T, S>
30where
31    T: Send + 'static,
32    S: SourceFunction<T> + Send,
33{
34    fn new(function: S) -> Self {
35        Self {
36            function,
37            task_id: 0,
38            input_cnt: 0,
39            output_cnt: 0,
40            last_input_cnt: 0,
41            last_output_cnt: 0,
42            _marker: PhantomData,
43        }
44    }
45}
46
47impl_chain_segment_lifecycle! {
48    impl<T, S>
49    ChainSegment for SourceChainSegment<T, S>,
50    with_metrics
51    where
52        T: Send + 'static,
53        S: SourceFunction<T> + Send,
54}
55
56#[async_trait]
57impl<T, S> ChainSourceSegment for SourceChainSegment<T, S>
58where
59    T: Send + 'static,
60    S: SourceFunction<T> + Send,
61{
62    async fn poll_next(&mut self) -> StreamResult<Option<Box<dyn Any + Send>>> {
63        let result = self.function.poll_next().await?;
64        if result.is_some() {
65            self.input_cnt += 1;
66            self.output_cnt += 1;
67        }
68        Ok(result.map(|v| Box::new(v) as Box<dyn Any + Send>))
69    }
70}
71
72pub fn source_chain_factory<T, S, F>(factory: Arc<Mutex<F>>) -> ChainSegmentFactory
73where
74    T: Send + 'static,
75    S: SourceFunction<T> + Send + 'static,
76    F: Fn() -> S + Send + 'static,
77{
78    Box::new(move || {
79        let source = factory.lock().expect("source factory lock poisoned")();
80        NamedStep {
81            name: "Source".into(),
82            kind: ChainStepKindInner::Source(Box::new(SourceChainSegment::new(source))),
83        }
84    })
85}
tinyflow_framework/runtime/chain_segment/source.rs

tinyflow_framework/runtime/chain_segment/
source.rs
