dynamo_bench/kv_router/common/
trace_gen.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4#[derive(Clone, Debug, Default)]
5pub struct WorkerTimelines<T> {
6    entries: Vec<Vec<T>>,
7}
8
9impl<T> WorkerTimelines<T> {
10    pub fn new(entries: Vec<Vec<T>>) -> Self {
11        Self { entries }
12    }
13
14    pub fn len(&self) -> usize {
15        self.entries.len()
16    }
17
18    pub fn is_empty(&self) -> bool {
19        self.entries.is_empty()
20    }
21
22    pub fn iter(&self) -> std::slice::Iter<'_, Vec<T>> {
23        self.entries.iter()
24    }
25
26    pub fn into_inner(self) -> Vec<Vec<T>> {
27        self.entries
28    }
29
30    pub fn into_rescaled_from_first<GetTimestamp, WithTimestamp>(
31        self,
32        benchmark_duration_ms: u64,
33        timestamp_of: GetTimestamp,
34        with_timestamp: WithTimestamp,
35    ) -> Self
36    where
37        GetTimestamp: Fn(&T) -> u64 + Copy,
38        WithTimestamp: Fn(T, u64) -> T + Copy,
39    {
40        let target_us = u128::from(benchmark_duration_ms) * 1000;
41        let entries = self
42            .entries
43            .into_iter()
44            .map(|worker_trace| {
45                let Some(first_timestamp_us) = worker_trace.first().map(timestamp_of) else {
46                    return Vec::new();
47                };
48                let span_us = worker_trace
49                    .last()
50                    .map(timestamp_of)
51                    .unwrap_or(first_timestamp_us)
52                    .saturating_sub(first_timestamp_us)
53                    .max(1);
54
55                worker_trace
56                    .into_iter()
57                    .map(|entry| {
58                        let relative_us = timestamp_of(&entry).saturating_sub(first_timestamp_us);
59                        let scaled_timestamp =
60                            u128::from(relative_us) * target_us / u128::from(span_us);
61                        with_timestamp(entry, scaled_timestamp.min(u128::from(u64::MAX)) as u64)
62                    })
63                    .collect()
64            })
65            .collect();
66
67        Self { entries }
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::WorkerTimelines;
74
75    #[derive(Clone, Debug, PartialEq, Eq)]
76    struct Entry {
77        timestamp_us: u64,
78        label: &'static str,
79    }
80
81    fn ts(entry: &Entry) -> u64 {
82        entry.timestamp_us
83    }
84
85    #[test]
86    fn worker_timelines_rescale_from_each_workers_first_entry() {
87        let timelines = WorkerTimelines::new(vec![
88            vec![
89                Entry {
90                    timestamp_us: 10,
91                    label: "a",
92                },
93                Entry {
94                    timestamp_us: 20,
95                    label: "b",
96                },
97            ],
98            vec![
99                Entry {
100                    timestamp_us: 1_000,
101                    label: "c",
102                },
103                Entry {
104                    timestamp_us: 1_010,
105                    label: "d",
106                },
107            ],
108        ]);
109
110        let scaled = timelines.into_rescaled_from_first(1_000, ts, |entry, timestamp_us| Entry {
111            timestamp_us,
112            label: entry.label,
113        });
114        let scaled = scaled.into_inner();
115
116        assert_eq!(scaled[0][0].timestamp_us, 0);
117        assert_eq!(scaled[0][1].timestamp_us, 1_000_000);
118        assert_eq!(scaled[1][0].timestamp_us, 0);
119        assert_eq!(scaled[1][1].timestamp_us, 1_000_000);
120    }
121}
dynamo_bench/kv_router/common/trace_gen.rs

dynamo_bench/kv_router/common/
trace_gen.rs
