oo_ide/

task_registry.rs

//! Queue-based task scheduler for the IDE.
//!
//! [`TaskRegistry`] is the single source of truth for all scheduled, running,
//! and recently finished tasks.  It enforces the following scheduling rules:
//!
//! 1. **One running task per queue** — a queue may have at most one task with
//!    status [`TaskStatus::Running`] at any time.
//! 2. **Same `(queue, target)` cancels previous** — scheduling a task whose
//!    [`TaskKey`] matches a running or queued task cancels the earlier one first.
//! 3. **Different targets are queued** — if another task is already running on
//!    the same queue but with a different target, the new task is enqueued.
//! 4. **Queue compaction** — if the same [`TaskKey`] already sits in the pending
//!    queue, it is replaced in-place rather than appended a second time.
//! 5. **FIFO execution** — tasks start in insertion order after deduplication.
//!
//! # Design constraints
//!
//! * **No disk I/O** — purely in-memory.
//! * **No async** — deterministic and synchronous.  Callers are responsible for
//!   spawning async work and calling [`TaskRegistry::mark_running`] /
//!   [`TaskRegistry::mark_finished`] on completion.
//! * **Single-threaded** — lives on the main thread as part of
//!   [`crate::app_state::AppState`].

use std::collections::{HashMap, VecDeque};
use std::time::Instant;

use tokio_util::sync::CancellationToken;

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Opaque, monotonically-increasing task identifier.  Never reused within a
/// single registry lifetime.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct TaskId(pub u64);

/// Identifies a named scheduling queue (e.g. `"build"`, `"lint"`, `"test"`).
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TaskQueueId(pub String);

/// The compound identity of a task for deduplication and cancellation purposes.
///
/// Two tasks with equal `TaskKey` values are considered interchangeable; the
/// registry ensures that at most one such task is pending or running at any
/// time.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TaskKey {
    pub queue: TaskQueueId,
    /// Arbitrary string that scopes the task within its queue — e.g. a crate
    /// path, a file path, or `"*"` for queue-wide singletons.
    pub target: String,
}

/// Lifecycle state of a task.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TaskStatus {
    /// Created and waiting for its turn in the queue.
    Pending,
    /// Currently executing.
    Running,
    /// Finished successfully.
    Success,
    /// Finished with warnings but no hard errors.
    Warning,
    /// Finished with at least one error.
    Error,
    /// Cancelled before or during execution.
    Cancelled,
}

/// What initiated the task.
#[derive(Clone, Debug)]
pub enum TaskTrigger {
    Manual,
    OnSave,
    OnFileChange,
    Extension(String),
}

/// A single scheduled, running, or finished task.
pub struct Task {
    pub id: TaskId,
    pub key: TaskKey,
    pub status: TaskStatus,
    /// The shell command string to execute (e.g. `"cargo build -p my-crate"`).
    pub command: String,
    pub created_at: Instant,
    pub started_at: Option<Instant>,
    pub finished_at: Option<Instant>,
    /// Signal used to stop the async work for this task.  Callers poll
    /// `cancellation_token.is_cancelled()` (or `await` it) in their async
    /// work loop.
    pub cancellation_token: CancellationToken,
    pub trigger: TaskTrigger,
}

// ---------------------------------------------------------------------------
// Registry
// ---------------------------------------------------------------------------

/// Central scheduler and store for all IDE tasks.
pub struct TaskRegistry {
    /// All tasks ever created in this session, keyed by [`TaskId`].
    tasks: HashMap<TaskId, Task>,
    /// The task currently running on each queue.
    running: HashMap<TaskQueueId, TaskId>,
    /// FIFO pending queue per [`TaskQueueId`].
    queues: HashMap<TaskQueueId, VecDeque<TaskId>>,
    next_id: u64,
    /// Ring buffer of recently finished task IDs (max 5), newest at back.
    recent_finished: VecDeque<TaskId>,
}

