tinyflow_framework/runtime/
checkpoint_manager.rs

1use std::sync::{Arc, Mutex};
2use tokio::sync::mpsc::error::TrySendError;
3
4use crate::runtime::control_msg::{ControlMsg, ControlSender};
5use std::collections::{HashMap, HashSet};
6use tinyflow_api::error::{StreamResult, task_failed};
7
8type TaskId = u32;
9
10pub struct CheckpointManager {
11    interval_ms: u64,
12    next_checkpoint_id: u64,
13    pending_acks: HashMap<u64, HashSet<TaskId>>,
14    registered_tasks: Arc<Mutex<HashMap<TaskId, ControlSender>>>,
15    ack_rx: tokio::sync::mpsc::Receiver<ControlMsg>,
16}
17
18impl CheckpointManager {
19    pub fn new(
20        interval_ms: u64,
21        ack_rx: tokio::sync::mpsc::Receiver<ControlMsg>,
22        registered_tasks: Arc<Mutex<HashMap<TaskId, ControlSender>>>,
23    ) -> Self {
24        log::info!("[checkpoint] manager created, interval={}ms", interval_ms);
25        Self {
26            interval_ms,
27            next_checkpoint_id: 1,
28            pending_acks: HashMap::new(),
29            registered_tasks,
30            ack_rx,
31        }
32    }
33
34    pub fn trigger_checkpoint(&mut self) -> StreamResult<()> {
35        let ckpt_id = self.next_checkpoint_id;
36        self.next_checkpoint_id += 1;
37        let mut tasks = self.registered_tasks.lock().unwrap();
38        log::info!(
39            "[checkpoint] triggering checkpoint {}, broadcasting to {} tasks",
40            ckpt_id,
41            tasks.len()
42        );
43
44        let msg = ControlMsg::Checkpoint(ckpt_id);
45        let mut dead = Vec::new();
46        for (id, tx) in tasks.iter() {
47            match tx.try_send(msg.clone()) {
48                Err(TrySendError::Full(_)) => {
49                    return Err(task_failed(
50                        *id,
51                        format!(
52                            "[checkpoint] checkpoint {}: task {} control channel full, failing job",
53                            ckpt_id, id,
54                        ),
55                    ));
56                }
57                Err(TrySendError::Closed(_)) => {
58                    dead.push(*id);
59                }
60                Ok(_) => {}
61            }
62        }
63        for id in &dead {
64            tasks.remove(id);
65        }
66
67        if tasks.is_empty() {
68            return Ok(());
69        }
70
71        let acks: HashSet<TaskId> = tasks.keys().copied().collect();
72        self.pending_acks.insert(ckpt_id, acks);
73        Ok(())
74    }
75
76    pub fn pending_acks_len(&self, checkpoint_id: u64) -> Option<usize> {
77        self.pending_acks.get(&checkpoint_id).map(|s| s.len())
78    }
79
80    pub fn handle_ack(&mut self, checkpoint_id: u64, task_id: TaskId) -> bool {
81        if let Some(acks) = self.pending_acks.get_mut(&checkpoint_id) {
82            acks.remove(&task_id);
83            let remaining = acks.len();
84            log::info!(
85                "[checkpoint] received Ack(ckpt={}, task={}), {} acks remaining",
86                checkpoint_id,
87                task_id,
88                remaining
89            );
90            if acks.is_empty() {
91                self.pending_acks.remove(&checkpoint_id);
92                return true;
93            }
94        }
95        false
96    }
97
98    pub async fn run(&mut self) -> StreamResult<()> {
99        if self.interval_ms == 0 {
100            return Ok(());
101        }
102        log::info!("[checkpoint] run loop started");
103        loop {
104            if self.registered_tasks.lock().unwrap().is_empty() {
105                log::info!("[checkpoint] all tasks finished, stopping checkpoint manager");
106                break;
107            }
108
109            let sleep = tokio::time::sleep(std::time::Duration::from_millis(self.interval_ms));
110            tokio::pin!(sleep);
111
112            tokio::select! {
113                _ = &mut sleep => {
114                    self.trigger_checkpoint()?;
115                    if self.registered_tasks.lock().unwrap().is_empty() {
116                        break;
117                    }
118                }
119                msg = self.ack_rx.recv() => {
120                    match msg {
121                        Some(ControlMsg::Ack(ckpt_id, task_id)) => {
122                            if self.handle_ack(ckpt_id, task_id) {
123                                self.broadcast_checkpoint_end(ckpt_id)?;
124                            }
125                        }
126                        Some(_) => {}
127                        None => break,
128                    }
129                }
130            }
131        }
132        log::info!("[checkpoint] run loop ended");
133        Ok(())
134    }
135
136    fn broadcast_checkpoint_end(&mut self, checkpoint_id: u64) -> StreamResult<()> {
137        let msg = ControlMsg::CheckpointEnd(checkpoint_id);
138        let mut tasks = self.registered_tasks.lock().unwrap();
139        let total_tasks = tasks.len();
140        let mut dead = Vec::new();
141        let mut sent = 0usize;
142        for (id, tx) in tasks.iter() {
143            match tx.try_send(msg.clone()) {
144                Err(TrySendError::Full(_)) => {
145                    log::warn!(
146                        "[checkpoint] CheckpointEnd({}) partial broadcast: {} of {} tasks \
147                         received the message; remaining task {} has full control channel. \
148                         Tasks that already received CheckpointEnd will proceed to commit \
149                         Kafka offset and (for sink subtask 0) Iceberg transaction, so the \
150                         remaining tasks will fall behind. This is a known data-loss window \
151                         documented in CLAUDE.md (Kafka offset committed before Iceberg commit).",
152                        checkpoint_id,
153                        sent,
154                        total_tasks,
155                        id,
156                    );
157                    return Err(task_failed(
158                        *id,
159                        format!(
160                            "[checkpoint] CheckpointEnd({}): task {} control channel full, failing job",
161                            checkpoint_id, id,
162                        ),
163                    ));
164                }
165                Err(TrySendError::Closed(_)) => {
166                    dead.push(*id);
167                }
168                Ok(_) => {
169                    sent += 1;
170                }
171            }
172        }
173        for id in &dead {
174            tasks.remove(id);
175        }
176        log::info!(
177            "[checkpoint] CheckpointEnd({}) sent to {} tasks ({} removed)",
178            checkpoint_id,
179            tasks.len(),
180            dead.len(),
181        );
182        Ok(())
183    }
184}
tinyflow_framework/runtime/checkpoint_manager.rs

tinyflow_framework/runtime/
checkpoint_manager.rs
