rchd 2.1.12

//! Background recovery of dead client leases that still own worker source
//! (bd-nalyr).
//!
//! The lease scan keeps a lease whose wrapper died while its recovery recipe
//! still names worker-side source (bd-dmg2k): that lease is the only authority
//! able to release the worker's source-authority claim. Keeping it is correct,
//! but until now only an operator running `rch jobs recover` finished it, so
//! claims piled up (~15 per half hour across the fleet) and fenced overlapping
//! builds. This task runs that same recovery for leases whose owner is provably
//! gone, a few per cycle, backing off per lease after a failure.

use crate::api::{is_process_alive, lease_blocks_restart, lease_owns_unretired_source};
use crate::events::EventBus;
use rch_common::job_identity::{DurableJobLease, default_job_lease_directory};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;
use tokio::time::Instant;
use tracing::{info, warn};

/// How often the lease directory is checked for recoverable leases.
const RECOVERY_INTERVAL: Duration = Duration::from_secs(5 * 60);
/// Recoveries per cycle; each is serial, so a bad cycle costs at most this many
/// timeouts.
const MAX_RECOVERIES_PER_CYCLE: usize = 4;
/// Deadline handed to `rch jobs recover`, plus slack before the child is killed.
const RECOVER_TIMEOUT_SECS: u64 = 300;
const RECOVER_KILL_AFTER: Duration = Duration::from_secs(RECOVER_TIMEOUT_SECS + 30);
/// Retry delay after the first failure; doubles per failure up to the cap.
const BACKOFF_BASE: Duration = Duration::from_secs(10 * 60);
const BACKOFF_MAX: Duration = Duration::from_secs(6 * 60 * 60);

/// Leases in `lease_dir` that the daemon may recover on the owner's behalf:
/// the owner is provably gone (the same evidence that stops the lease blocking
/// a restart), the recipe still owns worker source, and the wrapper never
/// acknowledged a terminal state (`rch jobs recover` returns early on those).
/// Unreadable entries are skipped; the restart scan already fails closed on
/// them.
pub(crate) fn recoverable_lease_ids(
    lease_dir: &Path,
    now_unix_ms: u64,
    alive: impl Fn(u32) -> bool,
) -> Vec<String> {
    let Ok(entries) = std::fs::read_dir(lease_dir) else {
        return Vec::new();
    };
    let mut ids: Vec<String> = entries
        .filter_map(Result::ok)
        .map(|entry| entry.path())
        .filter(|path| {
            path.extension()
                .is_some_and(|extension| extension == "json")
        })
        .filter_map(|path| std::fs::read(path).ok())
        .filter_map(|bytes| serde_json::from_slice::<DurableJobLease>(&bytes).ok())
        .filter(|lease| {
            !lease.terminal_acknowledged
                && lease_owns_unretired_source(lease)
                && !lease_blocks_restart(lease, now_unix_ms, || alive(lease.wrapper_pid))
        })
        .map(|lease| lease.identity.local_wrapper_id)
        .collect();
    ids.sort();
    ids
}

/// Delay before retrying a lease that has failed `failures` times.
fn backoff_after(failures: u32) -> Duration {
    BACKOFF_BASE
        .saturating_mul(1u32 << failures.saturating_sub(1).min(16))
        .min(BACKOFF_MAX)
}

#[derive(Default)]
struct Backoff {
    failures: HashMap<String, (u32, Instant)>,
}

impl Backoff {
    fn ready(&self, id: &str, now: Instant) -> bool {
        self.failures.get(id).is_none_or(|(_, next)| now >= *next)
    }

    fn record_failure(&mut self, id: &str, now: Instant) -> u32 {
        let failures = self.failures.get(id).map_or(1, |(count, _)| count + 1);
        self.failures
            .insert(id.to_string(), (failures, now + backoff_after(failures)));
        failures
    }

    /// Forget leases that are no longer candidates (recovered, or reaped).
    fn retain(&mut self, candidates: &[String]) {
        self.failures.retain(|id, _| candidates.contains(id));
    }
}

/// Start the recovery loop. `socket` is this daemon's socket; the child
/// `rch` is pointed at it so it never reconciles against another daemon.
pub(crate) fn start(events: EventBus, socket: PathBuf) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        let rch = match std::env::current_exe() {
            Ok(exe) => exe.with_file_name("rch"),
            Err(error) => {
                warn!("Lease auto-recovery disabled: cannot resolve rchd path: {error}");
                return;
            }
        };
        let mut backoff = Backoff::default();
        let mut ticker =
            tokio::time::interval_at(Instant::now() + RECOVERY_INTERVAL, RECOVERY_INTERVAL);
        loop {
            ticker.tick().await;
            if !rch.exists() {
                continue;
            }
            let now_unix_ms = u64::try_from(chrono::Utc::now().timestamp_millis()).unwrap_or(0);
            let candidates = tokio::task::spawn_blocking(move || {
                recoverable_lease_ids(
                    &default_job_lease_directory(),
                    now_unix_ms,
                    is_process_alive,
                )
            })
            .await
            .unwrap_or_default();
            backoff.retain(&candidates);
            let now = Instant::now();
            let due: Vec<&String> = candidates
                .iter()
                .filter(|id| backoff.ready(id, now))
                .take(MAX_RECOVERIES_PER_CYCLE)
                .collect();
            for id in due {
                match recover(&rch, &socket, id).await {
                    Ok(()) => {
                        info!(wrapper_id = %id, "Recovered dead client lease");
                        events.emit(
                            "lease_auto_recovered",
                            &serde_json::json!({ "wrapper_id": id }),
                        );
                    }
                    Err(error) => {
                        let failures = backoff.record_failure(id, Instant::now());
                        warn!(wrapper_id = %id, failures, "Dead client lease recovery failed: {error}");
                        events.emit(
                            "lease_auto_recovery_failed",
                            &serde_json::json!({
                                "wrapper_id": id,
                                "failures": failures,
                                "error": error,
                            }),
                        );
                    }
                }
            }
        }
    })
}

