tinyflow_framework/runtime/chained_stream_task/

mod.rs

//! In-process operator chain (framework-only; no connector channel types).

mod run_loop;

#[cfg(test)]
mod tests;

use std::sync::Arc;

use crate::runtime::chain_segment::{
    ChainSegmentFactory, ChainSourceSegment, ChainStepKindInner, ChainTailSegment, NamedStep,
    instantiate_chain_steps,
};
use crate::runtime::control_msg::{ControlMsg, ControlReceiver};
use tinyflow_api::checkpoint::CheckpointContext;
use tinyflow_api::context::RuntimeContext;
use tinyflow_api::error::StreamResult;
use tinyflow_api::state::StateStore;

use run_loop::run_through_chain;

/// In-process operator chain for a single subtask.
pub struct ChainedStreamTask {
    task_id: u32,
    task_name: String,
    subtask_index: u32,
    chain_id: u64,
    parallelism: u32,
    source: Option<Box<dyn ChainSourceSegment>>,
    tail_steps: Vec<Box<dyn ChainTailSegment>>,
    step_names: Vec<String>,
    last_checkpoint_metrics_at_ms: u64,
    metric_key: String,
    control_rx: ControlReceiver,
    ack_tx: tokio::sync::mpsc::Sender<ControlMsg>,
    job_state: Arc<dyn StateStore>,
    chain_state: Option<Arc<dyn StateStore>>,
}

impl ChainedStreamTask {
    #[allow(clippy::too_many_arguments)]
    pub fn from_factories(
        task_id: u32,
        task_name: impl Into<String>,
        subtask_index: u32,
        chain_id: u64,
        parallelism: u32,
        factories: &[ChainSegmentFactory],
        control_rx: ControlReceiver,
        ack_tx: tokio::sync::mpsc::Sender<ControlMsg>,
        job_state: Arc<dyn StateStore>,
        chain_state: Option<Arc<dyn StateStore>>,
    ) -> StreamResult<Self> {
        Self::from_steps(
            task_id,
            task_name,
            subtask_index,
            chain_id,
            parallelism,
            instantiate_chain_steps(factories),
            control_rx,
            ack_tx,
            job_state,
            chain_state,
        )
    }

    #[allow(clippy::too_many_arguments)]
    pub fn from_steps(
        task_id: u32,
        task_name: impl Into<String>,
        subtask_index: u32,
        chain_id: u64,
        parallelism: u32,
        steps: Vec<NamedStep>,
        control_rx: ControlReceiver,
        ack_tx: tokio::sync::mpsc::Sender<ControlMsg>,
        job_state: Arc<dyn StateStore>,
        chain_state: Option<Arc<dyn StateStore>>,
    ) -> StreamResult<Self> {
        let mut source = None;
        let mut tail_steps = Vec::new();
        let mut step_names = Vec::new();
        for step in steps {
            step_names.push(step.name);
            match step.kind {
                ChainStepKindInner::Source(s) => {
                    if source.is_some() {
                        return Err(tinyflow_api::error::StreamError::TaskFailed {
                            task_id,
                            reason: "chain has multiple source steps".into(),
                        });
                    }
                    source = Some(s);
                }
                ChainStepKindInner::Tail(o) => tail_steps.push(o),
            }
        }
        if source.is_none() {
            return Err(tinyflow_api::error::StreamError::TaskFailed {
                task_id,
                reason: "chain missing source head".into(),
            });
        }
        Ok(Self {
            task_id,
            task_name: task_name.into(),
            subtask_index,
            chain_id,
            parallelism,
            source,
            tail_steps,
            step_names,
            last_checkpoint_metrics_at_ms: 0,
            metric_key: format!("metric/{}", subtask_index),
            control_rx,
            ack_tx,
            job_state,
            chain_state,
        })
    }

    async fn open_all(&mut self, ctx: &mut RuntimeContext) -> StreamResult<()> {
        if let Some(s) = &mut self.source {
            s.open(ctx).await?;
        }
        for step in &mut self.tail_steps {
            step.open(ctx).await?;
        }
        Ok(())
    }

    async fn close_all(&mut self, ctx: &mut RuntimeContext) -> StreamResult<()> {
        for step in &mut self.tail_steps {
            step.close(ctx).await?;
        }
        if let Some(s) = &mut self.source {
            s.close(ctx).await?;
        }
        Ok(())
    }

    async fn on_checkpoint_all(
        &mut self,
        runtime: &mut RuntimeContext,
        ckpt: &CheckpointContext,
    ) -> StreamResult<()> {
        if let Some(s) = &mut self.source {
            s.on_checkpoint(runtime, ckpt).await?;
        }
        for step in &mut self.tail_steps {
            step.on_checkpoint(runtime, ckpt).await?;
        }
        Ok(())
    }

    async fn on_checkpoint_end_all(
        &mut self,
        runtime: &mut RuntimeContext,
        ckpt: &CheckpointContext,
    ) -> StreamResult<()> {
        if let Some(s) = &mut self.source {
            s.on_checkpoint_end(runtime, ckpt).await?;
        }
        for step in &mut self.tail_steps {
            step.on_checkpoint_end(runtime, ckpt).await?;
        }
        self.collect_and_store_metrics(runtime, ckpt).await?;
        Ok(())
    }

