tinyflow-framework 0.1.1

//! In-process operator chain steps (fused subtask, no inter-op channels).
//!
//! Each segment holds a user [`SourceFunction`](crate::api::functions::SourceFunction) /
//! [`MapFunction`](crate::api::functions::MapFunction) / … directly — no separate operator wrapper layer.

mod middle;
mod sink;
mod source;

pub use middle::{filter_chain_factory, flat_map_chain_factory, map_chain_factory};
pub use sink::sink_chain_factory;
pub use source::source_chain_factory;

use async_trait::async_trait;
use std::any::Any;
use std::collections::HashMap;

use tinyflow_api::checkpoint::CheckpointContext;
use tinyflow_api::context::RuntimeContext;
use tinyflow_api::error::{StreamError, StreamResult};

/// One step in an operator chain (same subtask, in-process).
#[async_trait]
pub trait ChainSegment: Send {
    async fn open(&mut self, ctx: &mut RuntimeContext) -> StreamResult<()>;
    async fn close(&mut self, ctx: &mut RuntimeContext) -> StreamResult<()>;
    async fn on_checkpoint(
        &mut self,
        runtime: &mut RuntimeContext,
        ckpt: &CheckpointContext,
    ) -> StreamResult<()>;
    async fn on_checkpoint_end(
        &mut self,
        runtime: &mut RuntimeContext,
        ckpt: &CheckpointContext,
    ) -> StreamResult<()>;

    /// Returns (input_delta, output_delta) and resets the internal last-pointers
    fn drain_counters(&mut self) -> (u64, u64) {
        (0, 0)
    }

    /// Delegates to the underlying operator for custom metrics
    fn collect_metrics(&self) -> HashMap<String, f64> {
        HashMap::new()
    }
}

/// Source head: pull one record; `None` when the source is idle or finished.
#[async_trait]
pub trait ChainSourceSegment: ChainSegment {
    async fn poll_next(&mut self) -> StreamResult<Option<Box<dyn Any + Send>>>;
}

/// Tail chain steps (Map / Filter / FlatMap / **Sink**): one record in, zero or more out.
/// The last step is always a terminal sink (no downstream records).
#[async_trait]
pub trait ChainTailSegment: ChainSegment {
    async fn process_record(
        &mut self,
        input: Box<dyn Any + Send>,
        runtime: &mut RuntimeContext,
    ) -> StreamResult<Vec<Box<dyn Any + Send>>>;
}

/// Erased chain step for heterogeneous fused chains.
pub enum ChainStepKindInner {
    Source(Box<dyn ChainSourceSegment>),
    Tail(Box<dyn ChainTailSegment>),
}

pub struct NamedStep {
    pub name: String,
    pub kind: ChainStepKindInner,
}

pub type ChainSegmentFactory = Box<dyn Fn() -> NamedStep + Send>;

pub fn instantiate_chain_steps(factories: &[ChainSegmentFactory]) -> Vec<NamedStep> {
    factories.iter().map(|f| f()).collect()
}

pub(super) fn downcast_input<T: Send + 'static>(
    input: Box<dyn Any + Send>,
    task_id: u32,
) -> StreamResult<T> {
    match input.downcast::<T>() {
        Ok(v) => Ok(*v),
        Err(_) => Err(StreamError::TaskFailed {
            task_id,
            reason: format!(
                "chain segment type mismatch: expected {}",
                std::any::type_name::<T>()
            ),
        }),
    }
}

