proc_tree/
default_store.rs

1//! Default storage implementations using standard library types.
2//!
3//! [`DefaultStore<V>`] is a generic `HashMap<Mutex>` store with optional
4//! TTL-based eviction. [`DefaultTree`] and [`DefaultCache`] are type aliases.
5//!
6//! # Example
7//!
8//! ```rust
9//! use proc_tree::{DefaultTree, DefaultCache, snapshot};
10//!
11//! let tree = DefaultTree::new(65536, 600);
12//! let cache = DefaultCache::new(65536, 600);
13//! snapshot(&tree, &cache);
14//! ```
15
16use std::collections::HashMap;
17use std::sync::{Arc, Mutex};
18use std::time::{Duration, Instant};
19
20use crate::{CacheStore, PidNode, ProcInfo, TreeStore};
21
22// ---- Internal entry with optional TTL ----
23
24struct Entry<V> {
25    value: V,
26    inserted_at: Instant,
27}
28
29impl<V: Clone> Clone for Entry<V> {
30    fn clone(&self) -> Self {
31        Self {
32            value: self.value.clone(),
33            inserted_at: self.inserted_at,
34        }
35    }
36}
37
38// ---- Shared inner ----
39
40type Inner<V> = Arc<Mutex<HashMap<u32, Entry<V>>>>;
41
42fn get_inner<V: Clone>(inner: &Inner<V>, pid: u32, ttl: Duration) -> Option<V> {
43    let mut map = inner.lock().unwrap();
44    let entry = map.get(&pid)?;
45    if !ttl.is_zero() && entry.inserted_at.elapsed() >= ttl {
46        map.remove(&pid);
47        return None;
48    }
49    Some(entry.value.clone())
50}
51
52fn insert_inner<V: Clone>(inner: &Inner<V>, pid: u32, value: V) {
53    let mut map = inner.lock().unwrap();
54    map.insert(
55        pid,
56        Entry {
57            value,
58            inserted_at: Instant::now(),
59        },
60    );
61}
62
63fn len_inner<V>(inner: &Inner<V>) -> usize {
64    inner.lock().unwrap().len()
65}
66
67// ---- DefaultStore<V> ----
68
69/// Generic store backed by `HashMap<Mutex>` with optional TTL eviction.
70///
71/// Thread-safe via `Arc<Mutex<...>>`. Cloning shares the same data.
72pub struct DefaultStore<V> {
73    inner: Inner<V>,
74    ttl: Duration,
75}
76
77/// Process tree store. See [`DefaultStore`].
78pub type DefaultTree = DefaultStore<PidNode>;
79
80/// Process info cache. See [`DefaultStore`].
81pub type DefaultCache = DefaultStore<ProcInfo>;
82
83impl<V: Clone> DefaultStore<V> {
84    /// Create a new store with the given capacity hint and TTL in seconds.
85    /// `ttl_secs = 0` means no expiration.
86    pub fn new(_capacity: u64, ttl_secs: u64) -> Self {
87        Self {
88            inner: Arc::new(Mutex::new(HashMap::new())),
89            ttl: Duration::from_secs(ttl_secs),
90        }
91    }
92
93    /// Number of entries (including possibly-expired ones not yet evicted).
94    pub fn len(&self) -> usize {
95        len_inner(&self.inner)
96    }
97
98    /// Returns `true` if the store contains no entries.
99    pub fn is_empty(&self) -> bool {
100        self.len() == 0
101    }
102
103    /// Check if a PID exists and is not expired.
104    pub fn contains_key(&self, pid: u32) -> bool {
105        get_inner(&self.inner, pid, self.ttl).is_some()
106    }
107}
108
109impl<V: Clone> Clone for DefaultStore<V> {
110    fn clone(&self) -> Self {
111        Self {
112            inner: Arc::clone(&self.inner),
113            ttl: self.ttl,
114        }
115    }
116}
117
118impl<V: Clone> Default for DefaultStore<V> {
119    /// Creates a store with capacity 100 and no TTL.
120    fn default() -> Self {
121        Self::new(100, 0)
122    }
123}
124
125impl TreeStore for DefaultTree {
126    fn get_node(&self, pid: u32) -> Option<PidNode> {
127        get_inner(&self.inner, pid, self.ttl)
128    }
129
130    fn insert_node(&self, pid: u32, node: PidNode) {
131        insert_inner(&self.inner, pid, node);
132    }
133
134    fn all_pids(&self) -> Vec<u32> {
135        self.inner.lock().unwrap().keys().copied().collect()
136    }
137}
138
139impl CacheStore for DefaultCache {
140    fn get_info(&self, pid: u32) -> Option<ProcInfo> {
141        get_inner(&self.inner, pid, self.ttl)
142    }
143
144    fn insert_info(&self, pid: u32, info: ProcInfo) {
145        insert_inner(&self.inner, pid, info);
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    #[test]
154    fn default_tree_insert_get() {
155        let tree = DefaultTree::new(100, 0);
156        tree.insert_node(
157            1,
158            PidNode {
159                ppid: 0,
160                cmd: "init".into(),
161            },
162        );
163        let node = tree.get_node(1).unwrap();
164        assert_eq!(node.ppid, 0);
165        assert_eq!(node.cmd, "init");
166    }
167
168    #[test]
169    fn default_tree_ttl_expired() {
170        let tree = DefaultTree::new(100, 0); // ttl=0 means no expiry
171        tree.insert_node(
172            1,
173            PidNode {
174                ppid: 0,
175                cmd: "init".into(),
176            },
177        );
178        assert!(tree.get_node(1).is_some());
179
180        // With ttl=1, entry expires after 1 second
181        let tree = DefaultTree::new(100, 1);
182        tree.insert_node(
183            1,
184            PidNode {
185                ppid: 0,
186                cmd: "init".into(),
187            },
188        );
189        assert!(tree.get_node(1).is_some());
190        std::thread::sleep(Duration::from_millis(1100));
191        assert!(tree.get_node(1).is_none());
192    }
193
194    #[test]
195    fn default_cache_insert_get() {
196        let cache = DefaultCache::new(100, 0);
197        cache.insert_info(
198            42,
199            ProcInfo {
200                cmd: "bash".into(),
201                user: "root".into(),
202                ppid: 1,
203                tgid: 42,
204                start_time_ns: 0,
205            },
206        );
207        let info = cache.get_info(42).unwrap();
208        assert_eq!(info.cmd, "bash");
209        assert_eq!(info.ppid, 1);
210    }
211
212    #[test]
213    fn clone_shares_data() {
214        let tree = DefaultTree::new(100, 0);
215        tree.insert_node(
216            1,
217            PidNode {
218                ppid: 0,
219                cmd: "init".into(),
220            },
221        );
222        let tree2 = tree.clone();
223        assert!(tree2.get_node(1).is_some());
224        tree2.insert_node(
225            2,
226            PidNode {
227                ppid: 1,
228                cmd: "bash".into(),
229            },
230        );
231        assert!(tree.get_node(2).is_some());
232    }
233
234    #[test]
235    fn len_and_contains() {
236        let cache = DefaultCache::new(100, 0);
237        assert_eq!(cache.len(), 0);
238        cache.insert_info(
239            1,
240            ProcInfo {
241                cmd: "a".into(),
242                user: "u".into(),
243                ppid: 0,
244                tgid: 1,
245                start_time_ns: 0,
246            },
247        );
248        assert_eq!(cache.len(), 1);
249        assert!(cache.contains_key(1));
250        assert!(!cache.contains_key(999));
251    }
252
253    #[test]
254    fn default_cache_ttl_expired() {
255        let cache = DefaultCache::new(100, 0);
256        cache.insert_info(
257            1,
258            ProcInfo {
259                cmd: "a".into(),
260                user: "u".into(),
261                ppid: 0,
262                tgid: 1,
263                start_time_ns: 0,
264            },
265        );
266        assert!(cache.get_info(1).is_some());
267
268        let cache = DefaultCache::new(100, 1);
269        cache.insert_info(
270            1,
271            ProcInfo {
272                cmd: "a".into(),
273                user: "u".into(),
274                ppid: 0,
275                tgid: 1,
276                start_time_ns: 0,
277            },
278        );
279        assert!(cache.get_info(1).is_some());
280        std::thread::sleep(Duration::from_millis(1100));
281        assert!(cache.get_info(1).is_none());
282    }
283
284    #[test]
285    fn is_empty_default() {
286        let tree = DefaultTree::new(100, 0);
287        assert!(tree.is_empty());
288        tree.insert_node(
289            1,
290            PidNode {
291                ppid: 0,
292                cmd: "init".into(),
293            },
294        );
295        assert!(!tree.is_empty());
296
297        let cache = DefaultCache::new(100, 0);
298        assert!(cache.is_empty());
299        cache.insert_info(
300            1,
301            ProcInfo {
302                cmd: "a".into(),
303                user: "u".into(),
304                ppid: 0,
305                tgid: 1,
306                start_time_ns: 0,
307            },
308        );
309        assert!(!cache.is_empty());
310    }
311
312    #[test]
313    fn all_pids_returns_keys() {
314        let tree = DefaultTree::new(100, 0);
315        tree.insert_node(
316            1,
317            PidNode {
318                ppid: 0,
319                cmd: "a".into(),
320            },
321        );
322        tree.insert_node(
323            2,
324            PidNode {
325                ppid: 1,
326                cmd: "b".into(),
327            },
328        );
329        tree.insert_node(
330            3,
331            PidNode {
332                ppid: 1,
333                cmd: "c".into(),
334            },
335        );
336        let mut pids = tree.all_pids();
337        pids.sort();
338        assert_eq!(pids, vec![1, 2, 3]);
339    }
340
341    #[test]
342    fn tree_ttl_contains_key_expires() {
343        let tree = DefaultTree::new(100, 1);
344        tree.insert_node(
345            1,
346            PidNode {
347                ppid: 0,
348                cmd: "a".into(),
349            },
350        );
351        assert!(tree.contains_key(1));
352        std::thread::sleep(Duration::from_millis(1100));
353        assert!(!tree.contains_key(1));
354    }
355}
proc_tree/default_store.rs

proc_tree/
default_store.rs