impl TaskRegistry {
    pub fn new() -> Self {
        Self {
            tasks: HashMap::new(),
            running: HashMap::new(),
            queues: HashMap::new(),
            next_id: 1,
            recent_finished: VecDeque::with_capacity(5),
        }
    }

    // -----------------------------------------------------------------------
    // Public API
    // -----------------------------------------------------------------------

    /// Schedule a new task on `key.queue`.
    ///
    /// Any existing running or queued task with the **same** [`TaskKey`] is
    /// cancelled first (Rule 2 / Rule 4).  If no task is currently running on
    /// the queue the new task starts immediately (status → `Running`); otherwise
    /// it is appended to the FIFO queue (status → `Pending`).
    ///
    /// Returns the [`TaskId`] of the newly created task.
    pub fn schedule_task(&mut self, key: TaskKey, trigger: TaskTrigger, command: String) -> TaskId {
        // Cancel any existing task with the same key (running or queued).
        self.cancel_by_key(&key);

        let id = self.alloc_id();
        let task = Task {
            id,
            key: key.clone(),
            status: TaskStatus::Pending,
            command,
            created_at: Instant::now(),
            started_at: None,
            finished_at: None,
            cancellation_token: CancellationToken::new(),
            trigger,
        };
        self.tasks.insert(id, task);

        let queue_id = &key.queue;
        if self.running.contains_key(queue_id) {
            // Another task is running on this queue — enqueue.
            self.queues.entry(queue_id.clone()).or_default().push_back(id);
        } else {
            // Queue is idle — start immediately.
            self.mark_running(id);
        }

        id
    }

    /// Cancel the task identified by `task_id`.
    ///
    /// * If the task is **Running** its [`CancellationToken`] is triggered and
    ///   `status` is set to [`TaskStatus::Cancelled`].  The `running` slot is
    ///   cleared and `start_next` is called for the queue.
    /// * If the task is **Pending** it is removed from the queue and its status
    ///   is set to [`TaskStatus::Cancelled`].
    /// * If the task is already finished or cancelled this is a no-op.
    ///
    /// Returns the [`TaskId`] of the next task that was started as a side-effect
    /// of releasing the running slot (only possible when cancelling a Running
    /// task that had a non-empty pending queue).
    pub fn cancel(&mut self, task_id: TaskId) -> Option<TaskId> {
        let (status, queue_id) = match self.tasks.get(&task_id) {
            Some(t) => (t.status.clone(), t.key.queue.clone()),
            None => return None,
        };

        match status {
            TaskStatus::Running => {
                if let Some(t) = self.tasks.get_mut(&task_id) {
                    t.cancellation_token.cancel();
                    t.status = TaskStatus::Cancelled;
                    t.finished_at = Some(Instant::now());
                }
                self.running.remove(&queue_id);
                self.start_next(&queue_id)
            }
            TaskStatus::Pending => {
                if let Some(queue) = self.queues.get_mut(&queue_id) {
                    queue.retain(|&id| id != task_id);
                }
                if let Some(t) = self.tasks.get_mut(&task_id) {
                    t.status = TaskStatus::Cancelled;
                }
                None
            }
            // Already in a terminal state — nothing to do.
            _ => None,
        }
    }

    /// Cancel all running and queued tasks whose [`TaskKey`] equals `key`.
    ///
    /// Returns the [`TaskId`] of any task that started as a side-effect of
    /// releasing a running slot.
    pub fn cancel_by_key(&mut self, key: &TaskKey) -> Option<TaskId> {
        let ids_to_cancel: Vec<TaskId> = self
            .tasks
            .values()
            .filter(|t| {
                &t.key == key
                    && matches!(t.status, TaskStatus::Pending | TaskStatus::Running)
            })
            .map(|t| t.id)
            .collect();

        let mut started = None;
        for id in ids_to_cancel {
            if let Some(next) = self.cancel(id) {
                started = Some(next);
            }
        }
        started
    }

