hyperhive/hive-agent/src/schedule_timer.rs

//! Fires this agent's swarm-scheduled prompts, in its own container.
//!
//! swarm-controller keeps the agent's list under its key in the
//! `agent-schedules` KV bucket. This task reads that key every
//! [`POLL_INTERVAL`] with the one subject the agent's own queue credential is
//! granted there (`swarm_queue_client::agent_schedules::read_own`), and fires
//! each due schedule as a `schedule`-subsystem todo, keyed `schedule:<id>`.
//! The wake path is the todo store's.
//!
//! **Fire times come from the schedule alone**: `anchor + k * interval`
//! ([`Schedule::latest_slot`]). Rewriting the list, or changing a body,
//! leaves them where they were.
//!
//! **No duplicates.** The last slot fired for each schedule id is kept in the
//! state db ([`FiredSlots`]), and a slot fires only when it is later than
//! that. After downtime the latest missed slot fires once and the earlier
//! ones not at all; a restart refires nothing already recorded. A one-shot
//! has one slot, so it fires once however long it stays listed. Records are
//! never dropped when a schedule leaves the list: ids are never reused, and
//! a list that reads empty for a moment must not erase what was fired.
//!
//! **Paused** schedules do not fire, and each slot that passes while one is
//! paused is recorded as if fired, so resuming does not fire it.
//!
//! The list itself is held in memory only: until the first successful read
//! after a start, nothing fires.

use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use anyhow::{Context as _, Result};
use rusqlite::{Connection, OptionalExtension as _, params};
use swarm_queue_client::agent_schedules::{Read, read_own};
use swarm_wire::agent_schedules::{AgentSchedules, Schedule};
use tokio::sync::Notify;

use crate::todos::Todos;

/// How often the key is read and the list checked for due slots. Also the
/// longest a fire can lag its slot, and a change can take to reach the agent.
const POLL_INTERVAL: Duration = Duration::from_secs(10);

/// The todo subsystem and source a schedule fire carries.
const SUBSYSTEM: &str = "schedule";
const SOURCE: &str = "scheduled";

const SCHEMA: &str = r"
CREATE TABLE IF NOT EXISTS schedule_fired_slots (
    schedule_id TEXT    PRIMARY KEY,
    slot_unix   INTEGER NOT NULL
);
";

/// The last slot fired, or passed while paused, per schedule id.
pub struct FiredSlots {
    conn: Mutex<Connection>,
}

impl FiredSlots {
    /// Open (creating if needed) the table in the db at `path`.
    ///
    /// # Errors
    ///
    /// Propagates sqlite open / schema-apply failures.
    pub fn open(path: &Path) -> Result<Self> {
        let conn = Connection::open(path)
            .with_context(|| format!("open schedule slots db {}", path.display()))?;
        conn.execute_batch(SCHEMA)
            .context("apply schedule slots schema")?;
        Ok(Self {
            conn: Mutex::new(conn),
        })
    }

    fn last(&self, id: &str) -> Result<Option<i64>> {
        let conn = self
            .conn
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        Ok(conn
            .query_row(
                "SELECT slot_unix FROM schedule_fired_slots WHERE schedule_id = ?1",
                params![id],
                |row| row.get(0),
            )
            .optional()?)
    }

    fn record(&self, id: &str, slot: i64) -> Result<()> {
        let conn = self
            .conn
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        conn.execute(
            "INSERT INTO schedule_fired_slots (schedule_id, slot_unix) VALUES (?1, ?2) \
             ON CONFLICT(schedule_id) DO UPDATE SET slot_unix = excluded.slot_unix",
            params![id, slot],
        )?;
        Ok(())
    }
}

/// What one schedule needs at `now`.
#[derive(Debug, PartialEq, Eq)]
enum Step {
    /// Fire, then record the slot.
    Fire(i64),
    /// Record the slot without firing: it passed while paused.
    Skip(i64),
    Nothing,
}

fn decide(schedule: &Schedule, last: Option<i64>, now: i64) -> Step {
    let Some(slot) = schedule.latest_slot(now) else {
        return Step::Nothing;
    };
    if last.is_some_and(|last| last >= slot) {
        return Step::Nothing;
    }
    if schedule.paused {
        Step::Skip(slot)
    } else {
        Step::Fire(slot)
    }
}

