Last active
September 12, 2024 19:23
-
-
Save coolaj86/43e2c1d74a2849934b6ad87a28e60126 to your computer and use it in GitHub Desktop.
POSIX.1-2024 timeout implemented in Rust
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| // generated by GPT4o using the command summary in this README as the prompt: | |
| // https://github.com/posix-utilities/timeout | |
| use std::process::{Command, ExitStatus}; | |
| use std::time::Duration; | |
| use std::os::unix::process::CommandExt; | |
| use std::sync::{Arc, atomic::{AtomicBool, Ordering}}; | |
| use std::thread; | |
| use nix::sys::signal::{self, Signal}; | |
| use std::env; | |
| use std::ffi::OsString; | |
| fn parse_duration(s: &str) -> Result<Duration, &'static str> { | |
| if s.ends_with('s') { | |
| s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v)).map_err(|_| "Invalid seconds") | |
| } else if s.ends_with('m') { | |
| s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v * 60.0)).map_err(|_| "Invalid minutes") | |
| } else if s.ends_with('h') { | |
| s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v * 3600.0)).map_err(|_| "Invalid hours") | |
| } else if s.ends_with('d') { | |
| s[..s.len()-1].parse::<f64>().map(|v| Duration::from_secs_f64(v * 86400.0)).map_err(|_| "Invalid days") | |
| } else { | |
| s.parse::<f64>().map(|v| Duration::from_secs_f64(v)).map_err(|_| "Invalid duration") | |
| } | |
| } | |
| fn main() { | |
| let args: Vec<OsString> = env::args_os().collect(); | |
| if args.len() < 3 { | |
| eprintln!("Usage: timeout [options] <duration> <command> [arguments...]"); | |
| std::process::exit(125); | |
| } | |
| let mut duration: Option<Duration> = None; | |
| let mut kill_after: Option<Duration> = None; | |
| let mut signal = Signal::SIGTERM; | |
| let mut preserve_exit_status = false; | |
| let mut for_process_only = false; | |
| let mut command_idx = None; | |
| let mut i = 1; | |
| while i < args.len() { | |
| let arg = &args[i]; | |
| match arg.to_str() { | |
| Some("-f") => for_process_only = true, | |
| Some("-k") => { | |
| i += 1; | |
| if i < args.len() { | |
| kill_after = Some(parse_duration(&args[i].to_string_lossy()).unwrap()); | |
| } | |
| } | |
| Some("-p") => preserve_exit_status = true, | |
| Some("-s") => { | |
| i += 1; | |
| if i < args.len() { | |
| signal = Signal::from_str(&args[i].to_string_lossy()).unwrap_or(Signal::SIGTERM); | |
| } | |
| } | |
| Some(s) if duration.is_none() => { | |
| duration = Some(parse_duration(s).unwrap()); | |
| } | |
| _ => { | |
| command_idx = Some(i); | |
| break; | |
| } | |
| } | |
| i += 1; | |
| } | |
| let duration = duration.unwrap_or_else(|| { | |
| eprintln!("Duration must be specified"); | |
| std::process::exit(125); | |
| }); | |
| let command_idx = command_idx.unwrap_or_else(|| { | |
| eprintln!("Command must be specified"); | |
| std::process::exit(125); | |
| }); | |
| let command = &args[command_idx]; | |
| let command_args = &args[command_idx + 1..]; | |
| let child = Command::new(command) | |
| .args(command_args) | |
| .spawn() | |
| .expect("Failed to spawn process"); | |
| let pid = child.id() as i32; | |
| let timeout_occurred = Arc::new(AtomicBool::new(false)); | |
| let timeout_occurred_clone = Arc::clone(&timeout_occurred); | |
| let handle = thread::spawn(move || { | |
| thread::sleep(duration); | |
| timeout_occurred_clone.store(true, Ordering::SeqCst); | |
| if for_process_only { | |
| signal::kill(pid, signal).unwrap(); | |
| } else { | |
| signal::killpg(pid, signal).unwrap(); | |
| } | |
| }); | |
| let status = child.wait().expect("Failed to wait on child"); | |
| if timeout_occurred.load(Ordering::SeqCst) { | |
| if let Some(kill_duration) = kill_after { | |
| thread::sleep(kill_duration); | |
| signal::kill(pid, Signal::SIGKILL).unwrap(); | |
| } | |
| } | |
| handle.join().unwrap(); | |
| let exit_status_code = if preserve_exit_status { | |
| match status.code() { | |
| Some(code) => code, | |
| None => 124, // Default timeout status | |
| } | |
| } else if timeout_occurred.load(Ordering::SeqCst) { | |
| 124 | |
| } else { | |
| status.code().unwrap_or(125) | |
| }; | |
| std::process::exit(exit_status_code); | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment