返回 CodeWhale
test_env_lock.rs
根目录 / crates / tui / src / test_env_lock.rs
1 //! Process-wide test-environment barrier owned by shell dispatch.
2 //!
3 //! `shell_dispatcher.rs` is also compiled directly by integration harnesses,
4 //! outside the main binary crate. Keeping the barrier under that module makes
5 //! shell detection self-contained while `crate::test_support` re-exports the
6 //! same instance to its existing environment-mutating callers in the main crate.
7
8 use std::sync::{Mutex, MutexGuard, OnceLock, TryLockError};
9 use std::thread::ThreadId;
10
11 fn env_lock() -> &'static Mutex<()> {
12 static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
13 LOCK.get_or_init(|| Mutex::new(()))
14 }
15
16 /// Who currently counts as "inside" the process-wide env lock.
17 ///
18 /// The owner is the thread holding [`TestEnvLock`]. `adopted` holds helper
19 /// threads that owner explicitly enrolled with [`join_env_scope`] — see that
20 /// function for why a worker thread of the current test must not be treated as
21 /// a foreign reader.
22 #[derive(Default)]
23 struct EnvScope {
24 /// Bumped on every acquisition, so a ticket minted by an earlier test can
25 /// never enroll a thread into a later test's environment.
26 generation: u64,
27 owner: Option<ThreadId>,
28 adopted: Vec<ThreadId>,
29 }
30
31 fn env_scope() -> &'static Mutex<EnvScope> {
32 static SCOPE: OnceLock<Mutex<EnvScope>> = OnceLock::new();
33 SCOPE.get_or_init(|| Mutex::new(EnvScope::default()))
34 }
35
36 fn lock_env_scope() -> MutexGuard<'static, EnvScope> {
37 match env_scope().lock() {
38 Ok(scope) => scope,
39 Err(poisoned) => poisoned.into_inner(),
40 }
41 }
42
43 fn open_env_scope() {
44 let mut scope = lock_env_scope();
45 scope.generation = scope.generation.wrapping_add(1);
46 scope.owner = Some(std::thread::current().id());
47 scope.adopted.clear();
48 }
49
50 fn current_thread_owns_contended_env_lock() -> bool {
51 let scope = lock_env_scope();
52 let current = std::thread::current().id();
53 scope.owner == Some(current) || scope.adopted.contains(&current)
54 }
55
56 /// Proof that the calling thread belongs to a live [`lock_test_env`] scope, handed to
57 /// a worker thread so it can join that scope with [`join_env_scope`].
58 ///
59 /// Owners and adopted workers can pass the same live generation to child
60 /// workers. Foreign threads cannot mint a ticket for a test they did not join.
61 #[derive(Clone, Copy, Debug)]
62 pub(crate) struct EnvScopeTicket {
63 generation: u64,
64 }
65
66 impl EnvScopeTicket {
67 /// Which sealed environment this ticket authorizes. Callers that gate real
68 /// disk writes on a live scope key their bookkeeping by this value, so a
69 /// straggler from generation N can never be mistaken for work belonging to
70 /// generation N+1.
71 pub(crate) fn generation(&self) -> u64 {
72 self.generation
73 }
74 }
75
76 /// The generation of the env scope the calling thread is currently inside, as
77 /// owner or as a [`join_env_scope`]-adopted worker; `None` when the thread is a
78 /// foreign reader with no sealed environment of its own.
79 ///
80 /// This is the authorization primitive for anything that must only touch disk
81 /// on behalf of a test that actually sealed `HOME`. A process-global "writes
82 /// are enabled" flag cannot distinguish unrelated parallel tests.
83 pub(crate) fn current_env_scope_generation() -> Option<u64> {
84 let scope = lock_env_scope();
85 let current = std::thread::current().id();
86 if scope.owner == Some(current) || scope.adopted.contains(&current) {
87 Some(scope.generation)
88 } else {
89 None
90 }
91 }
92
93 pub(crate) fn env_scope_ticket() -> Option<EnvScopeTicket> {
94 let scope = lock_env_scope();
95 let current = std::thread::current().id();
96 (scope.owner == Some(current) || scope.adopted.contains(&current)).then_some(EnvScopeTicket {
97 generation: scope.generation,
98 })
99 }
100
101 /// Enroll the calling thread in the ticket's env scope for as long as the
102 /// returned guard lives.
103 ///
104 /// [`with_test_env_lock`] stops a foreign test from resolving another test's
105 /// temporary `HOME`. A helper thread doing work for the sealing test must see
106 /// that same environment without blocking on the mutex its owner holds.
107 pub(crate) fn join_env_scope(ticket: Option<EnvScopeTicket>) -> Option<EnvScopeMembership> {
108 let ticket = ticket?;
109 let mut scope = lock_env_scope();
110 if scope.owner.is_none() || scope.generation != ticket.generation {
111 return None;
112 }
113 let thread = std::thread::current().id();
114 if !scope.adopted.contains(&thread) {
115 scope.adopted.push(thread);
116 }
117 Some(EnvScopeMembership {
118 generation: ticket.generation,
119 thread,
120 })
121 }
122
123 pub(crate) struct EnvScopeMembership {
124 generation: u64,
125 thread: ThreadId,
126 }
127
128 impl Drop for EnvScopeMembership {
129 fn drop(&mut self) {
130 let mut scope = lock_env_scope();
131 if scope.generation == self.generation {
132 scope.adopted.retain(|thread| *thread != self.thread);
133 }
134 }
135 }
136
137 /// Owned process-wide test-environment lock.
138 ///
139 /// Clearing the owner before the underlying mutex unlocks keeps re-entrant
140 /// reader detection exact. Closing the scope also evicts adopted workers, so
141 /// enrollment cannot outlive the test that granted it.
142 pub(crate) struct TestEnvLock {
143 _guard: MutexGuard<'static, ()>,
144 }
145
146 impl Drop for TestEnvLock {
147 fn drop(&mut self) {
148 let mut scope = lock_env_scope();
149 if scope.owner == Some(std::thread::current().id()) {
150 scope.owner = None;
151 scope.adopted.clear();
152 }
153 }
154 }
155
156 /// Acquire the process-wide env-var mutex.
157 ///
158 /// If a prior test panicked while holding the lock, recover the guard instead
159 /// of cascading failures across unrelated tests.
160 pub(crate) fn lock_test_env() -> TestEnvLock {
161 let guard = match env_lock().lock() {
162 Ok(guard) => guard,
163 Err(poisoned) => poisoned.into_inner(),
164 };
165 open_env_scope();
166 TestEnvLock { _guard: guard }
167 }
168
169 /// Read process-global test environment while respecting [`lock_test_env`].
170 ///
171 /// The owner check makes the barrier re-entrant for a test that reads its own
172 /// guarded override.
173 pub(crate) fn with_test_env_lock<T>(read: impl FnOnce() -> T) -> T {
174 if current_thread_owns_contended_env_lock() {
175 return read();
176 }
177
178 let _guard = lock_test_env();
179 read()
180 }
181
182 /// [`with_test_env_lock`] for readers that run inside a `LazyLock`/`OnceLock`
183 /// initializer: takes the barrier when it is free, and reads *without* it
184 /// rather than waiting when another thread holds it.
185 ///
186 /// Blocking is not an option there. The initializer's own lock is held for the
187 /// duration, so a second thread that already owns the env barrier and then
188 /// touches the same lazy value deadlocks against this one: A waits for the
189 /// initializer, B (this thread) waits for A's barrier. Both threads wedge, and
190 /// libtest has no per-test timeout, so the whole test binary hangs instead of
191 /// failing — which is how a single wedged run burns an entire CI job.
192 ///
193 /// The trade is deliberate and small: an uncontended read is still fully
194 /// serialized, and a contended one reads a variable mid-mutation at worst.
195 /// Callers must therefore only use this for values whose staleness is
196 /// tolerable, never to gate a disk write.
197 pub(crate) fn with_test_env_lock_if_uncontended<T>(read: impl FnOnce() -> T) -> T {
198 if current_thread_owns_contended_env_lock() {
199 return read();
200 }
201
202 match env_lock().try_lock() {
203 Ok(_guard) => read(),
204 Err(TryLockError::Poisoned(poisoned)) => {
205 let _guard = poisoned.into_inner();
206 read()
207 }
208 Err(TryLockError::WouldBlock) => read(),
209 }
210 }
211
212 pub(crate) fn current_thread_holds_test_env_lock() -> bool {
213 match env_lock().try_lock() {
214 Ok(guard) => {
215 drop(guard);
216 false
217 }
218 Err(TryLockError::Poisoned(poisoned)) => {
219 drop(poisoned.into_inner());
220 false
221 }
222 Err(TryLockError::WouldBlock) => current_thread_owns_contended_env_lock(),
223 }
224 }
225
226 #[cfg(test)]
227 mod tests {
228 use super::{lock_test_env, with_test_env_lock_if_uncontended};
229 use std::sync::mpsc;
230 use std::time::{Duration, Instant};
231
232 #[test]
233 fn adopted_worker_can_pass_only_its_live_environment_generation() {
234 let owner = lock_test_env();
235 let ticket = super::env_scope_ticket().unwrap();
236 let nested = std::thread::spawn(move || {
237 assert!(
238 super::env_scope_ticket().is_none(),
239 "foreign thread has no authority"
240 );
241 let _membership = super::join_env_scope(Some(ticket)).unwrap();
242 let child_ticket =
243 super::env_scope_ticket().expect("adopted worker can enroll its child");
244 std::thread::spawn(move || {
245 let _membership = super::join_env_scope(Some(child_ticket)).unwrap();
246 assert_eq!(
247 super::current_env_scope_generation(),
248 Some(child_ticket.generation())
249 );
250 assert_eq!(super::with_test_env_lock(|| 7), 7);
251 child_ticket
252 })
253 .join()
254 .unwrap()
255 })
256 .join()
257 .unwrap();
258 drop(owner);
259 let _next = lock_test_env();
260 assert!(
261 super::join_env_scope(Some(nested)).is_none(),
262 "old generation cannot join the next test"
263 );
264 }
265
266 /// The lock-order inversion this helper exists to break.
267 ///
268 /// A reader running inside a `LazyLock`/`OnceLock` initializer must not wait
269 /// for the env barrier: the thread holding the barrier may itself be waiting
270 /// on that initializer. Both threads then wedge, and libtest has no per-test
271 /// timeout, so the entire test binary hangs — reproduced locally as
272 /// `cargo test -p codewhale-tui --lib -- shell` never returning.
273 #[test]
274 fn uncontended_read_returns_while_another_thread_holds_the_barrier() {
275 let (held_tx, held_rx) = mpsc::channel();
276 let (release_tx, release_rx) = mpsc::channel::<()>();
277 let holder = std::thread::spawn(move || {
278 let _guard = lock_test_env();
279 held_tx.send(()).expect("announce that the barrier is held");
280 // Outlive the assertion below without depending on timing.
281 let _ = release_rx.recv_timeout(Duration::from_secs(30));
282 });
283 held_rx
284 .recv_timeout(Duration::from_secs(10))
285 .expect("holder thread must take the barrier");
286
287 let start = Instant::now();
288 let value = with_test_env_lock_if_uncontended(|| 7);
289 let elapsed = start.elapsed();
290
291 release_tx.send(()).ok();
292 holder.join().expect("holder thread");
293
294 assert_eq!(value, 7, "the read still runs, just unsynchronized");
295 assert!(
296 elapsed < Duration::from_secs(1),
297 "a contended read must not block; waited {elapsed:?}"
298 );
299 }
300 }
301
301 lines RUST