pi/core/tools/
mutation_queue.rs1use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use std::sync::{Arc, LazyLock, Mutex, PoisonError, Weak};
14
15use thiserror::Error;
16use tokio::sync::Semaphore;
17
18use super::path_utils::resolve_lexically_absolute;
19
20#[derive(Debug, Error)]
22pub enum MutationQueueError {
23 #[error("failed to resolve mutation queue key for {path}: {source}")]
26 ResolveKey {
27 path: PathBuf,
29 source: std::io::Error,
31 },
32 #[error("mutation queue for {path} is unavailable")]
35 QueueUnavailable {
36 path: PathBuf,
38 },
39}
40
41static REGISTRY: LazyLock<Mutex<HashMap<PathBuf, Weak<Semaphore>>>> =
45 LazyLock::new(|| Mutex::new(HashMap::new()));
46
47fn lock_registry() -> std::sync::MutexGuard<'static, HashMap<PathBuf, Weak<Semaphore>>> {
48 REGISTRY.lock().unwrap_or_else(PoisonError::into_inner)
49}
50
51async fn mutation_queue_key(file_path: &Path) -> Result<PathBuf, MutationQueueError> {
57 let resolved =
58 resolve_lexically_absolute(file_path).map_err(|source| MutationQueueError::ResolveKey {
59 path: file_path.to_path_buf(),
60 source,
61 })?;
62 match tokio::fs::canonicalize(&resolved).await {
63 Ok(canonical) => Ok(canonical),
64 Err(error)
65 if error.kind() == std::io::ErrorKind::NotFound
66 || error.kind() == std::io::ErrorKind::NotADirectory =>
67 {
68 Ok(resolved)
69 }
70 Err(source) => Err(MutationQueueError::ResolveKey {
71 path: file_path.to_path_buf(),
72 source,
73 }),
74 }
75}
76
77struct QueueRegistration {
80 key: PathBuf,
81 gate: Arc<Semaphore>,
82}
83
84impl QueueRegistration {
85 fn register(key: PathBuf) -> Self {
86 let mut map = lock_registry();
87 let gate = if let Some(existing) = map.get(&key).and_then(Weak::upgrade) {
88 existing
89 } else {
90 let gate = Arc::new(Semaphore::new(1));
91 map.insert(key.clone(), Arc::downgrade(&gate));
92 gate
93 };
94 Self { key, gate }
95 }
96}
97
98impl Drop for QueueRegistration {
99 fn drop(&mut self) {
100 let mut map = lock_registry();
101 let own_gate = Arc::downgrade(&self.gate);
102 let maps_to_this_gate = map
103 .get(&self.key)
104 .is_some_and(|mapped| Weak::ptr_eq(mapped, &own_gate));
105
106 if maps_to_this_gate && Arc::strong_count(&self.gate) == 1 {
110 map.remove(&self.key);
111 }
112 }
113}
114
115pub async fn with_file_mutation_queue<T, F, Fut>(
123 file_path: impl AsRef<Path>,
124 f: F,
125) -> Result<T, MutationQueueError>
126where
127 F: FnOnce() -> Fut,
128 Fut: Future<Output = T>,
129{
130 let path = file_path.as_ref();
131 let key = mutation_queue_key(path).await?;
132 let registration = QueueRegistration::register(key);
135 let permit = registration
136 .gate
137 .clone()
138 .acquire_owned()
139 .await
140 .map_err(|_| MutationQueueError::QueueUnavailable {
141 path: path.to_path_buf(),
142 })?;
143 let result = f().await;
147 drop(permit);
148 Ok(result)
149}
150
151#[cfg(test)]
152fn registry_holds_key(key: &Path) -> bool {
153 lock_registry().get(key).and_then(Weak::upgrade).is_some()
154}
155
156#[cfg(test)]
157mod tests {
158 use super::*;
159 use std::error::Error;
160 use std::sync::atomic::{AtomicUsize, Ordering};
161 use std::time::Duration;
162
163 use tokio::sync::Notify;
164
165 type TestResult = Result<(), Box<dyn Error + Send + Sync>>;
166
167 #[tokio::test]
168 async fn same_key_runs_serially() -> TestResult {
169 let dir = tempfile::tempdir()?;
170 let path = dir.path().join("target.txt");
171 std::fs::write(&path, b"seed")?;
172
173 let active = Arc::new(AtomicUsize::new(0));
174 let peak = Arc::new(AtomicUsize::new(0));
175 let completed = Arc::new(Mutex::new(Vec::new()));
176 let mut handles = Vec::new();
177 for index in 0..4 {
178 let path = path.clone();
179 let active = Arc::clone(&active);
180 let peak = Arc::clone(&peak);
181 let completed = Arc::clone(&completed);
182 handles.push(tokio::spawn(async move {
183 with_file_mutation_queue(&path, || async {
184 let current = active.fetch_add(1, Ordering::SeqCst) + 1;
185 peak.fetch_max(current, Ordering::SeqCst);
186 tokio::time::sleep(Duration::from_millis(5)).await;
187 completed
188 .lock()
189 .unwrap_or_else(PoisonError::into_inner)
190 .push(index);
191 active.fetch_sub(1, Ordering::SeqCst);
192 })
193 .await
194 }));
195 }
196
197 for handle in handles {
198 handle
199 .await
200 .map_err(|error| std::io::Error::other(error.to_string()))??;
201 }
202 assert_eq!(peak.load(Ordering::SeqCst), 1);
203 let mut recorded = completed
204 .lock()
205 .unwrap_or_else(PoisonError::into_inner)
206 .clone();
207 recorded.sort_unstable();
208 assert_eq!(recorded, vec![0, 1, 2, 3]);
209 Ok(())
210 }
211
212 #[tokio::test]
213 async fn distinct_keys_run_concurrently() -> TestResult {
214 let dir = tempfile::tempdir()?;
215 let path_a = dir.path().join("a.txt");
216 let path_b = dir.path().join("b.txt");
217 std::fs::write(&path_a, b"a")?;
218 std::fs::write(&path_b, b"b")?;
219
220 let a_entered = Arc::new(Notify::new());
221 let b_entered = Arc::new(Notify::new());
222 let a_signal = Arc::clone(&a_entered);
223 let b_wait = Arc::clone(&b_entered);
224 let b_signal = Arc::clone(&b_entered);
225
226 let a = tokio::spawn(async move {
227 with_file_mutation_queue(&path_a, || async {
228 a_signal.notify_one();
229 b_wait.notified().await;
232 "a"
233 })
234 .await
235 });
236 let b = tokio::spawn(async move {
237 a_entered.notified().await;
240 with_file_mutation_queue(&path_b, || async {
241 b_signal.notify_one();
242 "b"
243 })
244 .await
245 });
246
247 let a_result = tokio::time::timeout(Duration::from_secs(2), a)
248 .await
249 .map_err(|_| std::io::Error::other("A timed out; keys may be serialized"))?
250 .map_err(|error| std::io::Error::other(error.to_string()))??;
251 let b_result = tokio::time::timeout(Duration::from_secs(2), b)
252 .await
253 .map_err(|_| std::io::Error::other("B timed out; keys may be serialized"))?
254 .map_err(|error| std::io::Error::other(error.to_string()))??;
255 assert_eq!(a_result, "a");
256 assert_eq!(b_result, "b");
257 Ok(())
258 }
259
260 #[cfg(unix)]
261 #[tokio::test]
262 async fn symlink_and_realpath_share_one_key() -> TestResult {
263 let dir = tempfile::tempdir()?;
264 let real = dir.path().join("real.txt");
265 let link = dir.path().join("link.txt");
266 std::fs::write(&real, b"seed")?;
267 std::os::unix::fs::symlink(&real, &link)?;
268
269 let order = Arc::new(Mutex::new(Vec::new()));
270 let order_a = Arc::clone(&order);
271 let order_b = Arc::clone(&order);
272
273 let a = tokio::spawn(async move {
276 with_file_mutation_queue(&link, || async {
277 order_a
278 .lock()
279 .unwrap_or_else(PoisonError::into_inner)
280 .push("a-start");
281 tokio::time::sleep(Duration::from_millis(30)).await;
282 order_a
283 .lock()
284 .unwrap_or_else(PoisonError::into_inner)
285 .push("a-end");
286 })
287 .await
288 });
289 tokio::time::sleep(Duration::from_millis(5)).await;
291 let b = tokio::spawn(async move {
292 with_file_mutation_queue(&real, || async {
293 order_b
294 .lock()
295 .unwrap_or_else(PoisonError::into_inner)
296 .push("b-start");
297 order_b
298 .lock()
299 .unwrap_or_else(PoisonError::into_inner)
300 .push("b-end");
301 })
302 .await
303 });
304
305 a.await
306 .map_err(|error| std::io::Error::other(error.to_string()))??;
307 b.await
308 .map_err(|error| std::io::Error::other(error.to_string()))??;
309 let recorded = order.lock().unwrap_or_else(PoisonError::into_inner).clone();
310 assert_eq!(recorded, vec!["a-start", "a-end", "b-start", "b-end"]);
311 Ok(())
312 }
313
314 #[tokio::test]
315 async fn missing_path_uses_resolved_key_and_runs() -> TestResult {
316 let dir = tempfile::tempdir()?;
317 let path = dir.path().join("does-not-exist-yet.txt");
318 let result = with_file_mutation_queue(&path, || async { 42 }).await?;
319 assert_eq!(result, 42);
320 Ok(())
321 }
322
323 #[tokio::test]
324 async fn registry_is_cleaned_after_completion() -> TestResult {
325 let dir = tempfile::tempdir()?;
326 let path = dir.path().join("cleanup.txt");
327 std::fs::write(&path, b"seed")?;
328 let key = tokio::fs::canonicalize(&path).await?;
329
330 with_file_mutation_queue(&path, || async {
331 assert!(registry_holds_key(&key));
333 })
334 .await?;
335 assert!(!registry_holds_key(&key));
336
337 with_file_mutation_queue(&path, || async {}).await?;
339 with_file_mutation_queue(&path, || async {}).await?;
340 assert!(!registry_holds_key(&key));
341 Ok(())
342 }
343
344 #[test]
345 fn stale_last_registration_cannot_remove_replacement_gate() -> TestResult {
346 let key = PathBuf::from("replacement-race-key");
347 let stale = QueueRegistration::register(key.clone());
348 let replacement_gate = Arc::new(Semaphore::new(1));
349
350 lock_registry().insert(key.clone(), Arc::downgrade(&replacement_gate));
354 drop(stale);
355
356 let Some(mapped) = lock_registry().get(&key).and_then(Weak::upgrade) else {
357 return Err("stale drop removed the replacement gate".into());
358 };
359 assert!(Arc::ptr_eq(&mapped, &replacement_gate));
360 lock_registry().remove(&key);
361 Ok(())
362 }
363
364 #[tokio::test]
365 async fn concurrent_ops_on_same_key_leave_registry_empty() -> TestResult {
366 let dir = tempfile::tempdir()?;
367 let path = dir.path().join("shared.txt");
368 std::fs::write(&path, b"seed")?;
369 let key = tokio::fs::canonicalize(&path).await?;
370
371 let counter = Arc::new(AtomicUsize::new(0));
372 let mut handles = Vec::new();
373 for _ in 0..8 {
374 let path = path.clone();
375 let counter = Arc::clone(&counter);
376 handles.push(tokio::spawn(async move {
377 with_file_mutation_queue(&path, || async {
378 counter.fetch_add(1, Ordering::SeqCst);
379 })
380 .await
381 }));
382 }
383 for handle in handles {
384 handle
385 .await
386 .map_err(|error| std::io::Error::other(error.to_string()))??;
387 }
388 assert_eq!(counter.load(Ordering::SeqCst), 8);
389 assert!(!registry_holds_key(&key));
390 Ok(())
391 }
392}