marver 0.0.9

//! Deciding which queued tasks may start.
//!
//! One global cap on how many agents are live at once. The subtlety is not the
//! counting, it is which states count:
//!
//! | State | Holds a slot | Why |
//! |---|---|---|
//! | `running` | yes | obviously working |
//! | `blocked` | **yes** | resumes the instant the user answers, so a blocked task that freed its slot would let answering two prompts put you over the cap |
//! | `awaiting-review` | **no** | the turn is finished and consumes nothing; counting it would let unreviewed work starve the queue |
//! | `queued`, terminal states | no | nothing is running |
//!
//! Deciding *which* tasks start is separated from actually starting them.
//! [`Scheduler::plan`] is a pure query over the store, so the policy can be
//! tested without creating worktrees or tmux sessions; [`Scheduler::tick`] adds
//! the side effects through a [`Launch`] implementation.

use chrono::{DateTime, Utc};

use crate::domain::{Task, TaskState};
use crate::store::{Store, Transition};

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error(transparent)]
    Store(#[from] crate::store::Error),
}

pub type Result<T> = std::result::Result<T, Error>;

/// States that consume a concurrency slot.
pub const OCCUPYING: &[TaskState] = &[TaskState::Running, TaskState::Blocked];

/// How many agents may be live at once when nothing says otherwise.
pub const DEFAULT_CAP: usize = 3;

/// Starts a task for real: worktrees, tmux session, agent.
///
/// A trait so the scheduler's policy is testable on its own, and so the
/// launcher can be built independently.
pub trait Launch {
    /// Bring `task` to life. The error is recorded as the task's failure reason,
    /// so it should read as an explanation rather than a code.
    ///
    /// Takes the store because launching produces facts worth persisting — the
    /// tmux session name, the worktree paths — and losing them would leave a
    /// live agent marver could not find again.
    fn launch(&self, store: &mut Store, task: &Task) -> std::result::Result<(), String>;
}

/// What one pass of the scheduler did.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Tick {
    /// Tasks now running.
    pub started: Vec<i64>,
    /// Tasks that could not be started, with the reason. Each was moved to
    /// `failed`, not left queued.
    pub failed: Vec<(i64, String)>,
    /// Slots still free after this pass.
    pub remaining: usize,
}

pub struct Scheduler {
    cap: usize,
}

impl Scheduler {
    pub fn new(cap: usize) -> Self {
        Self { cap }
    }

    pub fn cap(&self) -> usize {
        self.cap
    }

    /// How many slots are in use.
    pub fn occupied(&self, store: &Store) -> Result<usize> {
        let mut total = 0;
        for state in OCCUPYING {
            total += store.list_tasks_in_state(*state)?.len();
        }
        Ok(total)
    }

    /// How many tasks could start right now.
    ///
    /// Saturating: lowering the cap below what is already live leaves zero free
    /// rather than underflowing, and running tasks are left alone.
    pub fn available(&self, store: &Store) -> Result<usize> {
        Ok(self.cap.saturating_sub(self.occupied(store)?))
    }

    /// Which queued tasks should start, oldest first.
    ///
    /// Pure: reads the store and decides. Nothing is created or transitioned.
    pub fn plan(&self, store: &Store) -> Result<Vec<Task>> {
        let free = self.available(store)?;
        if free == 0 {
            return Ok(Vec::new());
        }
        let mut queued = store.list_tasks_in_state(TaskState::Queued)?;
        // Oldest first: ids ascend with creation, and list_tasks_in_state
        // already orders by id, but say so rather than rely on it.
        queued.sort_by_key(|task| task.id);
        queued.truncate(free);
        Ok(queued)
    }

