Merge pull request #16 from appsignal/fix/graceful-shutdown-and-mount-leak

guifranchi · web-flow · commit 5cacc0f2650c · 2026-07-16T17:05:28.000-03:00
Graceful shutdown on SIGTERM + fix rootfs mount leak that wedges slots
diff --git a/manager/src/instance.rs b/manager/src/instance.rs
@@ -188,32 +188,43 @@ impl Instance {
         let rootfs_lv = format!("rootfs-{}-{}", self.role, self.idx);
         let base_lv = format!("base-{}", self.role);
 
+        // A mount leaked by a previous cycle (error between mount and unmount,
+        // or the whole manager getting SIGKILLed) keeps the rootfs LV open, so
+        // the lvremove/lvcreate below would fail on every retry — wedging this
+        // slot permanently. Clean it up first.
+        self.cleanup_stale_mount();
+
         // Remove stale snapshot and create fresh one from base
         let _ = lvm::lvremove(&self.lvm.volume_group, &rootfs_lv);
         lvm::lvcreate_snapshot(&self.lvm.volume_group, &base_lv, &rootfs_lv)?;
 
+        // Fetch the registration token before mounting: it is a network call
+        // and must not be able to fail while the rootfs is mounted.
+        let token = self.github.registration_token()?;
+        let runner_name = self.name();
+        let labels = self.labels();
+
         // Mount rootfs, inject config, unmount
         let mnt = Utf8PathBuf::from(self.work_dir.join("mnt").as_str());
         fs::create_dir_all(&mnt)?;
 
         let rootfs_device = self.rootfs_device_path();
         mount::mount_image(&rootfs_device, &mnt)?;
 
-        let token = self.github.registration_token()?;
-        let runner_name = self.name();
-        let labels = self.labels();
-
-        inject::inject_config(
+        let inject_result = inject::inject_config(
             &mnt,
             &self.github.org.clone(),
             &token,
             &runner_name,
             &labels,
             &self.network_allocation,
-        )?;
+        );
 
-        mount::unmount(&mnt)?;
+        // Always unmount, even when injection failed — see cleanup_stale_mount.
+        let unmount_result = mount::unmount(&mnt);
         let _ = fs::remove_dir(&mnt);
+        inject_result?;
+        unmount_result?;
 
         // Write Firecracker config.json
         let config_json = serde_json::to_string(&self.firecracker_config()?)?;
@@ -253,6 +264,21 @@ impl Instance {
         // Nothing to reset in new model; slot loop handles retry
     }
 
+    /// Unmount and remove this slot's `mnt` directory if a previous cycle
+    /// leaked it (error while mounted, or the manager was SIGKILLed). Mounts
+    /// are kernel state and survive process death; a held mount makes every
+    /// subsequent `lvremove`/`lvcreate_snapshot` for this slot fail with
+    /// "contains a filesystem in use".
+    pub fn cleanup_stale_mount(&self) {
+        let mnt = self.work_dir.join("mnt");
+        if mnt.as_std_path().exists() {
+            if mount::unmount(&mnt).is_ok() {
+                info!(role = %self.role, idx = self.idx, "unmounted stale rootfs mount");
+            }
+            let _ = fs::remove_dir(mnt.as_std_path());
+        }
+    }
+
     #[instrument(skip(self), fields(role = %self.role, idx = self.idx))]
     pub fn try_clear_cache(&mut self) -> Result<()> {
         match self.cache.usage_pct() {
diff --git a/manager/src/lib.rs b/manager/src/lib.rs
@@ -9,6 +9,7 @@ pub mod inject;
 pub mod instance;
 pub mod lvm;
 pub mod network;
+pub mod shutdown;
 pub mod slot;
 
 use network::Forwarding;
@@ -38,6 +39,10 @@ impl Manager {
             lvm::reconcile(&self.config)?;
         }
 
+        // Install the SIGTERM/SIGINT handler and VM-killing watcher before any
+        // slot exists, so an early stop is handled gracefully too.
+        shutdown::install()?;
+
         let config = Arc::new(self.config.clone());
         let mut handles = Vec::new();
 
@@ -63,30 +68,15 @@ impl Manager {
 
         info!(slots = handles.len(), "all slot threads started");
 
-        // Install SIGINT/SIGTERM handler
-        setup_signal_handler()?;
-
-        // Join all slot threads (they run forever unless killed)
+        // Join all slot threads. They run until a shutdown signal is received,
+        // at which point each slot tears down (unmount, PID file) and exits.
         for handle in handles {
             if let Err(e) = handle.join() {
                 warn!("slot thread panicked: {:?}", e);
             }
         }
 
-        info!("manager shutting down");
+        info!("all slots stopped, manager shutting down");
         Ok(())
     }
 }
-
-fn setup_signal_handler() -> Result<()> {
-    use nix::sys::signal::{sigaction, SaFlags, SigAction, SigHandler, SigSet, Signal};
-
-    // Install a no-op handler for SIGTERM and SIGINT so the process can shut down
-    // gracefully. Slot processes receive the signal via their own process group.
-    let action = SigAction::new(SigHandler::SigIgn, SaFlags::empty(), SigSet::empty());
-    unsafe {
-        sigaction(Signal::SIGTERM, &action)?;
-        sigaction(Signal::SIGINT, &action)?;
-    }
-    Ok(())
-}
diff --git a/manager/src/shutdown.rs b/manager/src/shutdown.rs
@@ -0,0 +1,84 @@
+//! Graceful shutdown coordination.
+//!
+//! SIGTERM/SIGINT set a global flag; a watcher thread then SIGKILLs every
+//! registered Firecracker process group so blocked `child.wait()` calls in the
+//! slot threads return. The slot loops observe the flag, run their teardown
+//! (unmount, PID file removal) and exit, letting the manager exit 0 instead of
+//! being SIGKILLed by systemd after `TimeoutStopSec`.
+
+use std::collections::BTreeMap;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Mutex;
+use std::time::Duration;
+use tracing::{info, warn};
+
+static SHUTDOWN: AtomicBool = AtomicBool::new(false);
+static CHILDREN: Mutex<BTreeMap<(String, usize), u32>> = Mutex::new(BTreeMap::new());
+
+/// Whether a shutdown signal has been received.
+pub fn requested() -> bool {
+    SHUTDOWN.load(Ordering::SeqCst)
+}
+
+extern "C" fn handle_signal(_: libc::c_int) {
+    // Only async-signal-safe work here: set the flag, nothing else.
+    SHUTDOWN.store(true, Ordering::SeqCst);
+}
+
+/// Install the SIGTERM/SIGINT handler and spawn the watcher thread that kills
+/// running Firecracker VMs once shutdown is requested. Call before spawning
+/// slot threads.
+pub fn install() -> anyhow::Result<()> {
+    use nix::sys::signal::{sigaction, SaFlags, SigAction, SigHandler, SigSet, Signal};
+
+    let action = SigAction::new(
+        SigHandler::Handler(handle_signal),
+        SaFlags::SA_RESTART,
+        SigSet::empty(),
+    );
+    unsafe {
+        sigaction(Signal::SIGTERM, &action)?;
+        sigaction(Signal::SIGINT, &action)?;
+    }
+
+    std::thread::Builder::new()
+        .name("shutdown-watcher".to_string())
+        .spawn(|| {
+            while !requested() {
+                std::thread::sleep(Duration::from_millis(250));
+            }
+            info!("shutdown requested, stopping running VMs");
+            kill_children();
+        })?;
+
+    Ok(())
+}
+
+/// Record a running Firecracker child for a slot so the watcher can stop it.
+pub fn register_child(role: &str, idx: usize, pid: u32) {
+    CHILDREN
+        .lock()
+        .unwrap()
+        .insert((role.to_string(), idx), pid);
+    // The watcher may have already swept before this child was registered.
+    if requested() {
+        kill_children();
+    }
+}
+
+/// Remove a slot's child after it has exited.
+pub fn unregister_child(role: &str, idx: usize) {
+    CHILDREN.lock().unwrap().remove(&(role.to_string(), idx));
+}
+
+fn kill_children() {
+    let children = CHILDREN.lock().unwrap();
+    for ((role, idx), pid) in children.iter() {
+        info!(role = %role, idx = *idx, pid = *pid, "killing Firecracker process group");
+        // Firecracker setsid()s into its own process group; kill the group.
+        let ret = unsafe { libc::kill(-(*pid as libc::pid_t), libc::SIGKILL) };
+        if ret != 0 {
+            warn!(role = %role, idx = *idx, pid = *pid, "kill failed (process already gone?)");
+        }
+    }
+}
diff --git a/manager/src/slot.rs b/manager/src/slot.rs
@@ -7,7 +7,7 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tracing::{error, info, instrument, warn};
 
-use crate::{instance::Instance, network::NetworkAllocation};
+use crate::{instance::Instance, network::NetworkAllocation, shutdown};
 
 /// Run the per-slot loop for a given role and slot index.
 /// This function is intended to run in its own OS process or thread.
@@ -42,6 +42,11 @@ pub fn run_slot(config: Arc<ManagerConfig>, role: &Role, idx: usize) -> Result<(
     let mut cycle_n: u64 = 0;
 
     loop {
+        if shutdown::requested() {
+            info!(role = %role.name, idx, "shutdown requested, exiting boot loop");
+            break;
+        }
+
         let pid_file = instance.pid_file_path();
 
         // Graceful restart: if Firecracker is already running for this slot, reattach
@@ -53,7 +58,10 @@ pub fn run_slot(config: Arc<ManagerConfig>, role: &Role, idx: usize) -> Result<(
                     pid,
                     "Firecracker already running, reattaching"
                 );
+                // Register so the shutdown watcher can stop the reattached VM too.
+                shutdown::register_child(&role.slug(), idx, pid);
                 wait_for_pid(pid);
+                shutdown::unregister_child(&role.slug(), idx);
                 remove_pid_file(&pid_file);
                 // Fall through to next cycle after the job finishes
                 continue;
@@ -71,6 +79,7 @@ pub fn run_slot(config: Arc<ManagerConfig>, role: &Role, idx: usize) -> Result<(
         match instance.start() {
             Ok(mut child) => {
                 let pid = child.id();
+                shutdown::register_child(&role.slug(), idx, pid);
                 if let Err(e) = write_pid_file(&pid_file, pid) {
                     error!(role = %role.name, idx, pid, error = %e, "failed to write PID file");
                 }
@@ -154,11 +163,18 @@ pub fn run_slot(config: Arc<ManagerConfig>, role: &Role, idx: usize) -> Result<(
                     }
                 }
 
+                shutdown::unregister_child(&role.slug(), idx);
                 remove_pid_file(&pid_file);
             }
             Err(e) => {
                 error!(role = %role.name, idx, cycle = cycle_n, error = %e, "boot cycle failed");
-                std::thread::sleep(Duration::from_secs(20));
+                // Interruptible retry back-off so shutdown isn't delayed by up to 20s.
+                for _ in 0..20 {
+                    if shutdown::requested() {
+                        break;
+                    }
+                    std::thread::sleep(Duration::from_secs(1));
+                }
                 instance.reset();
                 continue;
             }
@@ -169,6 +185,12 @@ pub fn run_slot(config: Arc<ManagerConfig>, role: &Role, idx: usize) -> Result<(
             error!(role = %role.name, idx, error = %e, "cache clear failed");
         }
     }
+
+    // Shutdown path: make sure nothing from the last cycle is left behind that
+    // would wedge the slot on the next manager start.
+    instance.cleanup_stale_mount();
+    info!(role = %role.name, idx, "slot shut down cleanly");
+    Ok(())
 }
 
 fn read_child_stderr(child: &mut std::process::Child) -> String {
