hypercall/shared/

task_leak.rs

//! Task leak detection utilities for testing.
//!
//! This module provides a mechanism to verify that all spawned tasks have exited
//! after a test completes, similar to [goleak](https://github.com/uber-go/goleak)
//! for Go goroutines.
//!
//! # Example
//! ```ignore
//! use crate::shared::task_leak::TaskTracker;
//! use crate::shared::shutdown::Shutdown;
//!
//! #[tokio::test]
//! async fn test_no_task_leaks() {
//!     let shutdown = Shutdown::new();
//!     let tracker = TaskTracker::new();
//!
//!     // Spawn tasks using the tracker
//!     let mut shutdown_rx = shutdown.subscribe();
//!     tracker.spawn("MyTask", async move {
//!         let _ = shutdown_rx.recv().await;
//!         Ok(())
//!     });
//!
//!     // Trigger shutdown
//!     shutdown.trigger();
//!
//!     // Verify all tasks complete within timeout
//!     tracker.verify_no_leaks(Duration::from_secs(5)).await.unwrap();
//! }
//! ```

use anyhow::{anyhow, Result};
use std::future::Future;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, Mutex};
use tokio::task::JoinHandle;
use tracing::{error, info, warn};

/// Tracks spawned tasks and detects leaks.
///
/// Similar to goleak, this helps detect tasks that don't properly
/// respond to shutdown signals and continue running indefinitely.
pub struct TaskTracker {
    tasks: Arc<Mutex<Vec<TrackedTask>>>,
    task_tx: mpsc::UnboundedSender<TrackedTask>,
    spawn_count: AtomicUsize,
    complete_count: AtomicUsize,
}

struct TrackedTask {
    name: &'static str,
    handle: JoinHandle<Result<()>>,
}

impl TaskTracker {
    /// Create a new task tracker wrapped in an Arc.
    pub fn new() -> Arc<Self> {
        let (task_tx, mut task_rx) = mpsc::unbounded_channel();
        let tasks = Arc::new(Mutex::new(Vec::new()));

        // Background task to process queued tasks
        let tasks_for_bg = tasks.clone();
        tokio::spawn(async move {
            while let Some(task) = task_rx.recv().await {
                tasks_for_bg.lock().await.push(task);
            }
        });

        Arc::new(Self {
            tasks,
            task_tx,
            spawn_count: AtomicUsize::new(0),
            complete_count: AtomicUsize::new(0),
        })
    }

