hypercall/shared/

task_group.rs

//! Task group for managing background tasks with graceful shutdown.
//!
//! This module provides a `TaskGroup` that tracks spawned tasks and allows
//! for coordinated shutdown with timeout handling.
//!
//! # Example
//! ```ignore
//! use crate::shared::task_group::TaskGroup;
//! use crate::shared::shutdown::Shutdown;
//! use std::time::Duration;
//!
//! let shutdown = Shutdown::new();
//! let mut group = TaskGroup::new();
//!
//! // Spawn a task that respects shutdown
//! let mut shutdown_rx = shutdown.subscribe();
//! group.spawn("MyTask", async move {
//!     loop {
//!         tokio::select! {
//!             _ = shutdown_rx.recv() => break,
//!             // ... do work
//!         }
//!     }
//!     Ok(())
//! });
//!
//! // Later, shutdown all tasks
//! group.shutdown_and_join(&shutdown, Duration::from_secs(5)).await?;
//! ```

use anyhow::{anyhow, Result};
use std::future::Future;
use std::time::Duration;
use tokio::task::JoinHandle;
use tracing::{error, info, warn};

use super::shutdown::Shutdown;

/// A handle to a spawned background task.
struct TaskHandle {
    name: &'static str,
    handle: JoinHandle<Result<()>>,
}

/// A group of background tasks that can be shut down together.
///
/// Use `TaskGroup` to spawn related background tasks and ensure they all
/// shut down gracefully with a timeout.
pub struct TaskGroup {
    tasks: Vec<TaskHandle>,
}

impl Default for TaskGroup {
    fn default() -> Self {
        Self::new()
    }
}

impl TaskGroup {
    /// Create a new empty task group.
    pub fn new() -> Self {
        Self { tasks: Vec::new() }
    }

    /// Spawn a new task and add it to the group.
    ///
    /// The task should respect shutdown signals (e.g., via `tokio::select!`
    /// on a shutdown receiver) to allow graceful termination.
    ///
    /// # Arguments
    /// * `name` - A static name for the task (used in error messages)
    /// * `fut` - The future to spawn
    pub fn spawn<F>(&mut self, name: &'static str, fut: F)
    where
        F: Future<Output = Result<()>> + Send + 'static,
    {
        let handle = tokio::spawn(fut);
        self.tasks.push(TaskHandle { name, handle });
    }

    /// Wait for all tasks to complete with a timeout.
    ///
    /// Returns `Ok(())` if all tasks complete successfully within the timeout.
    /// Returns `Err` if any task times out, panics, or returns an error.
    ///
    /// Note: Tasks are joined sequentially. Each task gets the full timeout duration,
    /// so the total wait time may exceed the timeout if multiple tasks are slow.
    /// Also note that timed-out tasks continue running in the background - the timeout
    /// means "gave up waiting" not "task was stopped".
    pub async fn join_all(self, timeout: Duration) -> Result<()> {
        let mut errors = Vec::new();
        let deadline = tokio::time::Instant::now() + timeout;

        for task in self.tasks {
            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
            match tokio::time::timeout(remaining, task.handle).await {
                Ok(Ok(Ok(()))) => {
                    info!("Task '{}' completed successfully", task.name);
                }
                Ok(Ok(Err(e))) => {
                    error!("Task '{}' returned error: {}", task.name, e);
                    errors.push(format!("{}: {}", task.name, e));
                }
                Ok(Err(join_error)) => {
                    if join_error.is_panic() {
                        error!("Task '{}' panicked", task.name);
                        errors.push(format!("{}: panicked", task.name));
                    } else if join_error.is_cancelled() {
                        warn!("Task '{}' was cancelled", task.name);
                        // Cancelled tasks are not treated as errors
                    }
                }
                Err(_) => {
                    error!("Task '{}' timed out after {:?}", task.name, timeout);
                    errors.push(format!("{}: timed out", task.name));
                }
            }
        }

        if errors.is_empty() {
            Ok(())
        } else {
            Err(anyhow!("Task group shutdown failed: {}", errors.join(", ")))
        }
    }

    /// Trigger shutdown and wait for all tasks to complete.
    ///
    /// This is a convenience method that calls `shutdown.trigger()` and then
    /// waits for all tasks to complete with the given timeout.
    pub async fn shutdown_and_join(self, shutdown: &Shutdown, timeout: Duration) -> Result<()> {
        info!("Triggering shutdown for task group");
        shutdown.trigger();
        self.join_all(timeout).await
    }

    /// Get the number of tasks in the group.
    pub fn len(&self) -> usize {
        self.tasks.len()
    }

    /// Check if the group is empty.
    pub fn is_empty(&self) -> bool {
        self.tasks.is_empty()
    }
}

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

    #[tokio::test]
    async fn test_task_group_basic() {
        let mut group = TaskGroup::new();
        let completed = Arc::new(AtomicBool::new(false));
        let completed_clone = completed.clone();

        group.spawn("TestTask", async move {
            completed_clone.store(true, Ordering::SeqCst);
            Ok(())
        });

        let result = group.join_all(Duration::from_secs(1)).await;
        assert!(result.is_ok());
        assert!(completed.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_task_group_with_shutdown() {
        let shutdown = Shutdown::new();
        let mut group = TaskGroup::new();

        // Task that waits for shutdown
        let mut shutdown_rx = shutdown.subscribe();
        group.spawn("ShutdownAwareTask", async move {
            let _ = shutdown_rx.recv().await;
            Ok(())
        });

        // Spawn a task that triggers shutdown after a short delay
        let shutdown_clone = shutdown.clone();
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(10)).await;
            shutdown_clone.trigger();
        });

        let result = group.join_all(Duration::from_secs(1)).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_shutdown_and_join() {
        let shutdown = Shutdown::new();
        let mut group = TaskGroup::new();

        let mut shutdown_rx = shutdown.subscribe();
        group.spawn("WaitForShutdown", async move {
            let _ = shutdown_rx.recv().await;
            Ok(())
        });

        // This should trigger and wait
        let result = group
            .shutdown_and_join(&shutdown, Duration::from_secs(1))
            .await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_task_error_propagation() {
        let mut group = TaskGroup::new();

        group.spawn(
            "FailingTask",
            async move { Err(anyhow!("intentional failure")) },
        );

        let result = group.join_all(Duration::from_secs(1)).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("FailingTask"));
    }

    #[tokio::test]
    async fn test_multiple_tasks() {
        let shutdown = Shutdown::new();
        let mut group = TaskGroup::new();

        // Spawn multiple tasks
        for i in 0..3 {
            let mut shutdown_rx = shutdown.subscribe();
            let name: &'static str = match i {
                0 => "Task0",
                1 => "Task1",
                2 => "Task2",
                _ => unreachable!(),
            };
            group.spawn(name, async move {
                let _ = shutdown_rx.recv().await;
                Ok(())
            });
        }

        assert_eq!(group.len(), 3);

        let result = group
            .shutdown_and_join(&shutdown, Duration::from_secs(1))
            .await;
        assert!(result.is_ok());
    }
}