soma_provider_adapters/python/
cache_prune.rs1use 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}