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,
})),
})
}