dotfiles/rust/snapshot_manager/src/main.rs

use std::{
    collections::HashMap,
    fs,
    path::{Path, PathBuf},
    process::ExitCode,
};

use anyhow::{Context, Result};
use chrono::{DateTime, Datelike, NaiveDate, NaiveTime, Timelike, Utc};
use clap::Parser;
use log::{error, info};
use serde::Deserialize;
use signal::Signal;

#[derive(Parser)]
#[command(version, about)]
struct Args {
    /// TOML file containing default and per-dataset snapshot retention counts.
    config_file: PathBuf,

    /// Send Signal failure alerts (requires both phone numbers).
    #[arg(long, env = "SIGNAL_ALERTS", default_value_t = false, action = clap::ArgAction::Set)]
    signal_alerts: bool,

    /// Signal account sending failure alerts.
    #[arg(
        long,
        env = "SIGNAL_ALERT_FROM_PHONE",
        hide_env_values = true,
        required_if_eq("signal_alerts", "true"),
        value_parser = clap::builder::NonEmptyStringValueParser::new()
    )]
    signal_alert_from_phone: Option<String>,

    /// Signal recipient for failure alerts.
    #[arg(
        long,
        env = "SIGNAL_ALERT_TO_PHONE",
        hide_env_values = true,
        required_if_eq("signal_alerts", "true"),
        value_parser = clap::builder::NonEmptyStringValueParser::new()
    )]
    signal_alert_to_phone: Option<String>,
}

#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct Retention {
    #[serde(rename = "15_min")]
    quarter_hourly: usize,
    hourly: usize,
    daily: usize,
    monthly: usize,
}

impl Default for Retention {
    fn default() -> Self {
        Self {
            quarter_hourly: 4,
            hourly: 12,
            daily: 0,
            monthly: 0,
        }
    }
}

#[derive(Debug, Deserialize)]
#[serde(transparent)]
struct Config(HashMap<String, Retention>);

impl Config {
    fn load(path: &Path) -> Result<Self> {
        let contents = fs::read_to_string(path)
            .with_context(|| format!("Failed to read configuration {}", path.display()))?;
        toml::from_str(&contents)
            .with_context(|| format!("Invalid configuration {}", path.display()))
    }

    fn retention_for(&self, dataset: &str) -> Retention {
        self.0
            .get(dataset)
            .or_else(|| self.0.get("default"))
            .copied()
            .unwrap_or_default()
    }
}

fn main() -> ExitCode {
    let args = Args::parse();
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
    info!("Starting snapshot_manager");
    let signal = if args.signal_alerts {
        args.signal_alert_from_phone
            .map(|from_phone| Signal { from_phone })
    } else {
        None
    };
    let alerts = signal.as_ref().zip(args.signal_alert_to_phone.as_deref());
    match run(&args.config_file, alerts) {
        Ok(()) => {
            info!("snapshot_manager completed");
            ExitCode::SUCCESS
        }
        Err(error) => {
            error!("{error:#}");
            ExitCode::FAILURE
        }
    }
}

fn run(config_file: &Path, alerts: Option<(&Signal, &str)>) -> Result<()> {
    let result = (|| {
        // Validate all retention settings before making any changes to ZFS.
        let config = Config::load(config_file)?;
        let snapshot_time = timestamp(Utc::now());
        for dataset_name in zfs::list_datasets()? {
            if let Err(error) = zfs::create_snapshot(&dataset_name, &snapshot_time) {
                let message = format!("{dataset_name} failed to create snapshot {snapshot_time}");
                error!("{message}: {error:#}");
                send_alert(alerts, &message)?;
                continue;
            }
            remove_old_snapshots(&config, alerts, &dataset_name)
                .context("Failed to remove old snapshots")?;
        }
        Ok(())
    })();
    if let Err(error) = &result {
        error!("snapshot_manager failed: {error:#}");
        if let Err(alert_error) = send_alert(alerts, "snapshot_manager failed") {
            error!("Failed to send failure alert: {alert_error:#}");
        }
    }
    result
}

fn timestamp(now: DateTime<Utc>) -> String {
    let minute = now.minute().saturating_sub(now.minute() % 15);
    format!("auto_{}{minute:02}", now.format("%Y%m%d%H"))
}

fn remove_old_snapshots(
    config: &Config,
    alerts: Option<(&Signal, &str)>,
    dataset_name: &str,
) -> Result<()> {
    let retention = config.retention_for(dataset_name);
    info!("Using {retention:?} for {dataset_name}");
    let snapshots = zfs::list_snapshots(dataset_name)?;
    for snapshot in snapshots_to_delete(&snapshots, retention) {
        info!("Deleting {dataset_name}@{snapshot}");
        if let Some(error) = zfs::delete_snapshot(dataset_name, snapshot)? {
            let message = format!("{dataset_name}@{snapshot} failed to delete: {error}");
            error!("{message}");
            send_alert(alerts, &message)?;
        }
    }
    Ok(())
}

fn snapshots_to_delete(snapshots: &[String], retention: Retention) -> Vec<&str> {
    let mut tiers: [Vec<&str>; 4] = Default::default();
    for snapshot in snapshots {
        if let Some(index) = bucket(snapshot)
            && let Some(tier) = tiers.get_mut(index)
        {
            tier.push(snapshot);
        }
    }

    let counts = [
        retention.quarter_hourly,
        retention.hourly,
        retention.daily,
        retention.monthly,
    ];
    let mut deletions = Vec::new();
    for (mut tier, count) in tiers.into_iter().zip(counts) {
        tier.sort_unstable();
        let delete_count = tier.len().saturating_sub(count);
        deletions.extend(tier.into_iter().take(delete_count));
    }
    deletions
}