    /// Transition a task from `Pending` to `Running` and record `started_at`.
    ///
    /// Also registers the task in the `running` map for its queue.
    /// Panics in debug builds if the task is not in `Pending` state.
    pub fn mark_running(&mut self, task_id: TaskId) {
        let queue_id = match self.tasks.get_mut(&task_id) {
            Some(t) => {
                debug_assert_eq!(
                    t.status,
                    TaskStatus::Pending,
                    "mark_running called on task {:?} with status {:?}",
                    task_id,
                    t.status
                );
                t.status = TaskStatus::Running;
                t.started_at = Some(Instant::now());
                t.key.queue.clone()
            }
            None => return,
        };
        self.running.insert(queue_id, task_id);
    }

    /// Record a terminal status for a finished task and start the next one.
    ///
    /// If the task was already `Cancelled` the status is **not** overwritten —
    /// a cancelled task stays cancelled regardless of the final process outcome.
    ///
    /// Returns the [`TaskId`] of the next task that was started as a
    /// side-effect, if the queue had a pending task waiting.
    pub fn mark_finished(&mut self, task_id: TaskId, status: TaskStatus) -> Option<TaskId> {
        let queue_id = match self.tasks.get_mut(&task_id) {
            Some(t) => {
                // A cancelled task must not be re-labelled.
                if t.status != TaskStatus::Cancelled {
                    t.status = status;
                    t.finished_at = Some(Instant::now());
                }
                t.key.queue.clone()
            }
            None => return None,
        };

        // Add to recently-finished ring (cap at 5, discard oldest).
        if self.recent_finished.len() >= 5 {
            self.recent_finished.pop_front();
        }
        self.recent_finished.push_back(task_id);

        // Clear the running slot only if this task still owns it.
        if self.running.get(&queue_id) == Some(&task_id) {
            self.running.remove(&queue_id);
            self.start_next(&queue_id)
        } else {
            None
        }
    }

    // -----------------------------------------------------------------------
    // Read accessors
    // -----------------------------------------------------------------------

    /// Returns an iterator over all tasks ever created (in arbitrary order).
    pub fn all_tasks(&self) -> impl Iterator<Item = &Task> {
        self.tasks.values()
    }

    /// Returns the total number of tasks ever created (running, pending, or finished).
    pub fn task_count(&self) -> usize {
        self.tasks.len()
    }

    /// Returns a reference to a task by ID, or `None` if it does not exist.
    pub fn get(&self, id: TaskId) -> Option<&Task> {
        self.tasks.get(&id)
    }

    /// Returns the [`TaskId`] of the currently running task on `queue`, if any.
    pub fn running_task(&self, queue: &TaskQueueId) -> Option<TaskId> {
        self.running.get(queue).copied()
    }

    /// Returns an iterator over all currently running tasks `(queue_id, task_id)`.
    pub fn running_tasks(&self) -> impl Iterator<Item = (&TaskQueueId, &TaskId)> {
        self.running.iter()
    }

