Add an in-memory cache for ManagedInstallations::find_all

crates/uv-python/src/managed.rs

@@ -1,9 +1,11 @@
 use core::fmt;
 use std::cmp::Reverse;
+use std::collections::HashMap;
 use std::ffi::OsStr;
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
+use std::sync::{Mutex, OnceLock};
 
 use fs_err as fs;
 use itertools::Itertools;
@@ -26,6 +28,9 @@ use crate::platform::Error as PlatformError;
 use crate::platform::{Arch, Libc, Os};
 use crate::python_version::PythonVersion;
 use crate::{PythonRequest, PythonVariant};
+
+static FIND_ALL_CACHE: OnceLock<Mutex<HashMap<PathBuf, Vec<PathBuf>>>> = OnceLock::new();
+
 #[derive(Error, Debug)]
 pub enum Error {
     #[error(transparent)]
@@ -181,31 +186,41 @@ impl ManagedPythonInstallations {
     pub fn find_all(
         &self,
     ) -> Result<impl DoubleEndedIterator<Item = ManagedPythonInstallation>, Error> {
-        let dirs = match fs_err::read_dir(&self.root) {
-            Ok(installation_dirs) => {
-                // Collect sorted directory paths; `read_dir` is not stable across platforms
-                let directories: Vec<_> = installation_dirs
-                    .filter_map(|read_dir| match read_dir {
-                        Ok(entry) => match entry.file_type() {
-                            Ok(file_type) => file_type.is_dir().then_some(Ok(entry.path())),
+        let find_all_dirs_cache = FIND_ALL_CACHE.get_or_init(|| Mutex::new(HashMap::new()));
+        let mut find_all_dirs_cache = find_all_dirs_cache.lock().unwrap();
+
+        let dirs = if let Some(cached) = find_all_dirs_cache.get(&self.root) {
+            // Skip the `read_dir` call on subsequent invocations in-process
+            cached.clone()
+        } else {
+            let dirs = match fs_err::read_dir(&self.root) {
+                Ok(installation_dirs) => {
+                    // Collect sorted directory paths; `read_dir` is not stable across platforms
+                    let directories: Vec<_> = installation_dirs
+                        .filter_map(|read_dir| match read_dir {
+                            Ok(entry) => match entry.file_type() {
+                                Ok(file_type) => file_type.is_dir().then_some(Ok(entry.path())),
+                                Err(err) => Some(Err(err)),
+                            },
                             Err(err) => Some(Err(err)),
-                        },
-                        Err(err) => Some(Err(err)),
-                    })
-                    .collect::<Result<_, io::Error>>()
-                    .map_err(|err| Error::ReadError {
+                        })
+                        .collect::<Result<_, io::Error>>()
+                        .map_err(|err| Error::ReadError {
+                            dir: self.root.clone(),
+                            err,
+                        })?;
+                    directories
+                }
+                Err(err) if err.kind() == io::ErrorKind::NotFound => vec![],
+                Err(err) => {
+                    return Err(Error::ReadError {
                         dir: self.root.clone(),
                         err,
-                    })?;
-                directories
-            }
-            Err(err) if err.kind() == io::ErrorKind::NotFound => vec![],
-            Err(err) => {
-                return Err(Error::ReadError {
-                    dir: self.root.clone(),
-                    err,
-                })
-            }
+                    })
+                }
+            };
+            find_all_dirs_cache.insert(self.root.clone(), dirs.clone());
+            dirs
         };
         let cache = self.cache();
         Ok(dirs