signal handling

src/executor/docker.rs

@@ -7,7 +7,11 @@ use bollard::{
     Docker,
 };
 use futures_util::StreamExt;
+use once_cell::sync::Lazy;
 use std::path::Path;
+use std::sync::Mutex;
+
+static RUNNING_CONTAINERS: Lazy<Mutex<Vec<String>>> = Lazy::new(|| Mutex::new(Vec::new()));
 
 pub struct DockerRuntime {
     docker: Docker,
@@ -37,6 +41,37 @@ pub fn is_available() -> bool {
     }
 }
 
+// Add container to tracking
+fn track_container(id: &str) {
+    if let Ok(mut containers) = RUNNING_CONTAINERS.lock() {
+        containers.push(id.to_string());
+    }
+}
+
+// Remove container from tracking
+fn untrack_container(id: &str) {
+    if let Ok(mut containers) = RUNNING_CONTAINERS.lock() {
+        containers.retain(|c| c != id);
+    }
+}
+
+// Clean up all tracked containers
+pub async fn cleanup_containers(docker: &Docker) {
+    let containers_to_cleanup = {
+        if let Ok(containers) = RUNNING_CONTAINERS.lock() {
+            containers.clone()
+        } else {
+            vec![]
+        }
+    };
+
+    for container_id in containers_to_cleanup {
+        let _ = docker.stop_container(&container_id, None).await;
+        let _ = docker.remove_container(&container_id, None).await;
+        untrack_container(&container_id);
+    }
+}
+
 #[async_trait]
 impl ContainerRuntime for DockerRuntime {
     async fn run_container(
@@ -105,6 +140,8 @@ impl ContainerRuntime for DockerRuntime {
             .await
             .map_err(|e| ContainerError::ContainerExecutionFailed(e.to_string()))?;
 
+        track_container(&container.id);
+
         // Wait for container to finish
         let wait_result = self
             .docker
@@ -143,6 +180,7 @@ impl ContainerRuntime for DockerRuntime {
 
         // Clean up container
         let _ = self.docker.remove_container(&container.id, None).await;
+        untrack_container(&container.id);
 
         Ok(ContainerOutput {
             stdout,

src/executor/mod.rs

@@ -7,3 +7,5 @@ mod environment;
 pub use engine::{
     execute_workflow, ExecutionResult, JobResult, JobStatus, RuntimeType, StepResult, StepStatus,
 };
+
+pub use docker::cleanup_containers;

src/main.rs

@@ -8,6 +8,7 @@ mod ui;
 mod utils;
 mod validators;
 
+use bollard::Docker;
 use clap::{Parser, Subcommand};
 use std::path::PathBuf;
 
@@ -55,10 +56,33 @@ enum Commands {
     },
 }
 
+async fn cleanup_on_exit() {
+    match Docker::connect_with_local_defaults() {
+        Ok(docker) => {
+            executor::cleanup_containers(&docker).await;
+        }
+        Err(_) => {
+            // Docker not available, nothing to clean up
+        }
+    }
+}
+
+async fn handle_signals() {
+    // This creates a future that completes when CTRL+C is received
+    tokio::signal::ctrl_c()
+        .await
+        .expect("Failed to listen for ctrl+c event");
+
+    println!("Received Ctrl+C, shutting down...");
+    cleanup_on_exit().await;
+    std::process::exit(0);
+}
+
 #[tokio::main]
 async fn main() {
     let cli = Wrkflw::parse();
     let verbose = cli.verbose;
+    tokio::spawn(handle_signals());
 
     match &cli.command {
         Some(Commands::Validate { path }) => {