From 195b29261b47bc0e9dedd6f24a21f0563554e1e4 Mon Sep 17 00:00:00 2001 From: KiriAky 107 Date: Wed, 9 Sep 2026 05:37:37 +0800 Subject: [PATCH] =?UTF-8?q?feat(extensions):=20=E6=94=B6=E5=88=B0=E5=8E=9F?= =?UTF-8?q?=E7=94=9F=20CPU=20=E5=8E=8B=E5=8A=9B=E9=80=9A=E7=9F=A5=E6=97=B6?= =?UTF-8?q?=E7=BB=88=E6=AD=A2=E4=BD=9C=E4=B8=9A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- frontend/src-tauri/src/extension_job.rs | 251 +++++++++++++++++++- frontend/src-tauri/src/extension_process.rs | 3 + 2 files changed, 251 insertions(+), 3 deletions(-) diff --git a/frontend/src-tauri/src/extension_job.rs b/frontend/src-tauri/src/extension_job.rs index 989e4ec..bca2936 100644 --- a/frontend/src-tauri/src/extension_job.rs +++ b/frontend/src-tauri/src/extension_job.rs @@ -11,10 +11,14 @@ use windows_sys::Win32::System::{ pub struct Job { handle: OwnedHandle, + state: std::sync::Arc, + _cpu: Option, } impl Job { pub(crate) fn clone_for_deadline(&self) -> Result { Ok(Self { + state: std::sync::Arc::clone(&self.state), + _cpu: None, handle: self .handle .try_clone() @@ -29,8 +33,10 @@ impl Job { if raw.is_null() { return Err(HostError::new("EXTENSION_RESOURCE_UNAVAILABLE")); } - let job = Self { + let mut job = Self { handle: unsafe { OwnedHandle::from_raw_handle(raw) }, + state: std::sync::Arc::new(std::sync::atomic::AtomicU8::new(0)), + _cpu: None, }; let limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION { BasicLimitInformation: JOBOBJECT_BASIC_LIMIT_INFORMATION { @@ -49,14 +55,26 @@ impl Job { return Err(HostError::new("EXTENSION_RESOURCE_UNAVAILABLE")); } let cpu = JOBOBJECT_CPU_RATE_CONTROL_INFORMATION { - ControlFlags: JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP, + ControlFlags: JOB_OBJECT_CPU_RATE_CONTROL_ENABLE + | JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP + | JOB_OBJECT_CPU_RATE_CONTROL_NOTIFY, Anonymous: JOBOBJECT_CPU_RATE_CONTROL_INFORMATION_0 { CpuRate: (10000 / processors).max(1), }, }; job.set(JobObjectCpuRateControlInformation, &cpu)?; + job._cpu = Some(CpuMonitor::arm(&job)?); Ok(job) } + pub fn check_resources(&self) -> Result<()> { + use std::sync::atomic::Ordering; + match self.state.load(Ordering::Acquire) { + 0 => Ok(()), + 1 => Err(HostError::new("EXTENSION_RESOURCE_CPU_EXCEEDED")), + 3 => Err(HostError::new("EXTENSION_RESOURCE_TERMINATE_FAILED")), + _ => Err(HostError::new("EXTENSION_RESOURCE_MONITOR_FAILED")), + } + } fn set(&self, class: JOBOBJECTINFOCLASS, value: &T) -> Result<()> { if unsafe { SetInformationJobObject( @@ -77,6 +95,7 @@ impl Job { /// Caller must own an unresumed CREATE_SUSPENDED process and terminate it on /// any error. Resume only after all AppContainer/handle/permission checks pass. pub unsafe fn assign_suspended(&self, process: BorrowedHandle<'_>) -> Result<()> { + self.check_resources()?; if unsafe { AssignProcessToJobObject(self.handle.as_raw_handle(), process.as_raw_handle()) } == 0 { @@ -108,6 +127,122 @@ impl Job { } } +/// The Windows notification uses a ten-second window and ToleranceHigh (60% +/// over budget). This is not a measurement of ten uninterrupted busy seconds. +/// Only the original Job owns this monitor; observation/deadline clones do not. +const JOB_NOTIFICATION_LIMIT: u32 = 11; // JOB_OBJECT_MSG_NOTIFICATION_LIMIT (Windows SDK) +struct CpuMonitor { + stop: std::sync::Arc, + worker: Option>, +} +impl CpuMonitor { + fn arm(job: &Job) -> Result { + use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }; + use windows_sys::Win32::{Foundation::INVALID_HANDLE_VALUE, System::IO::*}; + let raw = + unsafe { CreateIoCompletionPort(INVALID_HANDLE_VALUE, std::ptr::null_mut(), 0, 1) }; + if raw.is_null() { + return Err(HostError::new("EXTENSION_RESOURCE_MONITOR_UNAVAILABLE")); + } + let port = unsafe { OwnedHandle::from_raw_handle(raw) }; + job.set( + JobObjectAssociateCompletionPortInformation, + &JOBOBJECT_ASSOCIATE_COMPLETION_PORT { + CompletionKey: std::ptr::dangling_mut::().cast(), + CompletionPort: port.as_raw_handle(), + }, + )?; + job.set( + JobObjectNotificationLimitInformation, + &JOBOBJECT_NOTIFICATION_LIMIT_INFORMATION { + LimitFlags: JOB_OBJECT_LIMIT_RATE_CONTROL, + RateControlTolerance: ToleranceHigh, + RateControlToleranceInterval: ToleranceIntervalShort, + ..Default::default() + }, + )?; + let owned_job = job + .handle + .try_clone() + .map_err(|_| HostError::new("EXTENSION_RESOURCE_MONITOR_UNAVAILABLE"))?; + let state = Arc::clone(&job.state); + let stop = Arc::new(AtomicBool::new(false)); + let thread_stop = Arc::clone(&stop); + let worker = std::thread::Builder::new() + .name("extension-cpu".into()) + .spawn(move || { + // All exits, including a caught panic or completion-port failure, + // terminate the tree while this worker still owns a Job handle. + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + while !thread_stop.load(Ordering::Acquire) { + let (mut code, mut key, mut pointer) = (0, 0, std::ptr::null_mut()); + let ok = unsafe { + GetQueuedCompletionStatus( + port.as_raw_handle(), + &mut code, + &mut key, + &mut pointer, + 100, + ) + }; + if ok == 0 { + if unsafe { windows_sys::Win32::Foundation::GetLastError() } + == windows_sys::Win32::Foundation::WAIT_TIMEOUT + { + continue; + } + return 2; + } + if key != 1 { + return 2; + } + if code != JOB_NOTIFICATION_LIMIT { + continue; + } + let mut info = JOBOBJECT_LIMIT_VIOLATION_INFORMATION::default(); + if unsafe { + QueryInformationJobObject( + owned_job.as_raw_handle(), + JobObjectLimitViolationInformation, + (&mut info as *mut JOBOBJECT_LIMIT_VIOLATION_INFORMATION).cast(), + size_of::() as u32, + std::ptr::null_mut(), + ) + } == 0 + { + return 2; + } + if info.ViolationLimitFlags & JOB_OBJECT_LIMIT_RATE_CONTROL != 0 { + return 1; + } + } + 0 + })) + .unwrap_or(2); + state.store(outcome, Ordering::Release); + if unsafe { TerminateJobObject(owned_job.as_raw_handle(), 1) } == 0 { + state.store(3, Ordering::Release); + } + }) + .map_err(|_| HostError::new("EXTENSION_RESOURCE_MONITOR_UNAVAILABLE"))?; + Ok(Self { + stop, + worker: Some(worker), + }) + } +} +impl Drop for CpuMonitor { + fn drop(&mut self) { + self.stop.store(true, std::sync::atomic::Ordering::Release); + if let Some(worker) = self.worker.take() { + let _ = worker.join(); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -199,6 +334,114 @@ mod tests { assert!(excessive.try_reserve_exact(600 * 1024 * 1024).is_err()); } #[test] + #[ignore = "helper consumes CPU inside the parent controlled job"] + fn worker_cpu() { + let end = std::time::Instant::now() + Duration::from_secs(40); + std::thread::scope(|scope| { + for _ in 0..8 { + scope.spawn(move || { + let mut n = 1u64; + while std::time::Instant::now() < end { + for _ in 0..10000 { + n = std::hint::black_box( + n.wrapping_mul(6364136223846793005).wrapping_add(1), + ); + } + } + }); + } + }); + } + #[test] + #[ignore = "real ten-second CPU pressure and Host save acceptance; run explicitly"] + fn cpu_pressure_terminates_job_and_host_can_save() { + let idle_job = Job::new().unwrap(); + let idle = worker(); + unsafe { + idle_job.assign_suspended(idle.process.as_handle()).unwrap(); + assert_ne!(ResumeThread(idle.thread.as_raw_handle()), u32::MAX); + } + let job = Job::new().unwrap(); + let mut rate = JOBOBJECT_CPU_RATE_CONTROL_INFORMATION::default(); + assert_ne!( + unsafe { + QueryInformationJobObject( + job.handle.as_raw_handle(), + JobObjectCpuRateControlInformation, + (&mut rate as *mut JOBOBJECT_CPU_RATE_CONTROL_INFORMATION).cast(), + size_of::() as u32, + std::ptr::null_mut(), + ) + }, + 0 + ); + assert_eq!( + rate.ControlFlags, + JOB_OBJECT_CPU_RATE_CONTROL_ENABLE + | JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP + | JOB_OBJECT_CPU_RATE_CONTROL_NOTIFY + ); + assert_eq!( + unsafe { rate.Anonymous.CpuRate }, + (10000 / unsafe { GetActiveProcessorCount(ALL_PROCESSOR_GROUPS) }).max(1) + ); + let child = worker_named("worker_cpu"); + let start = std::time::Instant::now(); + unsafe { + job.assign_suspended(child.process.as_handle()).unwrap(); + assert_ne!(ResumeThread(child.thread.as_raw_handle()), u32::MAX); + } + let vault = tempfile::tempdir().unwrap(); + let mut workspace = crate::workspace::Workspace::open(vault.path()).unwrap(); + workspace + .write("cpu.md", "", b"while CPU is busy", "local") + .unwrap(); + while job.check_resources().is_ok() && start.elapsed() < Duration::from_secs(20) { + std::thread::sleep(Duration::from_millis(25)); + } + assert_eq!( + job.check_resources().unwrap_err().code, + "EXTENSION_RESOURCE_CPU_EXCEEDED" + ); + let detected = start.elapsed(); + let cleanup = std::time::Instant::now(); + while job.active_processes().unwrap() != 0 && cleanup.elapsed() < Duration::from_secs(5) { + std::thread::sleep(Duration::from_millis(10)); + } + assert_eq!(job.active_processes().unwrap(), 0); + let saved = workspace.read("cpu.md").unwrap(); + workspace + .write("cpu.md", &saved.entry.hash, b"after CPU cleanup", "local") + .unwrap(); + drop(workspace); + assert_eq!( + crate::workspace::Workspace::open(vault.path()) + .unwrap() + .read("cpu.md") + .unwrap() + .content, + "after CPU cleanup" + ); + eprintln!( + "CPU pressure detected at {detected:?}; Job empty after {:?}", + cleanup.elapsed() + ); + idle_job.check_resources().unwrap(); + assert!(idle_job.active_processes().unwrap() > 0); + // An observation handle must not keep the monitor alive after its owner + // is dropped, or keep the idle process running indefinitely. + let observation = idle_job.clone_for_deadline().unwrap(); + let stop = std::time::Instant::now(); + drop(idle_job); + assert!(stop.elapsed() < Duration::from_secs(2)); + while observation.active_processes().unwrap() != 0 + && stop.elapsed() < Duration::from_secs(5) + { + std::thread::sleep(Duration::from_millis(10)); + } + assert_eq!(observation.active_processes().unwrap(), 0); + } + #[test] fn actual_allocation_above_job_memory_budget_is_refused() { let job = Job::new().unwrap(); let child = worker_named("worker_memory"); @@ -306,7 +549,9 @@ mod tests { ); assert_eq!( cpu.ControlFlags, - JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP + JOB_OBJECT_CPU_RATE_CONTROL_ENABLE + | JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP + | JOB_OBJECT_CPU_RATE_CONTROL_NOTIFY ); let child = worker(); unsafe { diff --git a/frontend/src-tauri/src/extension_process.rs b/frontend/src-tauri/src/extension_process.rs index e9d85c5..5eafa73 100644 --- a/frontend/src-tauri/src/extension_process.rs +++ b/frontend/src-tauri/src/extension_process.rs @@ -239,6 +239,7 @@ impl<'a> Suspended<'a> { /// broker and all resource policy requirements immediately before this call. /// None of those authorization checks is supplied by this low-level module. pub unsafe fn resume(self) -> Result> { + self.0.job.check_resources()?; if unsafe { ResumeThread(self.0.handles.thread.as_raw_handle()) } != 1 { return Err(HostError::new("EXTENSION_PROCESS_RESUME_FAILED")); } @@ -261,6 +262,7 @@ impl<'a> Suspended<'a> { ) -> Result> { let watch = crate::extension_revocation::Watch::arm(&self.0.job, lease)?; watch.check()?; + self.0.job.check_resources()?; if unsafe { ResumeThread(self.0.handles.thread.as_raw_handle()) } != 1 { return Err(HostError::new("EXTENSION_PROCESS_RESUME_FAILED")); } @@ -291,6 +293,7 @@ impl Running<'_> { } pub fn check_authorization(&self) -> Result<()> { + self.process.job.check_resources()?; #[cfg(feature = "desktop")] if let Some(watch) = &self.revocation { return watch.check();