tinyflow_framework/stream/chain_execute/
spawn.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use crate::env::stream_env::StreamEnv;
5use crate::runtime::chain_job::ChainStep;
6use crate::runtime::chain_segment::NamedStep;
7use crate::runtime::chained_stream_task::ChainedStreamTask;
8use crate::runtime::control_msg::ControlMsg;
9use crate::runtime::job_runtime::ChainJobRuntime;
10use crate::runtime::pending_task::PendingTask;
11use crate::runtime::task_lifecycle::spawn_monitored_task;
12use crate::state::open_chain_state_db;
13use tinyflow_api::error::{StreamError, StreamResult};
14use tinyflow_api::state::StateStore;
15
16use super::validate::{resolve_step_name, validate_chain_job};
17
18/// Build one pending chain task per subtask from the env's [`ChainJob`](crate::runtime::chain_job::ChainJob).
19#[allow(clippy::too_many_arguments)]
20pub async fn spawn_chained_tasks(
21    env: &StreamEnv,
22    job_state: Arc<dyn StateStore>,
23    state_work_dir: Option<&Path>,
24    job_runtime: &ChainJobRuntime,
25    ack_tx: &tokio::sync::mpsc::Sender<ControlMsg>,
26    error_tx: &tokio::sync::mpsc::Sender<StreamError>,
27    chain_id: u64,
28    task_id_counter: &mut u32,
29) -> StreamResult<Vec<PendingTask>> {
30    let (par, step_count, step_names, steps) = {
31        let mut guard = env.chain_job.lock().unwrap();
32        let job = guard.as_mut().ok_or_else(|| StreamError::TaskFailed {
33            task_id: 0,
34            reason: "no ChainJob; call add_source first".into(),
35        })?;
36
37        validate_chain_job(env, job)?;
38
39        let par = job.parallelism.max(1);
40        let step_count = job.steps.len();
41        let step_names: Vec<String> = job
42            .steps
43            .iter()
44            .map(|s| resolve_step_name(s.kind, &s.user_name))
45            .collect();
46        let steps = std::mem::take(&mut job.steps);
47        (par, step_count, step_names, steps)
48    };
49
50    let mut pending = Vec::new();
51    for pi in 0..par {
52        let task_id = *task_id_counter;
53        *task_id_counter += 1;
54
55        let (ctrl_tx, ctrl_rx) = tokio::sync::mpsc::channel(64);
56        job_runtime.register_control(task_id, ctrl_tx);
57
58        let chain_label = if step_names.len() == 1 {
59            step_names[0].clone()
60        } else {
61            format!(
62                "{}..{}",
63                step_names.first().unwrap(),
64                step_names.last().unwrap()
65            )
66        };
67        let task_name = format!("{}-{}", chain_label, pi);
68
69        assert_eq!(steps.len(), step_count);
70        let mut substeps: Vec<NamedStep> =
71            steps.iter().map(|s: &ChainStep| (s.factory)()).collect();
72        for (i, step_info) in steps.iter().enumerate() {
73            substeps[i].name = resolve_step_name(step_info.kind, &step_info.user_name);
74        }
75
76        let chain_state = if env.enable_chain_state {
77            let work_dir = state_work_dir.ok_or_else(|| StreamError::TaskFailed {
78                task_id: 0,
79                reason: "enable_chain_state requires state_work_dir at spawn".into(),
80            })?;
81            Some(open_chain_state_db(work_dir, pi as u32).await?)
82        } else {
83            None
84        };
85
86        let task = ChainedStreamTask::from_steps(
87            task_id,
88            task_name.clone(),
89            pi as u32,
90            chain_id,
91            par as u32,
92            substeps,
93            ctrl_rx,
94            ack_tx.clone(),
95            Arc::clone(&job_state),
96            chain_state,
97        )?;
98
99        let error_tx = error_tx.clone();
100        pending.push(PendingTask::new(task_id, move || {
101            spawn_monitored_task(task_id, task_name, error_tx, move || async move {
102                task.run().await
103            })
104        }));
105    }
106
107    Ok(pending)
108}
tinyflow_framework/stream/chain_execute/spawn.rs

tinyflow_framework/stream/chain_execute/
spawn.rs
