1use std::collections::BTreeMap;
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7
8use laminar_core::checkpoint::{CheckpointManifest, StateFrameKey};
9
10use crate::db::LaminarDB;
11use crate::error::DbError;
12use crate::operator_graph::VnodeStateLoad;
13
14const MANIFEST_READ_TIMEOUT: Duration = Duration::from_secs(10);
15
16#[derive(Default)]
18pub(crate) struct VnodeLoadBoard(parking_lot::Mutex<Option<Arc<VnodeLoadSample>>>);
19
20pub(crate) struct VnodeLoadSample {
21 taken_at: Instant,
22 operators: Vec<(Arc<str>, Vec<VnodeStateLoad>)>,
23}
24
25impl VnodeLoadBoard {
26 pub(crate) fn publish(&self, operators: Vec<(Arc<str>, Vec<VnodeStateLoad>)>) {
27 let sample = Arc::new(VnodeLoadSample {
28 taken_at: Instant::now(),
29 operators,
30 });
31 *self.0.lock() = Some(sample);
32 }
33
34 fn latest(&self) -> Option<Arc<VnodeLoadSample>> {
35 self.0.lock().clone()
36 }
37
38 #[cfg(test)]
39 pub(crate) fn latest_for_test(&self, operator: &str) -> Vec<VnodeStateLoad> {
40 self.latest()
41 .and_then(|sample| {
42 sample
43 .operators
44 .iter()
45 .find(|(name, _)| &**name == operator)
46 .map(|(_, vnodes)| vnodes.clone())
47 })
48 .unwrap_or_default()
49 }
50}
51
52#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize)]
54pub struct VnodeLoadReport {
55 pub live_sample_age_ms: Option<u64>,
58 pub committed: Option<VnodeLoadCheckpoint>,
60 pub operators: Vec<OperatorVnodeLoad>,
62}
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
66pub struct VnodeLoadCheckpoint {
67 pub checkpoint_id: u64,
69 pub epoch: u64,
71}
72
73#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
75pub struct OperatorVnodeLoad {
76 pub operator: String,
78 pub vnodes: Vec<VnodeLoadEntry>,
80}
81
82#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
84pub struct VnodeLoadEntry {
85 pub vnode: u32,
87 pub live: Option<LiveVnodeLoad>,
89 pub committed_frames: Vec<CommittedVnodeFrame>,
91}
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
95pub struct LiveVnodeLoad {
96 pub state_bytes: u64,
98 pub groups: u64,
100 pub applied_rows: u64,
104}
105
106#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
108pub struct CommittedVnodeFrame {
109 pub participant_id: u64,
111 pub frame_bytes: u64,
113}
114
115impl LaminarDB {
116 pub async fn cluster_vnode_load(&self) -> Result<VnodeLoadReport, DbError> {
125 let sample = self
126 .engine_metrics
127 .lock()
128 .as_ref()
129 .and_then(|metrics| metrics.flow.vnode_loads.latest());
130 let views = self.mv_store.read().committed_views();
131 let manifests = match views {
132 Some(views) => tokio::time::timeout(MANIFEST_READ_TIMEOUT, views.latest_manifests())
133 .await
134 .map_err(|_| {
135 DbError::Checkpoint(
136 "timed out reading the committed checkpoint's manifests".into(),
137 )
138 })??,
139 None => None,
140 };
141 Ok(build_report(sample.as_deref(), manifests.as_deref()))
142 }
143}
144
145type OperatorEntries = BTreeMap<String, BTreeMap<u32, VnodeLoadEntry>>;
146
147fn entry<'a>(
148 operators: &'a mut OperatorEntries,
149 operator: &str,
150 vnode: u32,
151) -> &'a mut VnodeLoadEntry {
152 operators
153 .entry(operator.to_owned())
154 .or_default()
155 .entry(vnode)
156 .or_insert_with(|| VnodeLoadEntry {
157 vnode,
158 live: None,
159 committed_frames: Vec::new(),
160 })
161}
162
163fn build_report(
164 sample: Option<&VnodeLoadSample>,
165 manifests: Option<&[CheckpointManifest]>,
166) -> VnodeLoadReport {
167 let mut operators = OperatorEntries::new();
168 for (operator, vnodes) in sample.map_or(&[][..], |sample| &sample.operators) {
169 for load in vnodes {
170 entry(&mut operators, operator, load.vnode).live = Some(LiveVnodeLoad {
171 state_bytes: u64::try_from(load.state_bytes).unwrap_or(u64::MAX),
172 groups: u64::try_from(load.groups).unwrap_or(u64::MAX),
173 applied_rows: load.applied_rows,
174 });
175 }
176 }
177 let manifests = manifests.unwrap_or_default();
178 for manifest in manifests {
179 for frame in &manifest.state_frames {
180 let StateFrameKey::Vnode { operator_id, vnode } = &frame.key else {
181 continue;
182 };
183 let operator = operator_id.strip_prefix("graph:").unwrap_or(operator_id);
184 entry(&mut operators, operator, u32::from(*vnode))
185 .committed_frames
186 .push(CommittedVnodeFrame {
187 participant_id: manifest.participant_id,
188 frame_bytes: frame.range.length,
189 });
190 }
191 }
192 VnodeLoadReport {
193 live_sample_age_ms: sample
194 .map(|sample| u64::try_from(sample.taken_at.elapsed().as_millis()).unwrap_or(u64::MAX)),
195 committed: manifests.first().map(|manifest| VnodeLoadCheckpoint {
196 checkpoint_id: manifest.checkpoint_id,
197 epoch: manifest.epoch,
198 }),
199 operators: operators
200 .into_iter()
201 .map(|(operator, vnodes)| OperatorVnodeLoad {
202 operator,
203 vnodes: vnodes.into_values().collect(),
204 })
205 .collect(),
206 }
207}
208
209#[cfg(test)]
210mod tests {
211 use laminar_core::checkpoint::{ByteRange, StateChunkId, StateFrame};
212
213 use super::*;
214
215 fn manifest(participant_id: u64, frames: &[(StateFrameKey, u64)]) -> CheckpointManifest {
216 let mut manifest = CheckpointManifest::new(7, 7);
217 manifest.participant_id = participant_id;
218 manifest.state_frames = frames
219 .iter()
220 .map(|(key, length)| StateFrame {
221 key: key.clone(),
222 chunk: StateChunkId {
223 participant_id,
224 checkpoint_id: 7,
225 },
226 range: ByteRange {
227 offset: 0,
228 length: *length,
229 },
230 sha256: String::new(),
231 })
232 .collect();
233 manifest
234 }
235
236 fn vnode_frame(operator: &str, vnode: u16) -> StateFrameKey {
237 StateFrameKey::Vnode {
238 operator_id: format!("graph:{operator}"),
239 vnode,
240 }
241 }
242
243 #[test]
244 fn report_joins_live_vnodes_with_every_participants_committed_frames() {
245 let sample = VnodeLoadSample {
246 taken_at: Instant::now(),
247 operators: vec![(
248 Arc::from("agg"),
249 vec![VnodeStateLoad {
250 vnode: 5,
251 state_bytes: 300,
252 groups: 2,
253 applied_rows: 9,
254 }],
255 )],
256 };
257 let manifests = [
258 manifest(
259 1,
260 &[
261 (
262 StateFrameKey::OperatorWhole {
263 operator_id: "graph:agg".into(),
264 },
265 11,
266 ),
267 (vnode_frame("agg", 5), 120),
268 (vnode_frame("view", 1), 40),
269 ],
270 ),
271 manifest(2, &[(vnode_frame("agg", 2), 80)]),
272 ];
273
274 let report = build_report(Some(&sample), Some(&manifests));
275
276 assert!(report.live_sample_age_ms.is_some());
277 assert_eq!(
278 report.committed,
279 Some(VnodeLoadCheckpoint {
280 checkpoint_id: 7,
281 epoch: 7,
282 })
283 );
284 let frame = |participant_id, frame_bytes| CommittedVnodeFrame {
285 participant_id,
286 frame_bytes,
287 };
288 assert_eq!(
289 report.operators,
290 vec![
291 OperatorVnodeLoad {
292 operator: "agg".into(),
293 vnodes: vec![
294 VnodeLoadEntry {
295 vnode: 2,
296 live: None,
297 committed_frames: vec![frame(2, 80)],
298 },
299 VnodeLoadEntry {
300 vnode: 5,
301 live: Some(LiveVnodeLoad {
302 state_bytes: 300,
303 groups: 2,
304 applied_rows: 9,
305 }),
306 committed_frames: vec![frame(1, 120)],
307 },
308 ],
309 },
310 OperatorVnodeLoad {
311 operator: "view".into(),
312 vnodes: vec![VnodeLoadEntry {
313 vnode: 1,
314 live: None,
315 committed_frames: vec![frame(1, 40)],
316 }],
317 },
318 ]
319 );
320 }
321
322 #[test]
323 fn report_is_empty_before_a_sample_or_a_commit() {
324 assert_eq!(build_report(None, None), VnodeLoadReport::default());
325 }
326}