quickdid / src / task_manager.rs

//! Task management utilities for consistent background task handling
//!
//! This module provides helpers for spawning and managing background tasks with:
//! - Consistent start/stop logging
//! - Automatic shutdown on task failure
//! - Graceful shutdown support

use std::future::Future;
use tokio_util::{sync::CancellationToken, task::TaskTracker};
use tracing::{error, info};

/// Spawn a background task with cancellation support
///
/// This version allows the task to be cancelled via the token and handles
/// both graceful shutdown and unexpected failures
pub fn spawn_cancellable_task<F, Fut>(
    tracker: &TaskTracker,
    app_token: CancellationToken,
    task_name: &str,
    task_builder: F,
) where
    F: FnOnce(CancellationToken) -> Fut + Send + 'static,
    Fut: Future<Output = anyhow::Result<()>> + Send + 'static,
{
    info!(task = task_name, "Starting cancellable background task");

    let task_token = app_token.clone();
    let cancel_token = app_token.clone();

    let inner_task_name = task_name.to_string();

    tracker.spawn(async move {
        tokio::select! {
            result = task_builder(cancel_token.clone()) => {
                match result {
                    Ok(()) => {
                        info!(task = inner_task_name, "Background task completed successfully");
                    }
                    Err(e) => {
                        error!(error = ?e, task = inner_task_name, "Background task failed unexpectedly");
                        // Trigger application shutdown on task failure
                        task_token.cancel();
                    }
                }
            }
            () = task_token.cancelled() => {
                info!(task = inner_task_name, "Background task shutting down gracefully");
            }
        }
    });
}

/// Macro for consistent task error handling within a task
#[macro_export]
macro_rules! task_try {
    ($expr:expr, $task_name:expr) => {
        match $expr {
            Ok(val) => val,
            Err(e) => {
                tracing::error!(task = $task_name, error = ?e, "Task operation failed");
                return Err(e.into());
            }
        }
    };
}

/// Macro for logging task checkpoints
#[macro_export]
macro_rules! task_checkpoint {
    ($task_name:expr, $checkpoint:expr) => {
        tracing::debug!(
            task = $task_name,
            checkpoint = $checkpoint,
            "Task checkpoint reached"
        );
    };
}