vtcode_commons/
env_lock.rs

1//! Process-environment mutation lock.
2//!
3//! Rust 2024 marks [`std::env::set_var`] and [`std::env::remove_var`] as
4//! `unsafe` because POSIX `setenv`/`getenv` are not thread-safe. Each call site
5//! that needs to mutate the environment previously rolled its own
6//! `OnceLock<Mutex<()>>` plus duplicated `set_env_var` / `remove_env_var`
7//! helpers carrying their own `SAFETY:` comments. This module consolidates that
8//! invariant into a single sound wrapper.
9//!
10//! # Usage
11//!
12//! Acquire the guard once, then mutate the environment through its methods.
13//! The guard holds a process-wide [`Mutex`] so concurrent callers serialize
14//! automatically.
15//!
16//! ```no_run
17//! use vtcode_commons::env_lock;
18//!
19//! let env = env_lock::lock();
20//! env.set_var("MY_TEST_VAR", "1");
21//! // ... run code that reads MY_TEST_VAR ...
22//! env.remove_var("MY_TEST_VAR");
23//! ```
24//!
25//! All in-process code that mutates the environment **must** go through this
26//! module; direct `std::env::set_var` / `std::env::remove_var` calls bypass the
27//! lock and re-introduce the data race the wrapper is here to prevent.
28
29use std::ffi::OsStr;
30use std::sync::{Mutex, MutexGuard, OnceLock};
31
32static ENV_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
33
34fn raw_lock() -> MutexGuard<'static, ()> {
35    ENV_LOCK
36        .get_or_init(|| Mutex::new(()))
37        .lock()
38        .unwrap_or_else(|poisoned| poisoned.into_inner())
39}
40
41/// RAII guard that proves ownership of the process-wide environment lock.
42///
43/// Obtain one with [`lock`]; while it is alive, no other thread can enter the
44/// safe mutation methods on this type. Dropping the guard releases the lock.
45#[must_use = "EnvGuard releases the lock when dropped; bind it to a local"]
46pub struct EnvGuard(#[allow(dead_code)] MutexGuard<'static, ()>);
47
48impl EnvGuard {
49    /// Set a process environment variable.
50    ///
51    /// Safe because `self` proves the global env mutex is held, so no other
52    /// caller routed through this module is reading or writing the environment
53    /// concurrently.
54    #[expect(unsafe_code, reason = "guard serializes all env mutators, so no concurrent access")]
55    pub fn set_var(&self, key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) {
56        // SAFETY: `self` is the unique holder of the process-wide env mutex
57        // for the duration of this call; all set/remove calls in this module
58        // route through the same mutex, so no concurrent env access occurs.
59        unsafe {
60            std::env::set_var(key, value);
61        }
62    }
63
64    /// Remove a process environment variable.
65    ///
66    /// See [`Self::set_var`] for the safety argument.
67    #[expect(unsafe_code, reason = "see set_var — guard serializes all env mutators")]
68    pub fn remove_var(&self, key: impl AsRef<OsStr>) {
69        // SAFETY: see `set_var` — the guard serializes all mutators.
70        unsafe {
71            std::env::remove_var(key);
72        }
73    }
74
75    /// Restore a variable to its previous value, or remove it if there was none.
76    ///
77    /// Pairs with [`std::env::var_os`] snapshots taken before a mutation.
78    pub fn restore_var<T: AsRef<OsStr>>(&self, key: &str, previous: Option<T>) {
79        match previous {
80            Some(value) => self.set_var(key, value),
81            None => self.remove_var(key),
82        }
83    }
84}
85
86/// Acquire the process-wide environment lock.
87///
88/// Blocks until no other [`EnvGuard`] is alive. Poisoning is ignored — the
89/// inner `()` cannot become corrupted, so callers can safely recover.
90pub fn lock() -> EnvGuard {
91    EnvGuard(raw_lock())
92}
93
94/// One-shot safe wrapper around [`std::env::set_var`].
95///
96/// Acquires the process-wide environment lock for the duration of the call.
97/// Use [`lock`] when you need to perform multiple env operations atomically.
98pub fn set_var(key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) {
99    lock().set_var(key, value);
100}
101
102/// One-shot safe wrapper around [`std::env::remove_var`].
103///
104/// Acquires the process-wide environment lock for the duration of the call.
105/// Use [`lock`] when you need to perform multiple env operations atomically.
106pub fn remove_var(key: impl AsRef<OsStr>) {
107    lock().remove_var(key);
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn set_and_remove_roundtrip() {
116        let env = lock();
117        env.set_var("VTCODE_ENV_LOCK_TEST", "value-a");
118        assert_eq!(std::env::var("VTCODE_ENV_LOCK_TEST").as_deref(), Ok("value-a"));
119        env.remove_var("VTCODE_ENV_LOCK_TEST");
120        assert!(std::env::var("VTCODE_ENV_LOCK_TEST").is_err());
121    }
122
123    #[test]
124    fn restore_var_restores_previous_value() {
125        let env = lock();
126        env.set_var("VTCODE_ENV_LOCK_RESTORE", "original");
127        let previous = std::env::var_os("VTCODE_ENV_LOCK_RESTORE");
128        env.set_var("VTCODE_ENV_LOCK_RESTORE", "temporary");
129        env.restore_var("VTCODE_ENV_LOCK_RESTORE", previous);
130        assert_eq!(std::env::var("VTCODE_ENV_LOCK_RESTORE").as_deref(), Ok("original"));
131        env.remove_var("VTCODE_ENV_LOCK_RESTORE");
132    }
133
134    #[test]
135    fn restore_var_removes_when_previous_is_none() {
136        let env = lock();
137        env.set_var("VTCODE_ENV_LOCK_RESTORE_NONE", "temporary");
138        env.restore_var::<&str>("VTCODE_ENV_LOCK_RESTORE_NONE", None);
139        assert!(std::env::var("VTCODE_ENV_LOCK_RESTORE_NONE").is_err());
140    }
141}
vtcode_commons/env_lock.rs

vtcode_commons/
env_lock.rs
