Skip to main content

soma_provider_adapters/python/
cache_prune.rs

1//! Conservative planning and application of Python cache cleanup.
2
3use std::{
4    fs, io,
5    path::{Path, PathBuf},
6    sync::atomic::{AtomicU64, Ordering},
7};
8
9use serde::Serialize;
10use thiserror::Error;
11
12use super::{
13    PythonEnvironmentCache, PythonEnvironmentCacheEntry, PythonEnvironmentCacheError,
14    PythonEnvironmentCacheState, inspect_entry,
15};
16
17static PRUNE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
20pub struct PythonEnvironmentPrunePolicy {
21    pub stale_before_unix_seconds: u64,
22    pub remove_incomplete: bool,
23    pub remove_invalid: bool,
24    pub remove_staging: bool,
25}
26
27impl PythonEnvironmentPrunePolicy {
28    pub fn conservative(stale_before_unix_seconds: u64) -> Self {
29        Self {
30            stale_before_unix_seconds,
31            remove_incomplete: true,
32            remove_invalid: true,
33            remove_staging: true,
34        }
35    }
36
37    fn selects(self, entry: &PythonEnvironmentCacheEntry) -> bool {
38        let selected_state = match entry.state {
39            PythonEnvironmentCacheState::Ready => false,
40            PythonEnvironmentCacheState::Incomplete => self.remove_incomplete,
41            PythonEnvironmentCacheState::Invalid => self.remove_invalid,
42            PythonEnvironmentCacheState::Staging => self.remove_staging,
43        };
44        selected_state
45            && entry
46                .modified_unix_seconds
47                .is_some_and(|modified| modified <= self.stale_before_unix_seconds)
48    }
49}
50
51#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
52pub struct PythonEnvironmentPruneCandidate {
53    pub entry: PythonEnvironmentCacheEntry,
54    pub reason: String,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
58pub struct PythonEnvironmentPrunePlan {
59    pub root: PathBuf,
60    pub policy: PythonEnvironmentPrunePolicy,
61    pub candidates: Vec<PythonEnvironmentPruneCandidate>,
62    pub reclaimable_size_bytes: u64,
63    pub reclaimable_file_count: u64,
64}
65
66#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
67#[serde(rename_all = "snake_case")]
68pub enum PythonEnvironmentPruneOutcome {
69    Removed {
70        directory: PathBuf,
71        reclaimed_size_bytes: u64,
72        reclaimed_file_count: u64,
73    },
74    Missing {
75        directory: PathBuf,
76    },
77    Changed {
78        directory: PathBuf,
79    },
80}
81
82#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
83pub struct PythonEnvironmentPruneReport {
84    pub root: PathBuf,
85    pub outcomes: Vec<PythonEnvironmentPruneOutcome>,
86    pub removed: usize,
87    pub missing: usize,
88    pub changed: usize,
89    pub reclaimed_size_bytes: u64,
90    pub reclaimed_file_count: u64,
91}
92
93#[derive(Debug, Error)]
94pub enum PythonEnvironmentPruneError {
95    #[error(transparent)]
96    Inventory(#[from] PythonEnvironmentCacheError),
97    #[error("Python prune plan root {} does not match cache root {}", plan.display(), cache.display())]
98    RootMismatch { plan: PathBuf, cache: PathBuf },
99    #[error("Python prune candidate is outside the managed cache: {}", path.display())]
100    OutsideCache { path: PathBuf },
101    #[error("Python prune plans may never delete ready environments: {}", path.display())]
102    ReadyEnvironment { path: PathBuf },
103    #[error("Python prune I/O failed at {}: {source}", path.display())]
104    Io {
105        path: PathBuf,
106        #[source]
107        source: io::Error,
108    },
109}
110
111impl PythonEnvironmentCache {
112    pub fn plan_prune(
113        &self,
114        policy: PythonEnvironmentPrunePolicy,
115    ) -> Result<PythonEnvironmentPrunePlan, PythonEnvironmentPruneError> {
116        let inventory = self.inventory()?;
117        let candidates = inventory
118            .entries
119            .into_iter()
120            .filter(|entry| policy.selects(entry))
121            .map(|entry| PythonEnvironmentPruneCandidate {
122                reason: entry
123                    .issue
124                    .clone()
125                    .unwrap_or_else(|| "selected stale cache entry".to_owned()),
126                entry,
127            })
128            .collect::<Vec<_>>();
129        let reclaimable_size_bytes = candidates.iter().fold(0_u64, |total, candidate| {
130            total.saturating_add(candidate.entry.size_bytes)
131        });
132        let reclaimable_file_count = candidates.iter().fold(0_u64, |total, candidate| {
133            total.saturating_add(candidate.entry.file_count)
134        });
135        Ok(PythonEnvironmentPrunePlan {
136            root: self.root.clone(),
137            policy,
138            candidates,
139            reclaimable_size_bytes,
140            reclaimable_file_count,
141        })
142    }
143
144    pub fn apply_prune(
145        &self,
146        plan: &PythonEnvironmentPrunePlan,
147    ) -> Result<PythonEnvironmentPruneReport, PythonEnvironmentPruneError> {
148        if plan.root != self.root {
149            return Err(PythonEnvironmentPruneError::RootMismatch {
150                plan: plan.root.clone(),
151                cache: self.root.clone(),
152            });
153        }
154        for candidate in &plan.candidates {
155            validate_candidate_path(&self.root, &candidate.entry)?;
156            if candidate.entry.state == PythonEnvironmentCacheState::Ready {
157                return Err(PythonEnvironmentPruneError::ReadyEnvironment {
158                    path: candidate.entry.directory.clone(),
159                });
160            }
161        }
162
163        let current_inventory = self.inventory()?;
164        let mut outcomes = Vec::with_capacity(plan.candidates.len());
165        for candidate in &plan.candidates {
166            let path = &candidate.entry.directory;
167            let Some(current) = current_inventory
168                .entries
169                .iter()
170                .find(|entry| entry.directory == *path)
171            else {
172                outcomes.push(PythonEnvironmentPruneOutcome::Missing {
173                    directory: path.clone(),
174                });
175                continue;
176            };
177            if current != &candidate.entry || !plan.policy.selects(current) {
178                outcomes.push(PythonEnvironmentPruneOutcome::Changed {
179                    directory: path.clone(),
180                });
181                continue;
182            }
183            let revalidated = match fs::symlink_metadata(path) {
184                Ok(_) => inspect_entry(path, current.plan_directory_version)?,
185                Err(error) if error.kind() == io::ErrorKind::NotFound => {
186                    outcomes.push(PythonEnvironmentPruneOutcome::Missing {
187                        directory: path.clone(),
188                    });
189                    continue;
190                }
191                Err(source) => {
192                    return Err(PythonEnvironmentPruneError::Io {
193                        path: path.clone(),
194                        source,
195                    });
196                }
197            };
198            if revalidated != *current || !plan.policy.selects(&revalidated) {
199                outcomes.push(PythonEnvironmentPruneOutcome::Changed {
200                    directory: path.clone(),
201                });
202                continue;
203            }
204
205            let quarantine = quarantine_path(path)?;
206            match fs::rename(path, &quarantine) {
207                Ok(()) => {}
208                Err(error) if error.kind() == io::ErrorKind::NotFound => {
209                    outcomes.push(PythonEnvironmentPruneOutcome::Missing {
210                        directory: path.clone(),
211                    });
212                    continue;
213                }
214                Err(source) => {
215                    return Err(PythonEnvironmentPruneError::Io {
216                        path: path.clone(),
217                        source,
218                    });
219                }
220            }
221            remove_quarantine(&quarantine)?;
222            outcomes.push(PythonEnvironmentPruneOutcome::Removed {
223                directory: path.clone(),
224                reclaimed_size_bytes: current.size_bytes,
225                reclaimed_file_count: current.file_count,
226            });
227        }
228
229        let mut report = PythonEnvironmentPruneReport {
230            root: self.root.clone(),
231            outcomes,
232            removed: 0,
233            missing: 0,
234            changed: 0,
235            reclaimed_size_bytes: 0,
236            reclaimed_file_count: 0,
237        };
238        for outcome in &report.outcomes {
239            match outcome {
240                PythonEnvironmentPruneOutcome::Removed {
241                    reclaimed_size_bytes,
242                    reclaimed_file_count,
243                    ..
244                } => {
245                    report.removed += 1;
246                    report.reclaimed_size_bytes = report
247                        .reclaimed_size_bytes
248                        .saturating_add(*reclaimed_size_bytes);
249                    report.reclaimed_file_count = report
250                        .reclaimed_file_count
251                        .saturating_add(*reclaimed_file_count);
252                }
253                PythonEnvironmentPruneOutcome::Missing { .. } => report.missing += 1,
254                PythonEnvironmentPruneOutcome::Changed { .. } => report.changed += 1,
255            }
256        }
257        Ok(report)
258    }
259}
260
261fn validate_candidate_path(
262    root: &Path,
263    entry: &PythonEnvironmentCacheEntry,
264) -> Result<(), PythonEnvironmentPruneError> {
265    let path = &entry.directory;
266    let parent = path.parent();
267    let managed = parent == Some(root)
268        || parent
269            .and_then(Path::parent)
270            .is_some_and(|grandparent| grandparent == root);
271    if managed {
272        Ok(())
273    } else {
274        Err(PythonEnvironmentPruneError::OutsideCache { path: path.clone() })
275    }
276}
277
278fn quarantine_path(path: &Path) -> Result<PathBuf, PythonEnvironmentPruneError> {
279    let parent = path
280        .parent()
281        .ok_or_else(|| PythonEnvironmentPruneError::OutsideCache {
282            path: path.to_path_buf(),
283        })?;
284    let name = path
285        .file_name()
286        .and_then(|name| name.to_str())
287        .unwrap_or("environment");
288    loop {
289        let sequence = PRUNE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
290        let candidate = parent.join(format!(".{name}.prune-{}-{sequence}", std::process::id()));
291        match fs::symlink_metadata(&candidate) {
292            Ok(_) => continue,
293            Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(candidate),
294            Err(source) => {
295                return Err(PythonEnvironmentPruneError::Io {
296                    path: candidate,
297                    source,
298                });
299            }
300        }
301    }
302}
303
304fn remove_quarantine(path: &Path) -> Result<(), PythonEnvironmentPruneError> {
305    let metadata =
306        fs::symlink_metadata(path).map_err(|source| PythonEnvironmentPruneError::Io {
307            path: path.to_path_buf(),
308            source,
309        })?;
310    let result = if metadata.file_type().is_symlink() || !metadata.is_dir() {
311        fs::remove_file(path)
312    } else {
313        fs::remove_dir_all(path)
314    };
315    result.map_err(|source| PythonEnvironmentPruneError::Io {
316        path: path.to_path_buf(),
317        source,
318    })
319}