wsb 0.275.7622

//! The worktree lock: the tree held still while a snapshot of it is taken.
//!
//! A test run on another machine, or a ship, archives the worktree as it
//! stands. Anything that writes into the tree while that copy is being made —
//! a checkpoint, a render, a regeneration — gives a snapshot that is neither
//! the old tree nor the new one, and the run then tests something nobody has.
//! Nobody had to remember not to do that until it happened: a checkpoint
//! taken while a ship was staging broke the ship.
//!
//! Who takes it: only what snapshots or archives the worktree, and only for
//! as long as the copy takes — released once the state is archived, before it
//! is sent anywhere. Who respects it: every command that writes into the tree.
//! It waits, saying what it is waiting for, and gives up after
//! [`WAIT_CAP`] naming the holder. `--ignore-worktree-lock` proceeds
//! regardless, and says so.
//!
//! Both ways. A writer registers itself for as long as it runs
//! ([`begin_writing`], a file per process in [`WRITERS`]), and a snapshot,
//! once it holds the lock, waits for every registered writer to finish. It
//! did not: the lock only stopped writers from STARTING, so a test snapshot
//! taken in the middle of a checkpoint copied a tree half taken back from the
//! local render — `pyproject.toml` asking for one lungo release and the
//! generated code for another — and the run failed building it.
//!
//! The lock is a file in the worktree's local state
//! ([`crate::local_state::LocalState::lock`]), holding who took it and why.
//! Clearing it by hand is deleting that file; every message about it says so.

use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};

/// Set (to anything) to proceed despite the lock: what
/// `--ignore-worktree-lock` sets, so every wsb a command starts inherits it.
pub const IGNORE_ENV: &str = "WSB_IGNORE_WORKTREE_LOCK";

/// Set in a registered writer's environment, so the wsb processes it starts —
/// a checkpoint's renders, the hooks its commits run — are covered by its
/// registration rather than registering, or waiting, themselves. A child that
/// waited on a snapshot which waits on its parent would wait for ever.
pub const WRITING_ENV: &str = "WSB_WORKTREE_WRITER";

/// How long a snapshot waits for running writers. Longer than [`WAIT_CAP`]: a
/// writer is a live process doing its work — a checkpoint takes minutes — not
/// a lock somebody forgot.
pub const WRITERS_WAIT_CAP: Duration = Duration::from_secs(30 * 60);

/// How long a command waits for the lock before it gives up. A snapshot takes
/// seconds to a few minutes; a lock held longer is a snapshot that is stuck,
/// or a holder that died without clearing it.
pub const WAIT_CAP: Duration = Duration::from_secs(5 * 60);

const POLL: Duration = Duration::from_secs(2);

/// Who holds the lock, and why — the file's contents.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Holder {
    pub command: String,
    pub argv: Vec<String>,
    pub pid: u32,
    pub host: String,
    pub taken_at: String,
    pub reason: String,
    pub protects: String,
}

/// The lock's path for the worktree at `root`.
pub fn path(root: &Path) -> Result<PathBuf, String> {
    Ok(crate::local_state::LocalState::of(root).map_err(|e| format!("{e:#}"))?.lock())
}

/// The folder the worktree at `root`'s running writers register in.
fn writers(root: &Path) -> Result<PathBuf, String> {
    Ok(crate::local_state::LocalState::of(root).map_err(|e| format!("{e:#}"))?.writers())
}

/// The lock, held. Released when dropped — if it is still this process's.
#[derive(Debug)]
pub struct Held {
    path: PathBuf,
    pid: u32,
}

impl Held {
    /// The lock file, for a process that releases it on this one's behalf
    /// (`mg-test run --release-after-snapshot`).
    pub fn path(&self) -> &Path {
        &self.path
    }
}

impl Drop for Held {
    fn drop(&mut self) {
        // Only if it is still ours: the process it was handed to may have
        // released it already, and another snapshot may hold it now.
        if read(&self.path).ok().flatten().is_some_and(|holder| holder.pid == self.pid) {
            let _ = std::fs::remove_file(&self.path);
        }
    }
}

/// Take the lock to snapshot the workspace at `root`, for `reason`. A lock
/// another snapshot holds is waited for, as any writer waits; the file is
/// created only if nobody holds it, so two snapshots never both think they
/// do.
pub fn take(root: &Path, command: &str, reason: &str) -> Result<Held, String> {
    take_within(root, command, reason, WAIT_CAP)
}

