clankers_testing/
context.rs

1use std::path::PathBuf;
2use std::sync::Mutex;
3use std::time::Duration;
4
5use clankers_core::{LatencyStats, RobotResult};
6use clankers_data::{Replay, ReplayResult};
7
8/// Aggregated per-frame inference accounting across a replay run.
9#[derive(Debug, Clone, Default, PartialEq)]
10pub struct AggregatedInferenceStats {
11    pub frame_count: u32,
12    pub total_copies: usize,
13    pub total_allocations: usize,
14    pub total_bytes_copied: usize,
15    pub latency: LatencyStats,
16}
17
18impl AggregatedInferenceStats {
19    pub fn record_frame(
20        &mut self,
21        latency: Duration,
22        copies: usize,
23        allocations: usize,
24        bytes_copied: usize,
25    ) {
26        self.frame_count += 1;
27        self.total_copies += copies;
28        self.total_allocations += allocations;
29        self.total_bytes_copied += bytes_copied;
30        self.latency.record(latency);
31    }
32}
33
34/// Context for replay-based tests.
35pub struct ReplayContext {
36    pub mcap_path: PathBuf,
37    pub deterministic: bool,
38    /// Per-frame inference latencies recorded during a replay, aggregated by
39    /// [`latency`](ReplayContext::latency). Interior-mutable so the replay
40    /// handler (which borrows the context by shared reference) can feed it.
41    frame_latency: Mutex<LatencyStats>,
42    /// Per-frame copy/allocation accounting from inference runs.
43    frame_inference: Mutex<AggregatedInferenceStats>,
44}
45
46#[derive(Debug, Clone)]
47pub struct ReplayTestResult {
48    pub replay: ReplayResult,
49    pub panics: u32,
50    pub topics_seen: Vec<String>,
51}
52
53impl ReplayContext {
54    pub fn new(mcap_path: impl Into<PathBuf>) -> Self {
55        Self {
56            mcap_path: mcap_path.into(),
57            deterministic: true,
58            frame_latency: Mutex::new(LatencyStats::new()),
59            frame_inference: Mutex::new(AggregatedInferenceStats::default()),
60        }
61    }
62
63    /// Record one frame's inference latency (e.g. `InferenceStats::latency` from
64    /// `clankers_ml`) so [`latency`](Self::latency) can report an aggregate over
65    /// the whole replay.
66    pub fn record_frame_latency(&self, latency: Duration) {
67        self.frame_latency.lock().unwrap().record(latency);
68    }
69
70    /// Record per-frame inference accounting (e.g. fields from
71    /// `clankers_ml::InferenceStats`) alongside latency.
72    pub fn record_frame_inference(
73        &self,
74        latency: Duration,
75        copies: usize,
76        allocations: usize,
77        bytes_copied: usize,
78    ) {
79        self.frame_inference.lock().unwrap().record_frame(
80            latency,
81            copies,
82            allocations,
83            bytes_copied,
84        );
85    }
86
87    pub async fn run_replay<F, Fut>(&self, mut handler: F) -> RobotResult<ReplayTestResult>
88    where
89        F: FnMut(clankers_data::McapRecord) -> Fut,
90        Fut: std::future::Future<Output = RobotResult<()>>,
91    {
92        let replay = Replay::from_mcap(&self.mcap_path)?;
93        let mut topics_seen = std::collections::HashSet::new();
94
95        let result = replay
96            .run(|msg| {
97                topics_seen.insert(msg.topic.clone());
98                handler(msg)
99            })
100            .await?;
101
102        let mut topics: Vec<_> = topics_seen.into_iter().collect();
103        topics.sort();
104
105        Ok(ReplayTestResult {
106            replay: result,
107            panics: 0,
108            topics_seen: topics,
109        })
110    }
111
112    /// The aggregate of every latency passed to
113    /// [`record_frame_latency`](Self::record_frame_latency) during the replay.
114    ///
115    /// For end-to-end replay timing (per message, measured by the replay engine),
116    /// use `result.replay.latency` from [`run_replay`](Self::run_replay) instead.
117    pub fn latency(&self) -> LatencyStats {
118        self.frame_latency.lock().unwrap().clone()
119    }
120
121    /// Aggregated copy/allocation accounting from
122    /// [`record_frame_inference`](Self::record_frame_inference).
123    pub fn inference_stats(&self) -> AggregatedInferenceStats {
124        self.frame_inference.lock().unwrap().clone()
125    }
126}
clankers_testing/context.rs

clankers_testing/
context.rs