    async fn collect_and_store_metrics(
        &mut self,
        runtime: &mut RuntimeContext,
        ckpt: &CheckpointContext,
    ) -> StreamResult<()> {
        let now_ms = ckpt.timestamp_ms;
        let elapsed_ms = now_ms.saturating_sub(self.last_checkpoint_metrics_at_ms);
        self.last_checkpoint_metrics_at_ms = now_ms;

        if elapsed_ms == 0 {
            return Ok(());
        }

        let mut all_metrics = serde_json::Map::new();
        let mut idx = 0;

        let mut add_entry =
            |label: String,
             input_delta: u64,
             output_delta: u64,
             custom: std::collections::HashMap<String, f64>| {
                let input_rate = (input_delta as f64 / elapsed_ms as f64) * 1000.0;
                let output_rate = (output_delta as f64 / elapsed_ms as f64) * 1000.0;
                let mut entry = serde_json::json!({
                    "input_rate": input_rate,
                    "output_rate": output_rate,
                });
                for (k, v) in custom {
                    if v.is_nan() || v.is_infinite() {
                        log::warn!("metric value is {v}, replacing with 0.0");
                        entry[k] =
                            serde_json::Value::Number(serde_json::Number::from_f64(0.0).unwrap());
                    } else {
                        entry[k] = serde_json::Value::Number(
                            serde_json::Number::from_f64(v)
                                .unwrap_or_else(|| serde_json::Number::from_f64(0.0).unwrap()),
                        );
                    }
                }
                all_metrics.insert(label, entry);
            };

        if let Some(s) = &mut self.source {
            let label = if idx < self.step_names.len() {
                format!("{}/{}", self.step_names[idx], self.subtask_index)
            } else {
                format!("source/{}", self.subtask_index)
            };
            idx += 1;
            let (input_delta, output_delta) = s.drain_counters();
            add_entry(label, input_delta, output_delta, s.collect_metrics());
        }

        for step in &mut self.tail_steps {
            let label = if idx < self.step_names.len() {
                format!("{}/{}", self.step_names[idx], self.subtask_index)
            } else {
                format!("tail/{}", self.subtask_index)
            };
            idx += 1;
            let (input_delta, output_delta) = step.drain_counters();
            add_entry(label, input_delta, output_delta, step.collect_metrics());
        }

        let json_bytes = serde_json::to_vec(&all_metrics).map_err(|e| {
            tinyflow_api::error::StreamError::TaskFailed {
                task_id: self.task_id,
                reason: format!("serialize metrics: {e}"),
            }
        })?;

        let metric_key = &self.metric_key;
        runtime
            .job_state
            .put("job_status", metric_key.as_bytes(), &json_bytes)
            .await
            .map_err(|e| tinyflow_api::error::StreamError::TaskFailed {
                task_id: self.task_id,
                reason: format!("persist metrics: {e}"),
            })?;

        Ok(())
    }

    async fn handle_checkpoint(
        &mut self,
        ckpt_id: u64,
        runtime: &mut RuntimeContext,
    ) -> StreamResult<()> {
        let timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;
        let ckpt_ctx = CheckpointContext::new(ckpt_id, timestamp);
        self.on_checkpoint_all(runtime, &ckpt_ctx).await?;

        self.ack_tx
            .send(ControlMsg::Ack(ckpt_id, self.task_id))
            .await
            .map_err(|_| tinyflow_api::error::StreamError::ChannelClosed {
                reason: format!("ack_tx for task {} at checkpoint {}", self.task_id, ckpt_id),
            })?;
        Ok(())
    }

    async fn run_loop(&mut self, ctx: &mut RuntimeContext) -> StreamResult<()> {
        loop {
            tokio::select! {
                ctrl = self.control_rx.recv() => {
                    match ctrl {
                        Some(ControlMsg::Checkpoint(ckpt_id)) => {
                            self.handle_checkpoint(ckpt_id, ctx).await?;
                        }
                        Some(ControlMsg::CheckpointEnd(id)) => {
                            let timestamp = std::time::SystemTime::now()
                                .duration_since(std::time::UNIX_EPOCH)
                                .unwrap_or_default()
                                .as_millis() as u64;
                            let ckpt = CheckpointContext::new(id, timestamp);
                            self.on_checkpoint_end_all(ctx, &ckpt).await?;
                        }
                        Some(ControlMsg::Ack(_, _)) => {}
                        None => break,
                    }
                }
                record = async {
                    self.source
                        .as_mut()
                        .expect("chain source present")
                        .poll_next()
                        .await
                } => {
                    match record? {
                        Some(r) => {
                            run_through_chain(&mut self.tail_steps, r, ctx).await?;
                        }
                        None => {
                            tokio::task::yield_now().await;
                        }
                    }
                }
            }
        }
        Ok(())
    }

    pub async fn run(mut self) -> StreamResult<()> {
        log::info!(
            "[{}] chained task running (task_id={}, chain={}, parallelism={})",
            self.task_name,
            self.task_id,
            self.chain_id,
            self.parallelism
        );
        let mut ctx = RuntimeContext::new(
            self.task_id,
            self.task_name.clone(),
            self.subtask_index,
            self.parallelism,
            self.chain_id,
            Arc::clone(&self.job_state),
            self.chain_state.as_ref().map(Arc::clone),
        );
        self.open_all(&mut ctx).await?;
        let result = self.run_loop(&mut ctx).await;
        if let Err(e) = self.close_all(&mut ctx).await {
            log::warn!("[{}] close_all failed: {e}", self.task_name);
        }
        result?;
        log::info!("[{}] chained task finished", self.task_name);
        Ok(())
    }
}