Lumotia/crates/hotkey/src/supervisor.rs

//! Task supervisor for the evdev hotkey listener.
//!
//! Owns `JoinHandle`s for every task the listener spawns (scanner, hotplug
//! watcher, hotplug retry tasks, per-device listeners). Provides a
//! broadcast shutdown channel that every cooperating task subscribes to.
//!
//! Three concurrency leaks this fixes:
//! - **Race-1**: per-device listener tasks had no cancellation path, so
//!   after `stop()` returned they kept running and emitting events.
//! - **Race-2**: the forwarder task in `commands::hotkey` had no
//!   `JoinHandle` tracking, so a reconfigure leaked a permanent forwarder
//!   that received duplicated events.
//! - **Race-extra (TOCTOU)**: the `tracked` `HashSet` could miss-attach or
//!   double-attach because removal happened after the listener task
//!   exited, leaving a window where a concurrent hotplug saw stale state.
//!   The new `HashMap<PathBuf, ()>` keyed by canonical path, coupled with
//!   insert-before-spawn under one mutex hold, makes attachment atomic.
//!
//! See `tests/listener_lifecycle.rs` for regression coverage.

use std::sync::Arc;
use std::time::Duration;

use tokio::sync::{broadcast, Mutex};
use tokio::task::JoinHandle;
use tokio::time::timeout;

/// How long to wait for any single task to drain on `shutdown()` before
/// we give up on it. Two seconds is generous for cooperative shutdown
/// via the broadcast channel — anything slower is treated as a stuck
/// task and detached (NOT aborted: `timeout(d, handle).await` consumes
/// the `JoinHandle` by value and dropping a `JoinHandle` detaches the
/// task, so the task keeps running until the tokio runtime tears it
/// down). A warning is logged so the operator can investigate.
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);

/// Capacity of the broadcast shutdown channel. Eight is far more than we
/// expect to ever need (one slot per concurrent subscriber, drained
/// immediately), but cheap enough that a tighter bound buys us nothing.
const SHUTDOWN_CHANNEL_CAPACITY: usize = 8;

/// Inner state of the supervisor, behind a mutex so concurrent spawn
/// sites can register handles.
struct SupervisorInner {
    handles: Vec<(&'static str, JoinHandle<()>)>,
}

/// Shareable handle to a task supervisor. Hand a clone of this to every
/// task that needs to register child tasks (e.g. the hotplug watcher
/// needs to register retry tasks it spawns). The actual shutdown is
/// driven by [`SupervisorHandle::shutdown`], which is called once by the
/// listener's `stop()`.
#[derive(Clone)]
pub(crate) struct SupervisorHandle {
    shutdown_tx: broadcast::Sender<()>,
    inner: Arc<Mutex<SupervisorInner>>,
}

impl SupervisorHandle {
    pub(crate) fn new() -> Self {
        let (shutdown_tx, _) = broadcast::channel(SHUTDOWN_CHANNEL_CAPACITY);
        Self {
            shutdown_tx,
            inner: Arc::new(Mutex::new(SupervisorInner {
                handles: Vec::new(),
            })),
        }
    }

    /// Subscribe to the shutdown signal. Tasks should `tokio::select!`
    /// this receiver alongside their work so they exit promptly on
    /// `shutdown()`.
    pub(crate) fn subscribe(&self) -> broadcast::Receiver<()> {
        self.shutdown_tx.subscribe()
    }

    /// Register a spawned task. The `label` is logged only when the task
    /// has to be force-aborted, so concise tags like `"scanner"` are
    /// sufficient.
    pub(crate) async fn register(&self, label: &'static str, handle: JoinHandle<()>) {
        self.inner.lock().await.handles.push((label, handle));
    }

    /// Current number of registered tasks. Useful for logging.
    pub(crate) async fn task_count(&self) -> usize {
        self.inner.lock().await.handles.len()
    }

    /// Fire the shutdown signal WITHOUT awaiting any tasks. Used by
    /// `Drop` for paranoia — async drop is not available in stable Rust,
    /// so we cannot join handles here.
    pub(crate) fn signal_shutdown_nonblocking(&self) {
        let _ = self.shutdown_tx.send(());
    }

