Skip to main content

tinyflow_framework/stream/chain_execute/
mod.rs

1//! Spawn [`ChainedStreamTask`](crate::runtime::chained_stream_task::ChainedStreamTask) from a
2//! [`ChainJob`](crate::runtime::chain_job::ChainJob) (framework-only, no committer).
3
4mod spawn;
5mod validate;
6
7use std::sync::Arc;
8
9use crate::env::stream_env::StreamEnv;
10use crate::runtime::checkpoint_manager::CheckpointManager;
11use crate::runtime::job_runtime::{ChainJobRuntime, shutdown_chain_job};
12use crate::runtime::task_lifecycle::report_task_failure;
13use tinyflow_api::error::StreamResult;
14
15pub use spawn::spawn_chained_tasks;
16
17use spawn::spawn_chained_tasks as spawn_tasks;
18use validate::resolve_job_state_for_execute;
19
20/// Run the built operator chain until Ctrl+C, task failure, or checkpoint manager exit.
21pub async fn execute(env: StreamEnv) -> StreamResult<()> {
22    let (job_state, work_dir) = resolve_job_state_for_execute(&env).await?;
23    let work_dir_ref = work_dir.as_deref();
24
25    {
26        let pending = std::mem::take(&mut *env.pending_job_state.lock().unwrap());
27        for (namespace, key, value) in pending {
28            job_state.put(&namespace, key.as_bytes(), &value).await?;
29        }
30    }
31
32    let (ack_tx, ack_rx) = tokio::sync::mpsc::channel(64);
33    let (error_tx, mut error_rx) = tokio::sync::mpsc::channel(64);
34    let job_runtime = ChainJobRuntime::new();
35    let mut task_id_counter = 0u32;
36
37    let pending = spawn_tasks(
38        &env,
39        job_state,
40        work_dir_ref,
41        &job_runtime,
42        &ack_tx,
43        &error_tx,
44        0,
45        &mut task_id_counter,
46    )
47    .await?;
48
49    for p in pending {
50        job_runtime.register_join(p.spawn());
51    }
52
53    let cm_handle = if env.checkpoint_interval_ms > 0
54        && !job_runtime.control_senders.lock().unwrap().is_empty()
55    {
56        let interval = env.checkpoint_interval_ms;
57        log::info!(
58            "[execute] checkpoint manager interval={}ms, {} tasks",
59            interval,
60            job_runtime.control_senders.lock().unwrap().len()
61        );
62        let mut cm =
63            CheckpointManager::new(interval, ack_rx, Arc::clone(&job_runtime.control_senders));
64        let cm_error_tx = error_tx.clone();
65        Some(tokio::spawn(async move {
66            if let Err(e) = cm.run().await {
67                report_task_failure(&cm_error_tx, 0, "checkpoint-manager", e).await;
68            }
69        }))
70    } else {
71        drop(ack_rx);
72        None
73    };
74
75    drop(ack_tx);
76
77    let result = if let Some(handle) = cm_handle {
78        tokio::select! {
79            _ = tokio::signal::ctrl_c() => Ok(()),
80            msg = error_rx.recv() => match msg {
81                Some(e) => Err(e),
82                None => Ok(()),
83            },
84            _ = handle => match error_rx.try_recv() {
85                Ok(e) => Err(e),
86                Err(_) => Ok(()),
87            },
88        }
89    } else {
90        tokio::select! {
91            _ = tokio::signal::ctrl_c() => Ok(()),
92            msg = error_rx.recv() => match msg {
93                Some(e) => Err(e),
94                None => Ok(()),
95            },
96        }
97    };
98
99    shutdown_chain_job(&job_runtime).await?;
100    result
101}