    /// Returns an iterator over recently finished task IDs (newest last, max 5).
    ///
    /// Only tasks with a terminal status (Success / Warning / Error / Cancelled)
    /// are included.  The caller can use [`TaskRegistry::get`] to read the full task record.
    pub fn recently_finished_tasks(&self) -> impl Iterator<Item = TaskId> + '_ {
        // Iterate newest-first so the most recent result appears closest to the
        // running task on the right side of the status bar.
        self.recent_finished.iter().rev().copied()
    }

    /// Returns the ordered pending task IDs for `queue` (front = next to run).
    pub fn pending_tasks(&self, queue: &TaskQueueId) -> &[TaskId] {
        self.queues
            .get(queue)
            .map(|q| q.as_slices().0)
            .unwrap_or(&[])
    }

    // -----------------------------------------------------------------------
    // Private helpers
    // -----------------------------------------------------------------------

    fn alloc_id(&mut self) -> TaskId {
        let id = TaskId(self.next_id);
        self.next_id += 1;
        id
    }

    /// Dequeue the next pending task for `queue` and mark it as running.
    /// Returns `Some(TaskId)` if a task was started, `None` if the queue is
    /// empty.
    fn start_next(&mut self, queue: &TaskQueueId) -> Option<TaskId> {
        let next_id = self.queues.get_mut(queue)?.pop_front()?;
        self.mark_running(next_id);
        Some(next_id)
    }
}

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

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn key(queue: &str, target: &str) -> TaskKey {
        TaskKey {
            queue: TaskQueueId(queue.into()),
            target: target.into(),
        }
    }

    fn sched(reg: &mut TaskRegistry, queue: &str, target: &str) -> TaskId {
        reg.schedule_task(key(queue, target), TaskTrigger::Manual, format!("echo {target}"))
    }

    // 1. Schedule same (queue, target) twice → first cancelled, second is active
    #[test]
    fn same_key_cancels_previous() {
        let mut reg = TaskRegistry::new();
        let id1 = sched(&mut reg, "build", "crate:a");
        let id2 = sched(&mut reg, "build", "crate:a");

        assert_eq!(reg.get(id1).unwrap().status, TaskStatus::Cancelled);
        assert!(reg.get(id1).unwrap().cancellation_token.is_cancelled());
        assert_eq!(reg.get(id2).unwrap().status, TaskStatus::Running);
        assert_eq!(reg.running_task(&TaskQueueId("build".into())), Some(id2));
    }

    // 2. Schedule A then B → A running, B queued
    #[test]
    fn different_targets_are_queued() {
        let mut reg = TaskRegistry::new();
        let id_a = sched(&mut reg, "build", "crate:a");
        let id_b = sched(&mut reg, "build", "crate:b");

        assert_eq!(reg.get(id_a).unwrap().status, TaskStatus::Running);
        assert_eq!(reg.get(id_b).unwrap().status, TaskStatus::Pending);
        assert_eq!(reg.running_task(&TaskQueueId("build".into())), Some(id_a));
        assert_eq!(reg.pending_tasks(&TaskQueueId("build".into())), &[id_b]);
    }

    // 3. Schedule A, B, A again → deduplicated: only latest A remains in queue
    #[test]
    fn queue_compaction_deduplicates_key() {
        let mut reg = TaskRegistry::new();
        let id_a1 = sched(&mut reg, "build", "crate:a");
        let id_b = sched(&mut reg, "build", "crate:b");
        // Re-schedule A — must replace queued A, not append.
        let id_a2 = sched(&mut reg, "build", "crate:a");

        // A1 was running → cancelled; A2 queued (not running, because B is running).
        // Actually: A1 started immediately. Then B queued. Then A2 is scheduled:
        //   cancel_by_key(A) hits A1 (Running) → cancelled + start_next → B starts.
        //   new A2 is created: B is now running → A2 enqueued.
        assert_eq!(reg.get(id_a1).unwrap().status, TaskStatus::Cancelled);
        assert_eq!(reg.get(id_b).unwrap().status, TaskStatus::Running);
        assert_eq!(reg.get(id_a2).unwrap().status, TaskStatus::Pending);

        // Queue must contain exactly A2 (not id_b which is running, not id_a1).
        let pending = reg.pending_tasks(&TaskQueueId("build".into()));
        assert_eq!(pending, &[id_a2]);
    }

    // 4. Cancel a queued task → removed from queue, status Cancelled
    #[test]
    fn cancel_queued_task_removes_from_queue() {
        let mut reg = TaskRegistry::new();
        let _id_a = sched(&mut reg, "build", "crate:a");
        let id_b = sched(&mut reg, "build", "crate:b");

        reg.cancel(id_b);

        assert_eq!(reg.get(id_b).unwrap().status, TaskStatus::Cancelled);
        assert!(reg.pending_tasks(&TaskQueueId("build".into())).is_empty());
    }

    // 5. Cancel a running task → token triggered, status Cancelled
    #[test]
    fn cancel_running_task_triggers_token() {
        let mut reg = TaskRegistry::new();
        let id = sched(&mut reg, "build", "crate:a");
        let token = reg.get(id).unwrap().cancellation_token.clone();

        reg.cancel(id);

        assert_eq!(reg.get(id).unwrap().status, TaskStatus::Cancelled);
        assert!(token.is_cancelled());
        assert_eq!(reg.running_task(&TaskQueueId("build".into())), None);
    }

    // 6. mark_finished → next queued task starts automatically; returns next id
    #[test]
    fn finish_starts_next_queued_task() {
        let mut reg = TaskRegistry::new();
        let id_a = sched(&mut reg, "build", "crate:a");
        let id_b = sched(&mut reg, "build", "crate:b");

        let next = reg.mark_finished(id_a, TaskStatus::Success);

        assert_eq!(next, Some(id_b));
        assert_eq!(reg.get(id_a).unwrap().status, TaskStatus::Success);
        assert_eq!(reg.get(id_b).unwrap().status, TaskStatus::Running);
        assert_eq!(reg.running_task(&TaskQueueId("build".into())), Some(id_b));
    }

    // 7. Multiple queues operate independently
    #[test]
    fn independent_queues_do_not_interfere() {
        let mut reg = TaskRegistry::new();
        let build_id = sched(&mut reg, "build", "crate:a");
        let lint_id = sched(&mut reg, "lint", "crate:a");

        assert_eq!(reg.get(build_id).unwrap().status, TaskStatus::Running);
        assert_eq!(reg.get(lint_id).unwrap().status, TaskStatus::Running);

        reg.mark_finished(build_id, TaskStatus::Success);
        // Lint queue should be unaffected.
        assert_eq!(reg.get(lint_id).unwrap().status, TaskStatus::Running);
    }

    // 8. A finished cancelled task must remain Cancelled
    #[test]
    fn cancelled_task_stays_cancelled_on_finish() {
        let mut reg = TaskRegistry::new();
        let id = sched(&mut reg, "build", "crate:a");
        reg.cancel(id);

        // Simulate the async worker not noticing cancellation in time and
        // reporting success anyway.
        reg.mark_finished(id, TaskStatus::Success);

        assert_eq!(reg.get(id).unwrap().status, TaskStatus::Cancelled);
    }

    // 9. start_next on empty queue → no-op, returns None
    #[test]
    fn finish_on_empty_queue_is_noop() {
        let mut reg = TaskRegistry::new();
        let id = sched(&mut reg, "build", "crate:a");
        let next = reg.mark_finished(id, TaskStatus::Success);

        assert_eq!(next, None);
        assert_eq!(reg.running_task(&TaskQueueId("build".into())), None);
        assert!(reg.pending_tasks(&TaskQueueId("build".into())).is_empty());
    }

    // Bonus: rapid rescheduling leaves no duplicate key in queue
    #[test]
    fn rapid_reschedule_no_duplicate_key_in_queue() {
        let mut reg = TaskRegistry::new();
        // Occupy the queue with a long-running anchor task.
        let _anchor = sched(&mut reg, "build", "anchor");

        // Schedule target A three times rapidly.
        let id1 = sched(&mut reg, "build", "crate:a");
        let id2 = sched(&mut reg, "build", "crate:a");
        let id3 = sched(&mut reg, "build", "crate:a");

        assert_eq!(reg.get(id1).unwrap().status, TaskStatus::Cancelled);
        assert_eq!(reg.get(id2).unwrap().status, TaskStatus::Cancelled);
        assert_eq!(reg.get(id3).unwrap().status, TaskStatus::Pending);

        let pending = reg.pending_tasks(&TaskQueueId("build".into()));
        assert_eq!(pending.len(), 1);
        assert_eq!(pending[0], id3);
    }

    // cancel() on a running task returns the next started task
    #[test]
    fn cancel_running_returns_next_started() {
        let mut reg = TaskRegistry::new();
        let id_a = sched(&mut reg, "build", "crate:a");
        let id_b = sched(&mut reg, "build", "crate:b");

        let next = reg.cancel(id_a);

        assert_eq!(next, Some(id_b));
        assert_eq!(reg.get(id_b).unwrap().status, TaskStatus::Running);
    }
}