    /// Signal every subscriber to shut down, then await each registered
    /// task with a per-task timeout. Any task that doesn't drain inside
    /// the timeout is detached and logged via `tracing::warn!`.
    pub(crate) async fn shutdown(&self) {
        // `send` only errors when there are no live receivers, which is a
        // perfectly fine state — every subscriber already exited or none
        // was ever attached.
        let _ = self.shutdown_tx.send(());

        // Drain handles. We hold the lock only long enough to swap the
        // Vec out so tasks racing to register late don't block our wait.
        let handles: Vec<(&'static str, JoinHandle<()>)> = {
            let mut guard = self.inner.lock().await;
            std::mem::take(&mut guard.handles)
        };

        let task_count = handles.len();

        for (label, handle) in handles {
            match timeout(SHUTDOWN_TIMEOUT, handle).await {
                Ok(Ok(())) => {
                    // Clean exit.
                }
                Ok(Err(join_err)) => {
                    tracing::warn!(
                        target: "lumotia_hotkey",
                        task = label,
                        error = %join_err,
                        "supervised task ended abnormally"
                    );
                }
                Err(_elapsed) => {
                    // Timed out — task is stuck. We can't await again
                    // after the timeout future consumed the handle, so
                    // log and detach. Every task in this crate selects
                    // on the shutdown broadcast and sees senders drop,
                    // so reaching this branch indicates a genuine bug.
                    tracing::warn!(
                        target: "lumotia_hotkey",
                        task = label,
                        timeout_secs = SHUTDOWN_TIMEOUT.as_secs(),
                        "supervised task did not drain within timeout; detaching"
                    );
                }
            }
        }

        // Drain anything registered while we were awaiting (a late
        // hotplug retry, for instance). These tasks have already seen
        // the broadcast and should be on their way out.
        let late: Vec<(&'static str, JoinHandle<()>)> = {
            let mut guard = self.inner.lock().await;
            std::mem::take(&mut guard.handles)
        };
        let late_count = late.len();
        for (label, handle) in late {
            match timeout(SHUTDOWN_TIMEOUT, handle).await {
                Ok(Ok(())) => {}
                Ok(Err(e)) => {
                    tracing::warn!(
                        target: "lumotia_hotkey",
                        task = label,
                        error = %e,
                        "late-registered task ended abnormally"
                    );
                }
                Err(_) => {
                    tracing::warn!(
                        target: "lumotia_hotkey",
                        task = label,
                        "late-registered task did not drain"
                    );
                }
            }
        }

        tracing::info!(
            target: "lumotia_hotkey",
            task_count = task_count + late_count,
            "supervisor stopped"
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn shutdown_joins_all_registered_tasks() {
        let sup = SupervisorHandle::new();
        let counter = Arc::new(AtomicUsize::new(0));

        for i in 0..5 {
            let mut rx = sup.subscribe();
            let counter = counter.clone();
            let handle = tokio::spawn(async move {
                let _ = rx.recv().await;
                counter.fetch_add(1, Ordering::SeqCst);
                tracing::debug!(task_id = i, "task exiting");
            });
            sup.register("test-task", handle).await;
        }

        assert_eq!(sup.task_count().await, 5);

        sup.shutdown().await;

        assert_eq!(
            counter.load(Ordering::SeqCst),
            5,
            "every registered task should observe shutdown and run its exit path"
        );
    }

    /// A stuck task — one that does not subscribe to broadcast shutdown
    /// and would otherwise run forever — must not block `shutdown()`
    /// past the per-task timeout. Note: the supervisor does NOT abort
    /// the task; it detaches it. Verifying detach behaviour directly is
    /// not possible from this test because `register()` moves the
    /// `JoinHandle` into the supervisor's inner Vec. The bounded-elapsed
    /// assertion below is what guards against a regression that
    /// reintroduces an unbounded `handle.await` in `shutdown()`.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn shutdown_does_not_block_on_stuck_tasks_after_timeout() {
        let sup = SupervisorHandle::new();
        let handle = tokio::spawn(async move {
            // Sleep forever — does NOT subscribe to shutdown.
            tokio::time::sleep(Duration::from_secs(3600)).await;
        });
        sup.register("stuck", handle).await;

        let start = std::time::Instant::now();
        sup.shutdown().await;
        let elapsed = start.elapsed();

        assert!(
            elapsed < Duration::from_secs(4),
            "shutdown should not block past timeout * 2, took {elapsed:?}"
        );
    }
}