proc_daemon/

shutdown.rs

//! Shutdown coordination system for graceful daemon termination.
//!
//! This module provides a robust shutdown coordination system that allows
//! all subsystems to be notified simultaneously and coordinate their
//! graceful shutdown within configurable timeouts.

use arc_swap::ArcSwap;
use parking_lot::Mutex;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::{debug, info, warn};

use crate::error::{Error, Result};

/// Reason for shutdown initiation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShutdownReason {
    /// Shutdown requested via signal (SIGTERM, SIGINT, etc.)
    Signal(i32),
    /// Shutdown requested programmatically
    Requested,
    /// Shutdown due to critical error
    Error,
    /// Shutdown due to resource exhaustion
    ResourceExhausted,
    /// Forced shutdown (timeout exceeded)
    Forced,
}

impl std::fmt::Display for ShutdownReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Signal(sig) => write!(f, "Signal({sig})"),
            Self::Requested => write!(f, "Requested"),
            Self::Error => write!(f, "Error"),
            Self::ResourceExhausted => write!(f, "ResourceExhausted"),
            Self::Forced => write!(f, "Forced"),
        }
    }
}

/// Handle for subsystems to monitor shutdown state and coordinate graceful termination.
#[derive(Debug, Clone)]
pub struct ShutdownHandle {
    inner: Arc<ShutdownInner>,
    subsystem_id: u64,
}

impl ShutdownHandle {
    /// Create a new shutdown handle for a specific subsystem.
    const fn new(inner: Arc<ShutdownInner>, subsystem_id: u64) -> Self {
        Self {
            inner,
            subsystem_id,
        }
    }

    /// Check if shutdown has been initiated.
    #[must_use]
    pub fn is_shutdown(&self) -> bool {
        self.inner.is_shutdown()
    }

    /// Wait for shutdown to be initiated.
    /// This is the primary method subsystems should use in their main loops.
    pub async fn cancelled(&mut self) {
        // Fast path: check with relaxed ordering first (common case: not shutdown)
        if self.inner.shutdown_initiated.load(Ordering::Relaxed) {
            return;
        }

        // Use tokio or async-std depending on feature flags
        #[cfg(feature = "tokio")]
        {
            let mut rx = self.inner.shutdown_tx.subscribe();
            if self.is_shutdown() {
                return;
            }
            let _ = rx.recv().await;
        }

        #[cfg(all(feature = "async-std", not(feature = "tokio")))]
        {
            // For async-std, we'll use a different approach
            let shutdown_flag = &self.inner.shutdown_initiated;
            loop {
                if shutdown_flag.load(Ordering::Acquire) {
                    break;
                }
                async_std::task::sleep(Duration::from_millis(10)).await;
            }
        }
    }

    /// Get the reason for shutdown (if initiated).
    #[must_use]
    pub fn shutdown_reason(&self) -> Option<ShutdownReason> {
        if self.is_shutdown() {
            Some(**self.inner.shutdown_reason.load())
        } else {
            None
        }
    }

    /// Get the time when shutdown was initiated.
    #[must_use]
    pub fn shutdown_time(&self) -> Option<Instant> {
        *self.inner.shutdown_time.lock()
    }

    /// Check if this is a forced shutdown.
    #[must_use]
    pub fn is_forced(&self) -> bool {
        matches!(self.shutdown_reason(), Some(ShutdownReason::Forced))
    }

    /// Mark this subsystem as ready for shutdown.
    /// This should be called when the subsystem has completed its cleanup.
    pub fn ready(&self) {
        self.inner.mark_subsystem_ready(self.subsystem_id);
    }

    /// Get the remaining time before forced shutdown.
    #[must_use]
    pub fn time_remaining(&self) -> Option<Duration> {
        self.shutdown_time().and_then(|shutdown_time| {
            let elapsed = shutdown_time.elapsed();
            let timeout =
                Duration::from_millis(self.inner.graceful_timeout_ms.load(Ordering::Acquire));

            if elapsed < timeout {
                timeout.checked_sub(elapsed)
            } else {
                None
            }
        })
    }
}

