fix: make cpu timer compatible on multi-thread

Author: nyannyacha

crates/base/src/runtime/mod.rs

@@ -2076,15 +2076,43 @@ fn get_cpu_metrics_guard<'l>(
     return scopeguard::guard((), Box::new(|_| {}));
   };
 
+  #[derive(Clone)]
+  struct CurrentCPUTimer {
+    thread_id: std::thread::ThreadId,
+    timer: CPUTimer,
+  }
+
   let current_thread_id = std::thread::current().id();
   let send_cpu_metrics_fn = move |metric: CPUUsageMetrics| {
     let _ = cpu_usage_metrics_tx.send(metric);
   };
 
-  send_cpu_metrics_fn(CPUUsageMetrics::Enter(
-    current_thread_id,
-    CPUTimer::get_or_init().unwrap(),
-  ));
+  let mut state = op_state.borrow_mut();
+  let cpu_timer = if state.has::<CurrentCPUTimer>() {
+    let current_cpu_timer = state.borrow::<CurrentCPUTimer>();
+    if current_cpu_timer.thread_id != current_thread_id {
+      state.take::<CurrentCPUTimer>();
+      None
+    } else {
+      Some(current_cpu_timer.timer.clone())
+    }
+  } else {
+    None
+  };
+  let cpu_timer = if let Some(timer) = cpu_timer {
+    timer
+  } else {
+    let cpu_timer = CurrentCPUTimer {
+      thread_id: current_thread_id,
+      timer: CPUTimer::new().unwrap(),
+    };
+
+    state.put(cpu_timer.clone());
+    cpu_timer.timer
+  };
+
+  drop(state);
+  send_cpu_metrics_fn(CPUUsageMetrics::Enter(current_thread_id, cpu_timer));
 
   let current_cpu_time_ns = get_current_cpu_time_ns().unwrap();
 

crates/cpu_timer/src/lib.rs

@@ -80,26 +80,17 @@ pub struct CPUTimer {}
 
 impl CPUTimer {
   #[cfg(not(target_os = "linux"))]
-  pub fn get_or_init(_: u64) -> Result<Self, Error> {
+  pub fn new(_: u64) -> Result<Self, Error> {
     log::error!("CPU timer: not enabled (need Linux)");
     Ok(Self {})
   }
 
   #[cfg(target_os = "linux")]
-  pub fn get_or_init() -> Result<Self, Error> {
-    use std::cell::RefCell;
+  pub fn new() -> Result<Self, Error> {
     use std::sync::atomic::Ordering;
 
     use linux::*;
 
-    thread_local! {
-        static TIMER: RefCell<Option<CPUTimer>> = RefCell::new(Option::default());
-    }
-
-    if let Some(timer) = TIMER.with_borrow(|v| v.clone()) {
-      return Ok(timer);
-    }
-
     let id = TIMER_COUNTER.fetch_add(1, Ordering::SeqCst);
     let mut tid = TimerId(std::ptr::null_mut());
     let mut sigev: libc::sigevent = unsafe { std::mem::zeroed() };
@@ -135,10 +126,6 @@ impl CPUTimer {
       .send(SignalMsg::Add((id, this.clone())))
       .unwrap();
 
-    TIMER.with_borrow_mut(|v| {
-      assert!(v.replace(this.clone()).is_none());
-    });
-
     Ok(this)
   }
 
@@ -159,13 +146,9 @@ impl CPUTimer {
   pub fn reset(&self, initial_expiry: u64, interval: u64) -> Result<(), Error> {
     use anyhow::Context;
     use linux::*;
-    use log::error;
 
-    error!("try get lock: {:?}", std::thread::current());
     let timer = self.timer.try_lock().context("failed to get the lock")?;
 
-    error!("get lock: {:?}", std::thread::current());
-
     let initial_expiry_secs = initial_expiry / 1000;
     let initial_expiry_msecs = initial_expiry % 1000;
     let interval_secs = interval / 1000;
@@ -182,12 +165,9 @@ impl CPUTimer {
       libc::timer_settime(timer.tid.0, 0, &tmspec, std::ptr::null_mut())
     } < 0
     {
-      error!("leave lock: {:?}", std::thread::current());
       bail!(std::io::Error::last_os_error())
     }
 
-    error!("leave lock: {:?}", std::thread::current());
-
     Ok(())
   }
 
@@ -264,22 +244,22 @@ fn register_sigalrm() {
             Some(msg) = sig_msg_rx.recv() => {
               match msg {
                 SignalMsg::Alarm(ref timer_id) => {
-                    if let Some(cpu_timer) = registry.get(timer_id) {
-                      if let Some(tx) = (*cpu_timer.cpu_alarm_val.cpu_alarms_tx.lock().await).clone() {
-                        if tx.send(()).is_err() {
-                            debug!("failed to send cpu alarm to the provided channel");
-                        }
-                    }
-                    } else {
-                      // NOTE: Unix signals are being delivered asynchronously,
-                      // and there are no guarantees to cancel the signal after
-                      // a timer has been deleted, and after a signal is
-                      // received, there may no longer be a target to accept it.
-                      error!(
-                        "can't find the cpu alarm signal matched with the received timer id: {}",
-                        *timer_id
-                      );
+                  if let Some(cpu_timer) = registry.get(timer_id) {
+                    if let Some(tx) = (*cpu_timer.cpu_alarm_val.cpu_alarms_tx.lock().await).clone() {
+                      if tx.send(()).is_err() {
+                          debug!("failed to send cpu alarm to the provided channel");
+                      }
                     }
+                  } else {
+                    // NOTE: Unix signals are being delivered asynchronously,
+                    // and there are no guarantees to cancel the signal after
+                    // a timer has been deleted, and after a signal is
+                    // received, there may no longer be a target to accept it.
+                    error!(
+                      "can't find the cpu alarm signal matched with the received timer id: {}",
+                      *timer_id
+                    );
+                  }
                 }
 
                 SignalMsg::Add((timer_id, cpu_timer)) => {