/// Apply [`decide`] to every schedule in `list`. `fire` returns whether the
/// fire landed; a slot is recorded only after it did, so a failed fire is
/// retried on the next tick.
fn tick(
    schedules: &AgentSchedules,
    slots: &FiredSlots,
    now: i64,
    mut fire: impl FnMut(&Schedule) -> bool,
) {
    for schedule in &schedules.schedules {
        let previous = match slots.last(&schedule.id) {
            Ok(previous) => previous,
            Err(e) => {
                tracing::warn!(id = %schedule.id, error = ?e, "schedule: reading the fired slot failed");
                continue;
            }
        };
        let slot = match decide(schedule, previous, now) {
            Step::Nothing => continue,
            Step::Skip(slot) => slot,
            Step::Fire(slot) => {
                if !fire(schedule) {
                    continue;
                }
                tracing::info!(id = %schedule.id, slot, "schedule fired");
                slot
            }
        };
        if let Err(e) = slots.record(&schedule.id, slot) {
            tracing::warn!(id = %schedule.id, slot, error = ?e, "schedule: recording the slot failed");
        }
    }
}

/// Push `schedule` as a todo, and wake the turn
/// loop when it is new or reopened.
fn push_todo(todos: &Todos, wake: &Notify, schedule: &Schedule) -> bool {
    let key = format!("schedule:{}", schedule.id);
    // `true`: each fire is a new occurrence, so an acked earlier one reopens.
    match todos.upsert(SUBSYSTEM, Some(&key), &schedule.body, Some(SOURCE), true) {
        Ok((_, changed)) => {
            if changed {
                wake.notify_one();
            }
            true
        }
        Err(e) => {
            tracing::warn!(id = %schedule.id, error = ?e, "schedule: pushing the todo failed");
            false
        }
    }
}

/// Run the ticker for the harness's lifetime, if this agent has a label to
/// name its key with.
pub async fn run(todos: Arc<Todos>, wake: Arc<Notify>) {
    let agent = crate::identity::label();
    if agent.is_empty() {
        tracing::warn!("this agent has no label; not reading its swarm schedules");
        return;
    }
    let slots = match FiredSlots::open(&crate::paths::state_db()) {
        Ok(slots) => slots,
        Err(e) => {
            tracing::error!(error = ?e, "open schedule slots failed — swarm schedules disabled");
            return;
        }
    };
    let client = crate::swarm_queue::client().await;
    tracing::info!(%agent, "reading swarm schedules from the swarm queue");
    let mut list: Option<AgentSchedules> = None;
    loop {
        refresh(&client, &agent, &mut list).await;
        if let Some(list) = &list {
            let now = chrono::Utc::now().timestamp();
            tick(list, &slots, now, |s| push_todo(&todos, &wake, s));
        }
        tokio::time::sleep(POLL_INTERVAL).await;
    }
}

