Skip to main content

tinyflow_api/
functions.rs

1use crate::checkpoint::CheckpointContext;
2use crate::context::RuntimeContext;
3use crate::error::StreamResult;
4use async_trait::async_trait;
5use std::collections::HashMap;
6
7#[async_trait]
8pub trait LifecycleAware {
9    async fn open(&mut self, _ctx: &mut RuntimeContext) -> StreamResult<()> {
10        Ok(())
11    }
12
13    async fn close(&mut self, _ctx: &mut RuntimeContext) -> StreamResult<()> {
14        Ok(())
15    }
16
17    async fn on_checkpoint(
18        &mut self,
19        _runtime: &mut RuntimeContext,
20        _ckpt: &CheckpointContext,
21    ) -> StreamResult<()> {
22        Ok(())
23    }
24
25    async fn on_checkpoint_end(
26        &mut self,
27        _runtime: &mut RuntimeContext,
28        _ckpt: &CheckpointContext,
29    ) -> StreamResult<()> {
30        Ok(())
31    }
32
33    fn report_metrics(&self) -> HashMap<String, f64> {
34        HashMap::new()
35    }
36}
37
38#[macro_export]
39macro_rules! lifecycle {
40    ($t:ty) => {
41        impl $crate::functions::LifecycleAware for $t {}
42    };
43}
44
45#[async_trait]
46pub trait MapFunction<T, O>: LifecycleAware {
47    async fn map(&mut self, input: T, runtime: &mut RuntimeContext) -> StreamResult<O>;
48}
49
50#[async_trait]
51pub trait FilterFunction<T>: LifecycleAware {
52    async fn filter(&mut self, input: &T, runtime: &mut RuntimeContext) -> StreamResult<bool>;
53}
54
55#[async_trait]
56pub trait FlatMapFunction<T, O>: LifecycleAware {
57    async fn flat_map(&mut self, input: T, runtime: &mut RuntimeContext) -> StreamResult<Vec<O>>;
58}
59
60#[async_trait]
61pub trait SinkFunction<T>: LifecycleAware {
62    async fn accept(&mut self, input: T, runtime: &mut RuntimeContext) -> StreamResult<()>;
63}
64
65#[async_trait]
66pub trait SourceFunction<T: Send>: LifecycleAware {
67    async fn poll_next(&mut self) -> StreamResult<Option<T>>;
68}