Handle signals and clean up subprocesses

CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.24.0] - 2024-04-05
+
+### Fixed
+- Docuum now cleans up child processes when exiting due to a signal (`SIGHUP`, `SIGINT`, or `SIGTERM`).
+
 ## [0.23.1] - 2023-10-02
 
 ### Added

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "docuum"
-version = "0.23.1"
+version = "0.24.0"
 authors = ["Stephan Boyer <stephan@stephanboyer.com>"]
 edition = "2021"
 description = "LRU eviction of Docker images."
@@ -24,11 +24,11 @@ colored = "2"
 dirs = "3"
 env_logger = { version = "0.8", default-features = false, features = ["termcolor", "atty"] }
 log = "0.4"
-scopeguard = "1"
+regex = { version = "1.5.5", default-features = false, features = ["std", "unicode-perl"] }
 serde_json = "1.0"
 serde_yaml = "0.8"
+signal-hook = "0.3"
 tempfile = "3"
-regex = { version = "1.5.5", default-features = false, features = ["std", "unicode-perl"] }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 sysinfo = "0.23.5"

src/main.rs

@@ -11,12 +11,17 @@ use {
     env_logger::{fmt::Color, Builder},
     log::{Level, LevelFilter},
     regex::RegexSet,
+    signal_hook::{
+        consts::{SIGHUP, SIGINT, SIGTERM},
+        iterator::Signals,
+    },
     std::{
         env,
         io::{self, Write},
         process::exit,
         str::FromStr,
-        thread::sleep,
+        sync::{Arc, Mutex},
+        thread::{self, sleep},
         time::Duration,
     },
 };
@@ -231,8 +236,37 @@ fn settings() -> io::Result<Settings> {
     })
 }
 
+// This function consumes and runs all the registered destructors.
+#[allow(clippy::type_complexity)]
+fn run_destructors(destructors: &Arc<Mutex<Vec<Box<dyn FnOnce() + Send>>>>) {
+    let mut mutex_guard = destructors.lock().unwrap();
+    let destructor_fns = std::mem::take(&mut *mutex_guard);
+    for destructor in destructor_fns {
+        destructor();
+    }
+}
+
 // Let the fun begin!
 fn main() {
+    // React to signals in a separate thread.
+    let destructors = Arc::new(Mutex::new(Vec::<Box<dyn FnOnce() + Send>>::new()));
+    let destructors_clone = destructors.clone();
+    match Signals::new([SIGHUP, SIGINT, SIGTERM]) {
+        Ok(mut signals) => {
+            thread::spawn(move || {
+                if let Some(signal) = signals.forever().next() {
+                    println!("Received signal {signal:?}.");
+                    run_destructors(&destructors_clone);
+                    exit(1);
+                }
+            });
+        }
+        Err(error) => {
+            // Log the error and proceed anyway.
+            error!("{}", error);
+        }
+    }
+
     // Determine whether to print colored output.
     colored::control::set_override(atty::is(Stream::Stderr));
 
@@ -265,10 +299,13 @@ fn main() {
 
     // Stream Docker events and vacuum when necessary. Restart if an error occurs.
     loop {
-        if let Err(e) = run(&settings, &mut state, &mut first_run) {
-            error!("{}", e);
-            info!("Retrying in 5 seconds\u{2026}");
-            sleep(Duration::from_secs(5));
+        if let Err(error) = run(&settings, &mut state, &mut first_run, &destructors) {
+            error!("{}", error);
         }
+
+        run_destructors(&destructors);
+
+        info!("Retrying in 5 seconds\u{2026}");
+        sleep(Duration::from_secs(5));
     }
 }

src/run.rs

@@ -7,15 +7,14 @@ use {
     byte_unit::Byte,
     chrono::DateTime,
     regex::RegexSet,
-    scopeguard::guard,
     serde::{Deserialize, Serialize},
     std::{
         cmp::max,
         collections::{hash_map::Entry, HashMap, HashSet},
         io::{self, BufRead, BufReader},
-        mem::drop,
         ops::Deref,
         process::{Command, Stdio},
+        sync::{Arc, Mutex},
         time::{Duration, SystemTime, UNIX_EPOCH},
     },
 };
@@ -731,7 +730,13 @@ fn vacuum(
 }
 
 // Stream Docker events and vacuum when necessary.
-pub fn run(settings: &Settings, state: &mut State, first_run: &mut bool) -> io::Result<()> {
+#[allow(clippy::type_complexity)]
+pub fn run(
+    settings: &Settings,
+    state: &mut State,
+    first_run: &mut bool,
+    destructors: &Arc<Mutex<Vec<Box<dyn FnOnce() + Send>>>>,
+) -> io::Result<()> {
     // Determine the threshold in bytes.
     let threshold = match settings.threshold {
         Threshold::Absolute(b) => b,
@@ -764,27 +769,22 @@ pub fn run(settings: &Settings, state: &mut State, first_run: &mut bool) -> io::
     *first_run = false;
 
     // Spawn `docker events --format '{{json .}}'`.
-    let mut child = guard(
-        Command::new("docker")
-            .args(["events", "--format", "{{json .}}"])
-            .stdout(Stdio::piped())
-            .spawn()?,
-        |mut child| {
-            drop(child.kill());
-            drop(child.wait());
-        },
-    );
+    let mut child = Command::new("docker")
+        .args(["events", "--format", "{{json .}}"])
+        .stdout(Stdio::piped()) // [tag:stdout]
+        .spawn()?;
 
-    // Buffer the data as we read it line-by-line.
-    let reader = BufReader::new(child.stdout.as_mut().map_or_else(
-        || {
-            Err(io::Error::new(
-                io::ErrorKind::Other,
-                format!("Unable to read output from {}.", "docker events".code_str()),
-            ))
-        },
-        Ok,
-    )?);
+    // Buffer the data as we read it line-by-line. The `unwrap` is safe due to [ref:stdout].
+    let reader = BufReader::new(child.stdout.take().unwrap());
+
+    // When this run is done (either due to an error or a signal), terminate the child process.
+    destructors.lock().unwrap().push(Box::new(move || {
+        if let Err(error) = child.kill() {
+            error!("{}", error);
+        } else if let Err(error) = child.wait() {
+            error!("{}", error);
+        }
+    }));
 
     // Handle each incoming event.
     info!("Listening for Docker events\u{2026}");
@@ -854,7 +854,7 @@ pub fn run(settings: &Settings, state: &mut State, first_run: &mut bool) -> io::
     // The `for` loop above will only terminate if something happened to `docker events`.
     Err(io::Error::new(
         io::ErrorKind::Other,
-        format!("{} unexpectedly terminated.", "docker events".code_str()),
+        format!("{} terminated.", "docker events".code_str()),
     ))
 }