fix: coalesce locator cache in-flight lookups (Fixes #399)

Author: karthiknadig

crates/pet-conda/src/lib.rs

@@ -391,9 +391,9 @@ impl Locator for Conda {
                         reporter.report_environment(&env);
 
                         // Also check for a mamba/micromamba manager in the same directory and report it.
-                        // Reporting inside the closure minimizes the TOCTOU window compared to a
-                        // separate contains_key check, though concurrent threads may still
-                        // briefly both invoke the closure before the write-lock double-check.
+                        // LocatorCache coalesces concurrent lookups for this conda_dir, so mamba
+                        // discovery and its reporting side effect run at most once per in-flight
+                        // key.
                         let _ = self.mamba_managers.get_or_insert_with(conda_dir.clone(), || {
                             let mgr = get_mamba_manager(conda_dir);
                             if let Some(ref m) = mgr {

crates/pet-core/src/cache.rs

@@ -6,7 +6,12 @@
 //! Provides a thread-safe cache wrapper that consolidates common caching patterns
 //! used across multiple locators in the codebase.
 
-use std::{collections::HashMap, hash::Hash, path::PathBuf, sync::RwLock};
+use std::{
+    collections::HashMap,
+    hash::Hash,
+    path::PathBuf,
+    sync::{Arc, Condvar, Mutex, RwLock},
+};
 
 use crate::{manager::EnvManager, python_environment::PythonEnvironment};
 
@@ -17,13 +22,71 @@ use crate::{manager::EnvManager, python_environment::PythonEnvironment};
 /// returned from the cache.
 pub struct LocatorCache<K, V> {
     cache: RwLock<HashMap<K, V>>,
+    in_flight: Mutex<HashMap<K, Arc<InFlightEntry<V>>>>,
+}
+
+struct InFlightEntry<V> {
+    result: Mutex<Option<Option<V>>>,
+    changed: Condvar,
+}
+
+struct InFlightOwnerGuard<'a, K: Eq + Hash, V> {
+    key: Option<K>,
+    entry: Arc<InFlightEntry<V>>,
+    in_flight: &'a Mutex<HashMap<K, Arc<InFlightEntry<V>>>>,
+}
+
+enum InFlightClaim<'a, K: Eq + Hash, V> {
+    Owner(InFlightOwnerGuard<'a, K, V>),
+    Waiter(Arc<InFlightEntry<V>>),
+}
+
+impl<V> InFlightEntry<V> {
+    fn new() -> Self {
+        Self {
+            result: Mutex::new(None),
+            changed: Condvar::new(),
+        }
+    }
+}
+
+impl<K: Eq + Hash, V> InFlightOwnerGuard<'_, K, V> {
+    fn complete(mut self, result: Option<V>) {
+        self.publish_result(result);
+    }
+
+    fn publish_result(&mut self, result: Option<V>) {
+        *self
+            .entry
+            .result
+            .lock()
+            .expect("locator cache in-flight result lock poisoned") = Some(result);
+
+        if let Some(key) = self.key.take() {
+            self.in_flight
+                .lock()
+                .expect("locator cache in-flight lock poisoned")
+                .remove(&key);
+        }
+
+        self.entry.changed.notify_all();
+    }
+}
+
+impl<K: Eq + Hash, V> Drop for InFlightOwnerGuard<'_, K, V> {
+    fn drop(&mut self) {
+        if self.key.is_some() {
+            self.publish_result(None);
+        }
+    }
 }
 
 impl<K: Eq + Hash, V: Clone> LocatorCache<K, V> {
     /// Creates a new empty cache.
     pub fn new() -> Self {
         Self {
             cache: RwLock::new(HashMap::new()),
+            in_flight: Mutex::new(HashMap::new()),
         }
     }
 
@@ -68,35 +131,96 @@ impl<K: Eq + Hash, V: Clone> LocatorCache<K, V> {
     /// Returns a cloned value for the given key if it exists, otherwise computes
     /// and inserts the value using the provided closure.
     ///
-    /// This method first checks with a read lock, then upgrades to a write lock
-    /// if the value needs to be computed and inserted.
+    /// This method first checks with a read lock. If the key is missing, it
+    /// claims a per-key in-flight slot before computing the value so concurrent
+    /// callers for the same key wait for the first computation instead of
+    /// running duplicate closures with duplicate side effects. `None` results
+    /// are shared with current waiters but are not stored in the cache, so later
+    /// calls can retry the computation.
     #[must_use]
     pub fn get_or_insert_with<F>(&self, key: K, f: F) -> Option<V>
     where
         F: FnOnce() -> Option<V>,
         K: Clone,
     {
-        // First check with read lock
+        // First check with read lock.
         {
             let cache = self.cache.read().expect("locator cache lock poisoned");
             if let Some(value) = cache.get(&key) {
                 return Some(value.clone());
             }
         }
 
+        let in_flight = match self.claim_in_flight(&key) {
+            InFlightClaim::Owner(in_flight) => in_flight,
+            InFlightClaim::Waiter(entry) => return Self::wait_for_in_flight(entry),
+        };
+
+        // Check again after claiming the in-flight slot. Another thread may have
+        // completed the same key while this thread was waiting.
+        {
+            let cache = self.cache.read().expect("locator cache lock poisoned");
+            if let Some(value) = cache.get(&key) {
+                let result = Some(value.clone());
+                in_flight.complete(result.clone());
+                return result;
+            }
+        }
+
         // Compute the value (outside of any lock)
-        if let Some(value) = f() {
+        let result = if let Some(value) = f() {
             // Acquire write lock and insert
             let mut cache = self.cache.write().expect("locator cache lock poisoned");
             // Double-check in case another thread inserted while we were computing
             if let Some(existing) = cache.get(&key) {
-                return Some(existing.clone());
+                Some(existing.clone())
+            } else {
+                cache.insert(key, value.clone());
+                Some(value)
             }
-            cache.insert(key, value.clone());
-            Some(value)
         } else {
             None
+        };
+
+        in_flight.complete(result.clone());
+        result
+    }
+
+    fn claim_in_flight(&self, key: &K) -> InFlightClaim<'_, K, V>
+    where
+        K: Clone,
+    {
+        let mut in_flight = self
+            .in_flight
+            .lock()
+            .expect("locator cache in-flight lock poisoned");
+
+        if let Some(entry) = in_flight.get(key) {
+            return InFlightClaim::Waiter(entry.clone());
+        }
+
+        let entry = Arc::new(InFlightEntry::new());
+        in_flight.insert(key.clone(), entry.clone());
+        InFlightClaim::Owner(InFlightOwnerGuard {
+            key: Some(key.clone()),
+            entry,
+            in_flight: &self.in_flight,
+        })
+    }
+
+    fn wait_for_in_flight(entry: Arc<InFlightEntry<V>>) -> Option<V> {
+        let mut result = entry
+            .result
+            .lock()
+            .expect("locator cache in-flight result lock poisoned");
+        while result.is_none() {
+            result = entry
+                .changed
+                .wait(result)
+                .expect("locator cache in-flight condvar poisoned");
         }
+
+        result.clone().unwrap()
     }
 
     /// Clears all entries from the cache.
@@ -160,6 +284,12 @@ pub type ManagerCache = LocatorCache<PathBuf, EnvManager>;
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::sync::{
+        atomic::{AtomicUsize, Ordering},
+        mpsc, Arc, Barrier, Mutex,
+    };
+    use std::thread;
+    use std::time::Duration;
 
     #[test]
     fn test_cache_get_and_insert() {
@@ -192,6 +322,127 @@ mod tests {
         assert!(!cache.contains_key(&"key2".to_string()));
     }
 
+    #[test]
+    fn test_cache_get_or_insert_with_runs_one_closure_per_key() {
+        let cache: Arc<LocatorCache<String, i32>> = Arc::new(LocatorCache::new());
+        let barrier = Arc::new(Barrier::new(3));
+        let calls = Arc::new(AtomicUsize::new(0));
+        let (started_tx, started_rx) = mpsc::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let release_rx = Arc::new(Mutex::new(release_rx));
+        let mut handles = vec![];
+
+        for _ in 0..2 {
+            let cache = cache.clone();
+            let barrier = barrier.clone();
+            let calls = calls.clone();
+            let started_tx = started_tx.clone();
+            let release_rx = release_rx.clone();
+            handles.push(thread::spawn(move || {
+                barrier.wait();
+                cache.get_or_insert_with("key".to_string(), || {
+                    calls.fetch_add(1, Ordering::SeqCst);
+                    started_tx.send(()).unwrap();
+                    release_rx
+                        .lock()
+                        .unwrap()
+                        .recv_timeout(Duration::from_secs(5))
+                        .unwrap();
+                    Some(42)
+                })
+            }));
+        }
+
+        barrier.wait();
+        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
+        assert_eq!(calls.load(Ordering::SeqCst), 1);
+        assert!(started_rx.try_recv().is_err());
+
+        release_tx.send(()).unwrap();
+        release_tx.send(()).unwrap();
+
+        let mut results = handles
+            .into_iter()
+            .map(|handle| handle.join().unwrap())
+            .collect::<Vec<_>>();
+        results.sort();
+
+        assert_eq!(results, vec![Some(42), Some(42)]);
+        assert_eq!(calls.load(Ordering::SeqCst), 1);
+    }
+
+    #[test]
+    fn test_cache_get_or_insert_with_shares_concurrent_none_result() {
+        let cache: Arc<LocatorCache<String, i32>> = Arc::new(LocatorCache::new());
+        let barrier = Arc::new(Barrier::new(3));
+        let calls = Arc::new(AtomicUsize::new(0));
+        let (started_tx, started_rx) = mpsc::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let release_rx = Arc::new(Mutex::new(release_rx));
+        let mut handles = vec![];
+
+        for _ in 0..2 {
+            let cache = cache.clone();
+            let barrier = barrier.clone();
+            let calls = calls.clone();
+            let started_tx = started_tx.clone();
+            let release_rx = release_rx.clone();
+            handles.push(thread::spawn(move || {
+                barrier.wait();
+                cache.get_or_insert_with("key".to_string(), || {
+                    calls.fetch_add(1, Ordering::SeqCst);
+                    started_tx.send(()).unwrap();
+                    release_rx
+                        .lock()
+                        .unwrap()
+                        .recv_timeout(Duration::from_secs(5))
+                        .unwrap();
+                    None
+                })
+            }));
+        }
+
+        barrier.wait();
+        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
+        assert_eq!(calls.load(Ordering::SeqCst), 1);
+        assert!(started_rx.try_recv().is_err());
+
+        release_tx.send(()).unwrap();
+        release_tx.send(()).unwrap();
+
+        let results = handles
+            .into_iter()
+            .map(|handle| handle.join().unwrap())
+            .collect::<Vec<_>>();
+
+        assert_eq!(results, vec![None, None]);
+        assert_eq!(calls.load(Ordering::SeqCst), 1);
+        assert!(!cache.contains_key(&"key".to_string()));
+
+        assert_eq!(
+            cache.get_or_insert_with("key".to_string(), || Some(42)),
+            Some(42)
+        );
+    }
+
+    #[test]
+    fn test_cache_get_or_insert_with_panic_releases_in_flight_key() {
+        let cache: LocatorCache<String, i32> = LocatorCache::new();
+
+        let result = std::panic::catch_unwind(|| {
+            let _ = cache.get_or_insert_with("key".to_string(), || -> Option<i32> {
+                panic!("boom");
+            });
+        });
+
+        assert!(result.is_err());
+        assert!(!cache.contains_key(&"key".to_string()));
+        assert_eq!(
+            cache.get_or_insert_with("key".to_string(), || Some(42)),
+            Some(42)
+        );
+    }
+
     #[test]
     fn test_cache_clear() {
         let cache: LocatorCache<String, i32> = LocatorCache::new();