Plugins: time the CPU budget with the thread's CPU time on Windows too

crates/tw-plugin/Cargo.toml

@@ -29,6 +29,10 @@ rand = { workspace = true }
 # 量线程的 CPU 时间（CLOCK_THREAD_CPUTIME_ID）
 libc = { workspace = true }
 
+[target.'cfg(windows)'.dependencies]
+# 量线程的 CPU 时间（GetThreadTimes）
+windows-sys = { workspace = true, features = ["Win32_Foundation", "Win32_System_Threading"] }
+
 [build-dependencies]
 # 构建时的编译器：Cranelift 把快照编成目标平台的机器码；all-arch 让交叉编译
 # （比如在 x64 上编 aarch64 的 Windows）也能编

crates/tw-plugin/src/cpu.rs

@@ -4,36 +4,141 @@
 //! 照走，而插件其实没在跑 —— 按墙上时间算，一个 5 ms 的钩子在编译大工程时也会被
 //! 判超时，请求跟着被拒。
 //!
-//! unix（Linux、macOS）用 `CLOCK_THREAD_CPUTIME_ID`。Windows 的线程时间按时钟中断
-//! 记账（15.6 ms 一格），对 20 ms 的预算太粗，那里退回墙上时间：调用都跑在专用
-//! 的线程池上，平时两者差不多，只在机器很忙时偏严。
+//! unix（Linux、macOS）用 `CLOCK_THREAD_CPUTIME_ID`，精确到纳秒。
+//!
+//! **Windows 的线程时间按时钟中断记账**（`GetThreadTimes`，15.6 ms 一格），对 20 ms 的
+//! 预算单独用太粗：一个只跑了 5 ms 的钩子，碰上一次中断就被记成一整格。所以那里取两个
+//! 数里小的那个：墙上时间，和线程时间加一格。机器空闲时两者差不多，墙上时间更小、更准，
+//! 和以前一样；机器忙时墙上时间照走，线程时间不走，被抢占的那段不再算在插件头上 ——
+//! 至多多算一格。
+//!
+//! 起点和量的那一刻要在**同一个线程**上（调用都在专用的线程池上，一段执行不换线程）。
+
+use std::time::{Duration, Instant};
+
+/// 一段执行的起点。[`Mark::elapsed`] 是从这里起这个线程用掉的 CPU 时间
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct Mark {
+    wall: Instant,
+    /// 这个线程到这一刻的 CPU 时间。量不到时是 None，退回墙上时间
+    thread: Option<Duration>,
+}
+
+impl Mark {
+    pub(crate) fn now() -> Self {
+        Self {
+            wall: Instant::now(),
+            thread: thread_time(),
+        }
+    }
 
-use std::time::Duration;
+    pub(crate) fn elapsed(&self) -> Duration {
+        let wall = self.wall.elapsed();
+        match (self.thread, thread_time()) {
+            (Some(from), Some(to)) => charge(wall, to.saturating_sub(from)),
+            _ => wall,
+        }
+    }
+}
+
+/// Windows 上线程时间的一格：时钟中断的间隔
+#[cfg_attr(not(windows), allow(dead_code))]
+const TICK: Duration = Duration::from_micros(15_625);
+
+/// 这段执行记多少：线程时间按格记账的平台上，取墙上时间和「线程时间加一格」里小的那个
+#[cfg_attr(not(windows), allow(dead_code))]
+fn coarse(wall: Duration, thread: Duration) -> Duration {
+    wall.min(thread.saturating_add(TICK))
+}
+
+#[cfg(windows)]
+fn charge(wall: Duration, thread: Duration) -> Duration {
+    coarse(wall, thread)
+}
+
+#[cfg(not(windows))]
+fn charge(_wall: Duration, thread: Duration) -> Duration {
+    thread
+}
 
-/// 从某个任意起点算起的「这个线程的 CPU 时间」。只拿来相减
 #[cfg(unix)]
