Skip to main content

pi/core/tools/
mutation_queue.rs

1//! Per-realpath serialization for file mutations (edit / write).
2//!
3//! Ports `.references/pi/packages/coding-agent/src/core/tools/file-mutation-queue.ts`.
4//! Operations targeting the same canonical path run FIFO; operations on
5//! distinct paths run concurrently. The queue key is
6//! `realpath(resolve(path))`, falling back to the lexically resolved path
7//! when the target does not yet exist (`ENOENT` / `ENOTDIR`). Map entries
8//! are cleaned up when the last waiter for a key finishes, including on
9//! cancellation via a drop guard.
10
11use 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/// Errors produced while registering a path into the mutation queue.
21#[derive(Debug, Error)]
22pub enum MutationQueueError {
23    /// Resolving the mutation-queue key for a path failed (working-directory
24    /// lookup, or a `realpath` failure that is not "path missing").
25    #[error("failed to resolve mutation queue key for {path}: {source}")]
26    ResolveKey {
27        /// The path whose key could not be resolved.
28        path: PathBuf,
29        /// Underlying I/O failure.
30        source: std::io::Error,
31    },
32    /// The per-path semaphore was closed. Unreachable in practice: the
33    /// semaphore is never closed by this module.
34    #[error("mutation queue for {path} is unavailable")]
35    QueueUnavailable {
36        /// The path whose queue is unavailable.
37        path: PathBuf,
38    },
39}
40
41/// Global registry of per-key FIFO gates. Entries are held as [`Weak`] so a
42/// finished key can be reclaimed; the live [`Arc`] is owned by each
43/// outstanding waiter.
44static 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
51/// Resolve the queue key for `file_path` (TypeScript `getMutationQueueKey`).
52///
53/// Absolute-ize and lexically normalize first, then `realpath`. Missing
54/// targets (`ENOENT` / `ENOTDIR`) fall back to the resolved path so a
55/// create-on-write and a subsequent edit of the same path share a key.
56async 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
77/// Holds one waiter's Arc-cloned gate; on drop, removes the registry entry
78/// when this is the last remaining strong reference.
79struct 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        // Identity verification, last-reference detection, and removal must
107        // be one registry-locked operation. Otherwise a concurrent registrar
108        // can install a replacement gate between the count check and remove.
109        if maps_to_this_gate && Arc::strong_count(&self.gate) == 1 {
110            map.remove(&self.key);
111        }
112    }
113}
114
115/// Serialize `f` against every other mutation targeting the same realpath.
116/// Distinct keys run in parallel. TypeScript `withFileMutationQueue`.
117///
118/// # Errors
119///
120/// Returns [`MutationQueueError`] when the queue key cannot be resolved or
121/// the (never-closed) semaphore rejects the acquire.
122pub 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    // Registration is held for the whole operation so Drop cleans the map
133    // even when the future is cancelled.
134    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    // Declaration order: permit drops before registration (reverse of
144    // declaration), releasing the next waiter before the map entry is
145    // considered for cleanup.
146    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                // Wait for B to enter its critical section. If the two paths
230                // shared a key this would deadlock; the timeout below fails.
231                b_wait.notified().await;
232                "a"
233            })
234            .await
235        });
236        let b = tokio::spawn(async move {
237            // Wait until A holds its permit so the only way B starts is
238            // cross-key parallelism.
239            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        // A holds the realpath key via the symlink path; B targets the real
274        // path and must wait for A to finish.
275        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        // Give A a chance to register and acquire first.
290        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            // Entry is live while the operation holds the registration.
332            assert!(registry_holds_key(&key));
333        })
334        .await?;
335        assert!(!registry_holds_key(&key));
336
337        // Two sequential ops leave this key empty each time.
338        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        // Deterministically model the exact race state: a new registrar has
351        // installed a replacement after the stale registration observed
352        // itself as the final strong owner.
353        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}