/// Forwards [`LifecycleAware`](crate::api::functions::LifecycleAware) hooks; sets `task_id` on `open`.
/// Segment must have `function` + `task_id` fields.
///
/// Uses the fully-qualified `#[async_trait::async_trait]` path so that macro expansion
/// does not depend on the caller having imported `async_trait`,
/// making it safe to invoke inside other macros (e.g. `middle_tail_step!`).
macro_rules! impl_chain_segment_lifecycle {
    (
        $(#[$attr:meta])*
        impl<$($gen:ident),*>
        ChainSegment for $seg:ty
        where
            $($where:tt)*
    ) => {
        #[async_trait::async_trait]
        $(#[$attr])*
        impl<$($gen),*> $crate::runtime::chain_segment::ChainSegment for $seg
        where
            $($where)*
        {
            async fn open(
                &mut self,
                ctx: &mut tinyflow_api::context::RuntimeContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                self.task_id = ctx.task_id;
                tinyflow_api::functions::LifecycleAware::open(&mut self.function, ctx).await
            }
            async fn close(
                &mut self,
                ctx: &mut tinyflow_api::context::RuntimeContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                tinyflow_api::functions::LifecycleAware::close(&mut self.function, ctx).await
            }
            async fn on_checkpoint(
                &mut self,
                runtime: &mut tinyflow_api::context::RuntimeContext,
                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                tinyflow_api::functions::LifecycleAware::on_checkpoint(
                    &mut self.function,
                    runtime,
                    ckpt,
                )
                .await
            }
            async fn on_checkpoint_end(
                &mut self,
                runtime: &mut tinyflow_api::context::RuntimeContext,
                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                tinyflow_api::functions::LifecycleAware::on_checkpoint_end(
                    &mut self.function,
                    runtime,
                    ckpt,
                )
                .await
            }
        }
    };
    (
        $(#[$attr:meta])*
        impl<$($gen:ident),*>
        ChainSegment for $seg:ty,
        with_metrics
        where
            $($where:tt)*
    ) => {
        #[async_trait::async_trait]
        $(#[$attr])*
        impl<$($gen),*> $crate::runtime::chain_segment::ChainSegment for $seg
        where
            $($where)*
        {
            async fn open(
                &mut self,
                ctx: &mut tinyflow_api::context::RuntimeContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                self.task_id = ctx.task_id;
                tinyflow_api::functions::LifecycleAware::open(&mut self.function, ctx).await
            }
            async fn close(
                &mut self,
                ctx: &mut tinyflow_api::context::RuntimeContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                tinyflow_api::functions::LifecycleAware::close(&mut self.function, ctx).await
            }
            async fn on_checkpoint(
                &mut self,
                runtime: &mut tinyflow_api::context::RuntimeContext,
                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                tinyflow_api::functions::LifecycleAware::on_checkpoint(
                    &mut self.function,
                    runtime,
                    ckpt,
                )
                .await
            }
            async fn on_checkpoint_end(
                &mut self,
                runtime: &mut tinyflow_api::context::RuntimeContext,
                ckpt: &tinyflow_api::checkpoint::CheckpointContext,
            ) -> tinyflow_api::error::StreamResult<()> {
                tinyflow_api::functions::LifecycleAware::on_checkpoint_end(
                    &mut self.function,
                    runtime,
                    ckpt,
                )
                .await
            }
            fn drain_counters(&mut self) -> (u64, u64) {
                let i = self.input_cnt - self.last_input_cnt;
                let o = self.output_cnt - self.last_output_cnt;
                self.last_input_cnt = self.input_cnt;
                self.last_output_cnt = self.output_cnt;
                (i, o)
            }
            fn collect_metrics(&self) -> std::collections::HashMap<String, f64> {
                self.function.report_metrics()
            }
        }
    };
}
pub(crate) use impl_chain_segment_lifecycle;

/// Middle tail segment: struct + lifecycle + factory; caller supplies `process_record` body.
macro_rules! middle_tail_step {
    (
        struct $seg:ident<$($gen:ident),*>;
        marker: $marker:ty;
        where { $($bound:tt)* }
        process { |$this:ident, $data:ident, $runtime:ident| $($body:tt)* }
        factory $factory:ident
    ) => {
        struct $seg<$($gen),*> {
            function: F,
            task_id: u32,
            input_cnt: u64,
            output_cnt: u64,
            last_input_cnt: u64,
            last_output_cnt: u64,
            _marker: std::marker::PhantomData<$marker>,
        }

        // Reuse lifecycle implementation: delegate to impl_chain_segment_lifecycle! instead of hand-writing the four methods.
        $crate::runtime::chain_segment::impl_chain_segment_lifecycle! {
            impl<$($gen),*>
            ChainSegment for $seg<$($gen),*>,
            with_metrics
            where
                $($bound)*
        }

        #[async_trait::async_trait]
        impl<$($gen),*> $crate::runtime::chain_segment::ChainTailSegment for $seg<$($gen),*>
        where
            $($bound)*
        {
            async fn process_record(
                &mut self,
                input: Box<dyn std::any::Any + Send>,
                runtime: &mut tinyflow_api::context::RuntimeContext,
            ) -> tinyflow_api::error::StreamResult<Vec<Box<dyn std::any::Any + Send>>> {
                let $this = self;
                $this.input_cnt += 1;
                let task_id = $this.task_id;
                let $data = super::downcast_input(input, task_id)?;
                let $runtime = runtime;
                let result = { $($body)* };
                if let Ok(ref records) = result {
                    $this.output_cnt += records.len() as u64;
                }
                result
            }
        }

        pub fn $factory<$($gen),*>(f: F) -> super::ChainSegmentFactory
        where
            $($bound)*
        {
            let step_name = stringify!($seg);
            Box::new(move || {
                super::NamedStep {
                    name: step_name.to_string(),
                    kind: super::ChainStepKindInner::Tail(Box::new($seg {
                        function: f.clone(),
                        task_id: 0,
                        input_cnt: 0,
                        output_cnt: 0,
                        last_input_cnt: 0,
                        last_output_cnt: 0,
                        _marker: std::marker::PhantomData,
                    })),
                }
            })
        }
    };
}
pub(crate) use middle_tail_step;