Skip to content

Commit 1d30dd4

Browse files
Refactor: Consolidate caching patterns into shared LocatorCache abstraction (#301)
- [x] Create a new `LocatorCache<K, V>` generic caching abstraction in `pet-core` - [x] Add `get()`, `get_or_insert_with()`, `insert()`, `contains_key()`, `clear()`, `values()`, `len()`, `is_empty()`, and `clone_map()` methods - [x] Create type aliases `EnvironmentCache` and `ManagerCache` for common uses - [x] Refactor `pet-conda` to use the new cache abstraction - [x] Refactor `pet-linux-global-python` to use the new cache abstraction - [x] Run tests to verify the changes work correctly - [x] Request code review - [x] Run CodeQL security checker - [x] Address PR review feedback: - [x] Add `insert_many()` method for atomic batch inserts (addresses lock granularity concern) - [x] Add `#[must_use]` attribute to `get_or_insert_with()` - [x] Update `pet-linux-global-python` to use `insert_many()` for symlinks + executable ## Summary This PR introduces a generic `LocatorCache<K, V>` abstraction in the `pet-core` crate that consolidates the common caching patterns used across multiple locators. The cache provides: 1. Thread-safe operations using `RwLock` 2. Efficient read access with separate read/write locks 3. Double-check locking in `get_or_insert_with()` to prevent redundant computation 4. Atomic batch inserts via `insert_many()` for better lock efficiency 5. Type aliases for common use cases (`EnvironmentCache`, `ManagerCache`) ## Security Summary No security vulnerabilities were discovered by CodeQL analysis. <!-- START COPILOT ORIGINAL PROMPT --> <details> <summary>Original prompt</summary> > > ---- > > *This section details on the original issue you should resolve* > > <issue_title>Refactor: Consolidate caching patterns across locators into a shared abstraction</issue_title> > <issue_description>## Summary > Multiple locators implement nearly identical caching patterns for environments and managers. Consolidating these into a shared abstraction would reduce code duplication and ensure consistent behavior. > > ## Current Duplicate Patterns > > ### Pattern: Cache with mutex + contains_key check > Found in: > - `crates/pet-conda/src/lib.rs` > - `crates/pet-poetry/src/lib.rs` > - `crates/pet-windows-store/src/lib.rs` > - `crates/pet-linux-global-python/src/lib.rs` > > ```rust > // This pattern appears in multiple files: > pub struct SomeLocator { > environments: Arc<Mutex<HashMap<PathBuf, PythonEnvironment>>>, > } > > impl SomeLocator { > fn find_with_cache(&self) -> Option<...> { > let environments = self.environments.lock().unwrap(); > if environments.contains_key(&key) { > return Some(environments.get(&key).unwrap().clone()); > } > // ... compute and cache > } > > fn clear(&self) { > self.environments.lock().unwrap().clear(); > } > } > ``` > > ## Proposed Solution > > ### Create a generic caching wrapper > ```rust > // In pet-core or a new pet-cache crate > pub struct LocatorCache<K, V> { > cache: RwLock<HashMap<K, V>>, > } > > impl<K: Eq + Hash, V: Clone> LocatorCache<K, V> { > pub fn new() -> Self { > Self { cache: RwLock::new(HashMap::new()) } > } > > pub fn get(&self, key: &K) -> Option<V> { > self.cache.read().unwrap().get(key).cloned() > } > > pub fn get_or_insert_with<F>(&self, key: K, f: F) -> V > where > F: FnOnce() -> Option<V>, > { > // Check read lock first > if let Some(v) = self.cache.read().unwrap().get(&key) { > return v.clone(); > } > // Upgrade to write lock and compute > if let Some(value) = f() { > self.cache.write().unwrap().insert(key, value.clone()); > value > } > } > > pub fn clear(&self) { > self.cache.write().unwrap().clear(); > } > } > > // Type aliases for common uses > pub type EnvironmentCache = LocatorCache<PathBuf, PythonEnvironment>; > pub type ManagerCache = LocatorCache<PathBuf, EnvManager>; > ``` > > ### Usage in locators > ```rust > pub struct Conda { > pub environments: EnvironmentCache, > pub managers: ManagerCache, > // ... > } > > impl Conda { > fn try_from(&self, env: &PythonEnv) -> Option<PythonEnvironment> { > self.environments.get_or_insert_with(path.clone(), || { > // compute environment > }) > } > } > ``` > > ## Benefits > 1. **Consistency**: All locators behave the same way > 2. **Less code**: Remove duplicate boilerplate > 3. **Centralized improvements**: RwLock upgrade, better error handling, etc. benefit all locators > 4. **Testability**: Cache behavior can be tested once > > ## Priority > Low - Requires significant refactoring but improves maintainability long-term.</issue_description> > > ## Comments on the Issue (you are @copilot in this section) > > <comments> > </comments> > </details> <!-- START COPILOT CODING AGENT SUFFIX --> - Fixes #291 <!-- START COPILOT CODING AGENT TIPS --> --- 💬 We'd love your input! Share your thoughts on Copilot coding agent in our [2 minute survey](https://gh.io/copilot-coding-agent-survey). --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: karthiknadig <3840081+karthiknadig@users.noreply.github.com>
1 parent 525aa9f commit 1d30dd4

4 files changed

Lines changed: 263 additions & 95 deletions

File tree

crates/pet-conda/src/lib.rs

Lines changed: 34 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use environments::{get_conda_environment_info, CondaEnvironment};
1111
use log::error;
1212
use manager::CondaManager;
1313
use pet_core::{
14+
cache::LocatorCache,
1415
env::PythonEnv,
1516
os_environment::Environment,
1617
python_environment::{PythonEnvironment, PythonEnvironmentKind},
@@ -21,7 +22,6 @@ use pet_fs::path::norm_case;
2122
use rayon::prelude::*;
2223
use serde::{Deserialize, Serialize};
2324
use std::{
24-
collections::HashMap,
2525
path::{Path, PathBuf},
2626
sync::{Arc, RwLock},
2727
thread,
@@ -62,24 +62,24 @@ pub struct CondaTelemetryInfo {
6262
}
6363

6464
pub struct Conda {
65-
pub environments: Arc<RwLock<HashMap<PathBuf, PythonEnvironment>>>,
66-
pub managers: Arc<RwLock<HashMap<PathBuf, CondaManager>>>,
65+
pub environments: Arc<LocatorCache<PathBuf, PythonEnvironment>>,
66+
pub managers: Arc<LocatorCache<PathBuf, CondaManager>>,
6767
pub env_vars: EnvVariables,
6868
conda_executable: Arc<RwLock<Option<PathBuf>>>,
6969
}
7070

7171
impl Conda {
7272
pub fn from(env: &dyn Environment) -> Conda {
7373
Conda {
74-
environments: Arc::new(RwLock::new(HashMap::new())),
75-
managers: Arc::new(RwLock::new(HashMap::new())),
74+
environments: Arc::new(LocatorCache::new()),
75+
managers: Arc::new(LocatorCache::new()),
7676
env_vars: EnvVariables::from(env),
7777
conda_executable: Arc::new(RwLock::new(None)),
7878
}
7979
}
8080
fn clear(&self) {
81-
self.environments.write().unwrap().clear();
82-
self.managers.write().unwrap().clear();
81+
self.environments.clear();
82+
self.managers.clear();
8383
}
8484
}
8585

@@ -92,17 +92,17 @@ impl CondaLocator for Conda {
9292
// Look for environments that we couldn't find without spawning conda.
9393
let user_provided_conda_exe = conda_executable.is_some();
9494
let conda_info = CondaInfo::from(conda_executable)?;
95-
let environments = self.environments.read().unwrap().clone();
95+
let environments_map = self.environments.clone_map();
9696
let new_envs = conda_info
9797
.envs
9898
.clone()
9999
.into_iter()
100-
.filter(|p| !environments.contains_key(p))
100+
.filter(|p| !environments_map.contains_key(p))
101101
.collect::<Vec<PathBuf>>();
102102
if new_envs.is_empty() {
103103
return None;
104104
}
105-
let environments = environments
105+
let environments = environments_map
106106
.into_values()
107107
.collect::<Vec<PythonEnvironment>>();
108108

@@ -120,10 +120,7 @@ impl CondaLocator for Conda {
120120

121121
fn get_info_for_telemetry(&self, conda_executable: Option<PathBuf>) -> CondaTelemetryInfo {
122122
let can_spawn_conda = CondaInfo::from(conda_executable).is_some();
123-
let environments = self.environments.read().unwrap().clone();
124-
let environments = environments
125-
.into_values()
126-
.collect::<Vec<PythonEnvironment>>();
123+
let environments = self.environments.values();
127124
let (conda_rcs, env_dirs) = get_conda_rcs_and_env_dirs(&self.env_vars, &environments);
128125
let mut environments_txt = None;
129126
let mut environments_txt_exists = None;
@@ -160,19 +157,14 @@ impl CondaLocator for Conda {
160157
if let Some(conda_dir) = manager.conda_dir.clone() {
161158
// Keep track to search again later.
162159
// Possible we'll find environments in other directories created using this manager
163-
let mut managers = self.managers.write().unwrap();
164-
// Keep track to search again later.
165-
// Possible we'll find environments in other directories created using this manager
166-
managers.insert(conda_dir.clone(), manager.clone());
167-
drop(managers);
160+
self.managers.insert(conda_dir.clone(), manager.clone());
168161

169162
// Find all the environments in the conda install folder. (under `envs` folder)
170163
for conda_env in
171164
get_conda_environments(&get_environments(&conda_dir), &manager.clone().into())
172165
{
173166
// If reported earlier, no point processing this again.
174-
let mut environments = self.environments.write().unwrap();
175-
if environments.contains_key(&conda_env.prefix) {
167+
if self.environments.contains_key(&conda_env.prefix) {
176168
continue;
177169
}
178170

@@ -184,7 +176,8 @@ impl CondaLocator for Conda {
184176
.and_then(|p| CondaManager::from(&p))
185177
.unwrap_or(manager.clone());
186178
let env = conda_env.to_python_environment(Some(manager.to_manager()));
187-
environments.insert(conda_env.prefix.clone(), env.clone());
179+
self.environments
180+
.insert(conda_env.prefix.clone(), env.clone());
188181
reporter.report_manager(&manager.to_manager());
189182
reporter.report_environment(&env);
190183
}
@@ -195,22 +188,8 @@ impl CondaLocator for Conda {
195188

196189
impl Conda {
197190
fn get_manager(&self, conda_dir: &Path) -> Option<CondaManager> {
198-
// First try to read from cache
199-
{
200-
let managers = self.managers.read().unwrap();
201-
if let Some(mgr) = managers.get(conda_dir) {
202-
return Some(mgr.clone());
203-
}
204-
}
205-
206-
// If not found, acquire write lock and insert
207-
if let Some(manager) = CondaManager::from(conda_dir) {
208-
let mut managers = self.managers.write().unwrap();
209-
managers.insert(conda_dir.into(), manager.clone());
210-
Some(manager)
211-
} else {
212-
None
213-
}
191+
self.managers
192+
.get_or_insert_with(conda_dir.to_path_buf(), || CondaManager::from(conda_dir))
214193
}
215194
}
216195

@@ -251,29 +230,25 @@ impl Locator for Conda {
251230
return None;
252231
}
253232

254-
// First check cache with read lock
255-
{
256-
let environments = self.environments.read().unwrap();
257-
if let Some(env) = environments.get(path) {
258-
return Some(env.clone());
259-
}
233+
// Check cache first
234+
if let Some(cached_env) = self.environments.get(path) {
235+
return Some(cached_env);
260236
}
261-
// Not in cache, build the environment and insert with write lock
237+
238+
// Not in cache, build the environment and insert
262239
if let Some(env) = get_conda_environment_info(path, &None) {
263240
if let Some(conda_dir) = &env.conda_dir {
264241
if let Some(manager) = self.get_manager(conda_dir) {
265242
let env = env.to_python_environment(Some(manager.to_manager()));
266-
let mut environments = self.environments.write().unwrap();
267-
environments.insert(path.clone(), env.clone());
243+
self.environments.insert(path.clone(), env.clone());
268244
return Some(env);
269245
} else {
270246
// We will still return the conda env even though we do not have the manager.
271247
// This might seem incorrect, however the tool is about discovering environments.
272248
// The client can activate this env either using another conda manager or using the activation scripts
273249
error!("Unable to find Conda Manager for env (even though we have a conda_dir): {:?}", env);
274250
let env = env.to_python_environment(None);
275-
let mut environments = self.environments.write().unwrap();
276-
environments.insert(path.clone(), env.clone());
251+
self.environments.insert(path.clone(), env.clone());
277252
return Some(env);
278253
}
279254
} else {
@@ -282,8 +257,7 @@ impl Locator for Conda {
282257
// The client can activate this env either using another conda manager or using the activation scripts
283258
error!("Unable to find Conda Manager for env: {:?}", env);
284259
let env = env.to_python_environment(None);
285-
let mut environments = self.environments.write().unwrap();
286-
environments.insert(path.clone(), env.clone());
260+
self.environments.insert(path.clone(), env.clone());
287261
return Some(env);
288262
}
289263
}
@@ -316,8 +290,7 @@ impl Locator for Conda {
316290
error!("Unable to find Conda Manager for the Conda env: {:?}", env);
317291
let prefix = env.prefix.clone();
318292
let env = env.to_python_environment(None);
319-
let mut environments = self.environments.write().unwrap();
320-
environments.insert(prefix, env.clone());
293+
self.environments.insert(prefix, env.clone());
321294
reporter.report_environment(&env);
322295
return None;
323296
}
@@ -326,38 +299,23 @@ impl Locator for Conda {
326299
// We will try to get the manager for this conda_dir
327300
let prefix = env.clone().prefix.clone();
328301

329-
{
330-
// 3.1 Check if we have already reported this environment.
331-
// Closure to quickly release lock
332-
let environments = self.environments.read().unwrap();
333-
if environments.contains_key(&env.prefix) {
334-
return None;
335-
}
302+
// 3.1 Check if we have already reported this environment.
303+
if self.environments.contains_key(&env.prefix) {
304+
return None;
336305
}
337306

338-
339307
// 4 Get the manager for this env.
340308
let conda_dir = &env.conda_dir.clone()?;
341-
let managers = self.managers.read().unwrap();
342-
let mut manager = managers.get(conda_dir).cloned();
343-
drop(managers);
344-
345-
if manager.is_none() {
346-
// 4.1 Build the manager from the conda dir if we do not have it.
347-
if let Some(conda_manager) = CondaManager::from(conda_dir) {
348-
let mut managers = self.managers.write().unwrap();
349-
managers.insert(conda_dir.to_path_buf().clone(), conda_manager.clone());
350-
manager = Some(conda_manager);
351-
}
352-
}
309+
let manager = self.managers.get_or_insert_with(conda_dir.clone(), || {
310+
CondaManager::from(conda_dir)
311+
});
353312

354313
// 5. Report this env.
355314
if let Some(manager) = manager {
356315
let env = env.to_python_environment(
357316
Some(manager.to_manager()),
358317
);
359-
let mut environments = self.environments.write().unwrap();
360-
environments.insert(prefix.clone(), env.clone());
318+
self.environments.insert(prefix.clone(), env.clone());
361319
reporter.report_manager(&manager.to_manager());
362320
reporter.report_environment(&env);
363321
} else {
@@ -366,8 +324,7 @@ impl Locator for Conda {
366324
// The client can activate this env either using another conda manager or using the activation scripts
367325
error!("Unable to find Conda Manager for Conda env (even though we have a conda_dir {:?}): Env Details = {:?}", conda_dir, env);
368326
let env = env.to_python_environment(None);
369-
let mut environments = self.environments.write().unwrap();
370-
environments.insert(prefix.clone(), env.clone());
327+
self.environments.insert(prefix.clone(), env.clone());
371328
reporter.report_environment(&env);
372329
}
373330
Option::<()>::Some(())

0 commit comments

Comments
 (0)