/// Internal shutdown state shared between coordinator and handles.
#[derive(Debug)]
struct ShutdownInner {
    /// Flag indicating shutdown has been initiated
    shutdown_initiated: AtomicBool,
    /// Reason for shutdown
    shutdown_reason: ArcSwap<ShutdownReason>,
    /// Time when shutdown was initiated
    shutdown_time: Mutex<Option<Instant>>,
    /// Graceful shutdown timeout in milliseconds
    graceful_timeout_ms: AtomicU64,
    /// Force shutdown timeout in milliseconds
    force_timeout_ms: AtomicU64,
    /// Kill timeout in milliseconds (Unix only)
    kill_timeout_ms: AtomicU64,
    /// Registered subsystems
    subsystems: Mutex<Vec<SubsystemState>>,
    /// Broadcast channel for shutdown notifications
    #[cfg(feature = "tokio")]
    shutdown_tx: tokio::sync::broadcast::Sender<ShutdownReason>,
}

/// State of a registered subsystem.
#[derive(Debug)]
struct SubsystemState {
    id: u64,
    name: String,
    ready: AtomicBool,
    #[allow(dead_code)]
    registered_at: Instant,
}

impl ShutdownInner {
    fn new(graceful_timeout_ms: u64, force_timeout_ms: u64, kill_timeout_ms: u64) -> Self {
        #[cfg(feature = "tokio")]
        let (shutdown_tx, _) = tokio::sync::broadcast::channel(16);

        Self {
            shutdown_initiated: AtomicBool::new(false),
            shutdown_reason: ArcSwap::new(Arc::new(ShutdownReason::Requested)),
            shutdown_time: Mutex::new(None),
            graceful_timeout_ms: AtomicU64::new(graceful_timeout_ms),
            force_timeout_ms: AtomicU64::new(force_timeout_ms),
            kill_timeout_ms: AtomicU64::new(kill_timeout_ms),
            subsystems: Mutex::new(Vec::new()),
            #[cfg(feature = "tokio")]
            shutdown_tx,
        }
    }

    /// Check if shutdown has been initiated.
    #[must_use]
    pub fn is_shutdown(&self) -> bool {
        // Use Relaxed ordering for fast reads (single atomic variable)
        self.shutdown_initiated.load(Ordering::Relaxed)
    }

    /// Initiate shutdown with the given reason.
    /// Returns true if this call initiated shutdown, false if shutdown was already in progress.
    #[must_use]
    pub fn initiate_shutdown(&self, reason: ShutdownReason) -> bool {
        // Use compare_exchange to ensure we only initiate once
        if self
            .shutdown_initiated
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_ok()
        {
            // Update shutdown reason and time
            self.shutdown_reason.store(Arc::new(reason));
            *self.shutdown_time.lock() = Some(Instant::now());

            // Notify all waiting tasks
            #[cfg(feature = "tokio")]
            {
                let _ = self.shutdown_tx.send(reason);
            }

            info!("Shutdown initiated: {}", reason);
            true
        } else {
            debug!("Shutdown already initiated, ignoring additional request");
            false
        }
    }

    fn register_subsystem(&self, name: &str) -> u64 {
        let id = fastrand::u64(..);
        let state = SubsystemState {
            id,
            name: name.to_string(),
            ready: AtomicBool::new(false),
            registered_at: Instant::now(),
        };

        self.subsystems.lock().push(state);
        debug!("Registered subsystem '{}' with ID {}", name, id);
        id
    }

    fn mark_subsystem_ready(&self, subsystem_id: u64) {
        let subsystems = self.subsystems.lock();
        // Use Relaxed ordering - only one thread marks each subsystem ready
        if let Some(subsystem) = subsystems.iter().find(|s| s.id == subsystem_id) {
            subsystem.ready.store(true, Ordering::Relaxed);
            debug!(
                "Subsystem '{}' marked as ready for shutdown",
                subsystem.name
            );
        }
    }

    fn are_all_subsystems_ready(&self) -> bool {
        let subsystems = self.subsystems.lock();
        // Early exit optimization: if any subsystem is not ready, return false immediately
        subsystems.iter().all(|s| s.ready.load(Ordering::Relaxed))
    }

    fn get_subsystem_states(&self) -> Vec<(String, bool)> {
        let subsystems = self.subsystems.lock();
        subsystems
            .iter()
            .map(|s| (s.name.clone(), s.ready.load(Ordering::Relaxed)))
            .collect()
    }
}

