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(
55        unsafe_code,
56        reason = "guard serializes all env mutators, so no concurrent access"
57    )]
58    pub fn set_var(&self, key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) {
59        // SAFETY: `self` is the unique holder of the process-wide env mutex
60        // for the duration of this call; all set/remove calls in this module
61        // route through the same mutex, so no concurrent env access occurs.
62        unsafe {
63            std::env::set_var(key, value);
64        }
65    }
66
67    /// Remove a process environment variable.
68    ///
69    /// See [`Self::set_var`] for the safety argument.
70    #[expect(
71        unsafe_code,
72        reason = "see set_var — guard serializes all env mutators"
73    )]
74    pub fn remove_var(&self, key: impl AsRef<OsStr>) {
75        // SAFETY: see `set_var` — the guard serializes all mutators.
76        unsafe {
77            std::env::remove_var(key);
78        }
79    }
80
81    /// Restore a variable to its previous value, or remove it if there was none.
82    ///
83    /// Pairs with [`std::env::var_os`] snapshots taken before a mutation.
84    pub fn restore_var<T: AsRef<OsStr>>(&self, key: &str, previous: Option<T>) {
85        match previous {
86            Some(value) => self.set_var(key, value),
87            None => self.remove_var(key),
88        }
89    }
90}
91
92/// Acquire the process-wide environment lock.
93///
94/// Blocks until no other [`EnvGuard`] is alive. Poisoning is ignored — the
95/// inner `()` cannot become corrupted, so callers can safely recover.
96pub fn lock() -> EnvGuard {
97    EnvGuard(raw_lock())
98}
99
100/// One-shot safe wrapper around [`std::env::set_var`].
101///
102/// Acquires the process-wide environment lock for the duration of the call.
103/// Use [`lock`] when you need to perform multiple env operations atomically.
104pub fn set_var(key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) {
105    lock().set_var(key, value);
106}
107
108/// One-shot safe wrapper around [`std::env::remove_var`].
109///
110/// Acquires the process-wide environment lock for the duration of the call.
111/// Use [`lock`] when you need to perform multiple env operations atomically.
112pub fn remove_var(key: impl AsRef<OsStr>) {
113    lock().remove_var(key);
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn set_and_remove_roundtrip() {
122        let env = lock();
123        env.set_var("VTCODE_ENV_LOCK_TEST", "value-a");
124        assert_eq!(
125            std::env::var("VTCODE_ENV_LOCK_TEST").as_deref(),
126            Ok("value-a")
127        );
128        env.remove_var("VTCODE_ENV_LOCK_TEST");
129        assert!(std::env::var("VTCODE_ENV_LOCK_TEST").is_err());
130    }
131
132    #[test]
133    fn restore_var_restores_previous_value() {
134        let env = lock();
135        env.set_var("VTCODE_ENV_LOCK_RESTORE", "original");
136        let previous = std::env::var_os("VTCODE_ENV_LOCK_RESTORE");
137        env.set_var("VTCODE_ENV_LOCK_RESTORE", "temporary");
138        env.restore_var("VTCODE_ENV_LOCK_RESTORE", previous);
139        assert_eq!(
140            std::env::var("VTCODE_ENV_LOCK_RESTORE").as_deref(),
141            Ok("original")
142        );
143        env.remove_var("VTCODE_ENV_LOCK_RESTORE");
144    }
145
146    #[test]
147    fn restore_var_removes_when_previous_is_none() {
148        let env = lock();
149        env.set_var("VTCODE_ENV_LOCK_RESTORE_NONE", "temporary");
150        env.restore_var::<&str>("VTCODE_ENV_LOCK_RESTORE_NONE", None);
151        assert!(std::env::var("VTCODE_ENV_LOCK_RESTORE_NONE").is_err());
152    }
153}
vtcode_commons/env_lock.rs

vtcode_commons/
env_lock.rs