    /// Spawn a task and track it.
    ///
    /// The task should respect shutdown signals to allow proper cleanup.
    pub fn spawn<F>(self: &Arc<Self>, name: &'static str, fut: F)
    where
        F: Future<Output = Result<()>> + Send + 'static,
    {
        self.spawn_count.fetch_add(1, Ordering::SeqCst);
        let tracker = Arc::clone(self);

        let wrapped = async move {
            let result = fut.await;
            tracker.complete_count.fetch_add(1, Ordering::SeqCst);
            result
        };

        let handle = tokio::spawn(wrapped);
        let task = TrackedTask { name, handle };

        // Queue task for tracking (non-blocking)
        let _ = self.task_tx.send(task);
    }

    /// Get the number of tasks spawned.
    pub fn spawned_count(&self) -> usize {
        self.spawn_count.load(Ordering::SeqCst)
    }

    /// Get the number of tasks that have completed.
    pub fn completed_count(&self) -> usize {
        self.complete_count.load(Ordering::SeqCst)
    }

    /// Verify that all tracked tasks complete within the timeout.
    ///
    /// Returns `Ok(())` if all tasks complete successfully.
    /// Returns `Err` with details about any leaked (timed out) or failed tasks.
    pub async fn verify_no_leaks(self: Arc<Self>, timeout: Duration) -> Result<()> {
        // Give background task a moment to process any queued tasks
        tokio::time::sleep(Duration::from_millis(10)).await;
        let mut tasks = self.tasks.lock().await;
        let total = tasks.len();

        if total == 0 {
            return Ok(());
        }

        info!("Verifying {} tracked tasks complete...", total);

        let mut leaked = Vec::new();
        let mut failed = Vec::new();
        let mut completed = 0;

        for task in tasks.drain(..) {
            match tokio::time::timeout(timeout, task.handle).await {
                Ok(Ok(Ok(()))) => {
                    completed += 1;
                }
                Ok(Ok(Err(e))) => {
                    error!("Task '{}' returned error: {}", task.name, e);
                    failed.push(format!("{}: {}", task.name, e));
                }
                Ok(Err(join_error)) => {
                    if join_error.is_panic() {
                        error!("Task '{}' panicked", task.name);
                        failed.push(format!("{}: panicked", task.name));
                    } else if join_error.is_cancelled() {
                        warn!("Task '{}' was cancelled", task.name);
                        // Cancelled is not necessarily a leak
                        completed += 1;
                    }
                }
                Err(_) => {
                    error!(
                        "Task '{}' leaked (timed out after {:?})",
                        task.name, timeout
                    );
                    leaked.push(task.name);
                }
            }
        }

        info!(
            "Task verification: {} completed, {} leaked, {} failed",
            completed,
            leaked.len(),
            failed.len()
        );

        if !leaked.is_empty() || !failed.is_empty() {
            let mut msg = String::new();
            if !leaked.is_empty() {
                msg.push_str(&format!("Leaked tasks: {:?}. ", leaked));
            }
            if !failed.is_empty() {
                msg.push_str(&format!("Failed tasks: {}.", failed.join(", ")));
            }
            return Err(anyhow!("Task leak detected: {}", msg));
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::shared::shutdown::Shutdown;

    #[tokio::test]
    async fn test_tracker_detects_completed_tasks() {
        let tracker = TaskTracker::new();

        // Spawn a task that completes immediately
        tracker.spawn("ImmediateTask", async move { Ok(()) });

        // Verify it completed
        let result = tracker.verify_no_leaks(Duration::from_secs(1)).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_tracker_detects_shutdown_aware_tasks() {
        let shutdown = Shutdown::new();
        let tracker = TaskTracker::new();

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

        // Task is running
        assert_eq!(tracker.spawned_count(), 1);

        // Trigger shutdown
        shutdown.trigger();

        // Verify all tasks complete
        let result = tracker.verify_no_leaks(Duration::from_secs(1)).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_tracker_detects_leaked_tasks() {
        let tracker = TaskTracker::new();

        // Spawn a task that never exits
        tracker.spawn("LeakyTask", async move {
            // Simulate a task that doesn't respect shutdown
            tokio::time::sleep(Duration::from_secs(10)).await;
            Ok(())
        });

        // Verify with a short timeout - should detect the leak
        let result = tracker.verify_no_leaks(Duration::from_millis(100)).await;
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("LeakyTask"), "Error: {}", err_msg);
        assert!(err_msg.contains("Leaked tasks"), "Error: {}", err_msg);
    }

    #[tokio::test]
    async fn test_tracker_detects_failed_tasks() {
        let tracker = TaskTracker::new();

        // Spawn a task that returns an error
        tracker.spawn("FailingTask", async move {
            Err(anyhow!("intentional test failure"))
        });

        // Verify - should detect the failure
        let result = tracker.verify_no_leaks(Duration::from_secs(1)).await;
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("FailingTask"), "Error: {}", err_msg);
    }

    #[tokio::test]
    async fn test_tracker_multiple_tasks() {
        let shutdown = Shutdown::new();
        let tracker = TaskTracker::new();

        // Spawn multiple tasks
        for _ in 0..5 {
            let mut rx = shutdown.subscribe();
            tracker.spawn("Task", async move {
                let _ = rx.recv().await;
                Ok(())
            });
        }

        assert_eq!(tracker.spawned_count(), 5);

        // Shutdown and verify
        shutdown.trigger();
        let result = tracker.verify_no_leaks(Duration::from_secs(1)).await;
        assert!(result.is_ok());
    }
}