fn bucket(name: &str) -> Option<usize> {
    let digits = name.strip_prefix("auto_")?;
    if digits.len() != 12 || !digits.bytes().all(|byte| byte.is_ascii_digit()) {
        return None;
    }
    let date = NaiveDate::from_ymd_opt(
        digits.get(..4)?.parse().ok()?,
        digits.get(4..6)?.parse().ok()?,
        digits.get(6..8)?.parse().ok()?,
    )?;
    let time = NaiveTime::from_hms_opt(
        digits.get(8..10)?.parse().ok()?,
        digits.get(10..)?.parse().ok()?,
        0,
    )?;

    // The tiers are disjoint: midnight belongs to daily, except on the first
    // of the month, when it belongs to monthly.
    match (time.minute(), time.hour(), date.day()) {
        (15 | 30 | 45, _, _) => Some(0),
        (0, 1..=23, _) => Some(1),
        (0, 0, 2..=31) => Some(2),
        (0, 0, 1) => Some(3),
        _ => None,
    }
}

fn send_alert(alerts: Option<(&Signal, &str)>, body: &str) -> Result<()> {
    if let Some((signal, to_phone)) = alerts {
        signal.send(to_phone, body)?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use chrono::TimeZone;

    use super::*;

    #[test]
    fn dataset_overrides_default_and_other_datasets_fall_back() {
        let config: Config = toml::from_str(
            r#"
            [default]
            15_min = 8
            hourly = 24
            daily = 0
            monthly = 0

            ["pool/home"]
            15_min = 2
            hourly = 3
            daily = 4
            monthly = 5
            "#,
        )
        .unwrap();
        assert_eq!(
            config.retention_for("pool/home"),
            Retention {
                quarter_hourly: 2,
                hourly: 3,
                daily: 4,
                monthly: 5,
            }
        );
        assert_eq!(config.retention_for("pool/other").hourly, 24);
        assert_eq!(config.retention_for("pool/home/child").quarter_hourly, 8);
    }

    #[test]
    fn absent_default_uses_original_builtin_counts() {
        let config: Config = toml::from_str("").unwrap();
        assert_eq!(config.retention_for("pool/home"), Retention::default());
    }

    #[test]
    fn rejects_incomplete_negative_and_unknown_retention_settings() {
        for settings in [
            "15_min = 1\nhourly = 2\ndaily = 3",
            "15_min = -1\nhourly = 2\ndaily = 3\nmonthly = 4",
            "15_min = 1\nhourly = 2\ndaily = 3\nmonthly = 4\nweekly = 5",
            "15_min = 1.5\nhourly = 2\ndaily = 3\nmonthly = 4",
        ] {
            assert!(toml::from_str::<Config>(&format!("[default]\n{settings}")).is_err());
        }
    }

    #[test]
    fn timestamps_round_down_at_quarter_hour_boundaries() {
        for (minute, expected) in [
            (0, "00"),
            (14, "00"),
            (15, "15"),
            (29, "15"),
            (30, "30"),
            (44, "30"),
            (45, "45"),
            (59, "45"),
        ] {
            let now = Utc.with_ymd_and_hms(2026, 1, 1, 0, minute, 59).unwrap();
            assert_eq!(timestamp(now), format!("auto_2026010100{expected}"));
        }
    }

    #[test]
    fn tiers_are_disjoint_at_day_and_month_boundaries() {
        for (name, expected) in [
            ("auto_202602010015", 0),
            ("auto_202602010030", 0),
            ("auto_202602010045", 0),
            ("auto_202602010100", 1),
            ("auto_202602012300", 1),
            ("auto_202602020000", 2),
            ("auto_202402290000", 2),
            ("auto_202602010000", 3),
            ("auto_202601010000", 3),
        ] {
            assert_eq!(bucket(name), Some(expected), "{name}");
        }
    }

    #[test]
    fn only_exact_valid_manager_names_are_eligible() {
        for name in [
            "manual",
            "manual_auto_202602010015",
            "auto_202602010015_backup",
            "auto_20260201001",
            "auto_202602300000",
            "auto_202602012400",
            "auto_202602010060",
            "auto_202602010005",
            "auto_２０２６02010015",
        ] {
            assert_eq!(bucket(name), None, "{name}");
        }
    }

    #[test]
    fn keeps_newest_in_each_tier_regardless_of_input_order() {
        let snapshots = [
            "auto_202601010000",
            "auto_202602020000",
            "auto_202602020115",
            "auto_202602020215",
            "auto_202602020200",
            "auto_202602020100",
            "auto_202602030000",
            "auto_202602010000",
            "manual",
        ]
        .map(String::from);
        let retention = Retention {
            quarter_hourly: 1,
            hourly: 1,
            daily: 1,
            monthly: 1,
        };
        assert_eq!(
            snapshots_to_delete(&snapshots, retention),
            [
                "auto_202602020115",
                "auto_202602020100",
                "auto_202602020000",
                "auto_202601010000",
            ]
        );
    }

    #[test]
    fn zero_counts_delete_all_managed_snapshots_and_large_counts_keep_them() {
        let snapshots = [
            "auto_202601010015",
            "auto_202601010100",
            "auto_202601020000",
            "auto_202601010000",
            "manual",
        ]
        .map(String::from);
        let none = Retention {
            quarter_hourly: 0,
            hourly: 0,
            daily: 0,
            monthly: 0,
        };
        assert_eq!(snapshots_to_delete(&snapshots, none).len(), 4);
        let all = Retention {
            quarter_hourly: 100,
            hourly: 100,
            daily: 100,
            monthly: 100,
        };
        assert!(snapshots_to_delete(&snapshots, all).is_empty());
        assert!(snapshots_to_delete(&[], none).is_empty());
    }
}