byob_plugins/
pyenv.rs

1//! The Python a plugin's scripts and backend run with.
2//!
3//! Every plugin gets byob's helper module (`byob_plugin.py`, [`HELPER`]) on `PYTHONPATH`.
4//! A plugin that lists `python.requirements` gets its own virtual environment, cached under
5//! byob's data folder and keyed by a hash of the interpreter and the requirements (so two
6//! plugins with the same list share one, and changing the list makes a new one). It is
7//! built the first time something needs it, never on the engine thread: with `uv` when it
8//! is on `PATH` (`uv venv` + `uv pip install`), else `python -m venv` + `pip install`. A
9//! failure is kept (with the installer's output) until the plugin is reloaded, and shows
10//! as one of the plugin's problems. A build can be called off (the tool call stopped, the
11//! backend stopped); it isn't kept as a failure then.
12//!
13//! Two byob processes may share the data folder: a build holds a lock file next to the
14//! environment (`<venv>.lock`), and the other waits for it and uses what it built. A lock
15//! whose holder stopped refreshing it (a byob that died mid-build) is taken over.
16
17use std::collections::BTreeMap;
18use std::ffi::OsString;
19use std::fs;
20use std::hash::Hasher;
21use std::path::{Path, PathBuf};
22use std::process::Command;
23use std::sync::{Arc, Mutex, OnceLock, PoisonError, TryLockError};
24use std::time::{Duration, Instant, SystemTime};
25
26use byob_sandbox::RunOptions;
27
28use crate::authoring::Fnv;
29
30/// byob's Python helper for plugins (`import byob_plugin`).
31pub const HELPER: &str = include_str!("byob_plugin.py");
32/// Its file name.
33pub const HELPER_FILE: &str = "byob_plugin.py";
34/// Longest an environment may take to build (downloads included).
35const BUILD_TIMEOUT: Duration = Duration::from_secs(15 * 60);
36/// Installer output kept per stream.
37const MAX_BUILD_OUTPUT: usize = 64 * 1024;
38/// Characters of the installer's output in a problem message.
39const PROBLEM_OUTPUT: usize = 3000;
40/// The file that says an environment is complete.
41const READY: &str = ".byob-ready";
42/// What a build that was called off says.
43pub const BUILD_CALLED_OFF: &str = "the Python environment wasn't built: stopped";
44/// How often a build refreshes its lock file.
45const LOCK_BEAT: Duration = Duration::from_secs(20);
46/// A lock file not refreshed for this long is left over from a byob that died.
47const LOCK_STALE: Duration = Duration::from_secs(120);
48/// How often a waiting build looks at the lock again.
49const LOCK_POLL: Duration = Duration::from_millis(100);
50
51/// How to start Python for one plugin.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct Python {
54    /// Interpreter and leading arguments.
55    pub program: Vec<OsString>,
56    /// Folders a sandbox must let it read (the venv, the helper).
57    pub read: Vec<PathBuf>,
58    /// Variables to set (`PYTHONPATH`).
59    pub env: BTreeMap<String, String>,
60}
61
62/// Every plugin environment of one engine.
63pub struct PyEnvs {
64    root: PathBuf,
65    base: Vec<OsString>,
66    /// `None`: look for `uv` on `PATH`; `Some(None)`: never use it.
67    uv: Option<Option<PathBuf>>,
68    offline: bool,
69    helper: OnceLock<Result<PathBuf, String>>,
70    envs: Mutex<BTreeMap<String, Arc<PyEnv>>>,
71}
72
73impl PyEnvs {
74    /// Environments go under `root` (e.g. `<data>/plugin-python`); `base` is the
75    /// interpreter they are made from (and the one plugins without requirements use).
76    pub fn new(root: impl Into<PathBuf>, base: Vec<OsString>) -> Self {
77        Self { root: root.into(), base, uv: None, offline: false, helper: OnceLock::new(), envs: Mutex::new(BTreeMap::new()) }
78    }
79
80    /// Use this `uv` (`None`: pip only) instead of looking for one.
81    pub fn with_uv(mut self, uv: Option<PathBuf>) -> Self {
82        self.uv = Some(uv);
83        self
84    }
85
86    /// Install without a package index (local wheels and folders only; tests).
87    pub fn offline(mut self) -> Self {
88        self.offline = true;
89        self
90    }
91
92    pub fn base(&self) -> &[OsString] {
93        &self.base
94    }
95
96    /// The folder holding `byob_plugin.py`, written there the first time it is asked for.
97    pub fn helper_dir(&self) -> Result<PathBuf, String> {
98        self.helper
99            .get_or_init(|| {
100                let dir = self.root.join("lib");
101                let file = dir.join(HELPER_FILE);
102                if fs::read_to_string(&file).ok().as_deref() != Some(HELPER) {
103                    fs::create_dir_all(&dir).map_err(|e| format!("{}: {e}", dir.display()))?;
104                    let tmp = dir.join(format!(".{HELPER_FILE}.{}", std::process::id()));
105                    fs::write(&tmp, HELPER).and_then(|()| fs::rename(&tmp, &file)).map_err(|e| format!("{}: {e}", file.display()))?;
106                }
107                Ok(dir)
108            })
109            .clone()
110    }
111
112    /// The environment for a plugin in `dir` that needs `requirements` (relative paths,
113    /// like `./wheels/x.whl`, are taken from the plugin folder).
114    pub fn for_plugin(&self, requirements: &[String], dir: &Path) -> Arc<PyEnv> {
115        let reqs: Vec<String> =
116            requirements.iter().map(|r| if r.starts_with("./") || r.starts_with("../") { dir.join(r).display().to_string() } else { r.clone() }).collect();
117        let mut h = Fnv::default();
118        for p in &self.base {
119            h.write(p.as_encoded_bytes());
120            h.write_u8(0);
121        }
122        let mut sorted = reqs.clone();
123        sorted.sort();
124        for r in &sorted {
125            h.write(r.as_bytes());
126            h.write_u8(1);
127        }
128        let key = format!("{:016x}", h.finish());
129        let mut envs = self.envs.lock().unwrap_or_else(PoisonError::into_inner);
130        let env = envs.entry(key.clone()).or_insert_with(|| {
131            let uv = match &self.uv {
132                Some(u) => u.clone(),
133                None => byob_sandbox::which(std::ffi::OsStr::new("uv")),
134            };
135            Arc::new(PyEnv {
136                dir: self.root.join("venvs").join(&key),
137                reqs,
138                base: self.base.clone(),
139                uv,
140                offline: self.offline,
141                state: Mutex::new(State::Unknown),
142                build: Mutex::new(()),
143            })
144        });
145        Arc::clone(env)
146    }
147
148    /// Forget failed builds so they are tried again (a reload).
149    pub fn forget_failed(&self) {
150        self.envs.lock().unwrap_or_else(PoisonError::into_inner).retain(|_, e| e.failure().is_none());
151    }
152
153    /// How to start Python for a plugin with this environment (`None`: no requirements).
154    /// Builds the environment if needed, until `cancelled()`: call it off the engine thread.
155    pub fn python(&self, env: Option<&PyEnv>, cancelled: &dyn Fn() -> bool) -> Result<Python, String> {
156        let helper = self.helper_dir();
157        let mut read = Vec::new();
158        let mut vars = BTreeMap::new();
159        match &helper {
160            Ok(dir) => {
161                let mut path = dir.display().to_string();
162                if let Some(old) = std::env::var_os("PYTHONPATH").filter(|p| !p.is_empty()) {
163                    path.push(if cfg!(windows) { ';' } else { ':' });
164                    path.push_str(&old.to_string_lossy());
165                }
166                vars.insert("PYTHONPATH".to_owned(), path);
167                read.push(dir.clone());
168            }
169            Err(e) => log::warn!("the byob_plugin helper isn't available: {e}"),
170        }
171        vars.insert("PYTHONUNBUFFERED".to_owned(), "1".to_owned());
172        // UTF-8 on stdin/stdout and in files even where the locale says otherwise (Windows)
173        vars.insert("PYTHONUTF8".to_owned(), "1".to_owned());
174        vars.insert("PYTHONIOENCODING".to_owned(), "utf-8".to_owned());
175        let program = match env {
176            None => self.base.clone(),
177            Some(env) => {
178                let python = env.ensure(cancelled)?;
179                read.push(env.dir.clone());
180                vec![python.into_os_string()]
181            }
182        };
183        Ok(Python { program, read, env: vars })
184    }
185}
186
187/// One plugin's Python: byob's helper, and its own environment if it has requirements.
188#[derive(Clone)]
189pub struct PluginPython {
190    pub envs: Arc<PyEnvs>,
191    pub env: Option<Arc<PyEnv>>,
192}
193
194impl PluginPython {
195    /// For a plugin in `dir` with these requirements.
196    pub fn new(envs: &Arc<PyEnvs>, requirements: &[String], dir: &Path) -> Self {
197        let env = (!requirements.is_empty()).then(|| envs.for_plugin(requirements, dir));
198        Self { envs: Arc::clone(envs), env }
199    }
200
201    /// See [`PyEnvs::python`] (may build the environment: not on the engine thread).
202    pub fn python(&self) -> Result<Python, String> {
203        self.python_until(&|| false)
204    }
205
206    /// [`Self::python`] whose build gives up once `cancelled()`.
207    pub fn python_until(&self, cancelled: &dyn Fn() -> bool) -> Result<Python, String> {
208        self.envs.python(self.env.as_deref(), cancelled)
209    }
210
211    /// Why its environment couldn't be built (never blocks).
212    pub fn failure(&self) -> Option<String> {
213        self.env.as_ref().and_then(|e| e.failure())
214    }
215}
216
217#[derive(Debug, Clone)]
218enum State {
219    Unknown,
220    Ready(PathBuf),
221    Failed(String),
222}
223
224/// One virtual environment.
225pub struct PyEnv {
226    dir: PathBuf,
227    reqs: Vec<String>,
228    base: Vec<OsString>,
229    uv: Option<PathBuf>,
230    offline: bool,
231    state: Mutex<State>,
232    /// Held while building, so one build runs at a time (in this process; the lock file
233    /// covers others).
234    build: Mutex<()>,
235}
236
237impl PyEnv {
238    pub fn dir(&self) -> &Path {
239        &self.dir
240    }
241
242    pub fn requirements(&self) -> &[String] {
243        &self.reqs
244    }
245
246    /// Why the last build failed (never blocks).
247    pub fn failure(&self) -> Option<String> {
248        match &*self.state.lock().unwrap_or_else(PoisonError::into_inner) {
249            State::Failed(e) => Some(e.clone()),
250            _ => None,
251        }
252    }
253
254    /// It is built (never blocks).
255    pub fn ready(&self) -> bool {
256        matches!(&*self.state.lock().unwrap_or_else(PoisonError::into_inner), State::Ready(_)) || self.dir.join(READY).is_file()
257    }
258
259    fn set(&self, s: State) {
260        *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s;
261    }
262
263    fn get(&self) -> State {
264        self.state.lock().unwrap_or_else(PoisonError::into_inner).clone()
265    }
266
267    /// The environment's interpreter, building it first if needed (blocks while it builds,
268    /// or another byob does, until `cancelled()`: then [`BUILD_CALLED_OFF`]).
269    pub fn ensure(&self, cancelled: &dyn Fn() -> bool) -> Result<PathBuf, String> {
270        match self.get() {
271            State::Ready(p) => return Ok(p),
272            State::Failed(e) => return Err(e),
273            State::Unknown => {}
274        }
275        let _building = loop {
276            match self.build.try_lock() {
277                Ok(g) => break g,
278                Err(TryLockError::Poisoned(p)) => break p.into_inner(),
279                Err(TryLockError::WouldBlock) if cancelled() => return Err(BUILD_CALLED_OFF.into()),
280                Err(TryLockError::WouldBlock) => std::thread::sleep(LOCK_POLL),
281            }
282        };
283        // Someone else may have finished while we waited.
284        match self.get() {
285            State::Ready(p) => return Ok(p),
286            State::Failed(e) => return Err(e),
287            State::Unknown => {}
288        }
289        let python = venv_python(&self.dir);
290        if self.built(&python) {
291            self.set(State::Ready(python.clone()));
292            return Ok(python);
293        }
294        if let Some(parent) = self.dir.parent() {
295            fs::create_dir_all(parent).map_err(|e| format!("{}: {e}", parent.display()))?;
296        }
297        let lock = BuildLock::take(&self.dir.with_extension("lock"), cancelled)?;
298        // Another byob may have built it while we waited.
299        if self.built(&python) {
300            self.set(State::Ready(python.clone()));
301            return Ok(python);
302        }
303        log::info!("building a Python environment for {} in {}", self.reqs.join(" "), self.dir.display());
304        let called_off = || {
305            lock.beat();
306            cancelled()
307        };
308        match self.build_now(&python, &called_off) {
309            Ok(()) => {
310                self.set(State::Ready(python.clone()));
311                Ok(python)
312            }
313            Err(e) => {
314                let _ = fs::remove_dir_all(&self.dir);
315                if cancelled() {
316                    // not a failure: the next use builds it again
317                    return Err(BUILD_CALLED_OFF.into());
318                }
319                let e = format!("python.requirements couldn't be installed: {e}");
320                self.set(State::Failed(e.clone()));
321                Err(e)
322            }
323        }
324    }
325
326    fn built(&self, python: &Path) -> bool {
327        self.dir.join(READY).is_file() && python.exists()
328    }
329
330    fn build_now(&self, python: &Path, cancelled: &dyn Fn() -> bool) -> Result<(), String> {
331        let _ = fs::remove_dir_all(&self.dir);
332        let Some((base, lead)) = self.base.split_first() else { return Err("no Python interpreter configured".into()) };
333        match &self.uv {
334            Some(uv) => {
335                let mut c = Command::new(uv);
336                c.arg("venv").arg("--quiet").arg("--python").arg(base).arg(&self.dir);
337                run(c, "uv venv", cancelled)?;
338                let mut c = Command::new(uv);
339                c.args(["pip", "install", "--python"]).arg(python).args(self.offline.then_some("--offline")).args(&self.reqs);
340                run(c, "uv pip install", cancelled)?;
341            }
342            None => {
343                let mut c = Command::new(base);
344                c.args(lead).args(["-m", "venv"]).arg(&self.dir);
345                run(c, "python -m venv", cancelled)?;
346                let mut c = Command::new(python);
347                c.args(["-m", "pip", "install", "--disable-pip-version-check", "--no-input"]).args(self.offline.then_some("--no-index")).args(&self.reqs);
348                run(c, "pip install", cancelled)?;
349            }
350        }
351        fs::write(self.dir.join(READY), self.reqs.join("\n")).map_err(|e| format!("{}: {e}", self.dir.display()))
352    }
353}
354
355/// The lock file of one environment's build, shared by every byob using the data folder:
356/// created (exclusively) to build, refreshed while building, removed when done.
357struct BuildLock {
358    path: PathBuf,
359    file: Option<fs::File>,
360    beat: std::cell::Cell<Instant>,
361}
362
363impl BuildLock {
364    /// Take the lock at `path`, waiting while another build holds it (until `cancelled()`).
365    fn take(path: &Path, cancelled: &dyn Fn() -> bool) -> Result<Self, String> {
366        let mut waited = false;
367        loop {
368            match fs::OpenOptions::new().write(true).create_new(true).open(path) {
369                Ok(mut file) => {
370                    let _ = std::io::Write::write_all(&mut file, format!("{}\n", std::process::id()).as_bytes());
371                    return Ok(Self { path: path.to_path_buf(), file: Some(file), beat: std::cell::Cell::new(Instant::now()) });
372                }
373                Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
374                    let age = fs::metadata(path).and_then(|m| m.modified()).ok().and_then(|t| SystemTime::now().duration_since(t).ok());
375                    if age.is_some_and(|a| a > LOCK_STALE) {
376                        log::warn!("{} is left over from a build that stopped; taking it over", path.display());
377                        let _ = fs::remove_file(path);
378                        continue;
379                    }
380                    if !waited {
381                        log::info!("waiting for another byob to build the Python environment ({})", path.display());
382                        waited = true;
383                    }
384                }
385                Err(e) => return Err(format!("{}: {e}", path.display())),
386            }
387            if cancelled() {
388                return Err(BUILD_CALLED_OFF.into());
389            }
390            std::thread::sleep(LOCK_POLL);
391        }
392    }
393
394    /// Say the build is still going (at most every [`LOCK_BEAT`]).
395    fn beat(&self) {
396        if self.beat.get().elapsed() >= LOCK_BEAT {
397            if let Some(f) = &self.file {
398                let _ = f.set_modified(SystemTime::now());
399            }
400            self.beat.set(Instant::now());
401        }
402    }
403}
404
405impl Drop for BuildLock {
406    fn drop(&mut self) {
407        drop(self.file.take());
408        let _ = fs::remove_file(&self.path);
409    }
410}
411
412fn venv_python(dir: &Path) -> PathBuf {
413    if cfg!(windows) { dir.join("Scripts").join("python.exe") } else { dir.join("bin").join("python") }
414}
415
416fn run(cmd: Command, what: &str, cancelled: &dyn Fn() -> bool) -> Result<(), String> {
417    let opts = RunOptions { stdin: Vec::new(), timeout: BUILD_TIMEOUT, max_output: MAX_BUILD_OUTPUT };
418    let o = byob_sandbox::run(cmd, &opts, cancelled).map_err(|e| format!("{what}: {e}"))?;
419    if o.success() {
420        return Ok(());
421    }
422    if o.cancelled {
423        return Err(BUILD_CALLED_OFF.into());
424    }
425    let out = format!("{}\n{}", o.stdout.text().trim(), o.stderr.text().trim());
426    let out = out.trim();
427    let start = out.char_indices().rev().nth(PROBLEM_OUTPUT).map_or(0, |(i, _)| i);
428    let why = if o.timed_out { "took too long".to_owned() } else { format!("exit code {}", o.exit_code.map_or_else(|| "none".into(), |c| c.to_string())) };
429    Err(format!("{what} failed ({why}):\n{}", out.get(start..).unwrap_or(out)))
430}
431
432#[cfg(test)]
433mod tests {
434    use super::*;
435
436    #[test]
437    fn same_requirements_share_an_environment() {
438        let root = tempfile::tempdir().unwrap();
439        let envs = PyEnvs::new(root.path(), crate::default_python()).with_uv(None);
440        let a = envs.for_plugin(&["x==1".into(), "y".into()], Path::new("/p/a"));
441        let b = envs.for_plugin(&["y".into(), "x==1".into()], Path::new("/p/b"));
442        let c = envs.for_plugin(&["x==2".into()], Path::new("/p/a"));
443        assert_eq!(a.dir(), b.dir());
444        assert_ne!(a.dir(), c.dir());
445        let w = envs.for_plugin(&["./wheels/w.whl".into()], Path::new("/p/a"));
446        assert_eq!(w.requirements(), [Path::new("/p/a").join("./wheels/w.whl").display().to_string()]);
447    }
448
449    #[test]
450    fn helper_is_written_once_and_on_the_path() {
451        let root = tempfile::tempdir().unwrap();
452        let envs = PyEnvs::new(root.path(), crate::default_python());
453        let py = envs.python(None, &|| false).unwrap();
454        let dir = envs.helper_dir().unwrap();
455        assert_eq!(fs::read_to_string(dir.join(HELPER_FILE)).unwrap(), HELPER);
456        assert!(py.env["PYTHONPATH"].starts_with(&dir.display().to_string()));
457        assert_eq!(py.program, crate::default_python());
458    }
459
460    #[test]
461    fn python_speaks_utf8_whatever_the_locale() {
462        let root = tempfile::tempdir().unwrap();
463        let py = PyEnvs::new(root.path(), crate::default_python()).python(None, &|| false).unwrap();
464        assert_eq!(py.env.get("PYTHONUTF8").map(String::as_str), Some("1"));
465        assert_eq!(py.env.get("PYTHONIOENCODING").map(String::as_str), Some("utf-8"));
466    }
467
468    /// An environment whose interpreter doesn't exist: a build fails at once.
469    fn broken(root: &Path) -> Arc<PyEnv> {
470        PyEnvs::new(root, vec!["/definitely/not/python".into()]).with_uv(None).for_plugin(&["x".into()], Path::new("/p"))
471    }
472
473    #[test]
474    fn another_byobs_build_is_waited_for_and_used_not_clobbered() {
475        let root = tempfile::tempdir().unwrap();
476        let env = broken(root.path());
477        let lock = env.dir().with_extension("lock");
478        fs::create_dir_all(env.dir()).unwrap();
479        fs::write(&lock, "4242\n").unwrap();
480        // the other byob finishes its build, then lets go of the lock
481        let dir = env.dir().to_path_buf();
482        let other = std::thread::spawn(move || {
483            std::thread::sleep(Duration::from_millis(300));
484            let python = venv_python(&dir);
485            fs::create_dir_all(python.parent().unwrap()).unwrap();
486            fs::write(&python, "").unwrap();
487            fs::write(dir.join(READY), "x").unwrap();
488            fs::remove_file(dir.with_extension("lock")).unwrap();
489        });
490        assert_eq!(env.ensure(&|| false), Ok(venv_python(env.dir())));
491        other.join().unwrap();
492        assert!(env.dir().join(READY).is_file(), "its build was left alone");
493    }
494
495    #[test]
496    fn waiting_for_another_build_can_be_called_off() {
497        let root = tempfile::tempdir().unwrap();
498        let env = broken(root.path());
499        fs::create_dir_all(env.dir()).unwrap();
500        fs::write(env.dir().join("theirs"), "").unwrap();
501        fs::write(env.dir().with_extension("lock"), "4242\n").unwrap();
502        let t = Instant::now();
503        assert_eq!(env.ensure(&|| t.elapsed() > Duration::from_millis(200)), Err(BUILD_CALLED_OFF.to_owned()));
504        assert!(env.dir().join("theirs").exists(), "their half-built environment was left alone");
505        assert!(env.failure().is_none(), "not a failure: tried again next time");
506    }
507
508    #[test]
509    fn a_lock_left_by_a_byob_that_died_is_taken_over() {
510        let root = tempfile::tempdir().unwrap();
511        let env = broken(root.path());
512        let lock = env.dir().with_extension("lock");
513        fs::create_dir_all(env.dir()).unwrap();
514        let f = fs::File::create(&lock).unwrap();
515        f.set_modified(SystemTime::now() - Duration::from_secs(3600)).unwrap();
516        drop(f);
517        let e = env.ensure(&|| false).unwrap_err();
518        assert!(e.starts_with("python.requirements couldn't be installed"), "{e}");
519        assert!(!lock.exists(), "the lock goes with the build");
520    }
521
522    #[cfg(unix)]
523    #[test]
524    fn a_build_can_be_called_off() {
525        let root = tempfile::tempdir().unwrap();
526        // an "interpreter" whose `-m venv` takes half a minute
527        let base: Vec<OsString> = vec!["sh".into(), "-c".into(), "sleep 30".into(), "sh".into()];
528        let env = PyEnvs::new(root.path(), base).with_uv(None).for_plugin(&["x".into()], Path::new("/p"));
529        let t = Instant::now();
530        assert_eq!(env.ensure(&|| t.elapsed() > Duration::from_millis(300)), Err(BUILD_CALLED_OFF.to_owned()));
531        assert!(t.elapsed() < Duration::from_secs(10), "{:?}", t.elapsed());
532        assert!(env.failure().is_none());
533        assert!(!env.dir().exists());
534        assert!(!env.dir().with_extension("lock").exists());
535    }
536}
byob_plugins/pyenv.rs

byob_plugins/
pyenv.rs
