Thomas 9 months ago
parent
commit
82e6e06117
1 changed files with 41 additions and 31 deletions
  1. 41 31
      src/runners.rs

+ 41 - 31
src/runners.rs

@@ -9,6 +9,7 @@ use std::{
     thread,
 };
 
+use anyhow::Context;
 use bgzip::{BGZFWriter, Compression};
 use chrono::{DateTime, Utc};
 use log::{info, warn};
@@ -17,11 +18,26 @@ use uuid::Uuid;
 
 use crate::{config::Config, DOCKER_ID};
 
+/// Trait for running a command.
 pub trait Run {
+    /// Runs the command.
+    ///
+    /// This method spawns the child process and sets up threads to capture
+    /// its stdout and stderr output.
+    ///
+    /// # Returns
+    /// `Ok(())` if the command was successfully started, otherwise an error
     fn run(&mut self) -> anyhow::Result<()>;
 }
 
+/// Trait for waiting on a command to complete.
 pub trait Wait {
+    /// Waits for the command to complete.
+    ///
+    /// This method monitors the command's output and waits for the process to exit.
+    ///
+    /// # Returns
+    /// `Ok(())` if the command completed successfully, otherwise an error
     fn wait(&mut self) -> anyhow::Result<()>;
 }
 
@@ -62,18 +78,28 @@ impl RunReport {
 }
 
 pub trait Log {
+    /// Returns the accumulated log of the command's execution.
+    ///
+    /// # Returns
+    /// A string containing the command's execution log
     fn log(&self) -> String;
 }
 
+/// Represents a Docker command to be run, with facilities for execution, monitoring, and logging.
 #[derive(Debug, Default)]
 pub struct DockerRun {
+    /// Arguments for the Docker command
     pub args: Vec<String>,
+    /// Container ID of the running Docker instance
     pub container_id: String,
+    /// Start time of the Docker command execution
     pub start: DateTime<Utc>,
+    /// Accumulated logs of the Docker command execution
     pub logs: Arc<Mutex<String>>,
 }
 
 impl DockerRun {
+    /// Helper function to execute a Docker command and capture its output.
     pub fn new(args: &[&str]) -> Self {
         DockerRun {
             args: args.iter().map(|e| e.to_string()).collect(),
@@ -85,8 +111,8 @@ impl DockerRun {
 
 impl Run for DockerRun {
     fn run(&mut self) -> anyhow::Result<()> {
-        // Setup Ctrl-C handler cant be defined two times
-        let _ = ctrlc::try_set_handler(move || {
+        // Sets up a Ctrl-C handler to stop Docker containers on interrupt.
+        ctrlc::try_set_handler(move || {
             if let Ok(container_id) = DOCKER_ID.lock() {
                 for id in container_id.iter() {
                     warn!("Stopping Docker container {id}...");
@@ -94,14 +120,14 @@ impl Run for DockerRun {
                 }
             }
             std::process::exit(1);
-        });
+        }).context("Failed to set Ctrl-C handler")?;
 
-        // limit memory
+        // Configures memory limits for the Docker container.
         let c = Config::default();
         self.args
             .insert(1, format!("--memory={}g", c.docker_max_memory_go));
 
-        // Spawn the main command
+        // Spawns the Docker command and captures its output.
         let output = Command::new("docker")
             .args(&self.args)
             .stdout(Stdio::piped())
@@ -229,13 +255,6 @@ impl CommandRun {
 }
 
 impl Run for CommandRun {
-    /// Runs the command.
-    ///
-    /// This method spawns the child process and sets up threads to capture
-    /// its stdout and stderr output.
-    ///
-    /// # Returns
-    /// `Ok(())` if the command was successfully started, otherwise an error
     fn run(&mut self) -> anyhow::Result<()> {
         let info = format!("Running command: {} {}", &self.bin, &self.args.join(" "));
         info!("{info}");
@@ -254,10 +273,12 @@ impl Run for CommandRun {
         let stderr_reader = BufReader::new(stderr);
 
         let tx = self.tx.clone();
+        let bin = self.bin.clone();
+
         std::thread::scope(|s| {
             s.spawn(|| {
                 for line in stdout_reader.lines().map_while(Result::ok) {
-                    info!("[{}:stdout] {}", self.bin, line);
+                    info!("[{}:stdout] {}", &bin, line);
                     tx.send(("stdout".to_string(), line.to_string()))
                         .expect("Channel send failed");
                 }
@@ -265,7 +286,7 @@ impl Run for CommandRun {
 
             s.spawn(|| {
                 for line in stderr_reader.lines().map_while(Result::ok) {
-                    info!("[{}:stderr] {}", self.bin, line);
+                    info!("[{}:stderr] {}", &bin, line);
                     tx.send(("stderr".to_string(), line.to_string()))
                         .expect("Channel send failed");
                 }
@@ -278,12 +299,6 @@ impl Run for CommandRun {
 }
 
 impl Wait for CommandRun {
-    /// Waits for the command to complete.
-    ///
-    /// This method monitors the command's output and waits for the process to exit.
-    ///
-    /// # Returns
-    /// `Ok(())` if the command completed successfully, otherwise an error
     fn wait(&mut self) -> anyhow::Result<()> {
         if let Some(child) = &mut self.child {
             loop {
@@ -292,14 +307,13 @@ impl Wait for CommandRun {
                         self.log
                             .push_str(&format!("[{}] {}: {}\n", self.bin, stream, line));
                     }
-                    Err(TryRecvError::Empty) => match child.try_wait()? {
-                        Some(_) => {
+                    Err(TryRecvError::Empty) => {
+                        if let Some(status) = child.try_wait()? {
+                            info!("{} exited with status: {}", self.bin, status);
                             break;
                         }
-                        None => {
-                            std::thread::sleep(std::time::Duration::from_millis(100));
-                        }
-                    },
+                        std::thread::sleep(std::time::Duration::from_millis(100));
+                    }
                     Err(TryRecvError::Disconnected) => {
                         break;
                     }
@@ -308,17 +322,13 @@ impl Wait for CommandRun {
 
             // Ensure child process has exited
             let status = child.wait()?;
-            info!("{} {}", self.bin, status);
+            info!("{} exited with status: {}", self.bin, status);
         }
         Ok(())
     }
 }
 
 impl Log for CommandRun {
-    /// Returns the accumulated log of the command's execution.
-    ///
-    /// # Returns
-    /// A string containing the command's execution log
     fn log(&self) -> String {
         self.log.clone()
     }