ssh_mcp/background/
registry.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::time::{Duration, SystemTime};
4
5use tokio::sync::RwLock;
6
7use super::job::SharedJobState;
8
9/// Thread-safe job registry.
10///
11/// Stores running jobs and keeps completed jobs for a retention window.
12#[derive(Debug, Clone)]
13pub struct JobRegistry {
14    jobs: Arc<RwLock<HashMap<String, SharedJobState>>>,
15    completed_retention: Duration,
16}
17
18impl JobRegistry {
19    pub fn new(completed_retention: Duration) -> Self {
20        Self {
21            jobs: Arc::new(RwLock::new(HashMap::new())),
22            completed_retention,
23        }
24    }
25
26    pub fn completed_retention(&self) -> Duration {
27        self.completed_retention
28    }
29
30    pub async fn insert(&self, job_id: String, job: SharedJobState) {
31        let mut guard = self.jobs.write().await;
32        guard.insert(job_id, job);
33    }
34
35    pub async fn get(&self, job_id: &str) -> Option<SharedJobState> {
36        let guard = self.jobs.read().await;
37        guard.get(job_id).cloned()
38    }
39
40    pub async fn remove(&self, job_id: &str) -> Option<SharedJobState> {
41        let mut guard = self.jobs.write().await;
42        guard.remove(job_id)
43    }
44
45    /// Remove completed jobs that exceeded the retention window.
46    ///
47    /// Returns the number of pruned jobs.
48    pub async fn prune_expired(&self) -> usize {
49        let now = SystemTime::now();
50
51        let snapshot: Vec<(String, SharedJobState)> = {
52            let guard = self.jobs.read().await;
53            guard
54                .iter()
55                .map(|(id, job)| (id.clone(), Arc::clone(job)))
56                .collect()
57        };
58
59        let mut expired = Vec::new();
60        for (job_id, job) in snapshot {
61            let job_for_list = Arc::clone(&job);
62            let job_guard = job.lock().await;
63            if !job_guard.is_terminal() {
64                continue;
65            };
66
67            let Some(completed_at) = job_guard.completed_at else {
68                continue;
69            };
70
71            let Ok(age) = now.duration_since(completed_at) else {
72                continue;
73            };
74            if age > self.completed_retention {
75                expired.push((job_id, job_for_list, completed_at));
76            }
77        }
78
79        if expired.is_empty() {
80            return 0;
81        }
82
83        let mut guard = self.jobs.write().await;
84        let mut removed = 0;
85        for (job_id, job, completed_at) in expired {
86            let Some(current) = guard.get(&job_id).cloned() else {
87                continue;
88            };
89
90            // Avoid deleting a new job if a job_id is reused.
91            if !Arc::ptr_eq(&current, &job) {
92                continue;
93            }
94
95            // Re-validate terminal + age under the write lock before removing.
96            let Ok(job_guard) = current.try_lock() else {
97                continue;
98            };
99            if !job_guard.is_terminal() {
100                continue;
101            }
102            if job_guard.completed_at != Some(completed_at) {
103                continue;
104            }
105            let Ok(age) = now.duration_since(completed_at) else {
106                continue;
107            };
108            if age <= self.completed_retention {
109                continue;
110            }
111
112            if guard.remove(&job_id).is_some() {
113                removed += 1;
114            }
115        }
116        removed
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::background::job::{JobState, NewRunningJob};
124    use tokio::sync::Mutex;
125
126    fn make_running(job_id: &str) -> SharedJobState {
127        Arc::new(Mutex::new(JobState::new_running(NewRunningJob {
128            job_id: job_id.to_string(),
129            pid: 123,
130            log_path: std::path::PathBuf::from("/tmp/ssh-mcp/test.log"),
131            command: "echo test".to_string(),
132            connection_id: "test@localhost:22".to_string(),
133        })))
134    }
135
136    #[tokio::test]
137    async fn test_insert_get_remove() {
138        let reg = JobRegistry::new(Duration::from_secs(60));
139
140        let job = make_running("job_1");
141        reg.insert("job_1".to_string(), Arc::clone(&job)).await;
142
143        let got = reg.get("job_1").await.expect("job should exist");
144        assert!(Arc::ptr_eq(&got, &job));
145
146        let removed = reg.remove("job_1").await.expect("job should be removed");
147        assert!(Arc::ptr_eq(&removed, &job));
148        assert!(reg.get("job_1").await.is_none());
149    }
150
151    #[tokio::test]
152    async fn test_prune_expired_removes_terminal_jobs_past_retention() {
153        let reg = JobRegistry::new(Duration::from_millis(10));
154
155        let running = make_running("job_running");
156        reg.insert("job_running".to_string(), Arc::clone(&running))
157            .await;
158
159        let completed = make_running("job_completed");
160        {
161            let mut guard = completed.lock().await;
162            guard.mark_exit(0);
163            guard.completed_at = guard
164                .completed_at
165                .and_then(|t| t.checked_sub(Duration::from_secs(1)));
166        }
167        reg.insert("job_completed".to_string(), Arc::clone(&completed))
168            .await;
169
170        let removed = reg.prune_expired().await;
171        assert_eq!(removed, 1);
172        assert!(reg.get("job_completed").await.is_none());
173        assert!(reg.get("job_running").await.is_some());
174    }
175}
ssh_mcp/background/registry.rs

ssh_mcp/background/
registry.rs