/// Shutdown coordinator that manages the graceful shutdown process.
#[derive(Debug)]
pub struct ShutdownCoordinator {
    inner: Arc<ShutdownInner>,
}

impl ShutdownCoordinator {
    /// Create a new shutdown coordinator.
    #[must_use]
    pub fn new(graceful_timeout_ms: u64, force_timeout_ms: u64, kill_timeout_ms: u64) -> Self {
        Self {
            inner: Arc::new(ShutdownInner::new(
                graceful_timeout_ms,
                force_timeout_ms,
                kill_timeout_ms,
            )),
        }
    }

    /// Create a shutdown handle for a subsystem.
    pub fn create_handle<S: Into<String>>(&self, subsystem_name: S) -> ShutdownHandle {
        let name = subsystem_name.into();
        let subsystem_id = self.inner.register_subsystem(&name);
        ShutdownHandle::new(Arc::clone(&self.inner), subsystem_id)
    }

    /// Initiate graceful shutdown.
    #[must_use]
    pub fn initiate_shutdown(&self, reason: ShutdownReason) -> bool {
        self.inner.initiate_shutdown(reason)
    }

    /// Check if shutdown has been initiated.
    #[must_use]
    pub fn is_shutdown(&self) -> bool {
        self.inner.is_shutdown()
    }

    /// Get the reason for shutdown, if any.
    #[must_use]
    pub fn get_reason(&self) -> Option<ShutdownReason> {
        if self.is_shutdown() {
            Some(**self.inner.shutdown_reason.load())
        } else {
            None
        }
    }

    /// Wait for all subsystems to complete graceful shutdown.
    /// Will return when either all subsystems are ready or the timeout is reached.
    ///
    /// # Errors
    ///
    /// Returns an `Error::timeout` if the graceful shutdown timeout is reached.
    pub async fn wait_for_shutdown(&self) -> Result<()> {
        if !self.is_shutdown() {
            return Err(Error::invalid_state("Shutdown not initiated"));
        }

        let shutdown_time = *self.inner.shutdown_time.lock();
        if shutdown_time.is_none() {
            return Err(Error::invalid_state("Shutdown time not set"));
        }

        let graceful_timeout =
            Duration::from_millis(self.inner.graceful_timeout_ms.load(Ordering::Acquire));

        info!(
            "Waiting for subsystems to shutdown gracefully (timeout: {:?})",
            graceful_timeout
        );

        // Wait for all subsystems to be ready or timeout
        let start = Instant::now();

        // Fast path: check if already complete
        if self.inner.are_all_subsystems_ready() {
            info!("All subsystems already shut down gracefully");
            return Ok(());
        }

        // Use exponential backoff polling with wakeup hints
        let mut poll_interval = Duration::from_millis(1);
        let max_poll_interval = Duration::from_millis(50);

        while start.elapsed() < graceful_timeout {
            if self.inner.are_all_subsystems_ready() {
                info!(
                    "All subsystems shut down gracefully in {:?}",
                    start.elapsed()
                );
                return Ok(());
            }

            // Exponential backoff to reduce CPU usage under load
            #[cfg(feature = "tokio")]
            tokio::time::sleep(poll_interval).await;

            #[cfg(all(feature = "async-std", not(feature = "tokio")))]
            async_std::task::sleep(poll_interval).await;

            // Exponential backoff: double interval up to max
            poll_interval = (poll_interval * 2).min(max_poll_interval);
        }

        // Timeout exceeded, log which subsystems are not ready
        let states = self.inner.get_subsystem_states();
        let not_ready: Vec<String> = states
            .into_iter()
            .filter_map(|(name, ready)| if ready { None } else { Some(name) })
            .collect();

        warn!(
            "Graceful shutdown timeout exceeded. Subsystems not ready: {:?}",
            not_ready
        );

        // Initiate forced shutdown
        let _ = self.inner.initiate_shutdown(ShutdownReason::Forced);

        let timeout_ms = u64::try_from(graceful_timeout.as_millis()).unwrap_or(u64::MAX);
        Err(Error::timeout("Graceful shutdown", timeout_ms))
    }

