feat(snapshot-manager): migrate snapshot management to Rust

rust/snapshot_manager/src/main.rs

@@ -0,0 +1,383 @@
+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());
+    }
+}