Skip to main content

start_command/
store_lock.rs

1//! Cross-language file lock. Publish fully written data with an atomic hard link.
2use crate::local_hostname;
3use serde_json::{json, Value};
4use std::fs::{self, OpenOptions};
5use std::io::Write;
6use std::path::PathBuf;
7use std::thread;
8use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
9use uuid::Uuid;
10
11pub struct LockManager {
12    path: PathBuf,
13    token: Option<String>,
14}
15
16impl LockManager {
17    pub fn new(path: PathBuf) -> Self {
18        Self { path, token: None }
19    }
20
21    pub fn acquire(&mut self, timeout: u64) -> bool {
22        let started = Instant::now();
23        let debug = |message: &str| {
24            if std::env::var("START_DEBUG").as_deref() == Ok("1") {
25                eprintln!("[DEBUG] Store lock {}: {}", self.path.display(), message);
26            }
27        };
28        while started.elapsed() < Duration::from_millis(timeout) {
29            if let Ok(stat) = fs::metadata(&self.path) {
30                let data = fs::read_to_string(&self.path)
31                    .ok()
32                    .and_then(|s| serde_json::from_str::<Value>(&s).ok());
33                if Self::stale(data.as_ref(), &stat) {
34                    if let Ok(current) = fs::metadata(&self.path) {
35                        if Self::same_file(&stat, &current) {
36                            debug("reclaiming stale or malformed lock");
37                            let _ = fs::remove_file(&self.path);
38                        }
39                    }
40                }
41            }
42            let token = Uuid::new_v4().to_string();
43            let temp = self
44                .path
45                .with_extension(format!("lock.{}.{}", std::process::id(), token));
46            let result = (|| -> std::io::Result<()> {
47                let mut file = OpenOptions::new()
48                    .write(true)
49                    .create_new(true)
50                    .open(&temp)?;
51                let data = json!({"pid": std::process::id(), "timestamp":
52                    SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_millis(),
53                    "hostname": local_hostname::get().map(|h| h.to_string_lossy().to_string()).unwrap_or_default(),
54                    "token": token});
55                file.write_all(data.to_string().as_bytes())?;
56                file.sync_all()?;
57                fs::hard_link(&temp, &self.path)
58            })();
59            let _ = fs::remove_file(&temp);
60            match result {
61                Ok(()) => {
62                    debug(&format!(
63                        "acquired after {}ms",
64                        started.elapsed().as_millis()
65                    ));
66                    self.token = Some(token);
67                    return true;
68                }
69                Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
70                    thread::sleep(Duration::from_millis(100));
71                }
72                Err(error) => {
73                    debug(&format!("publication failed: {}", error));
74                    return false;
75                }
76            }
77        }
78        debug("acquisition timed out");
79        false
80    }
81
82    fn same_file(a: &fs::Metadata, b: &fs::Metadata) -> bool {
83        #[cfg(unix)]
84        {
85            use std::os::unix::fs::MetadataExt;
86            a.dev() == b.dev() && a.ino() == b.ino() && a.modified().ok() == b.modified().ok()
87        }
88        #[cfg(not(unix))]
89        {
90            a.modified().ok() == b.modified().ok() && a.len() == b.len()
91        }
92    }
93
94    fn stale(data: Option<&Value>, stat: &fs::Metadata) -> bool {
95        let timestamp = data
96            .and_then(|d| d.get("timestamp"))
97            .and_then(Value::as_u64);
98        let pid = data.and_then(|d| d.get("pid")).and_then(Value::as_u64);
99        let (Some(timestamp), Some(pid)) =
100            (timestamp, pid.filter(|p| *p > 0 && *p <= i32::MAX as u64))
101        else {
102            return stat
103                .modified()
104                .ok()
105                .and_then(|m| m.elapsed().ok())
106                .is_some_and(|age| age >= Duration::from_secs(3));
107        };
108        let hostname = crate::launch_owner::hostname();
109        if data.and_then(|d| d.get("hostname")).and_then(Value::as_str) == Some(&hostname) {
110            return !crate::launch_owner::process_alive(pid);
111        }
112        let now = SystemTime::now()
113            .duration_since(UNIX_EPOCH)
114            .unwrap_or_default()
115            .as_millis() as u64;
116        now.saturating_sub(timestamp) > 60000
117    }
118
119    pub fn release(&mut self) {
120        if let Some(token) = self.token.take() {
121            let data = fs::read_to_string(&self.path)
122                .ok()
123                .and_then(|s| serde_json::from_str::<Value>(&s).ok());
124            if data
125                .as_ref()
126                .and_then(|d| d.get("token"))
127                .and_then(Value::as_str)
128                == Some(&token)
129            {
130                let _ = fs::remove_file(&self.path);
131            }
132        }
133    }
134}
135
136impl Drop for LockManager {
137    fn drop(&mut self) {
138        self.release();
139    }
140}