fn take_within(root: &Path, command: &str, reason: &str, cap: Duration) -> Result<Held, String> {
    let lock = path(root)?;
    if let Some(parent) = lock.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| format!("cannot create {}: {e}", parent.display()))?;
    }
    let holder = Holder {
        command: command.to_string(),
        argv: std::env::args().collect(),
        pid: std::process::id(),
        host: host_name(),
        taken_at: chrono::Utc::now().to_rfc3339(),
        reason: reason.to_string(),
        protects: format!("the worktree at {}, while it is snapshotted", root.display()),
    };
    let text = serde_json::to_string_pretty(&holder).expect("a holder serializes");
    let started = Instant::now();
    loop {
        wait_until_free_since(root, command, started, cap)?;
        if create_exclusively(&lock, &text)? {
            // Held from here, so no writer starts; dropped on any error.
            let held = Held { path: lock, pid: holder.pid };
            wait_for_writers(root, command, Instant::now(), WRITERS_WAIT_CAP)?;
            return Ok(held);
        }
        // Taken between the wait and the create: wait again.
    }
}

/// A registered writer: removed from the writers' folder when dropped. A process that
/// exits without dropping it leaves a file naming a dead process, which a
/// snapshot does not wait for.
#[derive(Debug)]
pub struct Writing {
    file: PathBuf,
}

impl Drop for Writing {
    fn drop(&mut self) {
        let _ = std::fs::remove_file(&self.file);
    }
}

/// Register `operation` as writing into the worktree at `root`, for as long
/// as the returned guard lives: wait while a snapshot holds the lock, register,
/// and if a snapshot took the lock in between, step back and wait again — so a
/// snapshot never copies a tree a writer is half-way through. `None` for a
/// process a registered writer started ([`WRITING_ENV`]), which its parent's
/// registration covers.
pub fn begin_writing(root: &Path, operation: &str) -> Result<Option<Writing>, String> {
    if std::env::var_os(WRITING_ENV).is_some() {
        return Ok(None);
    }
    let ignore = std::env::var_os(IGNORE_ENV).is_some();
    let dir = writers(root)?;
    std::fs::create_dir_all(&dir).map_err(|e| format!("cannot create {}: {e}", dir.display()))?;
    let pid = std::process::id();
    let record = Holder {
        command: operation.to_string(),
        argv: std::env::args().collect(),
        pid,
        host: host_name(),
        taken_at: chrono::Utc::now().to_rfc3339(),
        reason: format!("{operation} is writing into the tree"),
        protects: format!("the worktree at {}, while it is written", root.display()),
    };
    let text = serde_json::to_string_pretty(&record).expect("a holder serializes");
    let file = dir.join(format!("{}-{pid}.json", host_name()));
    loop {
        wait_until_free_with(root, operation, ignore, WAIT_CAP)?;
        std::fs::write(&file, &text).map_err(|e| format!("cannot write {}: {e}", file.display()))?;
        let writing = Writing { file: file.clone() };
        if ignore || read(&path(root)?)?.is_none() {
            std::env::set_var(WRITING_ENV, pid.to_string());
            return Ok(Some(writing));
        }
        // A snapshot took the lock between the wait and the registration: it
        // is now waiting for this writer, so this writer steps back for it.
        drop(writing);
    }
}

/// Wait until no live writer is registered on the worktree at `root`.
fn wait_for_writers(root: &Path, command: &str, started: Instant, cap: Duration) -> Result<(), String> {
    let dir = writers(root)?;
    let mut said = false;
    loop {
        let writers = live_writers(&dir)?;
        let Some(first) = writers.first() else {
            return Ok(());
        };
        if started.elapsed() >= cap {
            return Err(format!(
                "{command}: {} (pid {}) is still writing into the worktree after {}s, so it cannot be \
                 snapshotted. Its registration is {}.",
                first.command,
                first.pid,
                started.elapsed().as_secs(),
                dir.display()
            ));
        }
        if !said {
            crate::outln!(
                "{command}: waiting for {} (pid {}) to finish writing into the worktree",
                first.command,
                first.pid
            );
            said = true;
        }
        std::thread::sleep(POLL);
    }
}

/// The writers registered in `dir` whose process is still running. Another
/// machine's are counted live: their processes cannot be asked about.
fn live_writers(dir: &Path) -> Result<Vec<Holder>, String> {
    let entries = match std::fs::read_dir(dir) {
        Ok(entries) => entries,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
        Err(e) => return Err(format!("cannot read {}: {e}", dir.display())),
    };
    let mut live = Vec::new();
    for entry in entries {
        let entry = entry.map_err(|e| format!("cannot read {}: {e}", dir.display()))?;
        let Some(writer) = read(&entry.path())? else {
            continue;
        };
        if writer.host == host_name() && !crate::ops::process::is_alive(writer.pid) {
            // Its process ended without unregistering — an `exit` skips drops.
            let _ = std::fs::remove_file(entry.path());
            continue;
        }
        live.push(writer);
    }
    Ok(live)
}

