Skip to content

Commit 9f01a4d

Browse files
Copilotkarthiknadig
Copilot
and
karthiknadig
committed
feat: Add LocatorCache abstraction and refactor pet-conda and pet-linux-global-python
Co-authored-by: karthiknadig <[email protected]>
1 parent e9d2527 commit 9f01a4d

File tree

4 files changed

+226
-92
lines changed

4 files changed

+226
-92
lines changed

‎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},
@@ -20,7 +21,6 @@ use pet_core::{
2021
use pet_fs::path::norm_case;
2122
use serde::{Deserialize, Serialize};
2223
use std::{
23-
collections::HashMap,
2424
path::{Path, PathBuf},
2525
sync::{Arc, RwLock},
2626
thread,
@@ -61,24 +61,24 @@ pub struct CondaTelemetryInfo {
6161
}
6262

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

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

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

@@ -119,10 +119,7 @@ impl CondaLocator for Conda {
119119

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

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

@@ -183,7 +175,8 @@ impl CondaLocator for Conda {
183175
.and_then(|p| CondaManager::from(&p))
184176
.unwrap_or(manager.clone());
185177
let env = conda_env.to_python_environment(Some(manager.to_manager()));
186-
environments.insert(conda_env.prefix.clone(), env.clone());
178+
self.environments
179+
.insert(conda_env.prefix.clone(), env.clone());
187180
reporter.report_manager(&manager.to_manager());
188181
reporter.report_environment(&env);
189182
}
@@ -194,22 +187,8 @@ impl CondaLocator for Conda {
194187

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

@@ -250,29 +229,25 @@ impl Locator for Conda {
250229
return None;
251230
}
252231

253-
// First check cache with read lock
254-
{
255-
let environments = self.environments.read().unwrap();
256-
if let Some(env) = environments.get(path) {
257-
return Some(env.clone());
258-
}
232+
// Check cache first
233+
if let Some(cached_env) = self.environments.get(path) {
234+
return Some(cached_env);
259235
}
260-
// Not in cache, build the environment and insert with write lock
236+
237+
// Not in cache, build the environment and insert
261238
if let Some(env) = get_conda_environment_info(path, &None) {
262239
if let Some(conda_dir) = &env.conda_dir {
263240
if let Some(manager) = self.get_manager(conda_dir) {
264241
let env = env.to_python_environment(Some(manager.to_manager()));
265-
let mut environments = self.environments.write().unwrap();
266-
environments.insert(path.clone(), env.clone());
242+
self.environments.insert(path.clone(), env.clone());
267243
return Some(env);
268244
} else {
269245
// We will still return the conda env even though we do not have the manager.
270246
// This might seem incorrect, however the tool is about discovering environments.
271247
// The client can activate this env either using another conda manager or using the activation scripts
272248
error!("Unable to find Conda Manager for env (even though we have a conda_dir): {:?}", env);
273249
let env = env.to_python_environment(None);
274-
let mut environments = self.environments.write().unwrap();
275-
environments.insert(path.clone(), env.clone());
250+
self.environments.insert(path.clone(), env.clone());
276251
return Some(env);
277252
}
278253
} else {
@@ -281,8 +256,7 @@ impl Locator for Conda {
281256
// The client can activate this env either using another conda manager or using the activation scripts
282257
error!("Unable to find Conda Manager for env: {:?}", env);
283258
let env = env.to_python_environment(None);
284-
let mut environments = self.environments.write().unwrap();
285-
environments.insert(path.clone(), env.clone());
259+
self.environments.insert(path.clone(), env.clone());
286260
return Some(env);
287261
}
288262
}
@@ -315,8 +289,7 @@ impl Locator for Conda {
315289
error!("Unable to find Conda Manager for the Conda env: {:?}", env);
316290
let prefix = env.prefix.clone();
317291
let env = env.to_python_environment(None);
318-
let mut environments = self.environments.write().unwrap();
319-
environments.insert(prefix, env.clone());
292+
self.environments.insert(prefix, env.clone());
320293
reporter.report_environment(&env);
321294
return None;
322295
}
@@ -325,38 +298,23 @@ impl Locator for Conda {
325298
// We will try to get the manager for this conda_dir
326299
let prefix = env.clone().prefix.clone();
327300

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

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

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

0 commit comments

Comments
 (0)