perf: move sampling to a background thread and cache static driver facts

src/main.rs

@@ -2,13 +2,13 @@ mod net;
 mod netlink;
 mod nvlink;
 mod rdma;
+mod sampler;
 mod stat;
 mod trace;
 mod tui;
 mod xgmi;
 
 use std::io;
-use std::time::Instant;
 
 fn run_tui() -> io::Result<()> {
     crossterm::terminal::enable_raw_mode()?;
@@ -19,23 +19,27 @@ fn run_tui() -> io::Result<()> {
     let backend = ratatui::backend::CrosstermBackend::new(stdout);
     let mut terminal = ratatui::Terminal::new(backend)?;
     let mut app = tui::app::App::new();
-
-    let mut last_refresh = Instant::now() - app.refresh_interval;
+    let sampler = sampler::Sampler::spawn(app.refresh_interval);
 
     loop {
-        if last_refresh.elapsed() >= app.refresh_interval {
-            app.refresh_stats();
-            last_refresh = Instant::now();
+        if let Some(snap) = sampler.try_latest() {
+            app.apply_snapshot(snap);
+        }
+        if app.sampler_error.is_none() {
+            app.sampler_error = sampler.death_reason();
         }
 
         terminal.draw(|frame| tui::ui::draw(frame, &mut app))?;
         tui::events::handle_events(&mut app)?;
+        // `<`/`>` change app.refresh_interval; mirror it to the thread.
+        sampler.set_interval(app.refresh_interval);
 
         if app.should_quit {
             break;
         }
     }
 
+    sampler.stop();
     crossterm::terminal::disable_raw_mode()?;
     crossterm::execute!(
         terminal.backend_mut(),

src/nvlink.rs

@@ -5,7 +5,9 @@
 //! throughput is read via the raw `nvmlDeviceGetFieldValues` entry point
 //! because `nvml-wrapper` does not expose a safe helper for those field IDs.
 
+use std::collections::HashMap;
 use std::io;
+use std::sync::{LazyLock, Mutex, OnceLock};
 
 use nvml_wrapper::enum_wrappers::nv_link::{ErrorCounter, IntDeviceType};
 use nvml_wrapper::enums::device::SampleValue;
@@ -82,6 +84,49 @@ pub struct LinkSnapshot {
     pub recovery_error_count: Option<u64>,
 }
 
+/// Static per-(gpu, link) facts that cannot change while the driver is
+/// loaded. Cached once fully resolved; partial replies retry next poll.
+/// (`remote_pci_bdf` is optional: some remotes report no BDF at all.)
+#[derive(Clone)]
+struct LinkMeta {
+    version: Option<u32>,
+    remote_device_type: RemoteDeviceType,
+    remote_pci_bdf: Option<String>,
+    speed_gbps: Option<f64>,
+}
+
+fn link_meta(
+    nvml: &Nvml,
+    device: &nvml_wrapper::Device<'_>,
+    gpu: u32,
+    link_id: u32,
+    common_speed_gbps: Option<f64>,
+) -> LinkMeta {
+    static CACHE: LazyLock<Mutex<HashMap<(u32, u32), LinkMeta>>> =
+        LazyLock::new(|| Mutex::new(HashMap::new()));
+    let mut cache = CACHE.lock().unwrap_or_else(|e| e.into_inner());
+    if let Some(meta) = cache.get(&(gpu, link_id)) {
+        return meta.clone();
+    }
+    let nvlink = device.link_wrapper_for(link_id);
+    let meta = LinkMeta {
+        version: nvlink.version().ok(),
+        remote_device_type: read_remote_device_type(nvml, device, link_id),
+        remote_pci_bdf: nvlink.remote_pci_info().ok().map(|p| p.bus_id),
+        speed_gbps: read_link_speed_gbps(device, link_id).or(common_speed_gbps),
+    };
+    // Cache only fully-resolved replies: a partial answer during link
+    // training (e.g. version ok, speed still 0) must retry next poll, or
+    // the missing fields would be frozen for the process lifetime.
+    if meta.version.is_some()
+        && meta.remote_device_type != RemoteDeviceType::Unknown
+        && meta.speed_gbps.is_some()
+    {
+        cache.insert((gpu, link_id), meta.clone());
+    }
+    meta
+}
+
 /// Snapshot of every NVLink attached to one GPU.
 #[derive(Clone, Debug)]
 pub struct NvLinkSnapshot {
@@ -120,16 +165,34 @@ const LINK_SPEED_FIELD_IDS: [u32; 12] = [
     NVML_FI_DEV_NVLINK_SPEED_MBPS_L11,
 ];
 
+/// dlopen + nvmlInit once, kept for the process lifetime (nvmlShutdown is
+/// never called; process exit tears it down). A failed init retries on the
+/// next poll; the lock serializes init so racing callers can never build
+/// and drop a second Nvml instance.
+fn nvml() -> Option<&'static Nvml> {
+    static INSTANCE: OnceLock<Nvml> = OnceLock::new();
+    static INIT: Mutex<()> = Mutex::new(());
+    if let Some(n) = INSTANCE.get() {
+        return Some(n);
+    }
+    let _guard = INIT.lock().unwrap_or_else(|e| e.into_inner());
+    if let Some(n) = INSTANCE.get() {
+        return Some(n);
+    }
+    let n = Nvml::init().ok()?;
+    Some(INSTANCE.get_or_init(|| n))
+}
+
 /// Read NVLink statistics for every GPU in the system.
 ///
 /// If NVML cannot be initialised (e.g. no NVIDIA driver loaded), an empty
 /// vector is returned so the caller can keep running with no NVLink data.
 /// Per-GPU or per-link failures are skipped silently and the remaining data
 /// is still returned.
 pub fn read_all_nvlink_stats() -> io::Result<Vec<NvLinkSnapshot>> {
-    let nvml = match Nvml::init() {
-        Ok(n) => n,
-        Err(_) => return Ok(Vec::new()),
+    let nvml = match nvml() {
+        Some(n) => n,
+        None => return Ok(Vec::new()),
     };
 
     let device_count = match nvml.device_count() {
@@ -139,7 +202,7 @@ pub fn read_all_nvlink_stats() -> io::Result<Vec<NvLinkSnapshot>> {
 
     let mut snapshots = Vec::with_capacity(device_count as usize);
     for idx in 0..device_count {
-        if let Some(snap) = read_device_snapshot(&nvml, idx) {
+        if let Some(snap) = read_device_snapshot(nvml, idx) {
             snapshots.push(snap);
         }
     }
@@ -159,29 +222,33 @@ fn read_device_snapshot(nvml: &Nvml, idx: u32) -> Option<NvLinkSnapshot> {
     for link_id in 0..link_count {
         let nvlink = device.link_wrapper_for(link_id);
         let is_active = nvlink.is_active().unwrap_or(false);
-        let version = nvlink.version().ok();
-        let remote_device_type = read_remote_device_type(nvml, &device, link_id);
-        let remote_pci_bdf = nvlink.remote_pci_info().ok().map(|p| p.bus_id);
+        // Static facts are cached; inactive links skip them entirely
+        // (matching the old behavior of speed, extended to all metadata).
+        let meta = if is_active {
+            link_meta(nvml, &device, idx, link_id, common_speed_gbps)
+        } else {
+            LinkMeta {
+                version: None,
+                remote_device_type: RemoteDeviceType::Unknown,
+                remote_pci_bdf: None,
+                speed_gbps: None,
+            }
+        };
         let crc_error_count = nvlink.error_counter(ErrorCounter::DlCrcFlit).ok();
         let replay_error_count = nvlink.error_counter(ErrorCounter::DlReplay).ok();
         let recovery_error_count = nvlink.error_counter(ErrorCounter::DlRecovery).ok();
         let (tx_bytes, rx_bytes) = read_link_throughput(nvml, &device, link_id);
-        let speed_gbps = if is_active {
-            read_link_speed_gbps(&device, link_id).or(common_speed_gbps)
-        } else {
-            None
-        };
-        if let Some(s) = speed_gbps {
+        if let Some(s) = meta.speed_gbps {
             total_speed_gbps += s;
         }
 
         links.push(LinkSnapshot {
             link_id,
             is_active,
-            version,
-            speed_gbps,
-            remote_device_type,
-            remote_pci_bdf,
+            version: meta.version,
+            speed_gbps: meta.speed_gbps,
+            remote_device_type: meta.remote_device_type,
+            remote_pci_bdf: meta.remote_pci_bdf,
             tx_bytes,
             rx_bytes,
             crc_error_count,

src/sampler.rs

@@ -0,0 +1,167 @@
+//! Background sampling thread. Collects every subsystem into a `Snapshot`
+//! and hands it to the UI through a single-slot mailbox, so slow driver
+//! calls (NVML/amdsmi init, netlink) never block rendering or input.
+
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use crate::{net, nvlink, stat, xgmi};
+
+/// One subsystem's reading stamped right at its own read, so rate math
+/// never absorbs another subsystem's latency (e.g. a driver init earlier
+/// in the same pass).
+pub struct Sample<T> {
+    pub data: T,
+    pub taken_at: Instant,
+}
+
+fn sample<T>(read: impl FnOnce() -> std::io::Result<T>) -> Option<Sample<T>> {
+    let taken_at = Instant::now();
+    read().ok().map(|data| Sample { data, taken_at })
+}
+
+/// Everything one sampling pass produces. Raw counters only; delta/rate
+/// math stays on the UI thread where the previous snapshot lives.
+/// Each field is None when its read failed, so one failing subsystem
+/// (e.g. a kernel without rdma netlink) never blocks the others.
+pub struct Snapshot {
+    pub stats: Option<Sample<Vec<stat::PortStat>>>,
+    pub ifstats: Option<Sample<Vec<net::IfStats>>>,
+    pub nvlink: Option<Sample<Vec<nvlink::NvLinkSnapshot>>>,
+    pub xgmi: Option<Sample<Vec<xgmi::XgmiSnapshot>>>,
+    pub processes: Option<Vec<stat::ProcessRdmaInfo>>,
+    /// Pass start; used for the duplicate guard, staleness, and trace ts.
+    pub taken_at: Instant,
+}
+
+fn collect() -> Snapshot {
+    let taken_at = Instant::now();
+    let processes = stat::read_all_qps()
+        .ok()
+        .map(|qps| stat::aggregate_by_process(&qps));
+    Snapshot {
+        stats: sample(stat::read_all_stats),
+        ifstats: sample(net::read_all_ifstats),
+        nvlink: sample(nvlink::read_all_nvlink_stats),
+        xgmi: sample(xgmi::read_all_xgmi_stats),
+        processes,
+        taken_at,
+    }
+}
+
+/// State shared between the sampling thread and the UI-side `Sampler`.
+struct Shared {
+    /// Latest snapshot; the thread overwrites, the UI takes. A single slot
+    /// caps memory at one snapshot even if the UI stalls for hours.
+    slot: Mutex<Option<Snapshot>>,
+    interval_ms: AtomicU64,
+    stop: AtomicBool,
+    /// Panic message when the thread died; the UI surfaces it, since the
+    /// default panic output is lost inside the alternate screen.
+    died: Mutex<Option<String>>,
+}
+
+pub struct Sampler {
+    shared: Arc<Shared>,
+}
+
+impl Sampler {
+    /// Spawn the sampling thread. It samples immediately (the baseline),
+    /// then keeps sampling at the current interval until stopped.
+    pub fn spawn(interval: Duration) -> Self {
+        let shared = Arc::new(Shared {
+            slot: Mutex::new(None),
+            interval_ms: AtomicU64::new(interval_to_ms(interval)),
+            stop: AtomicBool::new(false),
+            died: Mutex::new(None),
+        });
+        let thread_shared = shared.clone();
+        std::thread::spawn(move || run(&thread_shared));
+        Self { shared }
+    }
+
+    /// Latest snapshot, if a new one arrived since the last call.
+    pub fn try_latest(&self) -> Option<Snapshot> {
+        self.shared
+            .slot
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .take()
+    }
+
+    /// The captured panic message when the thread died; None while alive.
+    pub fn death_reason(&self) -> Option<String> {
+        self.shared
+            .died
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .clone()
+    }
+
+    pub fn set_interval(&self, interval: Duration) {
+        self.shared
+            .interval_ms
+            .store(interval_to_ms(interval), Ordering::Relaxed);
+    }
+
+    /// Ask the thread to exit. Detach, never join: a thread stuck inside a
+    /// driver call must not block process exit.
+    pub fn stop(&self) {
+        self.shared.stop.store(true, Ordering::Relaxed);
+    }
+}
+
+impl Drop for Sampler {
+    // The mailbox has no disconnect signal (unlike a channel), so stopping
+    // on drop is what keeps the thread from sampling forever on error paths.
+    fn drop(&mut self) {
+        self.stop();
+    }
+}
+
+/// Interval as stored millis, clamped to 1ms: a zero value would make the
+/// sampling loop spin with no sleep at all.
+fn interval_to_ms(interval: Duration) -> u64 {
+    (interval.as_millis() as u64).max(1)
+}
+
+/// Render a `catch_unwind` payload (typically &str or String) for the UI.
+fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
+    if let Some(s) = payload.downcast_ref::<&str>() {
+        (*s).to_string()
+    } else if let Some(s) = payload.downcast_ref::<String>() {
+        s.clone()
+    } else {
+        "unknown panic".to_string()
+    }
+}
+
+fn run(shared: &Shared) {
+    loop {
+        if shared.stop.load(Ordering::Relaxed) {
+            return;
+        }
+        let snap = match std::panic::catch_unwind(collect) {
+            Ok(s) => s,
+            Err(payload) => {
+                let mut died = shared.died.lock().unwrap_or_else(|e| e.into_inner());
+                *died = Some(panic_message(payload));
+                return;
+            }
+        };
+        *shared.slot.lock().unwrap_or_else(|e| e.into_inner()) = Some(snap);
+        // Sleep in short slices so interval changes and stop apply quickly.
+        let started = Instant::now();
+        loop {
+            if shared.stop.load(Ordering::Relaxed) {
+                return;
+            }
+            let interval = Duration::from_millis(shared.interval_ms.load(Ordering::Relaxed));
+            if started.elapsed() >= interval {
+                break;
+            }
+            std::thread::sleep(Duration::from_millis(50));
+        }
+    }
+}

src/stat.rs

@@ -95,6 +95,9 @@ fn parse_port_stat(nlmsg: &NlMsg) -> Option<PortStat> {
 }
 
 /// Parse the port line rate from sysfs, e.g. "400 Gb/sec (4X NDR)" -> 400.0.
+/// Read fresh every poll on purpose: a down port reports an SDR placeholder
+/// and links can retrain at a different speed, so caching would freeze a
+/// wrong utilization denominator for the process lifetime.
 fn read_port_link_gbps(dev_name: &str, port: u32) -> Option<f64> {
     let path = format!("/sys/class/infiniband/{}/ports/{}/rate", dev_name, port);
     let raw = std::fs::read_to_string(path).ok()?;

src/trace.rs

@@ -38,8 +38,11 @@ impl Recorder {
     }
 
     /// Record one interval's metrics, timestamped relative to record start.
-    pub fn push(&mut self, ports: Vec<PortMetrics>) {
-        let ts_us = self.start.elapsed().as_micros() as u64;
+    /// `taken_at` is the sampling pass start (sampler-side), so trace spacing
+    /// is immune to UI queue and event-poll latency; a sample taken before
+    /// recording started saturates to ts 0.
+    pub fn push(&mut self, taken_at: Instant, ports: Vec<PortMetrics>) {
+        let ts_us = taken_at.saturating_duration_since(self.start).as_micros() as u64;
         self.samples.push(TraceSample { ts_us, ports });
     }