tinyflow-framework 0.1.1

Streaming runtime engine — checkpoint, chained task execution, and state store
Documentation
use std::any::Any;
use std::marker::PhantomData;

use async_trait::async_trait;

use tinyflow_api::context::RuntimeContext;
use tinyflow_api::error::StreamResult;
use tinyflow_api::functions::SinkFunction;

use super::{
    ChainSegmentFactory, ChainStepKindInner, ChainTailSegment, NamedStep, downcast_input,
    impl_chain_segment_lifecycle,
};

struct SinkChainSegment<T, W> {
    function: W,
    task_id: u32,
    input_cnt: u64,
    output_cnt: u64,
    last_input_cnt: u64,
    last_output_cnt: u64,
    _marker: PhantomData<T>,
}

impl_chain_segment_lifecycle! {
    impl<T, W>
    ChainSegment for SinkChainSegment<T, W>,
    with_metrics
    where
        T: Send + 'static,
        W: SinkFunction<T> + Clone + Send + 'static,
}

#[async_trait]
impl<T, W> ChainTailSegment for SinkChainSegment<T, W>
where
    T: Send + 'static,
    W: SinkFunction<T> + Clone + Send + 'static,
{
    async fn process_record(
        &mut self,
        input: Box<dyn Any + Send>,
        runtime: &mut RuntimeContext,
    ) -> StreamResult<Vec<Box<dyn Any + Send>>> {
        let record = downcast_input::<T>(input, self.task_id)?;
        self.input_cnt += 1;
        self.function.accept(record, runtime).await?;
        Ok(vec![])
    }
}

pub fn sink_chain_factory<T, F>(f: F) -> ChainSegmentFactory
where
    T: Send + 'static,
    F: SinkFunction<T> + Clone + Send + 'static,
{
    Box::new(move || NamedStep {
        name: "Sink".into(),
        kind: ChainStepKindInner::Tail(Box::new(SinkChainSegment {
            function: f.clone(),
            task_id: 0,
            input_cnt: 0,
            output_cnt: 0,
            last_input_cnt: 0,
            last_output_cnt: 0,
            _marker: PhantomData,
        })),
    })
}