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    monitor_owner: bool,
15    /// Extra fields published in the lock file, for example who holds it.
16    metadata: serde_json::Map<String, Value>,
17}
18
19impl LockManager {
20    pub fn new(path: PathBuf) -> Self {
21        Self {
22            path,
23            token: None,
24            monitor_owner: false,
25            metadata: serde_json::Map::new(),
26        }
27    }
28
29    /// Publish `key` in the lock file while it is held.
30    pub fn with_metadata(mut self, key: &str, value: Value) -> Self {
31        self.metadata.insert(key.to_string(), value);
32        self
33    }
34
35    /// The published lock file, if one is readable.
36    pub fn read_lock_file(&self) -> Option<Value> {
37        fs::read_to_string(&self.path)
38            .ok()
39            .and_then(|s| serde_json::from_str::<Value>(&s).ok())
40    }
41
42    pub fn for_cpu_monitor(mut self) -> Self {
43        self.monitor_owner = true;
44        self
45    }
46
47    pub fn acquire(&mut self, timeout: u64) -> bool {
48        let started = Instant::now();
49        let debug = |message: &str| {
50            if std::env::var("START_DEBUG").as_deref() == Ok("1") {
51                eprintln!("[DEBUG] Store lock {}: {}", self.path.display(), message);
52            }
53        };
54        while started.elapsed() < Duration::from_millis(timeout) {
55            if let Ok(stat) = fs::metadata(&self.path) {
56                let data = fs::read_to_string(&self.path)
57                    .ok()
58                    .and_then(|s| serde_json::from_str::<Value>(&s).ok());
59                if self.stale(data.as_ref(), &stat) {
60                    if let Ok(current) = fs::metadata(&self.path) {
61                        if Self::same_file(&stat, &current) {
62                            debug("reclaiming stale or malformed lock");
63                            let _ = fs::remove_file(&self.path);
64                        }
65                    }
66                }
67            }
68            let token = Uuid::new_v4().to_string();
69            let temp = self
70                .path
71                .with_extension(format!("lock.{}.{}", std::process::id(), token));
72            let result = (|| -> std::io::Result<()> {
73                let mut file = OpenOptions::new()
74                    .write(true)
75                    .create_new(true)
76                    .open(&temp)?;
77                let owner = json!({"pid": std::process::id(), "timestamp":
78                    SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_millis(),
79                    "hostname": local_hostname::get().map(|h| h.to_string_lossy().to_string()).unwrap_or_default(),
80                    "token": token, "processStart": Self::process_info(std::process::id() as u64).map(|p| p.1)});
81                // Metadata goes first so it can never replace the owner fields.
82                let mut data = self.metadata.clone();
83                if let Value::Object(owner) = owner {
84                    data.extend(owner);
85                }
86                file.write_all(Value::Object(data).to_string().as_bytes())?;
87                file.sync_all()?;
88                fs::hard_link(&temp, &self.path)
89            })();
90            let _ = fs::remove_file(&temp);
91            match result {
92                Ok(()) => {
93                    debug(&format!(
94                        "acquired after {}ms",
95                        started.elapsed().as_millis()
96                    ));
97                    self.token = Some(token);
98                    return true;
99                }
100                Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
101                    thread::sleep(Duration::from_millis(100));
102                }
103                Err(error) => {
104                    debug(&format!("publication failed: {}", error));
105                    return false;
106                }
107            }
108        }
109        debug("acquisition timed out");
110        false
111    }
112
113    fn same_file(a: &fs::Metadata, b: &fs::Metadata) -> bool {
114        #[cfg(unix)]
115        {
116            use std::os::unix::fs::MetadataExt;
117            a.dev() == b.dev() && a.ino() == b.ino() && a.modified().ok() == b.modified().ok()
118        }
119        #[cfg(not(unix))]
120        {
121            a.modified().ok() == b.modified().ok() && a.len() == b.len()
122        }
123    }
124
125    fn stale(&self, data: Option<&Value>, stat: &fs::Metadata) -> bool {
126        let timestamp = data
127            .and_then(|d| d.get("timestamp"))
128            .and_then(Value::as_u64);
129        let pid = data.and_then(|d| d.get("pid")).and_then(Value::as_u64);
130        let (Some(timestamp), Some(pid)) =
131            (timestamp, pid.filter(|p| *p > 0 && *p <= i32::MAX as u64))
132        else {
133            return stat
134                .modified()
135                .ok()
136                .and_then(|m| m.elapsed().ok())
137                .is_some_and(|age| age >= Duration::from_secs(3));
138        };
139        let hostname = crate::launch_owner::hostname();
140        if data.and_then(|d| d.get("hostname")).and_then(Value::as_str) == Some(&hostname) {
141            if !crate::launch_owner::process_alive(pid) {
142                return true;
143            }
144            if let Some((state, start)) = Self::process_info(pid) {
145                if state == "Z"
146                    || state == "X"
147                    || state == "x"
148                    || data
149                        .and_then(|d| d["processStart"].as_str())
150                        .is_some_and(|s| s != start)
151                {
152                    return true;
153                }
154            }
155            #[cfg(target_os = "linux")]
156            if self.monitor_owner && data.and_then(|d| d["processStart"].as_str()).is_none() {
157                if let Ok(command) = fs::read_to_string(format!("/proc/{}/cmdline", pid)) {
158                    return !command.contains("cpu-penalty-monitor")
159                        && !command.contains("__start-cpu-monitor");
160                }
161            }
162            return false;
163        }
164        let now = SystemTime::now()
165            .duration_since(UNIX_EPOCH)
166            .unwrap_or_default()
167            .as_millis() as u64;
168        now.saturating_sub(timestamp) > 60000
169    }
170
171    fn process_info(pid: u64) -> Option<(String, String)> {
172        #[cfg(target_os = "linux")]
173        {
174            let stat = fs::read_to_string(format!("/proc/{}/stat", pid)).ok()?;
175            let fields: Vec<&str> = stat[stat.rfind(')')? + 2..].split_whitespace().collect();
176            Some((fields.first()?.to_string(), fields.get(19)?.to_string()))
177        }
178        #[cfg(not(target_os = "linux"))]
179        {
180            let _ = pid;
181            None
182        }
183    }
184
185    pub fn release(&mut self) {
186        if let Some(token) = self.token.take() {
187            let data = fs::read_to_string(&self.path)
188                .ok()
189                .and_then(|s| serde_json::from_str::<Value>(&s).ok());
190            if data
191                .as_ref()
192                .and_then(|d| d.get("token"))
193                .and_then(Value::as_str)
194                == Some(&token)
195            {
196                let _ = fs::remove_file(&self.path);
197            }
198        }
199    }
200}
201
202impl Drop for LockManager {
203    fn drop(&mut self) {
204        self.release();
205    }
206}
start_command/store_lock.rs

start_command/
store_lock.rs
