marver 0.0.27

//! Deciding which queued tasks may start.
//!
//! One global cap. The subtlety is which states count against it:
//!
//! | State | Holds a slot | Why |
//! |---|---|---|
//! | `running` | yes | obviously working |
//! | `blocked` | **yes** | resumes the instant the user answers, so freeing its slot would let answering two prompts put you over the cap |
//! | `awaiting-review` | **no** | the turn is finished; counting it would let unreviewed work starve the queue |
//! | `paused` | **no** | releasing the slot is the point of pausing |
//! | `queued`, terminal | no | nothing is running |
//!
//! A paused task is not `queued`, so it is never picked — which is why `paused`
//! is a state rather than a flag this has to remember to consult.
//!
//! [`Scheduler::plan`] is a pure query over the store, so the policy is testable
//! without worktrees or tmux; [`Scheduler::tick`] adds the effects through
//! [`Launch`].

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.
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.
    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)>,
    /// Tasks that left `queued` while they were being started, with where they
    /// went. Their agents are live and their state is not `running`.
    pub abandoned: Vec<(i64, TaskState)>,
    /// 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.
    pub fn available(&self, store: &Store) -> Result<usize> {
        Ok(self.cap.saturating_sub(self.occupied(store)?))
    }

    /// Which queued tasks should start, oldest first.
    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.
    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(()) => {
                    // `transition_from`, naming the state this decision was
                    // made in.
                    match store.transition_from(
                        task.id,
                        TaskState::Queued,
                        TaskState::Running,
                        Transition::Plain,
                        now,
                    ) {
                        Ok(_) => tick.started.push(task.id),
                        Err(crate::store::Error::IllegalTransition { from, .. }) => {
                            tick.abandoned.push((task.id, from));
                        }
                        Err(err) => return Err(err.into()),
                    }
                }
                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();
    }

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

    #[test]
    fn a_paused_task_is_passed_over_for_the_next_one() {
        // Pausing before a task starts is only useful if the queue moves on
        // without it.
        let mut store = store();
        let ids = queue(&mut store, 2);
        hold(&mut store, ids[0]);

        let scheduler = Scheduler::new(3);
        let launcher = Recorder::default();
        let pass = scheduler.tick(&mut store, &launcher, at(3)).unwrap();

        assert_eq!(pass.started, [ids[1]]);
        assert_eq!(
            store.get_task(ids[0]).unwrap().state,
            TaskState::Paused,
            "a held task stays held"
        );
    }

    #[test]
    fn pausing_a_working_task_frees_its_slot() {
        // The reason pausing a running agent is worth having at all: the cap
        // is full, and setting one aside lets the next task start.
        let mut store = store();
        let ids = queue(&mut store, 3);
        run(&mut store, ids[0]);
        run(&mut store, ids[1]);

        let scheduler = Scheduler::new(2);
        assert_eq!(scheduler.available(&store).unwrap(), 0);

        hold(&mut store, ids[0]);

        assert_eq!(scheduler.available(&store).unwrap(), 1);
        let pass = scheduler
            .tick(&mut store, &Recorder::default(), at(3))
            .unwrap();
        assert_eq!(pass.started, [ids[2]]);
    }

    #[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"
        );
    }
}

#[cfg(test)]
mod race_tests {
    use super::*;
    use std::cell::RefCell;
    use std::path::Path;

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

    /// Stands in for the user pausing a task during the window where the
    /// launcher is making worktrees and starting tmux.
    struct PausesWhileLaunching {
        seen: RefCell<Vec<i64>>,
    }

    impl Launch for PausesWhileLaunching {
        fn launch(&self, store: &mut Store, task: &Task) -> std::result::Result<(), String> {
            self.seen.borrow_mut().push(task.id);
            store
                .transition(task.id, TaskState::Paused, Transition::Plain, at(9))
                .expect("the user can pause a queued task at any moment");
            Ok(())
        }
    }

    #[test]
    fn a_task_paused_during_its_launch_stays_paused() {
        let mut store = Store::open_in_memory().unwrap();
        let id = store
            .create_task("held", "do it", Path::new("/tmp/tasks"), &[], at(0))
            .unwrap()
            .id;

        let launcher = PausesWhileLaunching {
            seen: RefCell::new(Vec::new()),
        };
        let _ = Scheduler::new(3).tick(&mut store, &launcher, at(10));

        assert_eq!(
            store.get_task(id).unwrap().state,
            TaskState::Paused,
            "the scheduler must not undo a pause that landed while it was starting the task"
        );
    }
}