-pub(crate) fn now() -> Duration {
+fn thread_time() -> Option<Duration> {
     let mut ts = libc::timespec {
         tv_sec: 0,
         tv_nsec: 0,
     };
     // SAFETY: 只写我们给的这个 timespec
     let rc = unsafe { libc::clock_gettime(libc::CLOCK_THREAD_CPUTIME_ID, &mut ts) };
-    if rc == 0 {
-        Duration::new(ts.tv_sec as u64, ts.tv_nsec as u32)
-    } else {
-        wall()
-    }
+    (rc == 0).then(|| Duration::new(ts.tv_sec as u64, ts.tv_nsec as u32))
+}
+
+#[cfg(windows)]
+fn thread_time() -> Option<Duration> {
+    use windows_sys::Win32::Foundation::FILETIME;
+    use windows_sys::Win32::System::Threading::{GetCurrentThread, GetThreadTimes};
+    let zero = FILETIME {
+        dwLowDateTime: 0,
+        dwHighDateTime: 0,
+    };
+    let (mut created, mut exited, mut kernel, mut user) = (zero, zero, zero, zero);
+    // SAFETY: GetCurrentThread 是不用关闭的伪句柄；只写我们给的这四个 FILETIME
+    let ok = unsafe {
+        GetThreadTimes(
+            GetCurrentThread(),
+            &mut created,
+            &mut exited,
+            &mut kernel,
+            &mut user,
+        )
+    };
+    let units = |t: FILETIME| (u64::from(t.dwHighDateTime) << 32) | u64::from(t.dwLowDateTime);
+    // 单位是 100 ns
+    (ok != 0).then(|| Duration::from_nanos((units(kernel) + units(user)).saturating_mul(100)))
 }
 
-#[cfg(not(unix))]
-pub(crate) fn now() -> Duration {
-    wall()
+#[cfg(not(any(unix, windows)))]
+fn thread_time() -> Option<Duration> {
+    None
 }
 
-fn wall() -> Duration {
-    use std::sync::OnceLock;
-    use std::time::Instant;
-    static ORIGIN: OnceLock<Instant> = OnceLock::new();
-    ORIGIN.get_or_init(Instant::now).elapsed()
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// 睡着的线程不花预算：被抢占、在等的时间不算在插件头上
+    #[test]
+    fn a_thread_that_sleeps_spends_almost_none_of_its_budget() {
+        let m = Mark::now();
+        std::thread::sleep(Duration::from_millis(80));
+        let spent = m.elapsed();
+        // unix 上几乎是 0；Windows 上一两格（醒来那一刻碰上时钟中断会多记一格），都远小于
+        // 睡掉的 80 ms
+        assert!(spent < Duration::from_millis(40), "{spent:?}");
+    }
+
+    /// 按格记账时，记的不超过墙上时间，也不超过线程时间加一格
+    #[test]
+    fn a_coarse_thread_clock_is_capped_by_the_wall_and_by_one_tick() {
+        let ms = Duration::from_millis;
+        // 空闲：墙上时间更小，用它
+        assert_eq!(coarse(ms(5), ms(0)), ms(5));
+        assert_eq!(coarse(ms(12), ms(15)), ms(12));
+        // 忙：墙上走了 200 ms，线程只跑了 2 格
+        assert_eq!(coarse(ms(200), Duration::from_micros(31_250)), TICK * 3);
+        // 失控的循环：两者一起涨，照样超过预算
+        assert!(coarse(ms(40), ms(37)) >= ms(20));
+    }
+
+    /// 真在跑的线程照样花预算
+    #[test]
+    fn a_busy_thread_spends_its_budget() {
+        let m = Mark::now();
+        let until = Instant::now() + Duration::from_millis(60);
+        let mut x = 0u64;
+        while Instant::now() < until {
+            x = std::hint::black_box(x.wrapping_add(1));
+        }
+        // 机器再忙，60 ms 的忙循环也至少拿到几毫秒的 CPU
+        assert!(m.elapsed() > Duration::ZERO, "{x}");
+    }
 }

crates/tw-plugin/src/lib.rs

@@ -608,14 +608,14 @@ impl Plugin {
         let view_json = to_json(&view);
         let ctx_json = to_json(&ctx);
         let _running = self.inner.rt.ticker.enter();
-        let start = cpu::now();
+        let start = cpu::Mark::now();
         let (mut sb, mut logs) = match self.setup(limits.request_memory) {
             Ok(x) => x,
             Err((e, logs)) => {
                 return Invocation {
                     result: Err(e),
                     logs,
-                    cpu: cpu::now().saturating_sub(start),
+                    cpu: start.elapsed(),
                 };
             }
         };

crates/tw-plugin/src/sandbox.rs

@@ -50,7 +50,7 @@ pub(crate) struct HostState {
     max_log_line: usize,
     log_overflow: bool,
     /// 这一段执行的起点（线程 CPU 时间）和预算
-    start: Duration,
+    start: cpu::Mark,
     budget: Duration,
     cpu_hit: bool,
 }
@@ -96,7 +96,7 @@ impl ResourceLimiter for HostState {
 
 impl HostState {
     fn exceeded(&self) -> bool {
-        cpu::now().saturating_sub(self.start) >= self.budget
+        self.start.elapsed() >= self.budget
     }
 }
 
@@ -235,7 +235,7 @@ impl Sandbox {
         memory_cap: usize,
         budget: Duration,
     ) -> Result<Sandbox, RunError> {
-        let start = cpu::now();
+        let start = cpu::Mark::now();
         let mut store = Store::new(
             pre.module().engine(),
             HostState {
@@ -283,15 +283,15 @@ impl Sandbox {
     /// 下一段执行的 CPU 预算从现在算起
     pub(crate) fn arm(&mut self, budget: Duration) {
         let s = self.store.data_mut();
-        s.start = cpu::now();
+        s.start = cpu::Mark::now();
         s.budget = budget;
         s.cpu_hit = false;
         self.store.set_epoch_deadline(1);
     }
 
     /// 这一段执行（上次 `arm` 或建实例以来）用掉的 CPU 时间
     pub(crate) fn elapsed(&self) -> Duration {
-        cpu::now().saturating_sub(self.store.data().start)
+        self.store.data().start.elapsed()
     }
 
     pub(crate) fn take_logs(&mut self) -> Vec<LogLine> {