nano_watchdog/

tracker.rs

use dashmap::DashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

/// Tracks active tasks for deadlock detection.
///
/// Thread-safe — all methods take `&self`. Typically stored in a `static`
/// or shared via `Arc`.
pub struct TaskTracker {
    active_tasks: DashMap<u64, TaskInfo>,
    next_task_id: AtomicU64,
    completed_count: AtomicU64,
}

#[derive(Clone)]
struct TaskInfo {
    description: String,
    start_time: Instant,
    phase: String,
}

/// Information about a task that has exceeded the stuck threshold.
#[derive(Debug, Clone)]
pub struct StuckTaskInfo {
    pub task_id: u64,
    pub description: String,
    pub phase: String,
    pub duration: Duration,
}

/// RAII guard that tracks a task's lifetime.
///
/// When dropped, the task is automatically unregistered from the tracker
/// and the completed counter is incremented.
pub struct TaskGuard<'a> {
    tracker: &'a TaskTracker,
    task_id: u64,
}

impl TaskTracker {
    /// Create a new task tracker.
    pub fn new() -> Self {
        Self {
            active_tasks: DashMap::new(),
            next_task_id: AtomicU64::new(1),
            completed_count: AtomicU64::new(0),
        }
    }

    /// Start tracking a task. Returns a guard that removes the task when dropped.
    pub fn track(&self, description: &str) -> TaskGuard<'_> {
        let task_id = self.next_task_id.fetch_add(1, Ordering::Relaxed);
        let info = TaskInfo {
            description: description.to_string(),
            start_time: Instant::now(),
            phase: "started".to_string(),
        };
        self.active_tasks.insert(task_id, info);

        TaskGuard {
            tracker: self,
            task_id,
        }
    }

    /// Update the phase of a tracked task.
    pub(crate) fn update_phase(&self, task_id: u64, phase: &str) {
        if let Some(mut entry) = self.active_tasks.get_mut(&task_id) {
            entry.phase = phase.to_string();
        }
    }

    /// Mark a task as completed (called by TaskGuard::drop).
    fn complete_task(&self, task_id: u64) {
        self.active_tasks.remove(&task_id);
        self.completed_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Number of currently active tasks.
    pub fn active_count(&self) -> usize {
        self.active_tasks.len()
    }

    /// Total number of completed tasks since creation.
    pub fn completed_count(&self) -> u64 {
        self.completed_count.load(Ordering::Relaxed)
    }

    /// Find tasks running longer than `max_duration`.
    pub fn check_stuck_tasks(&self, max_duration: Duration) -> Vec<StuckTaskInfo> {
        let now = Instant::now();
        let mut stuck = Vec::new();

        for entry in self.active_tasks.iter() {
            let duration = now.duration_since(entry.start_time);
            if duration > max_duration {
                stuck.push(StuckTaskInfo {
                    task_id: *entry.key(),
                    description: entry.description.clone(),
                    phase: entry.phase.clone(),
                    duration,
                });
            }
        }

        stuck
    }

    /// Dump all active tasks as a formatted string for diagnostics.
    pub fn dump_all_tasks(&self) -> String {
        let now = Instant::now();
        let mut output = String::new();

        output.push_str("========== WATCHDOG TASK DUMP ==========\n");
        output.push_str(&format!("Time: {:?}\n", std::time::SystemTime::now()));
        output.push_str(&format!("Active tasks: {}\n", self.active_tasks.len()));
        output.push_str(&format!(
            "Completed tasks: {}\n\n",
            self.completed_count.load(Ordering::Relaxed)
        ));

        let mut tasks: Vec<_> = self.active_tasks.iter().collect();
        tasks.sort_by(|a, b| a.start_time.cmp(&b.start_time));

        for entry in tasks {
            let duration = now.duration_since(entry.start_time);
            output.push_str(&format!(
                "[Task {}] {} - Phase: {} - Duration: {:.2}s\n",
                entry.key(),
                entry.description,
                entry.phase,
                duration.as_secs_f64()
            ));
        }

        output.push_str("========================================\n");
        output
    }
}

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

impl<'a> TaskGuard<'a> {
    /// Update the current phase of this task.
    pub fn set_phase(&self, phase: &str) {
        self.tracker.update_phase(self.task_id, phase);
    }

    /// Get the task ID.
    pub fn id(&self) -> u64 {
        self.task_id
    }
}

impl<'a> Drop for TaskGuard<'a> {
    fn drop(&mut self) {
        self.tracker.complete_task(self.task_id);
    }
}

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

    #[test]
    fn test_track_and_drop() {
        let tracker = TaskTracker::new();
        assert_eq!(tracker.active_count(), 0);
        assert_eq!(tracker.completed_count(), 0);

        {
            let _guard = tracker.track("test task");
            assert_eq!(tracker.active_count(), 1);
            assert_eq!(tracker.completed_count(), 0);
        }
        // Guard dropped
        assert_eq!(tracker.active_count(), 0);
        assert_eq!(tracker.completed_count(), 1);
    }

    #[test]
    fn test_set_phase() {
        let tracker = TaskTracker::new();
        let guard = tracker.track("phased task");

        guard.set_phase("processing");

        let stuck = tracker.check_stuck_tasks(Duration::ZERO);
        assert_eq!(stuck.len(), 1);
        assert_eq!(stuck[0].phase, "processing");
    }

    #[test]
    fn test_task_id_increments() {
        let tracker = TaskTracker::new();
        let g1 = tracker.track("first");
        let g2 = tracker.track("second");
        assert_ne!(g1.id(), g2.id());
        assert!(g2.id() > g1.id());
    }

    #[test]
    fn test_check_stuck_tasks() {
        let tracker = TaskTracker::new();
        let _guard = tracker.track("slow task");

        // With zero threshold, everything is stuck
        let stuck = tracker.check_stuck_tasks(Duration::ZERO);
        assert_eq!(stuck.len(), 1);
        assert_eq!(stuck[0].description, "slow task");

        // With large threshold, nothing is stuck
        let stuck = tracker.check_stuck_tasks(Duration::from_secs(9999));
        assert_eq!(stuck.len(), 0);
    }

    #[test]
    fn test_check_stuck_excludes_completed() {
        let tracker = TaskTracker::new();
        {
            let _guard = tracker.track("done task");
        }
        let stuck = tracker.check_stuck_tasks(Duration::ZERO);
        assert_eq!(stuck.len(), 0);
    }

    #[test]
    fn test_dump_all_tasks() {
        let tracker = TaskTracker::new();
        let guard = tracker.track("my task");
        guard.set_phase("waiting");

        let dump = tracker.dump_all_tasks();
        assert!(dump.contains("my task"));
        assert!(dump.contains("waiting"));
        assert!(dump.contains("Active tasks: 1"));
    }

    #[test]
    fn test_concurrent_tracking() {
        let tracker = TaskTracker::new();
        let tracker_ref = &tracker;

        std::thread::scope(|s| {
            for i in 0..10 {
                s.spawn(move || {
                    let _guard = tracker_ref.track(&format!("thread-{}", i));
                    std::thread::sleep(Duration::from_millis(5));
                });
            }
        });

        assert_eq!(tracker.active_count(), 0);
        assert_eq!(tracker.completed_count(), 10);
    }
}