/// Write `text` to `lock` only if no such file exists — the one step two
/// snapshots racing between their waits and their creates cannot both win.
/// Whether this call made it.
fn create_exclusively(lock: &Path, text: &str) -> Result<bool, String> {
    match std::fs::OpenOptions::new().write(true).create_new(true).open(lock) {
        Ok(mut file) => {
            use std::io::Write;
            file.write_all(text.as_bytes())
                .map_err(|e| format!("cannot write {}: {e}", lock.display()))?;
            Ok(true)
        }
        Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(false),
        Err(e) => Err(format!("cannot take the worktree lock {}: {e}", lock.display())),
    }
}

/// Wait until nobody holds the lock on the worktree at `root`, before
/// `operation` writes into it. Proceeds at once when [`IGNORE_ENV`] is set,
/// saying so.
pub fn wait_until_free(root: &Path, operation: &str) -> Result<(), String> {
    wait_until_free_with(root, operation, std::env::var_os(IGNORE_ENV).is_some(), WAIT_CAP)
}

fn wait_until_free_with(root: &Path, operation: &str, ignore: bool, cap: Duration) -> Result<(), String> {
    if ignore {
        if let Some(holder) = read(&path(root)?)? {
            crate::outln!(
                "{operation}: ignoring the worktree lock held by {} (pid {}, {}) — --ignore-worktree-lock",
                holder.command,
                holder.pid,
                holder.reason
            );
        }
        return Ok(());
    }
    wait_until_free_since(root, operation, Instant::now(), cap)
}

fn wait_until_free_since(root: &Path, operation: &str, started: Instant, cap: Duration) -> Result<(), String> {
    let lock = path(root)?;
    let mut said = false;
    loop {
        let Some(holder) = read(&lock)? else {
            return Ok(());
        };
        if holder.host == host_name() && !crate::ops::process::is_alive(holder.pid) {
            return Err(format!(
                "{operation}: the worktree lock is stale — {} (pid {}) took it at {} to {}, \
                 and that process is gone. Delete {} to clear it.",
                holder.command,
                holder.pid,
                holder.taken_at,
                holder.reason,
                lock.display()
            ));
        }
        if started.elapsed() >= cap {
            return Err(format!(
                "{operation}: the worktree is still locked after {}s — {} (pid {} on {}) took it at {} to {}. \
                 Delete {} to clear it, or pass --ignore-worktree-lock.",
                started.elapsed().as_secs(),
                holder.command,
                holder.pid,
                holder.host,
                holder.taken_at,
                holder.reason,
                lock.display()
            ));
        }
        if !said {
            crate::outln!(
                "{operation}: waiting on the worktree lock — {} (pid {}) is {}. \
                 Deleting {} clears it.",
                holder.command,
                holder.pid,
                holder.reason,
                lock.display()
            );
            said = true;
        }
        std::thread::sleep(POLL);
    }
}

/// The holder, if the lock is held. A lock file that is not a holder record
/// is refused rather than taken for absent or held.
fn read(lock: &Path) -> Result<Option<Holder>, String> {
    match std::fs::read_to_string(lock) {
        Ok(text) => serde_json::from_str(&text).map(Some).map_err(|e| {
            format!(
                "the workspace lock {} is not one wsb wrote ({e}); delete it to clear it",
                lock.display()
            )
        }),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(e) => Err(format!("cannot read the workspace lock {}: {e}", lock.display())),
    }
}

/// This machine's name, as a holder records it.
#[cfg(unix)]
fn host_name() -> String {
    let mut buffer = [0u8; 256];
    // SAFETY: the buffer is large enough for a host name and is written
    // NUL-terminated by gethostname on success.
    let ok = unsafe { libc::gethostname(buffer.as_mut_ptr().cast(), buffer.len()) } == 0;
    if !ok {
        return String::from("unknown");
    }
    let end = buffer.iter().position(|b| *b == 0).unwrap_or(buffer.len());
    String::from_utf8_lossy(&buffer[..end]).into_owned()
}

/// This machine's name, as a holder records it.
#[cfg(windows)]
fn host_name() -> String {
    std::env::var("COMPUTERNAME").unwrap_or_else(|_| String::from("unknown"))
}

#[cfg(test)]
#[path = "worktree_lock_tests.rs"]
mod tests;