AnyBot/src/lock.rs

// One AnyBot per data folder. Two copies on one folder each keep their own idea of the chats and
// overwrite each other's saves, and a second copy starting up would take the first one's open
// turns in the work log (oplog.rs) for a crash and close them. So the window, `--serve`, `--ask`
// and `anybot setup` (while it writes) hold `anybot.lock` in the data folder.
//
// It is an OS file lock (flock on Linux and macOS, LockFileEx on Windows), not a file that merely
// exists: it goes away with the process however it ends — closed, crashed, killed — so a copy
// that died never blocks the next one. The holder writes its process id in the file, only to say
// who it is.

use std::io::{Read, Seek, Write};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};

/// How long a new copy waits for the folder before giving up: long enough for a copy that is
/// closing (a relaunch, a service restart) to finish.
pub const WAIT: Duration = Duration::from_secs(5);

/// Held while this copy uses the data folder; letting go of it (or the process ending) frees it.
pub struct DataLock {
    _file: Option<std::fs::File>,
}

pub fn path() -> PathBuf {
    let state = crate::model::State::path();
    state.parent().map(|d| d.join("anybot.lock")).unwrap_or_else(|| PathBuf::from("anybot.lock"))
}

/// Take this data folder for this copy, waiting up to `wait` for another copy to let go.
/// Err: what to tell the person (another AnyBot is using the folder).
pub fn take(wait: Duration) -> Result<DataLock, String> {
    take_at(&path(), wait)
}

fn take_at(p: &Path, wait: Duration) -> Result<DataLock, String> {
    if let Some(d) = p.parent() {
        let _ = std::fs::create_dir_all(d);
    }
    let mut f = match std::fs::OpenOptions::new().create(true).read(true).write(true).truncate(false).open(p) {
        Ok(f) => f,
        Err(e) => {
            // A folder the lock can't be made in can't hold the chats either; say so, but don't
            // refuse to start over the lock alone.
            eprintln!("AnyBot：沒辦法建立 {}（{e}），這次不檢查是否有別的 AnyBot 在用同一個資料夾。", p.display());
            return Ok(DataLock { _file: None });
        }
    };
    let until = Instant::now() + wait;
    loop {
        match f.try_lock() {
            Ok(()) => {
                let _ = f.set_len(0);
                let _ = f.rewind();
                let _ = f.write_all(std::process::id().to_string().as_bytes());
                let _ = f.sync_data();
                return Ok(DataLock { _file: Some(f) });
            }
            Err(std::fs::TryLockError::WouldBlock) if Instant::now() < until => std::thread::sleep(Duration::from_millis(100)),
            Err(std::fs::TryLockError::WouldBlock) => return Err(busy_message(p, &mut f)),
            Err(std::fs::TryLockError::Error(e)) => {
                // A file system without locks (some network shares): go on unchecked.
                eprintln!("AnyBot：{} 不能上鎖（{e}），這次不檢查是否有別的 AnyBot 在用同一個資料夾。", p.display());
                return Ok(DataLock { _file: None });
            }
        }
    }
}

/// "Another AnyBot (process 123) is using <folder>": what a second copy says.
fn busy_message(p: &Path, f: &mut std::fs::File) -> String {
    let mut who = String::new();
    let _ = f.rewind();
    let _ = f.read_to_string(&mut who);
    let who = who.trim();
    let dir = p.parent().map(|d| d.display().to_string()).unwrap_or_default();
    let pid = if who.chars().all(|c| c.is_ascii_digit()) && !who.is_empty() { format!("（程序 {who}）") } else { String::new() };
    format!("另一個 AnyBot{pid}已經在使用這個資料夾：\n{dir}\n\n同一個資料夾一次只能有一個 AnyBot（視窗、--serve 或 anybot setup）。請先關掉另一個，再打開這個。")
}

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

    fn scratch(name: &str) -> PathBuf {
        let d = std::env::temp_dir().join(format!("anybot-lock-{name}-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&d);
        d.join("anybot.lock")
    }

    /// The folder is one copy's at a time: a second take is refused, saying who holds it and
    /// which folder; once the first lets go, the next take gets it. Another folder is separate.
    #[test]
    fn one_copy_per_folder() {
        let p = scratch("one");
        let first = take_at(&p, Duration::ZERO).expect("first copy");
        assert_eq!(std::fs::read_to_string(&p).unwrap(), std::process::id().to_string());
        let msg = take_at(&p, Duration::ZERO).err().expect("second copy refused");
        assert!(msg.contains(&format!("程序 {}", std::process::id())) && msg.contains(&p.parent().unwrap().display().to_string()), "{msg}");
        let other = take_at(&scratch("other"), Duration::ZERO);
        assert!(other.is_ok(), "another data folder has its own lock");
        drop(first);
        assert!(take_at(&p, Duration::ZERO).is_ok());
    }

    /// A copy that is closing (a relaunch) is waited for.
    #[test]
    fn a_closing_copy_is_waited_for() {
        let p = scratch("wait");
        let first = take_at(&p, Duration::ZERO).unwrap();
        let t = std::thread::spawn(move || {
            std::thread::sleep(Duration::from_millis(300));
            drop(first);
        });
        let t0 = Instant::now();
        assert!(take_at(&p, Duration::from_secs(5)).is_ok());
        assert!(t0.elapsed() >= Duration::from_millis(250));
        t.join().unwrap();
    }

    /// Not a test on its own: the process `a_killed_copy_leaves_no_lock` starts, which takes the
    /// lock at ANYBOT_LOCK_TEST and waits to be killed.
    #[test]
    #[ignore]
    fn a_helper_holds_the_lock() {
        let Ok(p) = std::env::var("ANYBOT_LOCK_TEST") else { return };
        let _held = take_at(Path::new(&p), Duration::ZERO).expect("helper takes the lock");
        println!("LOCK-HELD");
        std::thread::sleep(Duration::from_secs(60));
    }

    /// A copy that crashed or was killed holds nothing: the lock went with its process.
    #[test]
    fn a_killed_copy_leaves_no_lock() {
        use std::io::BufRead;
        let p = scratch("killed");
        std::fs::create_dir_all(p.parent().unwrap()).unwrap();
        let mut child = std::process::Command::new(std::env::current_exe().unwrap())
            .args(["--exact", "lock::tests::a_helper_holds_the_lock", "--ignored", "--nocapture", "--test-threads=1"])
            .env("ANYBOT_LOCK_TEST", &p)
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::null())
            .spawn()
            .unwrap();
        let out = std::io::BufReader::new(child.stdout.take().unwrap());
        let held = out.lines().map_while(Result::ok).any(|l| l.contains("LOCK-HELD"));
        assert!(held, "the helper never took the lock");
        let msg = take_at(&p, Duration::ZERO).err().expect("held by the other process");
        assert!(msg.contains(&format!("程序 {}", child.id())), "{msg}");
        child.kill().unwrap();
        child.wait().unwrap();
        assert!(take_at(&p, Duration::from_secs(2)).is_ok(), "a killed copy's lock is gone");
    }
}