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<'a> Default for LockOptions<'a> {
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 pub fn timeout_ms(mut self, timeout_ms: u64) -> Self {
52 self.timeout = Duration::from_millis(timeout_ms);
53 self
54 }
55
56 pub fn heartbeat_ms(mut self, heartbeat_ms: u64) -> Self {
57 self.heartbeat = Duration::from_millis(heartbeat_ms);
58 self
59 }
60
61 pub fn stale_grace_ms(mut self, stale_grace_ms: u64) -> Self {
62 self.stale_grace = Duration::from_millis(stale_grace_ms);
63 self
64 }
65
66 pub fn command(mut self, command: &'a str) -> Self {
67 self.command = Some(command);
68 self
69 }
70
71 pub fn force_stale(mut self, force: bool) -> Self {
72 self.force_stale = force;
73 self
74 }
75}
76
77#[allow(dead_code)]
78pub struct LockfileGuard {
79 lock_path: PathBuf,
80 #[allow(dead_code)]
81 file: File,
82 file_id: FileId,
83 record: LockRecord,
84 heartbeat_interval: Duration,
85}
86
87#[allow(dead_code)]
88impl LockfileGuard {
89 pub fn heartbeat(&mut self) -> Result<()> {
90 if self.heartbeat_interval.is_zero() {
91 return Ok(());
92 }
93 self.record.touch()?;
94 registry::write_record(&self.record)?;
95 Ok(())
96 }
97
98 pub fn file_id(&self) -> &FileId {
99 &self.file_id
100 }
101
102 pub fn owner_hint(&self) -> LockOwnerHint {
103 self.record.to_owner_hint()
104 }
105}
106
107impl Drop for LockfileGuard {
108 fn drop(&mut self) {
109 let _ = registry::remove_record(&self.file_id);
110 let _ = fs::remove_file(&self.lock_path);
111 }
112}
113
114pub fn acquire(path: &Path, options: LockOptions<'_>) -> Result<LockfileGuard> {
115 let lock_path = lockfile_path(path);
116 let file_id = registry::compute_file_id(path)?;
117 let command = options
118 .command
119 .map(std::borrow::ToOwned::to_owned)
120 .unwrap_or_else(default_command);
121 let heartbeat_ms = options
122 .heartbeat
123 .as_millis()
124 .try_into()
125 .unwrap_or(DEFAULT_HEARTBEAT_MS);
126 let record = LockRecord::new(&file_id, path, command, heartbeat_ms)?;
127 let start = Instant::now();
128
129 loop {
130 match OpenOptions::new()
131 .write(true)
132 .create_new(true)
133 .open(&lock_path)
134 {
135 Ok(file) => {
136 if let Err(err) = registry::write_record(&record) {
137 let _ = fs::remove_file(&lock_path);
138 return Err(err);
139 }
140 return Ok(LockfileGuard {
141 lock_path,
142 file,
143 file_id,
144 record,
145 heartbeat_interval: options.heartbeat,
146 });
147 }
148 Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
149 let existing = registry::read_record(&file_id)?;
150 let stale = existing
151 .as_ref()
152 .map(|rec| registry::is_stale(rec, options.stale_grace))
153 .unwrap_or(true);
154
155 if options.force_stale && stale {
156 let _ = registry::remove_record(&file_id);
157 match fs::remove_file(&lock_path) {
158 Ok(()) => continue,
159 Err(inner) if inner.kind() == std::io::ErrorKind::NotFound => continue,
160 Err(inner) => return Err(inner.into()),
161 }
162 }
163
164 if start.elapsed() >= options.timeout {
165 let hint = registry::to_owner_hint(existing.clone());
166 let message = existing
167 .as_ref()
168 .map(|rec| {
169 format!(
170 "memory locked by pid {} (cmd: {}) since {}",
171 rec.pid, rec.cmd, rec.started_at
172 )
173 })
174 .unwrap_or_else(|| "memory locked by another process".to_string());
175 return Err(LockedError::new(path.to_path_buf(), message, hint, stale).into());
176 }
177
178 let remaining = options
179 .timeout
180 .checked_sub(start.elapsed())
181 .unwrap_or_else(|| Duration::from_millis(SPIN_SLEEP_MS));
182 let sleep = Duration::from_millis(SPIN_SLEEP_MS).min(remaining);
183 thread::sleep(sleep);
184 }
185 Err(err) => return Err(err.into()),
186 }
187 }
188}
189
190pub fn current_owner(path: &Path) -> Result<Option<LockOwnerHint>> {
191 let file_id = match registry::compute_file_id(path) {
192 Ok(id) => id,
193 Err(crate::error::MemvidError::Io { source, .. })
194 if source.kind() == std::io::ErrorKind::NotFound =>
195 {
196 return Ok(None);
197 }
198 Err(err) => return Err(err),
199 };
200 let record = registry::read_record(&file_id)?;
201 Ok(registry::to_owner_hint(record))
202}