ssh_stamp/
mem_probe.rs

1// SPDX-FileCopyrightText: 2026 Marko Malenic <mmalenic1@gmail.com>
2//
3// SPDX-License-Identifier: GPL-3.0-or-later
4
5//! The benchmarking instrumentation, including checkpoints and replay. Timing is
6//! not related to the platform in any way, so it doesn't need to live inside the
7//! binrary crate.
8//!
9
10#[cfg(feature = "mem-probe")]
11use embassy_time::Instant;
12#[cfg(feature = "mem-probe")]
13use portable_atomic::{AtomicU64, Ordering};
14
15/// Emits a structured benchmark line.
16#[macro_export]
17macro_rules! bench_emit {
18    ($($arg:tt)*) => {
19        ::log::info!("@BENCH {}", ::core::format_args!($($arg)*))
20    };
21}
22
23/// One point in the `ssh_stamp` flow.
24#[derive(Clone, Copy)]
25pub enum Checkpoint {
26    /// Boot reached.
27    Boot,
28    /// Core peripherals ready.
29    PeripheralsReady,
30    /// Wifi is up.
31    WifiUp,
32    /// TCP listener is ready.
33    TcpListening,
34    /// A TCP connection was accepted.
35    TcpAccept,
36    /// SSH key exchange finished.
37    KexComplete,
38    /// Client authenticated.
39    AuthSuccess,
40    /// A session channel was opened.
41    ChannelOpen,
42}
43
44#[cfg(feature = "mem-probe")]
45impl Checkpoint {
46    /// The name emitted on the `@BENCH checkpoint=` line. The host uses these to
47    /// compute the benchmark calculations.
48    const fn name(self) -> &'static str {
49        match self {
50            Self::Boot => "bench_boot",
51            Self::PeripheralsReady => "bench_peripherals_ready",
52            Self::WifiUp => "bench_wifi_up",
53            Self::TcpListening => "bench_tcp_listening",
54            Self::TcpAccept => "bench_tcp_accept",
55            Self::KexComplete => "bench_kex_complete",
56            Self::AuthSuccess => "bench_auth_success",
57            Self::ChannelOpen => "bench_channel_open",
58        }
59    }
60}
61
62/// Every [`Checkpoint`] in order, used for replaying checkpoints to the host.
63#[cfg(feature = "mem-probe")]
64const ALL: [Checkpoint; 8] = [
65    Checkpoint::Boot,
66    Checkpoint::PeripheralsReady,
67    Checkpoint::WifiUp,
68    Checkpoint::TcpListening,
69    Checkpoint::TcpAccept,
70    Checkpoint::KexComplete,
71    Checkpoint::AuthSuccess,
72    Checkpoint::ChannelOpen,
73];
74
75/// Timestamps indexed by [`Checkpoint`], used for replaying checkpoints to the host.
76/// A timestamp of 0 indicates that nothing has been emitted yet.
77#[cfg(feature = "mem-probe")]
78static T_US: [AtomicU64; ALL.len()] = [const { AtomicU64::new(0) }; ALL.len()];
79
80/// Logs the checkpoint.
81#[cfg(feature = "mem-probe")]
82pub fn checkpoint(c: Checkpoint) {
83    let t_us = Instant::now().as_micros();
84    bench_emit!("checkpoint={} t_us={t_us}", c.name());
85    T_US[c as usize].store(t_us, Ordering::Relaxed);
86}
87
88#[cfg(not(feature = "mem-probe"))]
89pub fn checkpoint(_c: Checkpoint) {}
90
91/// Replays all checkpoints that have been recorded so far. This is needed because
92/// the host may miss a logged line when resetting/flashing the device because the
93/// console is not attached yet.
94#[cfg(feature = "mem-probe")]
95pub fn replay_checkpoints() {
96    for c in ALL {
97        // Zero is the "never fired" marker, not a timestamp.
98        if let t_us @ 1.. = T_US[c as usize].load(Ordering::Relaxed) {
99            crate::bench_emit!("checkpoint={} t_us={t_us}", c.name());
100        }
101    }
102}
103
104#[cfg(not(feature = "mem-probe"))]
105pub fn replay_checkpoints() {}
106
107#[cfg(feature = "mem-probe")]
108static KEX_START_TICKS: AtomicU64 = AtomicU64::new(0);
109
110/// Records the current instant as the KEX start point.
111#[cfg(feature = "mem-probe")]
112pub fn mark_kex_start() {
113    KEX_START_TICKS.store(Instant::now().as_ticks(), Ordering::Relaxed);
114}
115
116#[cfg(not(feature = "mem-probe"))]
117pub fn mark_kex_start() {}
118
119/// Logs the elapsed time since `mark_kex_start`.
120#[cfg(feature = "mem-probe")]
121pub fn log_kex_elapsed(label: &str) {
122    let start_ticks = KEX_START_TICKS.swap(0, Ordering::Relaxed);
123    if start_ticks == 0 {
124        return;
125    }
126    let elapsed = Instant::from_ticks(start_ticks).elapsed();
127    bench_emit!("kex={label} elapsed_us={}", elapsed.as_micros());
128}
129
130#[cfg(not(feature = "mem-probe"))]
131pub fn log_kex_elapsed(_label: &str) {}
ssh_stamp/mem_probe.rs

ssh_stamp/
mem_probe.rs