    /// Wait for forced shutdown after a timeout.
    /// This should be used as a fallback after `wait_for_shutdown`.
    ///
    /// # Errors
    ///
    /// Returns an `Error::timeout` if the force shutdown timeout is reached.
    pub async fn wait_for_force_shutdown(&self) -> Result<()> {
        let force_timeout =
            Duration::from_millis(self.inner.force_timeout_ms.load(Ordering::Acquire));

        warn!("Waiting for forced shutdown timeout: {:?}", force_timeout);

        let start = Instant::now();
        while start.elapsed() < force_timeout {
            if self.inner.are_all_subsystems_ready() {
                info!("All subsystems shut down during force phase");
                return Ok(());
            }

            #[cfg(feature = "tokio")]
            tokio::time::sleep(Duration::from_millis(50)).await;

            #[cfg(all(feature = "async-std", not(feature = "tokio")))]
            async_std::task::sleep(Duration::from_millis(50)).await;
        }

        let timeout_ms = u64::try_from(force_timeout.as_millis()).unwrap_or(u64::MAX);
        Err(Error::timeout("Force shutdown", timeout_ms))
    }

    /// Wait for kill shutdown timeout after force timeout expires.
    ///
    /// # Errors
    ///
    /// Returns an `Error::timeout` if the kill timeout is reached.
    pub async fn wait_for_kill_shutdown(&self) -> Result<()> {
        let kill_timeout =
            Duration::from_millis(self.inner.kill_timeout_ms.load(Ordering::Acquire));

        warn!("Waiting for kill shutdown timeout: {:?}", kill_timeout);

        #[cfg(feature = "tokio")]
        tokio::time::sleep(kill_timeout).await;

        #[cfg(all(feature = "async-std", not(feature = "tokio")))]
        async_std::task::sleep(kill_timeout).await;

        let timeout_ms = u64::try_from(kill_timeout.as_millis()).unwrap_or(u64::MAX);
        Err(Error::timeout("Kill shutdown", timeout_ms))
    }

    /// Get statistics about the shutdown process.
    #[must_use]
    pub fn get_stats(&self) -> ShutdownStats {
        let subsystems = self.inner.get_subsystem_states();
        let total_subsystems = subsystems.len();
        let ready_subsystems = subsystems.iter().filter(|(_, ready)| *ready).count();

        ShutdownStats {
            is_shutdown: self.is_shutdown(),
            reason: if self.is_shutdown() {
                Some(**self.inner.shutdown_reason.load())
            } else {
                None
            },
            shutdown_time: *self.inner.shutdown_time.lock(),
            total_subsystems,
            ready_subsystems,
            subsystem_states: subsystems,
        }
    }

    /// Update timeout configurations at runtime.
    pub fn update_timeouts(
        &self,
        graceful_timeout_ms: u64,
        force_timeout_ms: u64,
        kill_timeout_ms: u64,
    ) {
        self.inner
            .graceful_timeout_ms
            .store(graceful_timeout_ms, Ordering::Release);
        self.inner
            .force_timeout_ms
            .store(force_timeout_ms, Ordering::Release);
        self.inner
            .kill_timeout_ms
            .store(kill_timeout_ms, Ordering::Release);
        debug!(
            "Updated shutdown timeouts: graceful={}ms, force={}ms, kill={}ms",
            graceful_timeout_ms, force_timeout_ms, kill_timeout_ms
        );
    }
}

impl Clone for ShutdownCoordinator {
    fn clone(&self) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
        }
    }
}

/// Statistics about the shutdown process.
#[derive(Debug, Clone)]
pub struct ShutdownStats {
    /// Whether shutdown has been initiated
    pub is_shutdown: bool,
    /// Reason for shutdown
    pub reason: Option<ShutdownReason>,
    /// Time when shutdown was initiated
    pub shutdown_time: Option<Instant>,
    /// Total number of registered subsystems
    pub total_subsystems: usize,
    /// Number of subsystems ready for shutdown
    pub ready_subsystems: usize,
    /// Individual subsystem states
    pub subsystem_states: Vec<(String, bool)>,
}

impl ShutdownStats {
    /// Get the shutdown progress as a percentage (0.0 to 1.0).
    #[must_use]
    #[allow(clippy::cast_precision_loss)]
    pub fn progress(&self) -> f64 {
        // Return a value between 0.0 and 1.0 representing the progress
        if self.total_subsystems == 0 {
            1.0
        } else {
            // Use f64 to prevent loss of precision
            self.ready_subsystems as f64 / self.total_subsystems as f64
        }
    }

