address race condition

Signed-off-by: chang-ning <spiderpower02@gmail.com>

Author: crazyguitar

src/nvlink.rs

@@ -167,13 +167,20 @@ const LINK_SPEED_FIELD_IDS: [u32; 12] = [
 
 /// 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 so a driver that loads after startup is still picked up.
+/// 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);
     }
-    Nvml::init().ok().map(|n| INSTANCE.get_or_init(|| 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.

src/sampler.rs

@@ -72,7 +72,7 @@ impl Sampler {
     pub fn spawn(interval: Duration) -> Self {
         let shared = Arc::new(Shared {
             slot: Mutex::new(None),
-            interval_ms: AtomicU64::new(interval.as_millis() as u64),
+            interval_ms: AtomicU64::new(interval_to_ms(interval)),
             stop: AtomicBool::new(false),
             died: Mutex::new(None),
         });
@@ -102,7 +102,7 @@ impl Sampler {
     pub fn set_interval(&self, interval: Duration) {
         self.shared
             .interval_ms
-            .store(interval.as_millis() as u64, Ordering::Relaxed);
+            .store(interval_to_ms(interval), Ordering::Relaxed);
     }
 
     /// Ask the thread to exit. Detach, never join: a thread stuck inside a
@@ -120,6 +120,12 @@ impl Drop for Sampler {
     }
 }
 
+/// 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>() {

src/trace.rs

@@ -38,9 +38,9 @@ impl Recorder {
     }
 
     /// Record one interval's metrics, timestamped relative to record start.
-    /// `taken_at` is when the counters were read (sampler-side), so trace
-    /// spacing is immune to UI queue and event-poll latency; a sample taken
-    /// before recording started saturates to ts 0.
+    /// `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 });