Skip to main content

oxicode_agent/tools/
file_mutation_queue.rs

1/// File mutation queue - serializes concurrent writes to the same file.
2///
3/// Prevents race conditions when multiple edit operations target the same file.
4/// Operations on *different* files run in parallel; operations on the *same*
5/// file are serialized.
6///
7/// Includes automatic stale-entry cleanup to prevent unbounded memory growth:
8/// every `CLEANUP_INTERVAL` operations, entries for files that no longer exist
9/// on disk are removed. If the map exceeds `MAX_ENTRIES`, excess entries are
10/// evicted (oldest first via HashMap iteration order).
11use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14use std::sync::atomic::{AtomicUsize, Ordering};
15use tokio::fs;
16use tokio::sync::Mutex;
17
18/// Clean up stale entries every N operations.
19const CLEANUP_INTERVAL: usize = 128;
20
21/// Maximum number of per-file mutex entries before eviction kicks in.
22const MAX_ENTRIES: usize = 1024;
23
24/// Global file mutation queue.
25static QUEUE: std::sync::OnceLock<FileMutationQueue> = std::sync::OnceLock::new();
26
27/// Get the global file mutation queue.
28pub fn global_mutation_queue() -> &'static FileMutationQueue {
29    QUEUE.get_or_init(FileMutationQueue::new)
30}
31
32/// Serializes file mutation operations per canonical path.
33#[derive(Debug)]
34pub struct FileMutationQueue {
35    /// Map from canonical path to a mutex that serializes operations.
36    queues: Arc<Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>>,
37    /// Monotonic operation counter for triggering periodic cleanup.
38    op_counter: AtomicUsize,
39}
40
41impl FileMutationQueue {
42    /// Create a new, empty file mutation queue.
43    pub fn new() -> Self {
44        Self {
45            queues: Arc::new(Mutex::new(HashMap::new())),
46            op_counter: AtomicUsize::new(0),
47        }
48    }
49
50    /// Execute a mutation operation on a file, serialized per canonical path.
51    ///
52    /// If the file doesn't exist yet, uses the path as-is for the key.
53    /// Periodically triggers cleanup of stale entries to prevent memory leaks.
54    pub async fn with_queue<F, Fut, T>(&self, path: &Path, f: F) -> T
55    where
56        F: FnOnce() -> Fut,
57        Fut: Future<Output = T>,
58    {
59        let ops = self.op_counter.fetch_add(1, Ordering::Relaxed);
60
61        // Periodic stale-entry cleanup to prevent unbounded memory growth.
62        if ops.is_multiple_of(CLEANUP_INTERVAL) && ops > 0 {
63            self.cleanup_stale().await;
64        }
65
66        let canonical = fs::canonicalize(path)
67            .await
68            .unwrap_or_else(|_| path.to_path_buf());
69
70        // Get or create a mutex for this file.
71        let mutex = {
72            let mut queues = self.queues.lock().await;
73
74            // Enforce capacity limit: evict stale entries if over max.
75            if queues.len() >= MAX_ENTRIES {
76                // First pass: remove entries for non-existent files.
77                let keys: Vec<PathBuf> = queues.keys().cloned().collect();
78                drop(queues);
79                for key in &keys {
80                    if fs::metadata(key).await.is_err() {
81                        let mut q = self.queues.lock().await;
82                        q.remove(key);
83                    }
84                }
85                queues = self.queues.lock().await;
86
87                // Second pass: if still over capacity, evict arbitrary entries.
88                while queues.len() >= MAX_ENTRIES {
89                    if let Some(key) = queues.keys().next().cloned() {
90                        queues.remove(&key);
91                    } else {
92                        break;
93                    }
94                }
95            }
96
97            queues
98                .entry(canonical)
99                .or_insert_with(|| Arc::new(Mutex::new(())))
100                .clone()
101        };
102
103        // Lock the per-file mutex.
104        let _guard = mutex.lock().await;
105
106        // Execute the operation.
107        f().await
108    }
109
110    /// Remove entries for files that no longer exist on disk.
111    ///
112    /// This is called automatically every `CLEANUP_INTERVAL` operations, but
113    /// can also be called manually.
114    pub async fn cleanup_stale(&self) {
115        let queues = self.queues.lock().await;
116        let keys: Vec<PathBuf> = queues.keys().cloned().collect();
117        drop(queues); // Release lock during IO.
118
119        let mut to_remove = Vec::new();
120        for key in &keys {
121            if fs::metadata(key).await.is_err() {
122                to_remove.push(key.clone());
123            }
124        }
125
126        if !to_remove.is_empty() {
127            tracing::debug!(
128                stale_count = to_remove.len(),
129                "FileMutationQueue: cleaning up stale entries"
130            );
131            let mut queues = self.queues.lock().await;
132            for key in to_remove {
133                queues.remove(&key);
134            }
135        }
136    }
137
138    /// Clean up entries for a specific file.
139    pub async fn cleanup(&self, path: &Path) {
140        let canonical = fs::canonicalize(path)
141            .await
142            .unwrap_or_else(|_| path.to_path_buf());
143        let mut queues = self.queues.lock().await;
144        queues.remove(&canonical);
145    }
146
147    /// Returns the number of entries currently in the queue map.
148    #[allow(dead_code)]
149    pub async fn entry_count(&self) -> usize {
150        self.queues.lock().await.len()
151    }
152}
153
154impl Default for FileMutationQueue {
155    fn default() -> Self {
156        Self::new()
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163    use std::sync::atomic::{AtomicUsize, Ordering};
164
165    #[tokio::test]
166    async fn test_same_file_serialized() {
167        let queue = Arc::new(FileMutationQueue::new());
168        let counter = Arc::new(AtomicUsize::new(0));
169        let path = PathBuf::from("/tmp/test_mutation_queue_file");
170
171        let mut handles = Vec::new();
172
173        for _ in 0..10 {
174            let queue = queue.clone();
175            let counter = counter.clone();
176            let path = path.clone();
177
178            handles.push(tokio::spawn(async move {
179                queue
180                    .with_queue(&path, || async {
181                        let prev = counter.fetch_add(1, Ordering::SeqCst);
182                        // Simulate some work
183                        tokio::time::sleep(std::time::Duration::from_millis(1)).await;
184                        prev
185                    })
186                    .await
187            }));
188        }
189
190        // All operations should complete
191        for handle in handles {
192            let _ = handle.await.unwrap();
193        }
194
195        assert_eq!(counter.load(Ordering::SeqCst), 10);
196    }
197
198    #[tokio::test]
199    async fn test_different_files_parallel() {
200        let queue = Arc::new(FileMutationQueue::new());
201        let counter = Arc::new(AtomicUsize::new(0));
202
203        let path1 = PathBuf::from("/tmp/test_file_1");
204        let path2 = PathBuf::from("/tmp/test_file_2");
205
206        let q1 = queue.clone();
207        let q2 = queue.clone();
208        let counter1 = counter.clone();
209        let counter2 = counter.clone();
210
211        let h1 = tokio::spawn(async move {
212            q1.with_queue(&path1, || async {
213                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
214                counter1.fetch_add(1, Ordering::SeqCst)
215            })
216            .await
217        });
218
219        let h2 = tokio::spawn(async move {
220            q2.with_queue(&path2, || async { counter2.fetch_add(1, Ordering::SeqCst) })
221                .await
222        });
223
224        // Both should complete quickly (parallel)
225        let r1 = tokio::time::timeout(std::time::Duration::from_millis(100), h1).await;
226        let r2 = tokio::time::timeout(std::time::Duration::from_millis(100), h2).await;
227
228        assert!(r1.is_ok());
229        assert!(r2.is_ok());
230    }
231
232    #[tokio::test]
233    async fn test_auto_cleanup_removes_stale() {
234        let queue = FileMutationQueue::new();
235
236        // Create a temp file, use it, then delete it.
237        let temp_path = std::env::temp_dir().join("oxicode_test_stale_queue_file");
238        std::fs::write(&temp_path, "test").unwrap();
239
240        let _ = queue.with_queue(&temp_path, || async { 42 }).await;
241        assert_eq!(queue.entry_count().await, 1);
242
243        // Delete the file.
244        std::fs::remove_file(&temp_path).ok();
245
246        // Trigger cleanup.
247        queue.cleanup_stale().await;
248        assert_eq!(queue.entry_count().await, 0);
249    }
250
251    #[tokio::test]
252    async fn test_max_entries_enforced() {
253        let queue = Arc::new(FileMutationQueue::new());
254
255        // Fill beyond MAX_ENTRIES with non-existent paths.
256        let mut handles = Vec::new();
257        for i in 0..(MAX_ENTRIES + 10) {
258            let q = queue.clone();
259            handles.push(tokio::spawn(async move {
260                let path = PathBuf::from(format!("/tmp/oxicode_test_max_entries_{}", i));
261                q.with_queue(&path, || async { i }).await
262            }));
263        }
264
265        for h in handles {
266            let _ = h.await.unwrap();
267        }
268
269        // Entry count should be capped at or near MAX_ENTRIES.
270        let count = queue.entry_count().await;
271        assert!(
272            count <= MAX_ENTRIES + 10, // Allow slack since eviction is on next access
273            "Entry count {} should be bounded near {}",
274            count,
275            MAX_ENTRIES
276        );
277    }
278}