agent_kernel/

shutdown.rs

//! Graceful shutdown coordination for agent lifecycle.
//!
//! This module provides a state machine for coordinating graceful shutdown of agents,
//! ensuring in-flight work completes before termination.

use std::sync::Arc;
use std::sync::atomic::{AtomicU8, AtomicU32, Ordering};
use std::time::Duration;

use thiserror::Error;
use tokio::sync::broadcast;
use tokio::time::timeout;
use tracing::{debug, warn};

/// Shutdown state enumeration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShutdownState {
    /// Agent is running normally.
    Running = 0,
    /// Agent is draining in-flight work.
    Draining = 1,
    /// Agent has terminated.
    Terminated = 2,
}

impl ShutdownState {
    /// Convert from atomic representation.
    fn from_u8(value: u8) -> Self {
        match value {
            0 => ShutdownState::Running,
            1 => ShutdownState::Draining,
            _ => ShutdownState::Terminated,
        }
    }

    /// Convert to atomic representation.
    fn as_u8(self) -> u8 {
        self as u8
    }
}

/// Errors from shutdown operations.
#[derive(Debug, Error)]
pub enum ShutdownError {
    /// Shutdown was already initiated.
    #[error("shutdown already initiated")]
    AlreadyShuttingDown,

    /// Drain timeout expired before all work completed.
    #[error("drain timeout expired with {remaining} tasks still in flight")]
    DrainTimeout {
        /// Number of tasks still in flight when timeout expired.
        remaining: u32,
    },

    /// Broadcast channel error.
    #[error("broadcast error: {0}")]
    BroadcastError(String),
}

/// Result type for shutdown operations.
pub type ShutdownResult<T> = Result<T, ShutdownError>;

/// Coordinates graceful shutdown of agent workloads.
///
/// The shutdown coordinator manages a state machine with three states:
/// - `Running`: Normal operation, new work can be registered
/// - `Draining`: Shutdown initiated, no new work accepted, existing work completes
/// - `Terminated`: All work completed or timeout expired
///
/// # Example
///
/// ```ignore
/// let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_secs(30)));
/// let guard = coordinator.register_work();
/// // ... do work ...
/// drop(guard); // Work completes
///
/// coordinator.shutdown().await.unwrap();
/// assert_eq!(coordinator.state(), ShutdownState::Terminated);
/// ```
#[derive(Debug)]
pub struct ShutdownCoordinator {
    state: AtomicU8,
    drain_timeout: Duration,
    in_flight: AtomicU32,
    shutdown_signal: broadcast::Sender<()>,
}

impl ShutdownCoordinator {
    /// Creates a new shutdown coordinator with the specified drain timeout.
    ///
    /// # Arguments
    ///
    /// * `drain_timeout` - Maximum duration to wait for in-flight work to complete
    #[must_use]
    pub fn new(drain_timeout: Duration) -> Self {
        let (tx, _) = broadcast::channel(1);
        Self {
            state: AtomicU8::new(ShutdownState::Running.as_u8()),
            drain_timeout,
            in_flight: AtomicU32::new(0),
            shutdown_signal: tx,
        }
    }

    /// Registers a unit of in-flight work.
    ///
    /// Returns a `WorkGuard` that automatically decrements the in-flight counter
    /// when dropped. If shutdown is already in progress, returns an error.
    ///
    /// # Errors
    ///
    /// Returns `ShutdownError::AlreadyShuttingDown` if shutdown has been initiated.
    pub fn register_work(self: &Arc<Self>) -> ShutdownResult<WorkGuard> {
        let current_state = ShutdownState::from_u8(self.state.load(Ordering::SeqCst));

        if current_state != ShutdownState::Running {
            return Err(ShutdownError::AlreadyShuttingDown);
        }

        self.in_flight.fetch_add(1, Ordering::SeqCst);
        Ok(WorkGuard {
            coordinator: Arc::clone(self),
        })
    }

