tinyflow_framework/runtime/
job_runtime.rs

1//! Job-level task handles for graceful shutdown.
2
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5use std::time::Duration;
6
7use tokio::task::JoinHandle;
8
9use crate::runtime::control_msg::ControlSender;
10use tinyflow_api::error::{StreamResult, task_failed};
11
12pub struct ChainJobRuntime {
13    pub task_joins: Mutex<Vec<JoinHandle<()>>>,
14    pub control_senders: Arc<Mutex<HashMap<u32, ControlSender>>>,
15}
16
17impl ChainJobRuntime {
18    pub fn new() -> Self {
19        Self {
20            task_joins: Mutex::new(Vec::new()),
21            control_senders: Arc::new(Mutex::new(HashMap::new())),
22        }
23    }
24
25    pub fn register_control(&self, task_id: u32, ctrl_tx: ControlSender) {
26        self.control_senders
27            .lock()
28            .unwrap()
29            .insert(task_id, ctrl_tx);
30    }
31
32    pub fn register_join(&self, join: JoinHandle<()>) {
33        self.task_joins.lock().unwrap().push(join);
34    }
35}
36
37impl Default for ChainJobRuntime {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43pub async fn shutdown_chain_job(runtime: &ChainJobRuntime) -> StreamResult<()> {
44    runtime.control_senders.lock().unwrap().clear();
45    let joins = std::mem::take(&mut *runtime.task_joins.lock().unwrap());
46    for (i, j) in joins.into_iter().enumerate() {
47        match tokio::time::timeout(Duration::from_secs(10), j).await {
48            Ok(Ok(())) => {}
49            Ok(Err(e)) => {
50                return Err(task_failed(
51                    i as u32,
52                    format!("chain task join failed during shutdown: {e}"),
53                ));
54            }
55            Err(_elapsed) => {
56                log::error!("[shutdown] chain task {} join timed out after 10s", i);
57                return Err(task_failed(i as u32, "chain task shutdown timed out (10s)"));
58            }
59        }
60    }
61    Ok(())
62}
tinyflow_framework/runtime/job_runtime.rs

tinyflow_framework/runtime/
job_runtime.rs
