tinyflow-framework 0.1.0

//! Job-level task handles for graceful shutdown.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use tokio::task::JoinHandle;

use crate::runtime::control_msg::ControlSender;
use tinyflow_api::error::{StreamResult, task_failed};

pub struct ChainJobRuntime {
    pub task_joins: Mutex<Vec<JoinHandle<()>>>,
    pub control_senders: Arc<Mutex<HashMap<u32, ControlSender>>>,
}

impl ChainJobRuntime {
    pub fn new() -> Self {
        Self {
            task_joins: Mutex::new(Vec::new()),
            control_senders: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub fn register_control(&self, task_id: u32, ctrl_tx: ControlSender) {
        self.control_senders
            .lock()
            .unwrap()
            .insert(task_id, ctrl_tx);
    }

    pub fn register_join(&self, join: JoinHandle<()>) {
        self.task_joins.lock().unwrap().push(join);
    }
}

impl Default for ChainJobRuntime {
    fn default() -> Self {
        Self::new()
    }
}

pub async fn shutdown_chain_job(runtime: &ChainJobRuntime) -> StreamResult<()> {
    runtime.control_senders.lock().unwrap().clear();
    let joins = std::mem::take(&mut *runtime.task_joins.lock().unwrap());
    for (i, j) in joins.into_iter().enumerate() {
        match tokio::time::timeout(Duration::from_secs(10), j).await {
            Ok(Ok(())) => {}
            Ok(Err(e)) => {
                return Err(task_failed(
                    i as u32,
                    format!("chain task join failed during shutdown: {e}"),
                ));
            }
            Err(_elapsed) => {
                log::error!("[shutdown] chain task {} join timed out after 10s", i);
                return Err(task_failed(i as u32, "chain task shutdown timed out (10s)"));
            }
        }
    }
    Ok(())
}