Skip to main content

laminar_db/
vnode_load.rs

1//! The cluster vnode-load report: how each keyed operator's state spreads over vnodes, live on
2//! this node and in every participant's frames at the newest committed checkpoint.
3
4use 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/// The per-vnode state of every operator that keeps it, as of the graph's last cold sample.
17#[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/// How keyed operator state spreads over vnodes, as one cluster node sees it.
53#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize)]
54pub struct VnodeLoadReport {
55    /// Milliseconds since this node last sampled its live state; `None` before the first
56    /// sample. A running pipeline samples every 32 cycles.
57    pub live_sample_age_ms: Option<u64>,
58    /// The committed checkpoint the frame sizes come from; `None` before the first commit.
59    pub committed: Option<VnodeLoadCheckpoint>,
60    /// Operators in name order, each with its vnodes in ascending order.
61    pub operators: Vec<OperatorVnodeLoad>,
62}
63
64/// The committed checkpoint a [`VnodeLoadReport`]'s frame sizes come from.
65#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
66pub struct VnodeLoadCheckpoint {
67    /// Checkpoint id.
68    pub checkpoint_id: u64,
69    /// Checkpoint epoch.
70    pub epoch: u64,
71}
72
73/// One operator's state per vnode.
74#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
75pub struct OperatorVnodeLoad {
76    /// Operator (graph node) name.
77    pub operator: String,
78    /// Vnodes with live state on this node or a committed frame, in ascending order.
79    pub vnodes: Vec<VnodeLoadEntry>,
80}
81
82/// One vnode of one operator.
83#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
84pub struct VnodeLoadEntry {
85    /// Vnode id.
86    pub vnode: u32,
87    /// This node's live state for the vnode; `None` when it holds none.
88    pub live: Option<LiveVnodeLoad>,
89    /// The vnode's frames at the committed checkpoint, one per participant that wrote one.
90    pub committed_frames: Vec<CommittedVnodeFrame>,
91}
92
93/// One vnode's state on this node at the last sample.
94#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
95pub struct LiveVnodeLoad {
96    /// Accounted state bytes, a lower-bound charge.
97    pub state_bytes: u64,
98    /// Groups; a window key counts once per open window or session.
99    pub groups: u64,
100    /// Input rows applied to the vnode since the operator was built on this node (start,
101    /// recovery or topology change). A row counts once per window it falls in; a global
102    /// aggregate counts none.
103    pub applied_rows: u64,
104}
105
106/// One vnode's state frame in a committed checkpoint.
107#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
108pub struct CommittedVnodeFrame {
109    /// Participant that wrote the frame.
110    pub participant_id: u64,
111    /// Frame length in bytes.
112    pub frame_bytes: u64,
113}
114
115impl LaminarDB {
116    /// How keyed operator state spreads over vnodes: this node's live state at the graph's last
117    /// metrics sample, and every participant's frame sizes at the newest committed checkpoint
118    /// this node knows. Live state needs [`Self::set_engine_metrics`]; frame sizes need
119    /// checkpointing.
120    ///
121    /// # Errors
122    ///
123    /// Fails when the committed checkpoint's manifests cannot be read within 10 seconds.
124    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}