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;
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(())
}
}