tinyflow_framework/runtime/chain_segment/
mod.rs

1//! In-process operator chain steps (fused subtask, no inter-op channels).
2//!
3//! Each segment holds a user [`SourceFunction`](crate::api::functions::SourceFunction) /
4//! [`MapFunction`](crate::api::functions::MapFunction) / … directly — no separate operator wrapper layer.
5
6mod middle;
7mod sink;
8mod source;
9
10pub use middle::{filter_chain_factory, flat_map_chain_factory, map_chain_factory};
11pub use sink::sink_chain_factory;
12pub use source::source_chain_factory;
13
14use async_trait::async_trait;
15use std::any::Any;
16use std::collections::HashMap;
17
18use tinyflow_api::checkpoint::CheckpointContext;
19use tinyflow_api::context::RuntimeContext;
20use tinyflow_api::error::{StreamError, StreamResult};
21
22/// One step in an operator chain (same subtask, in-process).
23#[async_trait]
24pub trait ChainSegment: Send {
25    async fn open(&mut self, ctx: &mut RuntimeContext) -> StreamResult<()>;
26    async fn close(&mut self, ctx: &mut RuntimeContext) -> StreamResult<()>;
27    async fn on_checkpoint(
28        &mut self,
29        runtime: &mut RuntimeContext,
30        ckpt: &CheckpointContext,
31    ) -> StreamResult<()>;
32    async fn on_checkpoint_end(
33        &mut self,
34        runtime: &mut RuntimeContext,
35        ckpt: &CheckpointContext,
36    ) -> StreamResult<()>;
37
38    /// Returns (input_delta, output_delta) and resets the internal last-pointers
39    fn drain_counters(&mut self) -> (u64, u64) {
40        (0, 0)
41    }
42
43    /// Delegates to the underlying operator for custom metrics
44    fn collect_metrics(&self) -> HashMap<String, f64> {
45        HashMap::new()
46    }
47}
48
49/// Source head: pull one record; `None` when the source is idle or finished.
50#[async_trait]
51pub trait ChainSourceSegment: ChainSegment {
52    async fn poll_next(&mut self) -> StreamResult<Option<Box<dyn Any + Send>>>;
53}
54
55/// Tail chain steps (Map / Filter / FlatMap / **Sink**): one record in, zero or more out.
56/// The last step is always a terminal sink (no downstream records).
57#[async_trait]
58pub trait ChainTailSegment: ChainSegment {
59    async fn process_record(
60        &mut self,
61        input: Box<dyn Any + Send>,
62        runtime: &mut RuntimeContext,
63    ) -> StreamResult<Vec<Box<dyn Any + Send>>>;
64}
65
66/// Erased chain step for heterogeneous fused chains.
67pub enum ChainStepKindInner {
68    Source(Box<dyn ChainSourceSegment>),
69    Tail(Box<dyn ChainTailSegment>),
70}
71
72pub struct NamedStep {
73    pub name: String,
74    pub kind: ChainStepKindInner,
75}
76
77pub type ChainSegmentFactory = Box<dyn Fn() -> NamedStep + Send>;
78
79pub fn instantiate_chain_steps(factories: &[ChainSegmentFactory]) -> Vec<NamedStep> {
80    factories.iter().map(|f| f()).collect()
81}
82
83pub(super) fn downcast_input<T: Send + 'static>(
84    input: Box<dyn Any + Send>,
85    task_id: u32,
86) -> StreamResult<T> {
87    match input.downcast::<T>() {
88        Ok(v) => Ok(*v),
89        Err(_) => Err(StreamError::TaskFailed {
90            task_id,
91            reason: format!(
92                "chain segment type mismatch: expected {}",
93                std::any::type_name::<T>()
94            ),
95        }),
96    }
97}
98
99/// Forwards [`LifecycleAware`](crate::api::functions::LifecycleAware) hooks; sets `task_id` on `open`.
100/// Segment must have `function` + `task_id` fields.
101///
102/// Uses the fully-qualified `#[async_trait::async_trait]` path so that macro expansion
103/// does not depend on the caller having imported `async_trait`,
104/// making it safe to invoke inside other macros (e.g. `middle_tail_step!`).
105macro_rules! impl_chain_segment_lifecycle {
106    (
107        $(#[$attr:meta])*
108        impl<$($gen:ident),*>
109        ChainSegment for $seg:ty
110        where
111            $($where:tt)*
112    ) => {
113        #[async_trait::async_trait]
114        $(#[$attr])*
115        impl<$($gen),*> $crate::runtime::chain_segment::ChainSegment for $seg
116        where
117            $($where)*
118        {
119            async fn open(
120                &mut self,
121                ctx: &mut tinyflow_api::context::RuntimeContext,
122            ) -> tinyflow_api::error::StreamResult<()> {
123                self.task_id = ctx.task_id;
124                tinyflow_api::functions::LifecycleAware::open(&mut self.function, ctx).await
125            }
126            async fn close(
127                &mut self,
128                ctx: &mut tinyflow_api::context::RuntimeContext,
129            ) -> tinyflow_api::error::StreamResult<()> {
130                tinyflow_api::functions::LifecycleAware::close(&mut self.function, ctx).await
131            }
132            async fn on_checkpoint(
133                &mut self,
134                runtime: &mut tinyflow_api::context::RuntimeContext,
135                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
136            ) -> tinyflow_api::error::StreamResult<()> {
137                tinyflow_api::functions::LifecycleAware::on_checkpoint(
138                    &mut self.function,
139                    runtime,
140                    ckpt,
141                )
142                .await
143            }
144            async fn on_checkpoint_end(
145                &mut self,
146                runtime: &mut tinyflow_api::context::RuntimeContext,
147                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
148            ) -> tinyflow_api::error::StreamResult<()> {
149                tinyflow_api::functions::LifecycleAware::on_checkpoint_end(
150                    &mut self.function,
151                    runtime,
152                    ckpt,
153                )
154                .await
155            }
156        }
157    };
158    (
159        $(#[$attr:meta])*
160        impl<$($gen:ident),*>
161        ChainSegment for $seg:ty,
162        with_metrics
163        where
164            $($where:tt)*
165    ) => {
166        #[async_trait::async_trait]
167        $(#[$attr])*
168        impl<$($gen),*> $crate::runtime::chain_segment::ChainSegment for $seg
169        where
170            $($where)*
171        {
172            async fn open(
173                &mut self,
174                ctx: &mut tinyflow_api::context::RuntimeContext,
175            ) -> tinyflow_api::error::StreamResult<()> {
176                self.task_id = ctx.task_id;
177                tinyflow_api::functions::LifecycleAware::open(&mut self.function, ctx).await
178            }
179            async fn close(
180                &mut self,
181                ctx: &mut tinyflow_api::context::RuntimeContext,
182            ) -> tinyflow_api::error::StreamResult<()> {
183                tinyflow_api::functions::LifecycleAware::close(&mut self.function, ctx).await
184            }
185            async fn on_checkpoint(
186                &mut self,
187                runtime: &mut tinyflow_api::context::RuntimeContext,
188                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
189            ) -> tinyflow_api::error::StreamResult<()> {
190                tinyflow_api::functions::LifecycleAware::on_checkpoint(
191                    &mut self.function,
192                    runtime,
193                    ckpt,
194                )
195                .await
196            }
197            async fn on_checkpoint_end(
198                &mut self,
199                runtime: &mut tinyflow_api::context::RuntimeContext,
200                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
201            ) -> tinyflow_api::error::StreamResult<()> {
202                tinyflow_api::functions::LifecycleAware::on_checkpoint_end(
203                    &mut self.function,
204                    runtime,
205                    ckpt,
206                )
207                .await
208            }
209            fn drain_counters(&mut self) -> (u64, u64) {
210                let i = self.input_cnt - self.last_input_cnt;
211                let o = self.output_cnt - self.last_output_cnt;
212                self.last_input_cnt = self.input_cnt;
213                self.last_output_cnt = self.output_cnt;
214                (i, o)
215            }
216            fn collect_metrics(&self) -> std::collections::HashMap<String, f64> {
217                self.function.report_metrics()
218            }
219        }
220    };
221}
222pub(crate) use impl_chain_segment_lifecycle;
223
224/// Middle tail segment: struct + lifecycle + factory; caller supplies `process_record` body.
225macro_rules! middle_tail_step {
226    (
227        struct $seg:ident<$($gen:ident),*>;
228        marker: $marker:ty;
229        where { $($bound:tt)* }
230        process { |$this:ident, $data:ident, $runtime:ident| $($body:tt)* }
231        factory $factory:ident
232    ) => {
233        struct $seg<$($gen),*> {
234            function: F,
235            task_id: u32,
236            input_cnt: u64,
237            output_cnt: u64,
238            last_input_cnt: u64,
239            last_output_cnt: u64,
240            _marker: std::marker::PhantomData<$marker>,
241        }
242
243        // Reuse lifecycle implementation: delegate to impl_chain_segment_lifecycle! instead of hand-writing the four methods.
244        $crate::runtime::chain_segment::impl_chain_segment_lifecycle! {
245            impl<$($gen),*>
246            ChainSegment for $seg<$($gen),*>,
247            with_metrics
248            where
249                $($bound)*
250        }
251
252        #[async_trait::async_trait]
253        impl<$($gen),*> $crate::runtime::chain_segment::ChainTailSegment for $seg<$($gen),*>
254        where
255            $($bound)*
256        {
257            async fn process_record(
258                &mut self,
259                input: Box<dyn std::any::Any + Send>,
260                runtime: &mut tinyflow_api::context::RuntimeContext,
261            ) -> tinyflow_api::error::StreamResult<Vec<Box<dyn std::any::Any + Send>>> {
262                let $this = self;
263                $this.input_cnt += 1;
264                let task_id = $this.task_id;
265                let $data = super::downcast_input(input, task_id)?;
266                let $runtime = runtime;
267                let result = { $($body)* };
268                if let Ok(ref records) = result {
269                    $this.output_cnt += records.len() as u64;
270                }
271                result
272            }
273        }
274
275        pub fn $factory<$($gen),*>(f: F) -> super::ChainSegmentFactory
276        where
277            $($bound)*
278        {
279            let step_name = stringify!($seg);
280            Box::new(move || {
281                super::NamedStep {
282                    name: step_name.to_string(),
283                    kind: super::ChainStepKindInner::Tail(Box::new($seg {
284                        function: f.clone(),
285                        task_id: 0,
286                        input_cnt: 0,
287                        output_cnt: 0,
288                        last_input_cnt: 0,
289                        last_output_cnt: 0,
290                        _marker: std::marker::PhantomData,
291                    })),
292                }
293            })
294        }
295    };
296}
297pub(crate) use middle_tail_step;
tinyflow_framework/runtime/chain_segment/mod.rs

tinyflow_framework/runtime/chain_segment/
mod.rs
