diff --git a/crates/session/src/runtime/process_tree.rs b/crates/session/src/runtime/process_tree.rs index 50eca30..04bc9b5 100644 --- a/crates/session/src/runtime/process_tree.rs +++ b/crates/session/src/runtime/process_tree.rs @@ -1,7 +1,10 @@ use std::{ collections::{BTreeMap, BTreeSet, HashMap}, io, - sync::{Arc, OnceLock}, + sync::{ + Arc, OnceLock, + atomic::{AtomicU64, Ordering}, + }, time::{Duration, Instant}, }; @@ -12,11 +15,13 @@ const SNAPSHOT_CACHE_TTL: Duration = Duration::from_millis(20); #[derive(Clone)] struct CachedSnapshot { + sequence: u64, captured_at: Instant, records: Arc<[ProcessRecord]>, } static PROCESS_SNAPSHOT_CACHE: OnceLock>> = OnceLock::new(); +static NEXT_SNAPSHOT_SEQUENCE: AtomicU64 = AtomicU64::new(1); #[derive(Clone, Debug, Eq, PartialEq)] struct ProcessIdentity { @@ -40,6 +45,7 @@ pub(super) struct TrackedProcess { pub(super) struct ProcessTree { known: BTreeMap, + latest_snapshot_sequence: u64, scan_timeout: Duration, max_tracked_processes: usize, } @@ -50,32 +56,45 @@ impl ProcessTree { scan_timeout: Duration, max_tracked_processes: usize, ) -> io::Result { - let records = process_snapshot(scan_timeout, true)?; + let snapshot = process_snapshot(scan_timeout, true)?; let root_pid = root_pid.as_raw(); let mut tree = Self { known: BTreeMap::new(), + latest_snapshot_sequence: snapshot.sequence, scan_timeout, max_tracked_processes, }; - if let Some(root) = records + if let Some(root) = snapshot .records .iter() .find(|record| record.identity.pid == root_pid && !record.zombie) { tree.known.insert(root_pid, root.identity.clone()); } - let _ = tree.absorb(&records.records)?; + let _ = tree.absorb_snapshot(&snapshot)?; Ok(tree) } pub fn refresh(&mut self) -> io::Result> { - let snapshot = process_snapshot(self.scan_timeout, false)?; - self.absorb(&snapshot.records) + // A global scan can begin before this tree's latest evidence and + // publish afterward. Never let that older view prune a proven process. + let snapshot = process_snapshot_after( + self.scan_timeout, + false, + Some(self.latest_snapshot_sequence), + )?; + self.absorb_snapshot(&snapshot) } pub fn refresh_fresh(&mut self) -> io::Result> { let snapshot = process_snapshot(self.scan_timeout, true)?; - self.absorb(&snapshot.records) + self.absorb_snapshot(&snapshot) + } + + fn absorb_snapshot(&mut self, snapshot: &CachedSnapshot) -> io::Result> { + let processes = self.absorb(&snapshot.records)?; + self.latest_snapshot_sequence = self.latest_snapshot_sequence.max(snapshot.sequence); + Ok(processes) } fn absorb(&mut self, records: &[ProcessRecord]) -> io::Result> { @@ -144,6 +163,7 @@ impl ProcessTree { fn with_root_for_test(root: ProcessRecord, max_tracked_processes: usize) -> Self { Self { known: BTreeMap::from([(root.identity.pid, root.identity)]), + latest_snapshot_sequence: 1, scan_timeout: Duration::from_secs(1), max_tracked_processes, } @@ -151,7 +171,16 @@ impl ProcessTree { } fn process_snapshot(timeout: Duration, force: bool) -> io::Result { + process_snapshot_after(timeout, force, None) +} + +fn process_snapshot_after( + timeout: Duration, + force: bool, + minimum_sequence: Option, +) -> io::Result { let started = Instant::now(); + let sequence = next_snapshot_sequence()?; let cache = PROCESS_SNAPSHOT_CACHE.get_or_init(|| Mutex::new(None)); if !force { let Some(cached) = cache.try_lock_for(timeout) else { @@ -161,7 +190,7 @@ fn process_snapshot(timeout: Duration, force: bool) -> io::Result io::Result io::Result io::Result { + // The counter is only an ordering token; the cache mutex publishes the + // snapshot data, so no cross-thread memory ordering is required here. + NEXT_SNAPSHOT_SEQUENCE + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |sequence| { + sequence.checked_add(1) + }) + .map_err(|_| io::Error::other("process-tree snapshot sequence exhausted")) +} + +fn cached_snapshot_is_usable(snapshot: &CachedSnapshot, minimum_sequence: Option) -> bool { + snapshot.captured_at.elapsed() <= SNAPSHOT_CACHE_TTL + && minimum_sequence.is_none_or(|minimum| snapshot.sequence >= minimum) +} + +fn publish_snapshot(cache: &mut Option, snapshot: CachedSnapshot) { + let should_publish = cache + .as_ref() + .is_none_or(|cached| snapshot.sequence > cached.sequence); + if should_publish { + *cache = Some(snapshot); + } +} + #[cfg(target_os = "linux")] fn scan_processes_uncached(timeout: Duration) -> io::Result> { const MAX_PROCESS_SNAPSHOT_RECORDS: usize = 262_144; @@ -365,6 +421,57 @@ mod tests { } } + fn snapshot( + sequence: u64, + captured_at: Instant, + records: impl Into>, + ) -> CachedSnapshot { + CachedSnapshot { + sequence, + captured_at, + records: records.into(), + } + } + + #[test] + fn out_of_order_scan_cannot_replace_a_newer_cached_snapshot() { + let newer = snapshot(2, Instant::now(), vec![record(200, 1, 200, "newer")]); + let older = snapshot(1, Instant::now(), vec![record(100, 1, 100, "older")]); + let mut cache = Some(newer); + + publish_snapshot(&mut cache, older); + + let cached = cache.expect("newer cache entry should be retained"); + assert_eq!(cached.sequence, 2); + assert_eq!(cached.records[0].identity.pid, 200); + } + + #[test] + fn cached_snapshot_from_before_latest_tree_evidence_is_not_usable() { + let stale = snapshot(1, Instant::now(), Vec::::new()); + + assert!(!cached_snapshot_is_usable(&stale, Some(2))); + } + + #[test] + fn fresh_snapshot_advances_the_tree_cache_floor() { + let root = record(100, 1, 100, "root"); + let child = record(101, 100, 101, "child"); + let mut tree = ProcessTree::with_root_for_test(root.clone(), 8); + let newer = snapshot(3, Instant::now(), vec![root.clone(), child]); + + tree.absorb_snapshot(&newer) + .expect("newer process-tree evidence should be accepted"); + + let stale = snapshot(2, Instant::now(), vec![root]); + assert_eq!(tree.latest_snapshot_sequence, 3); + assert!(tree.known.contains_key(&101)); + assert!(!cached_snapshot_is_usable( + &stale, + Some(tree.latest_snapshot_sequence) + )); + } + #[test] fn only_proven_descendants_are_retained_across_group_changes() { let root = record(100, 1, 100, "root");