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, /// 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, } #[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); impl Config { fn load(path: &Path) -> Result { 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) -> 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 { 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::(&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_202602010015", ] { 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()); } }