    /// Check if all subsystems are ready.
    #[must_use]
    pub const fn is_complete(&self) -> bool {
        self.total_subsystems > 0 && self.ready_subsystems == self.total_subsystems
    }

    /// Get the elapsed time since shutdown was initiated.
    #[must_use]
    pub fn elapsed(&self) -> Option<Duration> {
        self.shutdown_time.map(|t| t.elapsed())
    }
}

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

    #[cfg(feature = "tokio")]
    #[cfg_attr(miri, ignore)]
    #[tokio::test]
    async fn test_shutdown_coordination() {
        // Add a test timeout to prevent freezing
        let test_result = tokio::time::timeout(Duration::from_secs(5), async {
            // Use shorter timeouts for testing
            let coordinator = ShutdownCoordinator::new(100, 200, 300);

            // Create handles for subsystems
            let handle1 = coordinator.create_handle("subsystem1");
            let handle2 = coordinator.create_handle("subsystem2");

            // Initially not shutdown
            assert!(!coordinator.is_shutdown());
            assert!(!handle1.is_shutdown());

            // Initiate shutdown
            assert!(coordinator.initiate_shutdown(ShutdownReason::Requested));

            // Should be shutdown now
            assert!(coordinator.is_shutdown());
            assert!(handle1.is_shutdown());

            // Instead of trying to listen for the cancelled() notification,
            // we'll just verify that the handle is properly marked as shutdown
            assert!(handle1.is_shutdown());
            assert!(handle2.is_shutdown());

            // Mark subsystems as ready
            handle1.ready();
            handle2.ready();

            // All should be ready now
            let stats = coordinator.get_stats();
            assert!(stats.is_complete());
            // Use a more reasonable epsilon for floating point comparisons
            let epsilon: f64 = 1e-6;
            assert!((stats.progress() - 1.0).abs() < epsilon);
        })
        .await;

        assert!(test_result.is_ok(), "Test timed out after 5 seconds");
    }

    #[cfg(feature = "tokio")]
    #[cfg_attr(miri, ignore)]
    #[tokio::test]
    async fn test_shutdown_timeout() {
        // Add a test timeout to prevent the test itself from hanging
        let test_result = tokio::time::timeout(Duration::from_secs(5), async {
            let coordinator = ShutdownCoordinator::new(100, 200, 300); // Very short timeout

            let _handle1 = coordinator.create_handle("slow_subsystem");

            // Initiate shutdown but don't mark as ready
            let _ = coordinator.initiate_shutdown(ShutdownReason::Requested);

            // Wait for shutdown should timeout
            let result = coordinator.wait_for_shutdown().await;
            assert!(result.is_err());
            assert!(result.unwrap_err().is_timeout());
        })
        .await;

        assert!(test_result.is_ok(), "Test timed out after 5 seconds");
    }

    #[cfg(all(feature = "async-std", not(feature = "tokio")))]
    #[async_std::test]
    async fn test_shutdown_timeout() {
        // Add a test timeout to prevent the test itself from hanging
        let test_result = async_std::future::timeout(Duration::from_secs(5), async {
            let coordinator = ShutdownCoordinator::new(100, 200, 300); // Very short timeout

            let _handle1 = coordinator.create_handle("slow_subsystem");

            // Initiate shutdown but don't mark as ready
            let _ = coordinator.initiate_shutdown(ShutdownReason::Requested);

            // Wait for shutdown should timeout
            let result = coordinator.wait_for_shutdown().await;
            assert!(result.is_err());
            assert!(result.unwrap_err().is_timeout());
        })
        .await;

        assert!(test_result.is_ok(), "Test timed out after 5 seconds");
    }

    #[test]
    fn test_shutdown_reason_display() {
        assert_eq!(format!("{}", ShutdownReason::Signal(15)), "Signal(15)");
        assert_eq!(format!("{}", ShutdownReason::Requested), "Requested");
        assert_eq!(format!("{}", ShutdownReason::Error), "Error");
    }

    #[cfg(feature = "tokio")]
    #[cfg_attr(miri, ignore)]
    #[tokio::test]
    async fn test_multiple_shutdown_initiation() {
        // Add a test timeout to prevent freezing
        let test_result = tokio::time::timeout(Duration::from_secs(5), async {
            let coordinator = ShutdownCoordinator::new(5000, 10000, 15000);

            // First initiation should succeed
            assert!(coordinator.initiate_shutdown(ShutdownReason::Requested));

            // Subsequent initiations should be ignored
            assert!(!coordinator.initiate_shutdown(ShutdownReason::Signal(15)));
            assert!(!coordinator.initiate_shutdown(ShutdownReason::Error));

            // Reason should remain the first one
            assert_eq!(coordinator.get_reason(), Some(ShutdownReason::Requested));
        })
        .await;

        assert!(test_result.is_ok(), "Test timed out after 5 seconds");
    }

    #[cfg(all(feature = "async-std", not(feature = "tokio")))]
    #[async_std::test]
    async fn test_multiple_shutdown_initiation() {
        // Add a test timeout to prevent freezing
        let test_result = async_std::future::timeout(Duration::from_secs(5), async {
            let coordinator = ShutdownCoordinator::new(5000, 10000, 15000);

            // First initiation should succeed
            assert!(coordinator.initiate_shutdown(ShutdownReason::Requested));

            // Subsequent initiations should be ignored
            assert!(!coordinator.initiate_shutdown(ShutdownReason::Signal(15)));
            assert!(!coordinator.initiate_shutdown(ShutdownReason::Error));

            // Reason should remain the first one
            let stats = coordinator.get_stats();
            assert_eq!(stats.reason, Some(ShutdownReason::Requested));
        })
        .await;

        assert!(test_result.is_ok(), "Test timed out after 5 seconds");
    }

    #[test]
    fn test_shutdown_stats() {
        let coordinator = ShutdownCoordinator::new(5000, 10000, 15000);
        let handle1 = coordinator.create_handle("test1");
        let handle2 = coordinator.create_handle("test2");

        let stats = coordinator.get_stats();
        assert_eq!(stats.total_subsystems, 2);
        assert_eq!(stats.ready_subsystems, 0);
        assert!(!stats.is_complete());

        // Use a more reasonable epsilon for floating point comparisons
        let epsilon: f64 = 1e-6;
        assert!((stats.progress() - 0.0).abs() < epsilon);

        handle1.ready();
        let stats = coordinator.get_stats();
        assert_eq!(stats.ready_subsystems, 1);

        assert!((stats.progress() - 0.5).abs() < epsilon);

        handle2.ready();
        let stats = coordinator.get_stats();
        assert!(stats.is_complete());

        assert!((stats.progress() - 1.0).abs() < epsilon);
    }
}

