Rust · 8652 bytes Raw Blame History
1 use nix::sys::wait::WaitStatus;
2 use nix::unistd::Pid;
3 use std::collections::HashMap;
4
5 pub type JobId = u32;
6
7 /// State of a job
8 #[derive(Debug, Clone, PartialEq, Eq)]
9 pub enum JobState {
10 /// Job is currently running
11 Running,
12 /// Job has been stopped (Ctrl-Z)
13 Stopped,
14 /// Job has completed
15 Done(i32), // exit code
16 }
17
18 /// Represents a job (pipeline or command)
19 #[derive(Debug)]
20 pub struct Job {
21 /// Unique job ID
22 pub id: JobId,
23 /// Process group ID
24 pub pgid: Pid,
25 /// Command string for display
26 pub command: String,
27 /// Process IDs in this job (for pipelines)
28 pub pids: Vec<Pid>,
29 /// Current state
30 pub state: JobState,
31 /// Is this a foreground job?
32 pub foreground: bool,
33 }
34
35 impl Job {
36 /// Create a new job
37 pub fn new(
38 id: JobId,
39 pgid: Pid,
40 command: String,
41 pids: Vec<Pid>,
42 foreground: bool,
43 ) -> Self {
44 Self {
45 id,
46 pgid,
47 command,
48 pids,
49 state: JobState::Running,
50 foreground,
51 }
52 }
53
54 /// Check if all processes in the job have completed
55 pub fn is_completed(&self) -> bool {
56 matches!(self.state, JobState::Done(_))
57 }
58
59 /// Check if the job is stopped
60 pub fn is_stopped(&self) -> bool {
61 matches!(self.state, JobState::Stopped)
62 }
63
64 /// Check if the job is running
65 pub fn is_running(&self) -> bool {
66 matches!(self.state, JobState::Running)
67 }
68
69 /// Get a status string for display
70 pub fn status_string(&self) -> &str {
71 match &self.state {
72 JobState::Running => "Running",
73 JobState::Stopped => "Stopped",
74 JobState::Done(_) => "Done",
75 }
76 }
77 }
78
79 /// Manages all jobs
80 #[derive(Debug)]
81 pub struct JobList {
82 /// Map of job ID to Job
83 jobs: HashMap<JobId, Job>,
84 /// Next job ID to assign
85 next_job_id: JobId,
86 /// Shell's process group ID
87 shell_pgid: Pid,
88 /// Current foreground job (if any)
89 current_job: Option<JobId>,
90 }
91
92 impl JobList {
93 /// Create a new job list
94 pub fn new(shell_pgid: Pid) -> Self {
95 Self {
96 jobs: HashMap::new(),
97 next_job_id: 1,
98 shell_pgid,
99 current_job: None,
100 }
101 }
102
103 /// Add a new job
104 pub fn add_job(
105 &mut self,
106 pgid: Pid,
107 command: String,
108 pids: Vec<Pid>,
109 foreground: bool,
110 ) -> JobId {
111 let id = self.next_job_id;
112 self.next_job_id += 1;
113
114 let job = Job::new(id, pgid, command, pids, foreground);
115 self.jobs.insert(id, job);
116
117 if foreground {
118 self.current_job = Some(id);
119 }
120
121 id
122 }
123
124 /// Get a job by ID
125 pub fn get_job(&self, id: JobId) -> Option<&Job> {
126 self.jobs.get(&id)
127 }
128
129 /// Get a mutable reference to a job by ID
130 pub fn get_job_mut(&mut self, id: JobId) -> Option<&mut Job> {
131 self.jobs.get_mut(&id)
132 }
133
134 /// Remove a job
135 pub fn remove_job(&mut self, id: JobId) -> Option<Job> {
136 if self.current_job == Some(id) {
137 self.current_job = None;
138 }
139 self.jobs.remove(&id)
140 }
141
142 /// Get the current (most recent) job
143 pub fn current_job(&self) -> Option<&Job> {
144 self.current_job.and_then(|id| self.get_job(id))
145 }
146
147 /// Get all jobs
148 pub fn jobs(&self) -> impl Iterator<Item = &Job> {
149 self.jobs.values()
150 }
151
152 /// Get all jobs as a sorted list
153 pub fn jobs_sorted(&self) -> Vec<&Job> {
154 let mut jobs: Vec<_> = self.jobs.values().collect();
155 jobs.sort_by_key(|job| job.id);
156 jobs
157 }
158
159 /// Update job state based on wait status
160 pub fn update_job_status(&mut self, pid: Pid, status: WaitStatus) -> Option<JobId> {
161 // Find the job containing this PID
162 let job_id = self.jobs.iter().find_map(|(id, job)| {
163 if job.pids.contains(&pid) {
164 Some(*id)
165 } else {
166 None
167 }
168 })?;
169
170 if let Some(job) = self.get_job_mut(job_id) {
171 match status {
172 WaitStatus::Exited(_, code) => {
173 job.state = JobState::Done(code);
174 }
175 WaitStatus::Signaled(_, signal, _) => {
176 // Terminated by signal, use 128 + signal number as exit code
177 job.state = JobState::Done(128 + signal as i32);
178 }
179 WaitStatus::Stopped(_, _) => {
180 job.state = JobState::Stopped;
181 }
182 WaitStatus::Continued(_) => {
183 job.state = JobState::Running;
184 }
185 _ => {}
186 }
187 }
188
189 Some(job_id)
190 }
191
192 /// Remove all completed jobs
193 pub fn clean_completed(&mut self) -> Vec<Job> {
194 let completed: Vec<_> = self
195 .jobs
196 .iter()
197 .filter(|(_, job)| job.is_completed())
198 .map(|(id, _)| *id)
199 .collect();
200
201 completed
202 .into_iter()
203 .filter_map(|id| self.remove_job(id))
204 .collect()
205 }
206
207 /// Get shell's process group ID
208 pub fn shell_pgid(&self) -> Pid {
209 self.shell_pgid
210 }
211
212 /// Get the current job ID
213 pub fn current_job_id(&self) -> Option<JobId> {
214 self.current_job
215 }
216
217 /// Get the previous job ID (second most recent)
218 pub fn previous_job(&self) -> Option<JobId> {
219 // Return the second highest job ID that isn't the current one
220 let mut job_ids: Vec<_> = self.jobs.keys().copied().collect();
221 job_ids.sort_by(|a, b| b.cmp(a)); // Descending order
222
223 job_ids.into_iter().find(|&id| Some(id) != self.current_job)
224 }
225
226 /// Disown a specific job (remove from job control without killing it)
227 pub fn disown_job(&mut self, id: JobId) {
228 self.remove_job(id);
229 }
230
231 /// Disown all jobs, optionally only running jobs
232 pub fn disown_all(&mut self, running_only: bool) {
233 let to_disown: Vec<_> = self
234 .jobs
235 .iter()
236 .filter(|(_, job)| !running_only || job.is_running())
237 .map(|(id, _)| *id)
238 .collect();
239
240 for id in to_disown {
241 self.remove_job(id);
242 }
243 }
244
245 /// Send SIGHUP to all running jobs
246 /// Called when the shell is exiting to notify background jobs
247 pub fn send_hup_to_all(&self) {
248 use nix::sys::signal::{killpg, Signal};
249
250 for job in self.jobs.values() {
251 if job.is_running() || job.is_stopped() {
252 // Send SIGHUP to the process group
253 let _ = killpg(job.pgid, Signal::SIGHUP);
254 // Also send SIGCONT in case the job was stopped
255 if job.is_stopped() {
256 let _ = killpg(job.pgid, Signal::SIGCONT);
257 }
258 }
259 }
260 }
261
262 /// Check if there are any running or stopped jobs
263 pub fn has_active_jobs(&self) -> bool {
264 self.jobs.values().any(|job| job.is_running() || job.is_stopped())
265 }
266 }
267
268 #[cfg(test)]
269 mod tests {
270 use super::*;
271
272 #[test]
273 fn test_job_creation() {
274 let shell_pgid = Pid::from_raw(1000);
275 let mut job_list = JobList::new(shell_pgid);
276
277 let pgid = Pid::from_raw(2000);
278 let pids = vec![Pid::from_raw(2000), Pid::from_raw(2001)];
279 let id = job_list.add_job(pgid, "ls | grep test".to_string(), pids, false);
280
281 assert_eq!(id, 1);
282 assert_eq!(job_list.jobs().count(), 1);
283
284 let job = job_list.get_job(id).unwrap();
285 assert_eq!(job.id, 1);
286 assert_eq!(job.pgid, pgid);
287 assert!(job.is_running());
288 }
289
290 #[test]
291 fn test_job_state_transitions() {
292 let shell_pgid = Pid::from_raw(1000);
293 let mut job_list = JobList::new(shell_pgid);
294
295 let pgid = Pid::from_raw(2000);
296 let pid = Pid::from_raw(2000);
297 let id = job_list.add_job(pgid, "sleep 10".to_string(), vec![pid], false);
298
299 // Initially running
300 assert!(job_list.get_job(id).unwrap().is_running());
301
302 // Update to stopped
303 job_list.update_job_status(pid, WaitStatus::Stopped(pid, nix::sys::signal::SIGTSTP));
304 assert!(job_list.get_job(id).unwrap().is_stopped());
305
306 // Update to continued
307 job_list.update_job_status(pid, WaitStatus::Continued(pid));
308 assert!(job_list.get_job(id).unwrap().is_running());
309
310 // Update to done
311 job_list.update_job_status(pid, WaitStatus::Exited(pid, 0));
312 assert!(job_list.get_job(id).unwrap().is_completed());
313 }
314 }
315