    /// Initiates graceful shutdown.
    ///
    /// Transitions the coordinator to `Draining` state and waits for in-flight work
    /// to complete up to the configured drain timeout. If the timeout expires,
    /// returns an error but still transitions to `Terminated`.
    ///
    /// # Errors
    ///
    /// Returns `ShutdownError::DrainTimeout` if in-flight work doesn't complete
    /// within the drain timeout.
    pub async fn shutdown(&self) -> ShutdownResult<()> {
        // Transition to Draining
        let old_state = self.state.compare_exchange(
            ShutdownState::Running.as_u8(),
            ShutdownState::Draining.as_u8(),
            Ordering::SeqCst,
            Ordering::SeqCst,
        );

        if old_state.is_err() {
            return Err(ShutdownError::AlreadyShuttingDown);
        }

        debug!("shutdown initiated, draining in-flight work");

        // Broadcast shutdown signal
        let _ = self.shutdown_signal.send(());

        // Wait for in-flight work to complete or timeout
        let result = timeout(self.drain_timeout, self.wait_for_drain()).await;

        // Transition to Terminated
        self.state
            .store(ShutdownState::Terminated.as_u8(), Ordering::SeqCst);

        if result.is_ok() {
            debug!("shutdown complete, all work drained");
            Ok(())
        } else {
            let remaining = self.in_flight.load(Ordering::SeqCst);
            warn!(remaining, "drain timeout expired");
            Err(ShutdownError::DrainTimeout { remaining })
        }
    }

    /// Returns the current shutdown state.
    #[must_use]
    pub fn state(&self) -> ShutdownState {
        ShutdownState::from_u8(self.state.load(Ordering::SeqCst))
    }

    /// Subscribes to shutdown signal.
    ///
    /// Returns a receiver that will be notified when shutdown is initiated.
    #[must_use]
    pub fn subscribe(&self) -> broadcast::Receiver<()> {
        self.shutdown_signal.subscribe()
    }

    /// Returns the number of in-flight work items.
    #[must_use]
    pub fn in_flight_count(&self) -> u32 {
        self.in_flight.load(Ordering::SeqCst)
    }

    /// Waits for all in-flight work to complete.
    async fn wait_for_drain(&self) {
        loop {
            if self.in_flight.load(Ordering::SeqCst) == 0 {
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    }
}

/// RAII guard for tracking in-flight work.
///
/// When dropped, automatically decrements the in-flight work counter.
pub struct WorkGuard {
    coordinator: Arc<ShutdownCoordinator>,
}

impl Drop for WorkGuard {
    fn drop(&mut self) {
        self.coordinator.in_flight.fetch_sub(1, Ordering::SeqCst);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn initial_state_is_running() {
        let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_secs(1)));
        assert_eq!(coordinator.state(), ShutdownState::Running);
        assert_eq!(coordinator.in_flight_count(), 0);
    }

    #[test]
    fn register_work_increments_counter() {
        let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_secs(1)));
        let _guard = coordinator.register_work().unwrap();
        assert_eq!(coordinator.in_flight_count(), 1);
    }

    #[test]
    fn work_guard_drop_decrements_counter() {
        let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_secs(1)));
        {
            let _guard = coordinator.register_work().unwrap();
            assert_eq!(coordinator.in_flight_count(), 1);
        }
        assert_eq!(coordinator.in_flight_count(), 0);
    }

    #[test]
    fn register_work_fails_during_shutdown() {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_secs(1)));
            let _guard = coordinator.register_work().unwrap();

            let coordinator_clone = Arc::clone(&coordinator);
            tokio::spawn(async move {
                let _ = coordinator_clone.shutdown().await;
            });

            tokio::time::sleep(Duration::from_millis(50)).await;
            let result = coordinator.register_work();
            assert!(result.is_err());
        });
    }

    #[tokio::test]
    async fn shutdown_waits_for_work_to_complete() {
        let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_secs(5)));
        let guard = coordinator.register_work().unwrap();

        let coordinator_clone = Arc::clone(&coordinator);
        let shutdown_task = tokio::spawn(async move { coordinator_clone.shutdown().await });

        tokio::time::sleep(Duration::from_millis(100)).await;
        drop(guard);

        let result = shutdown_task.await.unwrap();
        assert!(result.is_ok());
        assert_eq!(coordinator.state(), ShutdownState::Terminated);
    }

    #[tokio::test]
    async fn shutdown_timeout_with_pending_work() {
        let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_millis(100)));
        let _guard = coordinator.register_work().unwrap();

        let result = coordinator.shutdown().await;
        assert!(result.is_err());
        assert_eq!(coordinator.state(), ShutdownState::Terminated);
    }

    #[tokio::test]
    async fn shutdown_signal_broadcast() {
        let coordinator = Arc::new(ShutdownCoordinator::new(Duration::from_secs(1)));
        let mut rx = coordinator.subscribe();

        let coordinator_clone = Arc::clone(&coordinator);
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(50)).await;
            let _ = coordinator_clone.shutdown().await;
        });

        let result = tokio::time::timeout(Duration::from_secs(1), rx.recv()).await;
        assert!(result.is_ok());
    }
}