/// Run `rch jobs recover` for one lease; the error is the tail of its stderr.
async fn recover(rch: &Path, socket: &Path, wrapper_id: &str) -> Result<(), String> {
    let child = tokio::process::Command::new(rch)
        .args(["jobs", "recover", wrapper_id, "--timeout-secs"])
        .arg(RECOVER_TIMEOUT_SECS.to_string())
        .env("RCH_SOCKET_PATH", socket)
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::piped())
        .kill_on_drop(true)
        .output();
    let output = match tokio::time::timeout(RECOVER_KILL_AFTER, child).await {
        Ok(Ok(output)) => output,
        Ok(Err(error)) => return Err(format!("cannot run {}: {error}", rch.display())),
        Err(_) => return Err(format!("timed out after {}s", RECOVER_KILL_AFTER.as_secs())),
    };
    if output.status.success() {
        return Ok(());
    }
    let stderr = String::from_utf8_lossy(&output.stderr);
    let tail: String = stderr
        .trim()
        .chars()
        .rev()
        .take(400)
        .collect::<Vec<_>>()
        .into_iter()
        .rev()
        .collect();
    Err(format!("{}: {tail}", output.status))
}

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

    fn lease(heartbeat_unix_ms: u64, pid: u32, recipe: serde_json::Value) -> DurableJobLease {
        let mut identity = JobIdentity::new_local();
        identity.admit(7);
        let mut lease = DurableJobLease::new(
            identity,
            pid,
            None,
            None,
            heartbeat_unix_ms,
            false,
            true,
            "blake3:test".to_string(),
        );
        lease.admit(7, "worker1".to_string(), heartbeat_unix_ms);
        lease.recovery = Some(recipe);
        lease
    }

    #[test]
    fn only_dead_stale_unacknowledged_source_owners_are_recoverable() {
        let dir = tempfile::tempdir().unwrap();
        let now = 100 * 60 * 60 * 1000;
        let stale = now - 60 * 60 * 1000;
        let owning = serde_json::json!({ "source_roots": ["/p"], "pair": null, "retire_root": null, "retired": false });
        let retired = serde_json::json!({ "source_roots": ["/p"], "pair": null, "retire_root": null, "retired": true });
        let (dead, live) = (11, 22);
        let mut acknowledged = lease(stale, dead, owning.clone());
        acknowledged.acknowledge_terminal(stale);
        let leases = [
            ("dead", lease(stale, dead, owning.clone())),
            ("live", lease(stale, live, owning.clone())),
            ("fresh", lease(now - 1000, dead, owning.clone())),
            ("retired", lease(stale, dead, retired)),
            ("no-pid", lease(stale, 0, owning)),
            ("acked", acknowledged),
        ];
        for (name, lease) in &leases {
            std::fs::write(
                dir.path().join(format!("{name}.json")),
                serde_json::to_vec(lease).unwrap(),
            )
            .unwrap();
        }
        std::fs::write(dir.path().join("junk.json"), b"{").unwrap();

        let ids = recoverable_lease_ids(dir.path(), now, |pid| pid == live);

        assert_eq!(ids, vec![leases[0].1.identity.local_wrapper_id.clone()]);
    }

    #[test]
    fn missing_lease_directory_has_no_candidates() {
        let dir = tempfile::tempdir().unwrap();
        assert!(recoverable_lease_ids(&dir.path().join("absent"), 0, |_| false).is_empty());
    }

    #[test]
    fn backoff_doubles_to_a_cap_and_forgets_resolved_leases() {
        assert_eq!(backoff_after(1), BACKOFF_BASE);
        assert_eq!(backoff_after(2), BACKOFF_BASE * 2);
        assert_eq!(backoff_after(40), BACKOFF_MAX);

        let mut backoff = Backoff::default();
        let now = Instant::now();
        assert!(backoff.ready("a", now));
        assert_eq!(backoff.record_failure("a", now), 1);
        assert!(!backoff.ready("a", now));
        assert!(backoff.ready("a", now + BACKOFF_BASE));
        assert_eq!(backoff.record_failure("a", now), 2);
        backoff.retain(&[]);
        assert!(backoff.ready("a", now));
    }
}