1use std::path::{Path, PathBuf};
2use std::thread;
3use std::time::{Duration, Instant};
4
5use fs_err::{self as fs, File, OpenOptions};
6
7use crate::error::{LockOwnerHint, LockedError, Result};
8use crate::registry::{self, FileId, LockRecord};
9
10const DEFAULT_TIMEOUT_MS: u64 = 250;
11const DEFAULT_HEARTBEAT_MS: u64 = 2_000;
12const DEFAULT_STALE_GRACE_MS: u64 = 10_000;
13const SPIN_SLEEP_MS: u64 = 10;
14
15fn default_command() -> String {
16 std::env::args().collect::<Vec<_>>().join(" ")
17}
18
19fn lockfile_path(path: &Path) -> PathBuf {
20 let mut lock_path = path.to_path_buf();
21 let suffix = match path.extension().and_then(|ext| ext.to_str()) {
22 Some(ext) if !ext.is_empty() => format!("{ext}.lock"),
23 _ => "lock".to_string(),
24 };
25 lock_path.set_extension(suffix);
26 lock_path
27}
28
29#[derive(Debug, Clone)]
30pub struct LockOptions<'a> {
31 pub timeout: Duration,
32 pub heartbeat: Duration,
33 pub stale_grace: Duration,
34 pub command: Option<&'a str>,
35 pub force_stale: bool,
36}
37
38impl Default for LockOptions<'_> {
39 fn default() -> Self {
40 Self {
41 timeout: Duration::from_millis(DEFAULT_TIMEOUT_MS),
42 heartbeat: Duration::from_millis(DEFAULT_HEARTBEAT_MS),
43 stale_grace: Duration::from_millis(DEFAULT_STALE_GRACE_MS),
44 command: None,
45 force_stale: false,
46 }
47 }
48}
49
50impl<'a> LockOptions<'a> {
51 #[must_use]
52 pub fn timeout_ms(mut self, timeout_ms: u64) -> Self {
53 self.timeout = Duration::from_millis(timeout_ms);
54 self
55 }
56
57 #[must_use]
58 pub fn heartbeat_ms(mut self, heartbeat_ms: u64) -> Self {
59 self.heartbeat = Duration::from_millis(heartbeat_ms);
60 self
61 }
62
63 #[must_use]
64 pub fn stale_grace_ms(mut self, stale_grace_ms: u64) -> Self {
65 self.stale_grace = Duration::from_millis(stale_grace_ms);
66 self
67 }
68
69 #[must_use]
70 pub fn command(mut self, command: &'a str) -> Self {
71 self.command = Some(command);
72 self
73 }
74
75 #[must_use]
76 pub fn force_stale(mut self, force: bool) -> Self {
77 self.force_stale = force;
78 self
79 }
80}
81
82#[allow(dead_code)]
83pub struct LockfileGuard {
84 lock_path: PathBuf,
85 #[allow(dead_code)]
86 file: File,
87 file_id: FileId,
88 record: LockRecord,
89 heartbeat_interval: Duration,
90}
91
92#[allow(dead_code)]
93impl LockfileGuard {
94 pub fn heartbeat(&mut self) -> Result<()> {
95 if self.heartbeat_interval.is_zero() {
96 return Ok(());
97 }
98 self.record.touch()?;
99 registry::write_record(&self.record)?;
100 Ok(())
101 }
102
103 #[must_use]
104 pub fn file_id(&self) -> &FileId {
105 &self.file_id
106 }
107
108 #[must_use]
109 pub fn owner_hint(&self) -> LockOwnerHint {
110 self.record.to_owner_hint()
111 }
112}
113
114impl Drop for LockfileGuard {
115 fn drop(&mut self) {
116 let _ = registry::remove_record(&self.file_id);
117 let _ = fs::remove_file(&self.lock_path);
118 }
119}
120
121pub fn acquire(path: &Path, options: LockOptions<'_>) -> Result<LockfileGuard> {
122 let lock_path = lockfile_path(path);
123 let file_id = registry::compute_file_id(path)?;
124 let command = options
125 .command
126 .map_or_else(default_command, std::borrow::ToOwned::to_owned);
127 let heartbeat_ms = options
128 .heartbeat
129 .as_millis()
130 .try_into()
131 .unwrap_or(DEFAULT_HEARTBEAT_MS);
132 let record = LockRecord::new(&file_id, path, command, heartbeat_ms)?;
133 let start = Instant::now();
134
135 loop {
136 match OpenOptions::new()
137 .write(true)
138 .create_new(true)
139 .open(&lock_path)
140 {
141 Ok(file) => {
142 if let Err(err) = registry::write_record(&record) {
143 let _ = fs::remove_file(&lock_path);
144 return Err(err);
145 }
146 return Ok(LockfileGuard {
147 lock_path,
148 file,
149 file_id,
150 record,
151 heartbeat_interval: options.heartbeat,
152 });
153 }
154 Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
155 let existing = registry::read_record(&file_id)?;
156 let stale = existing
157 .as_ref()
158 .is_none_or(|rec| registry::is_stale(rec, options.stale_grace));
159
160 if options.force_stale && stale {
161 let _ = registry::remove_record(&file_id);
162 match fs::remove_file(&lock_path) {
163 Ok(()) => continue,
164 Err(inner) if inner.kind() == std::io::ErrorKind::NotFound => continue,
165 Err(inner) => return Err(inner.into()),
166 }
167 }
168
169 if start.elapsed() >= options.timeout {
170 let hint = registry::to_owner_hint(existing.clone());
171 let message = existing
172 .as_ref()
173 .map(|rec| {
174 format!(
175 "memory locked by pid {} (cmd: {}) since {}",
176 rec.pid, rec.cmd, rec.started_at
177 )
178 })
179 .unwrap_or_else(|| "memory locked by another process".to_string());
180 return Err(Box::new(LockedError::new(
181 path.to_path_buf(),
182 message,
183 hint,
184 stale,
185 ))
186 .into());
187 }
188
189 let remaining = options
190 .timeout
191 .checked_sub(start.elapsed())
192 .unwrap_or_else(|| Duration::from_millis(SPIN_SLEEP_MS));
193 let sleep = Duration::from_millis(SPIN_SLEEP_MS).min(remaining);
194 thread::sleep(sleep);
195 }
196 Err(err) => return Err(err.into()),
197 }
198 }
199}
200
201pub fn current_owner(path: &Path) -> Result<Option<LockOwnerHint>> {
202 let file_id = match registry::compute_file_id(path) {
203 Ok(id) => id,
204 Err(crate::error::MemvidError::Io { source, .. })
205 if source.kind() == std::io::ErrorKind::NotFound =>
206 {
207 return Ok(None);
208 }
209 Err(err) => return Err(err),
210 };
211 let record = registry::read_record(&file_id)?;
212 Ok(registry::to_owner_hint(record))
213}