/// Replace `list` with the key's current value. On any failure the previous
/// list stays, so schedules keep firing through a queue outage.
async fn refresh(client: &async_nats::Client, agent: &str, list: &mut Option<AgentSchedules>) {
    // An unconnected client does not fail a JetStream request, it hangs on it.
    if swarm_queue_client::ensure_connected(client).is_err() {
        return;
    }
    match read_own(client, agent).await {
        Ok(Read::Current(current)) => *list = Some(current),
        Ok(Read::NoBucket) => {
            tracing::debug!("schedule: the agent-schedules bucket does not answer yet");
        }
        Err(e) => tracing::warn!(error = %e, "schedule: reading this agent's schedules failed"),
    }
}

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

    fn schedule(id: &str, anchor: i64, interval: Option<u64>) -> Schedule {
        Schedule {
            id: id.to_owned(),
            body: format!("{id} body"),
            anchor_unix: anchor,
            interval_seconds: interval,
            paused: false,
        }
    }

    fn list(schedules: Vec<Schedule>) -> AgentSchedules {
        AgentSchedules { schedules }
    }

    fn slots() -> FiredSlots {
        FiredSlots::open(Path::new(":memory:")).unwrap()
    }

    /// The bodies fired by one tick at `now`.
    fn fired(list: &AgentSchedules, slots: &FiredSlots, now: i64) -> Vec<String> {
        let mut out = Vec::new();
        tick(list, slots, now, |s| {
            out.push(s.body.clone());
            true
        });
        out
    }

    #[test]
    fn a_body_edit_does_not_move_the_cadence() {
        let slots = slots();
        let mut l = list(vec![schedule("a", 1000, Some(100))]);
        assert_eq!(fired(&l, &slots, 1000), ["a body"]);
        "edited".clone_into(&mut l.schedules[0].body);
        assert!(fired(&l, &slots, 1050).is_empty());
        assert!(fired(&l, &slots, 1099).is_empty());
        assert_eq!(fired(&l, &slots, 1100), ["edited"]);
    }

    #[test]
    fn an_offline_gap_of_several_slots_fires_once_then_rejoins_the_grid() {
        let slots = slots();
        let l = list(vec![schedule("a", 1000, Some(100))]);
        assert_eq!(fired(&l, &slots, 1000), ["a body"]);
        // Slots 1100 through 1500 pass unseen.
        assert_eq!(fired(&l, &slots, 1550), ["a body"]);
        assert!(fired(&l, &slots, 1560).is_empty());
        assert!(fired(&l, &slots, 1599).is_empty());
        assert_eq!(fired(&l, &slots, 1600), ["a body"]);
    }

    #[test]
    fn a_restart_after_a_fire_does_not_refire() {
        let dir = tempfile::tempdir().unwrap();
        let db = dir.path().join("state.sqlite");
        let l = list(vec![schedule("a", 1000, Some(100))]);
        let before = FiredSlots::open(&db).unwrap();
        assert_eq!(fired(&l, &before, 1010), ["a body"]);
        drop(before);
        let after = FiredSlots::open(&db).unwrap();
        assert!(fired(&l, &after, 1020).is_empty());
        // Control: the reopened store still fires the next slot.
        assert_eq!(fired(&l, &after, 1100), ["a body"]);
    }

    /// An agent dropped from a schedule and added back gets the same id, so
    /// the slot it already fired stays fired.
    #[test]
    fn leaving_and_rejoining_a_schedule_does_not_refire() {
        let slots = slots();
        let l = list(vec![schedule("a", 1000, Some(100))]);
        assert_eq!(fired(&l, &slots, 1000), ["a body"]);
        assert!(fired(&list(Vec::new()), &slots, 1020).is_empty());
        assert!(fired(&l, &slots, 1040).is_empty());
        assert_eq!(fired(&l, &slots, 1100), ["a body"]);
    }

    #[test]
    fn a_list_refresh_does_not_refire() {
        let slots = slots();
        let l = list(vec![
            schedule("a", 1000, Some(100)),
            schedule("b", 900, None),
        ]);
        assert_eq!(fired(&l, &slots, 1000), ["a body", "b body"]);
        let refreshed: AgentSchedules =
            serde_json::from_slice(&serde_json::to_vec(&l).unwrap()).unwrap();
        assert!(fired(&refreshed, &slots, 1001).is_empty());
        // An empty read in between erases nothing.
        assert!(fired(&list(Vec::new()), &slots, 1002).is_empty());
        assert!(fired(&refreshed, &slots, 1003).is_empty());
    }

    #[test]
    fn a_one_shot_fires_once() {
        let slots = slots();
        let l = list(vec![schedule("once", 1000, None)]);
        assert!(fired(&l, &slots, 999).is_empty());
        assert_eq!(fired(&l, &slots, 1000), ["once body"]);
        assert!(fired(&l, &slots, 1001).is_empty());
        assert!(fired(&l, &slots, 1_000_000).is_empty());
    }

    #[test]
    fn a_paused_schedule_does_not_fire_or_catch_up() {
        let slots = slots();
        let mut l = list(vec![schedule("a", 1000, Some(100))]);
        l.schedules[0].paused = true;
        assert!(fired(&l, &slots, 1000).is_empty());
        assert!(fired(&l, &slots, 1250).is_empty());
        l.schedules[0].paused = false;
        // Slot 1200 passed while paused, so resuming does not fire it.
        assert!(fired(&l, &slots, 1260).is_empty());
        assert_eq!(fired(&l, &slots, 1300), ["a body"]);
    }

    #[test]
    fn a_failed_fire_is_retried_and_not_recorded() {
        let slots = slots();
        let l = list(vec![schedule("a", 1000, Some(100))]);
        tick(&l, &slots, 1000, |_| false);
        assert_eq!(fired(&l, &slots, 1005), ["a body"]);
    }

    /// A fire is subsystem `schedule`, key `schedule:<id>`, source
    /// `scheduled`, and reopens once the agent acked the previous fire.
    #[test]
    fn a_fire_is_the_schedule_todo() {
        let todos = Todos::open(Path::new(":memory:")).unwrap();
        let wake = Notify::new();
        let s = schedule("abcdefgh", 1000, Some(100));
        assert!(push_todo(&todos, &wake, &s));
        let listed = todos.list(Some(SUBSYSTEM)).unwrap();
        assert_eq!(listed.len(), 1);
        assert_eq!(
            listed[0].subsystem_key.as_deref(),
            Some("schedule:abcdefgh")
        );
        assert_eq!(listed[0].summary, "abcdefgh body");
        assert_eq!(listed[0].source.as_deref(), Some(SOURCE));
        todos.mark_done(listed[0].id).unwrap();
        assert!(push_todo(&todos, &wake, &s));
        assert_eq!(todos.list(Some(SUBSYSTEM)).unwrap().len(), 1);
    }
}