use std::sync::{Arc, Mutex};
use tokio::sync::mpsc::error::TrySendError;
use crate::runtime::control_msg::{ControlMsg, ControlSender};
use std::collections::{HashMap, HashSet};
use tinyflow_api::error::{StreamResult, task_failed};
type TaskId = u32;
pub struct CheckpointManager {
interval_ms: u64,
next_checkpoint_id: u64,
pending_acks: HashMap<u64, HashSet<TaskId>>,
registered_tasks: Arc<Mutex<HashMap<TaskId, ControlSender>>>,
ack_rx: tokio::sync::mpsc::Receiver<ControlMsg>,
}
impl CheckpointManager {
pub fn new(
interval_ms: u64,
ack_rx: tokio::sync::mpsc::Receiver<ControlMsg>,
registered_tasks: Arc<Mutex<HashMap<TaskId, ControlSender>>>,
) -> Self {
log::info!("[checkpoint] manager created, interval={}ms", interval_ms);
Self {
interval_ms,
next_checkpoint_id: 1,
pending_acks: HashMap::new(),
registered_tasks,
ack_rx,
}
}
pub fn trigger_checkpoint(&mut self) -> StreamResult<()> {
let ckpt_id = self.next_checkpoint_id;
self.next_checkpoint_id += 1;
let mut tasks = self.registered_tasks.lock().unwrap();
log::info!(
"[checkpoint] triggering checkpoint {}, broadcasting to {} tasks",
ckpt_id,
tasks.len()
);
let msg = ControlMsg::Checkpoint(ckpt_id);
let mut dead = Vec::new();
for (id, tx) in tasks.iter() {
match tx.try_send(msg.clone()) {
Err(TrySendError::Full(_)) => {
return Err(task_failed(
*id,
format!(
"[checkpoint] checkpoint {}: task {} control channel full, failing job",
ckpt_id, id,
),
));
}
Err(TrySendError::Closed(_)) => {
dead.push(*id);
}
Ok(_) => {}
}
}
for id in &dead {
tasks.remove(id);
}
if tasks.is_empty() {
return Ok(());
}
let acks: HashSet<TaskId> = tasks.keys().copied().collect();
self.pending_acks.insert(ckpt_id, acks);
Ok(())
}
pub fn pending_acks_len(&self, checkpoint_id: u64) -> Option<usize> {
self.pending_acks.get(&checkpoint_id).map(|s| s.len())
}
pub fn handle_ack(&mut self, checkpoint_id: u64, task_id: TaskId) -> bool {
if let Some(acks) = self.pending_acks.get_mut(&checkpoint_id) {
acks.remove(&task_id);
let remaining = acks.len();
log::info!(
"[checkpoint] received Ack(ckpt={}, task={}), {} acks remaining",
checkpoint_id,
task_id,
remaining
);
if acks.is_empty() {
self.pending_acks.remove(&checkpoint_id);
return true;
}
}
false
}
pub async fn run(&mut self) -> StreamResult<()> {
if self.interval_ms == 0 {
return Ok(());
}
log::info!("[checkpoint] run loop started");
loop {
if self.registered_tasks.lock().unwrap().is_empty() {
log::info!("[checkpoint] all tasks finished, stopping checkpoint manager");
break;
}
let sleep = tokio::time::sleep(std::time::Duration::from_millis(self.interval_ms));
tokio::pin!(sleep);
tokio::select! {
_ = &mut sleep => {
self.trigger_checkpoint()?;
if self.registered_tasks.lock().unwrap().is_empty() {
break;
}
}
msg = self.ack_rx.recv() => {
match msg {
Some(ControlMsg::Ack(ckpt_id, task_id)) => {
if self.handle_ack(ckpt_id, task_id) {
self.broadcast_checkpoint_end(ckpt_id)?;
}
}
Some(_) => {}
None => break,
}
}
}
}
log::info!("[checkpoint] run loop ended");
Ok(())
}
fn broadcast_checkpoint_end(&mut self, checkpoint_id: u64) -> StreamResult<()> {
let msg = ControlMsg::CheckpointEnd(checkpoint_id);
let mut tasks = self.registered_tasks.lock().unwrap();
let total_tasks = tasks.len();
let mut dead = Vec::new();
let mut sent = 0usize;
for (id, tx) in tasks.iter() {
match tx.try_send(msg.clone()) {
Err(TrySendError::Full(_)) => {
log::warn!(
"[checkpoint] CheckpointEnd({}) partial broadcast: {} of {} tasks \
received the message; remaining task {} has full control channel. \
Tasks that already received CheckpointEnd will proceed to commit \
Kafka offset and (for sink subtask 0) Iceberg transaction, so the \
remaining tasks will fall behind. This is a known data-loss window \
documented in CLAUDE.md (Kafka offset committed before Iceberg commit).",
checkpoint_id,
sent,
total_tasks,
id,
);
return Err(task_failed(
*id,
format!(
"[checkpoint] CheckpointEnd({}): task {} control channel full, failing job",
checkpoint_id, id,
),
));
}
Err(TrySendError::Closed(_)) => {
dead.push(*id);
}
Ok(_) => {
sent += 1;
}
}
}
for id in &dead {
tasks.remove(id);
}
log::info!(
"[checkpoint] CheckpointEnd({}) sent to {} tasks ({} removed)",
checkpoint_id,
tasks.len(),
dead.len(),
);
Ok(())
}
}