diff --git a/src/executor/docker.rs b/src/executor/docker.rs index be0389a..fb7c1af 100644 --- a/src/executor/docker.rs +++ b/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>> = 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, diff --git a/src/executor/mod.rs b/src/executor/mod.rs index d6f39ac..43b9d56 100644 --- a/src/executor/mod.rs +++ b/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; diff --git a/src/main.rs b/src/main.rs index 26b4daa..2a3491e 100644 --- a/src/main.rs +++ b/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 }) => {