#[cfg(all(feature = "async-std", not(feature = "tokio")))]
#[async_std::test]
async fn test_shutdown_coordination() {
    // Add a test timeout to prevent freezing
    let test_result = async_std::future::timeout(Duration::from_secs(5), async {
        // Use shorter timeouts for testing
        let coordinator = ShutdownCoordinator::new(100, 200, 300);

        // Create handles for subsystems
        let handle1 = coordinator.create_handle("subsystem1");
        let handle2 = coordinator.create_handle("subsystem2");

        // Initially not shutdown
        assert!(!coordinator.is_shutdown());
        assert!(!handle1.is_shutdown());

        // Initiate shutdown
        assert!(coordinator.initiate_shutdown(ShutdownReason::Requested));

        // Should be shutdown now
        assert!(coordinator.is_shutdown());
        assert!(handle1.is_shutdown());

        // Instead of trying to listen for the cancelled() notification,
        // we'll just verify that the handle is properly marked as shutdown
        assert!(handle1.is_shutdown());
        assert!(handle2.is_shutdown());

        // Mark subsystems as ready
        handle1.ready();
        handle2.ready();

        // All should be ready now
        let stats = coordinator.get_stats();
        assert!(stats.is_complete());
        // Use a more reasonable epsilon for floating point comparisons
        let epsilon: f64 = 1e-6;
        assert!((stats.progress() - 1.0).abs() < epsilon);
    })
    .await;

    assert!(test_result.is_ok(), "Test timed out after 5 seconds");
}