proc_tree/
default_store.rs

1//! Default storage implementation using standard library types.
2//!
3//! [`DefaultStore`] is a `HashMap<Mutex>` store with optional TTL-based eviction.
4//!
5//! # Example
6//!
7//! ```rust
8//! use proc_tree::{DefaultStore, snapshot};
9//!
10//! let store = DefaultStore::new(600);
11//! snapshot(&store).expect("failed to read /proc");
12//! ```
13
14use std::collections::HashMap;
15use std::sync::{Arc, Mutex};
16use std::time::{Duration, Instant};
17
18use crate::traits::ProcessStore;
19use crate::types::ProcessInfo;
20
21// ---- Internal entry with optional TTL ----
22
23struct Entry {
24    value: ProcessInfo,
25    inserted_at: Instant,
26}
27
28impl Clone for Entry {
29    fn clone(&self) -> Self {
30        Self {
31            value: self.value.clone(),
32            inserted_at: self.inserted_at,
33        }
34    }
35}
36
37// ---- Shared inner ----
38
39type Inner = Arc<Mutex<HashMap<u32, Entry>>>;
40
41fn get_inner(
42    inner: &Inner,
43    pid: u32,
44    ttl: Duration,
45    children_index: &Arc<Mutex<HashMap<u32, Vec<u32>>>>,
46) -> Option<ProcessInfo> {
47    let mut map = inner.lock().expect("lock poisoned");
48    let entry = map.get(&pid)?;
49    if !ttl.is_zero() && entry.inserted_at.elapsed() >= ttl {
50        let info = map.remove(&pid).expect("entry should exist").value;
51        // Update children index
52        let mut index = children_index.lock().expect("lock poisoned");
53        if let Some(children) = index.get_mut(&info.ppid()) {
54            children.retain(|&c| c != pid);
55        }
56        // Clean up this process's own children index entry
57        index.remove(&pid);
58        return None;
59    }
60    Some(entry.value.clone())
61}
62
63fn insert_inner(inner: &Inner, pid: u32, value: ProcessInfo) {
64    let mut map = inner.lock().expect("lock poisoned");
65    map.insert(
66        pid,
67        Entry {
68            value,
69            inserted_at: Instant::now(),
70        },
71    );
72}
73
74// ---- DefaultStore ----
75
76/// Process tree store backed by `HashMap<Mutex>` with optional TTL eviction.
77///
78/// Thread-safe via `Arc<Mutex<...>>`. Cloning shares the same data.
79pub struct DefaultStore {
80    inner: Inner,
81    children_index: Arc<Mutex<HashMap<u32, Vec<u32>>>>,
82    ttl: Duration,
83}
84
85impl DefaultStore {
86    /// Create a new store with the given TTL in seconds.
87    /// `ttl_secs = 0` means no expiration.
88    pub fn new(ttl_secs: u64) -> Self {
89        Self {
90            inner: Arc::new(Mutex::new(HashMap::new())),
91            children_index: Arc::new(Mutex::new(HashMap::new())),
92            ttl: Duration::from_secs(ttl_secs),
93        }
94    }
95
96    /// Number of entries (including possibly-expired ones not yet evicted).
97    pub fn len(&self) -> usize {
98        self.inner.lock().expect("lock poisoned").len()
99    }
100
101    /// Returns `true` if the store contains no entries.
102    pub fn is_empty(&self) -> bool {
103        self.len() == 0
104    }
105
106    /// Check if a PID exists in the store.
107    ///
108    /// Does not trigger TTL eviction — use [`get_process`](ProcessStore::get_process)
109    /// to check existence with TTL enforcement.
110    pub fn contains_key(&self, pid: u32) -> bool {
111        self.inner.lock().expect("lock poisoned").contains_key(&pid)
112    }
113}
114
115impl Clone for DefaultStore {
116    fn clone(&self) -> Self {
117        Self {
118            inner: Arc::clone(&self.inner),
119            children_index: Arc::clone(&self.children_index),
120            ttl: self.ttl,
121        }
122    }
123}
124
125impl std::fmt::Debug for DefaultStore {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        f.debug_struct("DefaultStore")
128            .field("len", &self.len())
129            .field("ttl", &self.ttl)
130            .finish()
131    }
132}
133
134impl Default for DefaultStore {
135    /// Creates a store with no TTL.
136    fn default() -> Self {
137        Self::new(0)
138    }
139}
140
141impl ProcessStore for DefaultStore {
142    fn get_process(&self, pid: u32) -> Option<ProcessInfo> {
143        get_inner(&self.inner, pid, self.ttl, &self.children_index)
144    }
145
146    fn insert_process(&self, pid: u32, info: ProcessInfo) {
147        let new_ppid = info.ppid();
148
149        // Check if process already exists with different ppid
150        let old_ppid = {
151            let map = self.inner.lock().expect("lock poisoned");
152            map.get(&pid).map(|e| e.value.ppid())
153        };
154
155        // Insert the process
156        insert_inner(&self.inner, pid, info);
157
158        // Update children index
159        let mut index = self.children_index.lock().expect("lock poisoned");
160
161        // Remove from old parent's index if ppid changed
162        if let Some(old_ppid) = old_ppid
163            && old_ppid != new_ppid
164            && let Some(children) = index.get_mut(&old_ppid)
165        {
166            children.retain(|&c| c != pid);
167        }
168
169        // Add to new parent's index (avoid duplicates)
170        let children = index.entry(new_ppid).or_default();
171        if !children.contains(&pid) {
172            children.push(pid);
173        }
174    }
175
176    fn remove_process(&self, pid: u32) -> Option<ProcessInfo> {
177        // Remove from inner
178        let info = {
179            let mut map = self.inner.lock().expect("lock poisoned");
180            map.remove(&pid).map(|e| e.value)
181        };
182
183        // Remove from children index
184        if let Some(ref p) = info {
185            let mut index = self.children_index.lock().expect("lock poisoned");
186            // Remove from parent's index
187            if let Some(children) = index.get_mut(&p.ppid()) {
188                children.retain(|&c| c != pid);
189            }
190            // Clean up this process's own children index entry
191            index.remove(&pid);
192        }
193
194        info
195    }
196
197    fn all_pids(&self) -> Vec<u32> {
198        self.inner
199            .lock()
200            .expect("lock poisoned")
201            .keys()
202            .copied()
203            .collect()
204    }
205
206    fn for_each_child(&self, pid: u32, f: &mut dyn FnMut(u32)) {
207        // Collect children first, then release lock before calling f.
208        // f may call insert_process which also locks children_index.
209        let children: Vec<u32> = {
210            let index = self.children_index.lock().expect("lock poisoned");
211            index.get(&pid).cloned().unwrap_or_default()
212        };
213        for child in children {
214            f(child);
215        }
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn default_store_insert_get() {
225        let store = DefaultStore::new(0);
226        store.insert_process(
227            1,
228            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
229        );
230        let info = store.get_process(1).unwrap();
231        assert_eq!(info.ppid(), 0);
232        assert_eq!(info.cmd(), "init");
233    }
234
235    #[test]
236    fn default_store_ttl_expired() {
237        let store = DefaultStore::new(0); // ttl=0 means no expiry
238        store.insert_process(
239            1,
240            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
241        );
242        assert!(store.get_process(1).is_some());
243
244        // With ttl=1, entry expires after 1 second
245        let store = DefaultStore::new(1);
246        store.insert_process(
247            1,
248            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
249        );
250        assert!(store.get_process(1).is_some());
251        std::thread::sleep(Duration::from_millis(1100));
252        assert!(store.get_process(1).is_none());
253    }
254
255    #[test]
256    fn clone_shares_data() {
257        let store = DefaultStore::new(0);
258        store.insert_process(
259            1,
260            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
261        );
262        let store2 = store.clone();
263        assert!(store2.get_process(1).is_some());
264        store2.insert_process(
265            2,
266            ProcessInfo::new("bash".into(), "bash".into(), "root".into(), 1, 2, 0),
267        );
268        assert!(store.get_process(2).is_some());
269    }
270
271    #[test]
272    fn len_and_contains() {
273        let store = DefaultStore::new(0);
274        assert_eq!(store.len(), 0);
275        store.insert_process(
276            1,
277            ProcessInfo::new("a".into(), "a".into(), "u".into(), 0, 1, 0),
278        );
279        assert_eq!(store.len(), 1);
280        assert!(store.contains_key(1));
281        assert!(!store.contains_key(999));
282    }
283
284    #[test]
285    fn is_empty_default() {
286        let store = DefaultStore::new(0);
287        assert!(store.is_empty());
288        store.insert_process(
289            1,
290            ProcessInfo::new("a".into(), "a".into(), "u".into(), 0, 1, 0),
291        );
292        assert!(!store.is_empty());
293    }
294
295    #[test]
296    fn all_pids_returns_keys() {
297        let store = DefaultStore::new(0);
298        store.insert_process(
299            1,
300            ProcessInfo::new("a".into(), "a".into(), "u".into(), 0, 1, 0),
301        );
302        store.insert_process(
303            2,
304            ProcessInfo::new("b".into(), "b".into(), "u".into(), 1, 2, 0),
305        );
306        store.insert_process(
307            3,
308            ProcessInfo::new("c".into(), "c".into(), "u".into(), 1, 3, 0),
309        );
310        let mut pids = store.all_pids();
311        pids.sort();
312        assert_eq!(pids, vec![1, 2, 3]);
313    }
314
315    #[test]
316    fn ttl_contains_key_does_not_expire() {
317        let store = DefaultStore::new(1);
318        store.insert_process(
319            1,
320            ProcessInfo::new("a".into(), "a".into(), "u".into(), 0, 1, 0),
321        );
322        assert!(store.contains_key(1));
323        std::thread::sleep(Duration::from_millis(1100));
324        // contains_key does not trigger TTL eviction
325        assert!(store.contains_key(1));
326        // but get_process does
327        assert!(store.get_process(1).is_none());
328    }
329
330    #[test]
331    fn remove_process() {
332        let store = DefaultStore::new(0);
333        store.insert_process(
334            1,
335            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
336        );
337        store.insert_process(
338            2,
339            ProcessInfo::new("bash".into(), "bash".into(), "root".into(), 1, 2, 0),
340        );
341
342        assert_eq!(store.len(), 2);
343        assert!(store.contains_key(2));
344
345        let removed = store.remove_process(2);
346        assert!(removed.is_some());
347        assert_eq!(store.len(), 1);
348        assert!(!store.contains_key(2));
349    }
350
351    #[test]
352    fn children_of() {
353        let store = DefaultStore::new(0);
354        store.insert_process(
355            1,
356            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
357        );
358        store.insert_process(
359            100,
360            ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0),
361        );
362        store.insert_process(
363            200,
364            ProcessInfo::new("b".into(), "b".into(), "root".into(), 1, 200, 0),
365        );
366        store.insert_process(
367            300,
368            ProcessInfo::new("c".into(), "c".into(), "root".into(), 100, 300, 0),
369        );
370
371        let mut kids = store.children_of(1);
372        kids.sort();
373        assert_eq!(kids, vec![100, 200]);
374
375        let kids_100 = store.children_of(100);
376        assert_eq!(kids_100, vec![300]);
377
378        assert!(store.children_of(999).is_empty());
379    }
380
381    #[test]
382    fn insert_ppid_change_removes_from_old_parent() {
383        let store = DefaultStore::new(0);
384        store.insert_process(
385            1,
386            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
387        );
388        store.insert_process(
389            100,
390            ProcessInfo::new("other".into(), "other".into(), "root".into(), 0, 100, 0),
391        );
392        store.insert_process(
393            200,
394            ProcessInfo::new("child".into(), "child".into(), "root".into(), 100, 200, 0),
395        );
396
397        // child 200 is under parent 100
398        assert_eq!(store.children_of(100), vec![200]);
399        assert!(store.children_of(1).is_empty());
400
401        // Re-parent 200 from 100 to 1
402        store.insert_process(
403            200,
404            ProcessInfo::new("child".into(), "child".into(), "root".into(), 1, 200, 0),
405        );
406
407        // 200 should be removed from old parent's index
408        assert!(
409            store.children_of(100).is_empty(),
410            "old parent should have no children"
411        );
412        // 200 should be in new parent's index
413        assert_eq!(store.children_of(1), vec![200]);
414    }
415
416    #[test]
417    fn insert_same_ppid_no_duplicate() {
418        let store = DefaultStore::new(0);
419        store.insert_process(
420            1,
421            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
422        );
423        store.insert_process(
424            100,
425            ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0),
426        );
427        // Insert same pid with same ppid again
428        store.insert_process(
429            100,
430            ProcessInfo::new("a".into(), "a".into(), "root".into(), 1, 100, 0),
431        );
432        // Should not duplicate
433        assert_eq!(store.children_of(1), vec![100]);
434    }
435
436    #[test]
437    fn remove_process_cleans_own_children_index() {
438        let store = DefaultStore::new(0);
439        store.insert_process(
440            1,
441            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
442        );
443        store.insert_process(
444            100,
445            ProcessInfo::new("parent".into(), "parent".into(), "root".into(), 1, 100, 0),
446        );
447        store.insert_process(
448            200,
449            ProcessInfo::new("child".into(), "child".into(), "root".into(), 100, 200, 0),
450        );
451
452        assert_eq!(store.children_of(100), vec![200]);
453
454        // Remove parent 100
455        store.remove_process(100);
456
457        // children_of(100) should return empty, not stale [200]
458        assert!(
459            store.children_of(100).is_empty(),
460            "removed process should have no children index"
461        );
462        // Child 200 still exists in store but parent is gone
463        assert!(store.get_process(200).is_some());
464    }
465
466    #[test]
467    fn ttl_expiration_cleans_own_children_index() {
468        let store = DefaultStore::new(1); // 1 second TTL
469        store.insert_process(
470            1,
471            ProcessInfo::new("init".into(), "init".into(), "root".into(), 0, 1, 0),
472        );
473        store.insert_process(
474            100,
475            ProcessInfo::new("parent".into(), "parent".into(), "root".into(), 1, 100, 0),
476        );
477        store.insert_process(
478            200,
479            ProcessInfo::new("child".into(), "child".into(), "root".into(), 100, 200, 0),
480        );
481
482        assert_eq!(store.children_of(100), vec![200]);
483
484        // Wait for parent to expire
485        std::thread::sleep(Duration::from_millis(1100));
486
487        // Accessing expired parent triggers eviction
488        assert!(store.get_process(100).is_none());
489
490        // children_of(100) should return empty, not stale [200]
491        assert!(
492            store.children_of(100).is_empty(),
493            "expired process should have no children index"
494        );
495    }
496}
proc_tree/default_store.rs

proc_tree/
default_store.rs
