metrale_bench/
plugin.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: [`Plugin`], the abstraction a [`crate::Benchmark`]
4//! specialises, with its [`PluginHandle`] and the [`TargetEndpoint`] it
5//! drives.
6//!
7//! [`Plugin::load`] receives the handle: status line, log, progress and glow
8//! as [`PluginEvent`]s sent over a channel, plus the artifact store, the
9//! endpoint and the cancellation flag.
10//!
11//! Owner: bench (plugin API).
12//! Invariants:
13//! - `TargetEndpoint::base_url` never ends in `/` when built by `new` or
14//!   `local`.
15//! - Sending an event never fails the plugin: a closed channel is ignored.
16
17use std::future::Future;
18use std::sync::Arc;
19use std::sync::atomic::{AtomicBool, Ordering};
20use std::sync::mpsc::Sender;
21
22use anyhow::{Result, bail};
23
24use crate::artifacts::ArtifactStore;
25use crate::metadata::PluginMetadata;
26use crate::result::{LogLevel, LogLine};
27
28/// 2026-09-26: The served endpoint a benchmark drives. In the TUI it starts
29/// as the local server, and the Benchmarks pane can point it at another URL
30/// or model.
31///
32/// `Default` is the empty endpoint: `host_port` rejects it, so
33/// `http::probe` fails on it.
34#[derive(Clone, Debug, Default, PartialEq, Eq)]
35pub struct TargetEndpoint {
36    /// 2026-09-26: Base URL with no trailing slash, e.g.
37    /// `http://127.0.0.1:8888`.
38    pub base_url: String,
39    /// 2026-09-26: The `model` field sent in requests.
40    pub model: String,
41    /// 2026-09-26: The serve overrides the caller started this target under.
42    /// The headless run record takes its `serve_overrides` from here, and the
43    /// concurrency sweep reads `ssm_cache_slots` from here. Empty when the
44    /// caller did not start the server, which means "not stated", never "no
45    /// overrides".
46    pub serve_overrides: std::collections::BTreeMap<String, String>,
47}
48
49impl TargetEndpoint {
50    pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Self {
51        let base = base_url.into();
52        Self {
53            base_url: base.trim_end_matches('/').to_string(),
54            model: model.into(),
55            serve_overrides: std::collections::BTreeMap::new(),
56        }
57    }
58
59    /// 2026-09-26: The same endpoint, with the overrides it was started under.
60    #[must_use]
61    pub fn with_serve_overrides(
62        mut self,
63        serve_overrides: std::collections::BTreeMap<String, String>,
64    ) -> Self {
65        self.serve_overrides = serve_overrides;
66        self
67    }
68
69    /// 2026-09-26: An override the server was started with, parsed as an
70    /// integer. `None` when it was not stated or does not parse.
71    pub fn serve_override_usize(&self, key: &str) -> Option<usize> {
72        self.serve_overrides.get(key)?.trim().parse().ok()
73    }
74
75    /// 2026-09-26: `http://127.0.0.1:<port>`.
76    pub fn local(port: u16, model: impl Into<String>) -> Self {
77        Self::new(format!("http://127.0.0.1:{port}"), model)
78    }
79
80    /// 2026-09-26: Is this endpoint served on this box?
81    ///
82    /// The energy meter records the local GPU rail only for a loopback
83    /// target; for any other host the local reading describes another
84    /// machine. An unparseable target is not local.
85    pub fn is_loopback(&self) -> bool {
86        let Ok((host, _)) = self.host_port() else {
87            return false;
88        };
89        host == "localhost"
90            || host
91                .trim_start_matches('[')
92                .trim_end_matches(']')
93                .parse::<std::net::IpAddr>()
94                .is_ok_and(|ip| ip.is_loopback())
95    }
96
97    /// 2026-09-26: Split into `(host, port)` for a raw TCP connect. Only
98    /// `http://` is accepted; the port defaults to 80 when absent.
99    pub fn host_port(&self) -> Result<(String, u16)> {
100        let rest = self.base_url.strip_prefix("http://").ok_or_else(|| {
101            anyhow::anyhow!("only http:// targets are supported: {}", self.base_url)
102        })?;
103        let authority = rest.split('/').next().unwrap_or(rest);
104        match authority.rsplit_once(':') {
105            Some((host, port)) => {
106                let port: u16 = port
107                    .parse()
108                    .map_err(|_| anyhow::anyhow!("bad port in {}", self.base_url))?;
109                Ok((host.to_string(), port))
110            }
111            None => Ok((authority.to_string(), 80)),
112        }
113    }
114}
115
116/// 2026-09-26: A message from a running plugin to the terminal.
117#[derive(Clone, Debug)]
118pub enum PluginEvent {
119    Log(LogLine),
120    /// 2026-09-26: Replace the one-line status.
121    Status(String),
122    Progress {
123        done: u64,
124        total: u64,
125    },
126    /// 2026-09-26: Turn the terminal's work-in-flight glow on or off.
127    Glow(bool),
128}
129
130/// 2026-09-26: The plugin's view of its host.
131#[derive(Clone)]
132pub struct PluginHandle {
133    /// 2026-09-26: Distinguishes this run from the other runs of the same
134    /// executor, whose counter assigns it.
135    ///
136    /// Cold TTFT and the serve-matrix probes fold it, with the process id,
137    /// into their prompt prefix tags (`benchmarks::unique_prefix_tag`), so
138    /// two runs do not share a cached prefix.
139    run_id: u64,
140    target: TargetEndpoint,
141    artifacts: ArtifactStore,
142    events: Sender<PluginEvent>,
143    cancel: Arc<AtomicBool>,
144}
145
146impl PluginHandle {
147    /// 2026-09-26: This run's id; see the field doc.
148    pub fn run_id(&self) -> u64 {
149        self.run_id
150    }
151
152    pub fn new(
153        run_id: u64,
154        target: TargetEndpoint,
155        artifacts: ArtifactStore,
156        events: Sender<PluginEvent>,
157        cancel: Arc<AtomicBool>,
158    ) -> Self {
159        Self {
160            run_id,
161            target,
162            artifacts,
163            events,
164            cancel,
165        }
166    }
167
168    pub fn target(&self) -> &TargetEndpoint {
169        &self.target
170    }
171
172    pub fn artifacts(&self) -> &ArtifactStore {
173        &self.artifacts
174    }
175
176    /// 2026-09-26: Emit an event. A closed channel means the receiver is
177    /// gone; that is not an error for the plugin.
178    fn emit(&self, event: PluginEvent) {
179        let _ = self.events.send(event);
180    }
181
182    pub fn log(&self, level: LogLevel, text: impl Into<String>) {
183        self.emit(PluginEvent::Log(LogLine {
184            level,
185            text: text.into(),
186        }));
187    }
188    pub fn info(&self, text: impl Into<String>) {
189        self.log(LogLevel::Info, text);
190    }
191    pub fn warn(&self, text: impl Into<String>) {
192        self.log(LogLevel::Warn, text);
193    }
194    pub fn status(&self, text: impl Into<String>) {
195        self.emit(PluginEvent::Status(text.into()));
196    }
197    pub fn progress(&self, done: u64, total: u64) {
198        self.emit(PluginEvent::Progress { done, total });
199    }
200    pub fn set_glow(&self, on: bool) {
201        self.emit(PluginEvent::Glow(on));
202    }
203
204    pub fn is_cancelled(&self) -> bool {
205        self.cancel.load(Ordering::Relaxed)
206    }
207
208    /// 2026-09-26: `Err("cancelled")` once the run is cancelled. Call it at
209    /// every await point a benchmark controls: otherwise cancellation takes
210    /// effect only when the executor checks the flag, after each frame.
211    pub fn check_cancelled(&self) -> Result<()> {
212        if self.is_cancelled() {
213            bail!("cancelled");
214        }
215        Ok(())
216    }
217}
218
219/// 2026-09-26: The plugin abstraction. Every implementor today is a
220/// benchmark.
221pub trait Plugin {
222    /// 2026-09-26: Who wrote this plugin, where it came from, and where to
223    /// report it.
224    fn metadata(&self) -> &'static PluginMetadata;
225
226    /// 2026-09-26: First step. Acquire resources, provision artifacts, verify
227    /// the host has what this plugin needs. An `Err` ends the run with a
228    /// failed `setup` frame, so its message must name what is missing.
229    fn load(&mut self, handle: PluginHandle) -> impl Future<Output = Result<()>> + Send;
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn base_url_loses_its_trailing_slash() {
238        let t = TargetEndpoint::new("http://box:8888/", "m");
239        assert_eq!(t.base_url, "http://box:8888");
240    }
241
242    #[test]
243    fn host_port_splits_ipv4_and_defaults_to_80() {
244        assert_eq!(
245            TargetEndpoint::local(8888, "m").host_port().unwrap(),
246            ("127.0.0.1".to_string(), 8888)
247        );
248        assert_eq!(
249            TargetEndpoint::new("http://dgx3", "m").host_port().unwrap(),
250            ("dgx3".to_string(), 80)
251        );
252        assert!(TargetEndpoint::new("https://x", "m").host_port().is_err());
253    }
254
255    #[test]
256    fn cancellation_is_visible_through_the_handle() {
257        let (tx, rx) = std::sync::mpsc::channel();
258        let cancel = Arc::new(AtomicBool::new(false));
259        let h = PluginHandle::new(
260            1,
261            TargetEndpoint::local(8888, "m"),
262            ArtifactStore::with_root("/tmp/metrale-test"),
263            tx,
264            cancel.clone(),
265        );
266        h.check_cancelled().unwrap();
267        h.status("warming up");
268        cancel.store(true, Ordering::Relaxed);
269        assert!(h.check_cancelled().is_err());
270        assert!(matches!(rx.try_recv(), Ok(PluginEvent::Status(s)) if s == "warming up"));
271    }
272}
metrale_bench/plugin.rs

metrale_bench/
plugin.rs