    /// Plan, launch, and record the results.
    ///
    /// A task that fails to start is moved to `failed` rather than left queued.
    /// Leaving it would make it retry forever and hold a slot in every future
    /// pass; failing it makes the problem visible and frees the queue.
    pub fn tick(
        &self,
        store: &mut Store,
        launcher: &impl Launch,
        now: DateTime<Utc>,
    ) -> Result<Tick> {
        let mut tick = Tick::default();

        for task in self.plan(store)? {
            match launcher.launch(store, &task) {
                Ok(()) => {
                    store.transition(task.id, TaskState::Running, Transition::Plain, now)?;
                    tick.started.push(task.id);
                }
                Err(reason) => {
                    store.transition(
                        task.id,
                        TaskState::Failed,
                        Transition::Failed(reason.clone()),
                        now,
                    )?;
                    tick.failed.push((task.id, reason));
                }
            }
        }

        tick.remaining = self.available(store)?;
        Ok(tick)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::domain::BlockedKind;
    use crate::store::BlockedInfo;
    use std::cell::RefCell;
    use std::path::Path;

    fn at(secs: i64) -> DateTime<Utc> {
        DateTime::from_timestamp(secs, 0).expect("valid timestamp")
    }

    /// Records what it was asked to start; optionally refuses named tasks.
    #[derive(Default)]
    struct Recorder {
        seen: RefCell<Vec<i64>>,
        refuse: Vec<i64>,
    }

    impl Recorder {
        fn refusing(ids: &[i64]) -> Self {
            Self {
                seen: RefCell::new(Vec::new()),
                refuse: ids.to_vec(),
            }
        }

        fn launched(&self) -> Vec<i64> {
            self.seen.borrow().clone()
        }
    }

    impl Launch for Recorder {
        fn launch(&self, _store: &mut Store, task: &Task) -> std::result::Result<(), String> {
            self.seen.borrow_mut().push(task.id);
            if self.refuse.contains(&task.id) {
                return Err(format!("no worktree for task {}", task.id));
            }
            Ok(())
        }
    }

    fn store() -> Store {
        Store::open_in_memory().expect("store")
    }

    fn queue(store: &mut Store, count: usize) -> Vec<i64> {
        (0..count)
            .map(|i| {
                store
                    .create_task(
                        &format!("task {i}"),
                        "do it",
                        Path::new("/tmp/tasks"),
                        &[],
                        at(0),
                    )
                    .unwrap()
                    .id
            })
            .collect()
    }

    fn run(store: &mut Store, id: i64) {
        store
            .transition(id, TaskState::Running, Transition::Plain, at(1))
            .unwrap();
    }

    fn block(store: &mut Store, id: i64) {
        run(store, id);
        store
            .transition(
                id,
                TaskState::Blocked,
                Transition::Blocked(BlockedInfo::new(BlockedKind::Question)),
                at(2),
            )
            .unwrap();
    }

    fn review(store: &mut Store, id: i64) {
        run(store, id);
        store
            .transition(id, TaskState::AwaitingReview, Transition::Plain, at(2))
            .unwrap();
    }

    #[test]
    fn an_empty_queue_starts_nothing() {
        let store = store();
        let tick = Scheduler::new(3).plan(&store).unwrap();
        assert!(tick.is_empty());
    }

    #[test]
    fn starts_up_to_the_cap_oldest_first() {
        let mut store = store();
        let ids = queue(&mut store, 5);
        let launcher = Recorder::default();

        let tick = Scheduler::new(2)
            .tick(&mut store, &launcher, at(5))
            .unwrap();

        assert_eq!(tick.started, ids[..2]);
        assert_eq!(launcher.launched(), ids[..2], "oldest queued go first");
        assert_eq!(tick.remaining, 0);
        assert_eq!(
            store.list_tasks_in_state(TaskState::Queued).unwrap().len(),
            3
        );
    }

    #[test]
    fn a_cap_larger_than_the_queue_starts_everything() {
        let mut store = store();
        let ids = queue(&mut store, 2);
        let tick = Scheduler::new(10)
            .tick(&mut store, &Recorder::default(), at(5))
            .unwrap();
        assert_eq!(tick.started, ids);
        assert_eq!(tick.remaining, 8);
    }

    #[test]
    fn a_cap_of_zero_starts_nothing() {
        let mut store = store();
        queue(&mut store, 3);
        let launcher = Recorder::default();
        let tick = Scheduler::new(0)
            .tick(&mut store, &launcher, at(5))
            .unwrap();
        assert!(tick.started.is_empty());
        assert!(launcher.launched().is_empty(), "nothing should be touched");
    }

    #[test]
    fn a_blocked_task_still_holds_its_slot() {
        let mut store = store();
        let ids = queue(&mut store, 3);
        block(&mut store, ids[0]);

        let scheduler = Scheduler::new(1);
        assert_eq!(scheduler.occupied(&store).unwrap(), 1);
        assert!(
            scheduler.plan(&store).unwrap().is_empty(),
            "a blocked agent resumes the moment it is answered, so its slot is \
             not free"
        );
    }

    #[test]
    fn a_task_awaiting_review_frees_its_slot() {
        let mut store = store();
        let ids = queue(&mut store, 3);
        review(&mut store, ids[0]);

        let scheduler = Scheduler::new(1);
        assert_eq!(scheduler.occupied(&store).unwrap(), 0);
        let planned = scheduler.plan(&store).unwrap();
        assert_eq!(
            planned.iter().map(|t| t.id).collect::<Vec<_>>(),
            [ids[1]],
            "unreviewed work must not starve the queue"
        );
    }

    #[test]
    fn terminal_tasks_never_hold_a_slot() {
        let mut store = store();
        let ids = queue(&mut store, 4);
        review(&mut store, ids[0]);
        store
            .transition(ids[0], TaskState::Committed, Transition::Plain, at(3))
            .unwrap();
        store
            .transition(ids[1], TaskState::Cancelled, Transition::Plain, at(3))
            .unwrap();
        run(&mut store, ids[2]);
        store
            .transition(
                ids[2],
                TaskState::Failed,
                Transition::Failed("died".into()),
                at(3),
            )
            .unwrap();

        assert_eq!(Scheduler::new(2).occupied(&store).unwrap(), 0);
    }

    #[test]
    fn a_full_scheduler_is_a_no_op() {
        let mut store = store();
        let ids = queue(&mut store, 3);
        run(&mut store, ids[0]);
        run(&mut store, ids[1]);

        let launcher = Recorder::default();
        let tick = Scheduler::new(2)
            .tick(&mut store, &launcher, at(5))
            .unwrap();

        assert!(tick.started.is_empty());
        assert!(launcher.launched().is_empty());
        assert_eq!(store.get_task(ids[2]).unwrap().state, TaskState::Queued);
    }

    #[test]
    fn lowering_the_cap_below_what_is_live_does_not_underflow() {
        let mut store = store();
        let ids = queue(&mut store, 3);
        run(&mut store, ids[0]);
        run(&mut store, ids[1]);

        let scheduler = Scheduler::new(1);
        assert_eq!(scheduler.available(&store).unwrap(), 0);
        assert!(
            scheduler.plan(&store).unwrap().is_empty(),
            "over-capacity must not start more, nor panic"
        );
        // Nothing already running is disturbed.
        assert_eq!(store.get_task(ids[0]).unwrap().state, TaskState::Running);
    }

    #[test]
    fn a_task_that_cannot_start_fails_rather_than_requeueing() {
        let mut store = store();
        let ids = queue(&mut store, 2);
        let launcher = Recorder::refusing(&[ids[0]]);

        let tick = Scheduler::new(2)
            .tick(&mut store, &launcher, at(5))
            .unwrap();

        assert_eq!(tick.started, [ids[1]]);
        assert_eq!(tick.failed.len(), 1);
        assert_eq!(tick.failed[0].0, ids[0]);

        let failed = store.get_task(ids[0]).unwrap();
        assert_eq!(failed.state, TaskState::Failed);
        assert_eq!(
            failed.failure_reason.as_deref(),
            Some(format!("no worktree for task {}", ids[0]).as_str()),
            "the launcher's explanation must survive"
        );
    }

    #[test]
    fn a_failed_launch_frees_its_slot_for_the_next_pass() {
        let mut store = store();
        let ids = queue(&mut store, 3);

        // First pass: the only slot goes to a task that cannot start.
        let scheduler = Scheduler::new(1);
        scheduler
            .tick(&mut store, &Recorder::refusing(&[ids[0]]), at(5))
            .unwrap();
        assert_eq!(store.get_task(ids[0]).unwrap().state, TaskState::Failed);

        // Second pass: the slot is free again, so the queue advances rather
        // than being wedged behind the broken task.
        let launcher = Recorder::default();
        let tick = scheduler.tick(&mut store, &launcher, at(6)).unwrap();
        assert_eq!(tick.started, [ids[1]]);
    }

    #[test]
    fn repeated_ticks_do_not_double_start() {
        let mut store = store();
        queue(&mut store, 5);
        let scheduler = Scheduler::new(2);
        let launcher = Recorder::default();

        let first = scheduler.tick(&mut store, &launcher, at(5)).unwrap();
        let second = scheduler.tick(&mut store, &launcher, at(6)).unwrap();

        assert_eq!(first.started.len(), 2);
        assert!(second.started.is_empty());
        assert_eq!(launcher.launched().len(), 2, "no task launched twice");
    }

    #[test]
    fn finishing_a_task_lets_the_queue_advance() {
        let mut store = store();
        let ids = queue(&mut store, 3);
        let scheduler = Scheduler::new(1);
        let launcher = Recorder::default();

        scheduler.tick(&mut store, &launcher, at(5)).unwrap();
        assert_eq!(launcher.launched(), [ids[0]]);

        // The agent finishes; its slot is released.
        store
            .transition(ids[0], TaskState::AwaitingReview, Transition::Plain, at(6))
            .unwrap();
        let tick = scheduler.tick(&mut store, &launcher, at(7)).unwrap();

        assert_eq!(tick.started, [ids[1]]);
    }

    #[test]
    fn answering_a_blocked_task_does_not_exceed_the_cap() {
        let mut store = store();
        let ids = queue(&mut store, 4);
        let scheduler = Scheduler::new(2);
        let launcher = Recorder::default();

        scheduler.tick(&mut store, &launcher, at(5)).unwrap();
        // Both live agents get stuck on a prompt.
        for id in &ids[..2] {
            store
                .transition(
                    *id,
                    TaskState::Blocked,
                    Transition::Blocked(BlockedInfo::new(BlockedKind::PermissionPrompt)),
                    at(6),
                )
                .unwrap();
        }

        let tick = scheduler.tick(&mut store, &launcher, at(7)).unwrap();
        assert!(
            tick.started.is_empty(),
            "starting here would mean four live agents once both prompts are \
             answered, with a cap of two"
        );
    }
}