hyperhive/swarm-ci-metrics/src/cursor.rs

//! Which runs of each repo are already pushed, kept in one file in the unit's
//! state directory.
//!
//! Per repo, a run number `low` below which every run is handled, and the
//! handled run numbers at or above it. A run is handled once it is pushed or
//! skipped. A run still going holds `low` at or below it, so a run that
//! finishes after a later one is still pushed.
//!
//! The file is replaced whole after every push the collector accepted, and
//! never before. A crash between that push and the write re-pushes those runs
//! on the next pass; nothing is pushed only to be forgotten, and nothing is
//! marked handled before it is pushed. A missing file is a cursor at run 0,
//! so the first pass is the backfill.

use std::collections::{BTreeMap, BTreeSet};
use std::io::Write as _;
use std::path::Path;

use anyhow::{Context as _, Result};
use serde::{Deserialize, Serialize};

/// One repo's cursor.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RepoCursor {
    /// Every run numbered below this is handled.
    pub low: u64,
    /// Handled runs numbered `low` or above.
    pub handled: BTreeSet<u64>,
}

impl RepoCursor {
    /// Whether the walk down the newest-first run list has passed every run
    /// it still has to look at.
    pub fn is_below(&self, run: u64) -> bool {
        run < self.low
    }

    pub fn is_handled(&self, run: u64) -> bool {
        self.is_below(run) || self.handled.contains(&run)
    }

    pub fn mark(&mut self, runs: impl IntoIterator<Item = u64>) {
        self.handled.extend(runs);
    }

    /// Move `low` up once the pass has handled what it could: to the lowest
    /// run still going, or past the newest run seen when none is. Never down.
    pub fn advance(&mut self, lowest_pending: Option<u64>, newest_seen: Option<u64>) {
        let to = match (lowest_pending, newest_seen) {
            (Some(pending), _) => pending,
            (None, Some(newest)) => newest + 1,
            (None, None) => return,
        };
        self.low = self.low.max(to);
        let low = self.low;
        self.handled.retain(|&run| run >= low);
    }
}

/// Every repo's cursor, keyed by `owner/name`.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct State {
    pub repos: BTreeMap<String, RepoCursor>,
}

impl State {
    /// The state at `path`, or an empty one when there is no file.
    pub fn load(path: &Path) -> Result<Self> {
        match std::fs::read(path) {
            Ok(bytes) => serde_json::from_slice(&bytes)
                .with_context(|| format!("parsing the cursor at {}", path.display())),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
            Err(e) => Err(e).with_context(|| format!("reading the cursor at {}", path.display())),
        }
    }

    /// Replace the file at `path`: a reader sees the old state or the new one,
    /// never a mix.
    pub fn save(&self, path: &Path) -> Result<()> {
        let dir = path.parent().context("the cursor path has no directory")?;
        let tmp = path.with_extension("tmp");
        let mut file =
            std::fs::File::create(&tmp).with_context(|| format!("creating {}", tmp.display()))?;
        file.write_all(&serde_json::to_vec(self)?)?;
        file.sync_all()?;
        std::fs::rename(&tmp, path).with_context(|| format!("replacing {}", path.display()))?;
        std::fs::File::open(dir)?.sync_all()?;
        Ok(())
    }

    pub fn repo(&mut self, repo: &str) -> &mut RepoCursor {
        self.repos.entry(repo.to_owned()).or_default()
    }
}

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

    #[test]
    fn a_missing_file_is_a_cursor_at_run_zero() {
        let dir = tempfile::tempdir().expect("a scratch dir");
        let mut state = State::load(&dir.path().join("cursor.json")).expect("no file is fine");
        let c = state.repo("o/r");
        assert_eq!(c.low, 0);
        assert!(!c.is_handled(1));
    }

    #[test]
    fn a_run_still_going_holds_low_and_one_finished_above_it_stays_handled() {
        let mut c = RepoCursor::default();
        // 7 is running; 8 and 9 finished first.
        c.mark([5, 6, 8, 9]);
        c.advance(Some(7), Some(9));
        assert_eq!(c.low, 7);
        assert!(c.is_handled(6));
        assert!(!c.is_handled(7));
        assert!(c.is_handled(8));
        // 7 finishes.
        c.mark([7]);
        c.advance(None, Some(9));
        assert_eq!(c.low, 10);
        assert!(c.handled.is_empty());
    }

    #[test]
    fn low_never_moves_down() {
        let mut c = RepoCursor {
            low: 10,
            handled: BTreeSet::new(),
        };
        c.advance(Some(4), Some(12));
        assert_eq!(c.low, 10);
        c.advance(None, None);
        assert_eq!(c.low, 10);
    }

    #[test]
    fn a_saved_state_reads_back_and_leaves_no_temp_file() {
        let dir = tempfile::tempdir().expect("a scratch dir");
        let path = dir.path().join("cursor.json");
        let mut state = State::default();
        state.repo("o/r").mark([3, 4]);
        state.repo("o/r").advance(Some(3), Some(4));
        state.save(&path).expect("save");
        assert_eq!(State::load(&path).expect("load"), state);
        assert!(!path.with_extension("tmp").exists());
    }

    #[test]
    fn a_corrupt_file_is_an_error_not_a_fresh_backfill() {
        let dir = tempfile::tempdir().expect("a scratch dir");
        let path = dir.path().join("cursor.json");
        std::fs::write(&path, b"{").expect("write");
        assert!(State::load(&path).is_err());
    }
}