tinyflow_framework/env/
stream_env.rs

1use std::path::PathBuf;
2
3use crate::stream::data_stream::DataStream;
4use std::sync::{Arc, Mutex};
5use tinyflow_api::functions::*;
6
7use crate::runtime::chain_job::ChainJob;
8
9type PendingInitialState = Vec<(String, String, Vec<u8>)>;
10
11/// Job environment: DSL chain construction panics on programmer errors.
12/// Runtime entry: [`crate::stream::chain_execute::execute`].
13#[derive(Clone)]
14pub struct StreamEnv {
15    pub(crate) default_parallelism: usize,
16    pub(crate) checkpoint_interval_ms: u64,
17    pub chain_job: Arc<Mutex<Option<ChainJob>>>,
18    /// Root work directory for `job/state.db`; required when checkpoint or chain_state is enabled.
19    pub(crate) state_work_dir: Option<PathBuf>,
20    /// When true, each chain subtask opens `chains/{i}/chain.db` (requires a state work dir at run).
21    /// Default: **true** — call [`Self::set_enable_chain_state(false)`] to skip per-subtask DB I/O.
22    pub(crate) enable_chain_state: bool,
23    pub(crate) pending_job_state: Arc<Mutex<PendingInitialState>>,
24}
25
26impl Default for StreamEnv {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl StreamEnv {
33    pub fn new() -> Self {
34        Self {
35            default_parallelism: 1,
36            checkpoint_interval_ms: 0,
37            chain_job: Arc::new(Mutex::new(None)),
38            state_work_dir: None,
39            enable_chain_state: true,
40            pending_job_state: Arc::new(Mutex::new(Vec::new())),
41        }
42    }
43
44    pub fn set_state_work_dir(&mut self, path: impl Into<PathBuf>) {
45        self.state_work_dir = Some(path.into());
46    }
47
48    /// Enable per-subtask `chain_state` (`chains/{subtask}/chain.db`) at task spawn.
49    pub fn set_enable_chain_state(&mut self, enable: bool) {
50        self.enable_chain_state = enable;
51    }
52
53    /// Buffer an initial job-state entry that will be flushed into `job_state` before the stream starts.
54    pub fn add_initial_job_state(&self, namespace: &str, key: &str, value: &[u8]) {
55        self.pending_job_state.lock().unwrap().push((
56            namespace.to_string(),
57            key.to_string(),
58            value.to_vec(),
59        ));
60    }
61
62    pub fn set_parallelism(&mut self, p: usize) {
63        self.default_parallelism = p;
64    }
65
66    pub fn set_checkpoint_interval(&mut self, ms: u64) {
67        self.checkpoint_interval_ms = ms;
68    }
69
70    /// Run the built chain until Ctrl+C, task failure, or checkpoint manager exit.
71    pub async fn execute(self) -> tinyflow_api::error::StreamResult<()> {
72        crate::stream::chain_execute::execute(self).await
73    }
74
75    /// Test-only: whether per-subtask `chain_state` is enabled.
76    #[doc(hidden)]
77    pub fn test_enable_chain_state(&self) -> bool {
78        self.enable_chain_state
79    }
80
81    /// Test-only: inspect built chain step count.
82    #[doc(hidden)]
83    pub fn test_chain_step_count(&self) -> usize {
84        self.chain_job
85            .lock()
86            .unwrap()
87            .as_ref()
88            .map(|j| j.steps.len())
89            .unwrap_or(0)
90    }
91
92    #[doc(hidden)]
93    pub fn test_chain_frozen(&self) -> bool {
94        self.chain_job
95            .lock()
96            .unwrap()
97            .as_ref()
98            .map(|j| j.frozen)
99            .unwrap_or(false)
100    }
101
102    /// Test-only: simulate terminal sink freezing the chain.
103    #[doc(hidden)]
104    pub fn test_freeze_chain(&self) {
105        if let Some(job) = self.chain_job.lock().unwrap().as_mut() {
106            job.frozen = true;
107        }
108    }
109
110    pub fn add_source<T: Send + 'static, S: SourceFunction<T> + Clone + Send + 'static>(
111        &self,
112        source: S,
113    ) -> DataStream<T> {
114        self.add_source_with_factory(move || source.clone())
115    }
116
117    pub fn add_source_with_factory<T: Send + 'static, S: SourceFunction<T> + Send + 'static>(
118        &self,
119        factory: impl Fn() -> S + Send + 'static,
120    ) -> DataStream<T> {
121        let mut job_guard = self.chain_job.lock().unwrap();
122        if job_guard.is_some() {
123            panic!("StreamEnv supports only one ChainJob; add_source already called");
124        }
125        let par = self.default_parallelism.max(1);
126        let factory = std::sync::Arc::new(std::sync::Mutex::new(factory));
127        let chain_factory =
128            crate::runtime::chain_segment::source_chain_factory(std::sync::Arc::clone(&factory));
129        let mut job = ChainJob::new(par);
130        job.push_step_with_types(
131            crate::runtime::chain_job::ChainStep {
132                kind: crate::runtime::chain_job::ChainOperatorKind::Source,
133                user_name: None,
134                factory: chain_factory,
135            },
136            None,
137            std::any::TypeId::of::<T>(),
138        );
139        *job_guard = Some(job);
140        DataStream::new(self.clone(), 0)
141    }
142}
tinyflow_framework/env/stream_env.rs

tinyflow_framework/env/
stream_env.rs
