| 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(¤t) |
| 54 | } |
| 55 | |
| 56 | /// Proof that the calling thread owns a live [`lock_test_env`] scope, handed to |
| 57 | /// a worker thread so it can join that scope with [`join_env_scope`]. |
| 58 | /// |
| 59 | /// Returns `None` when the caller is not the owner, so a ticket can never be |
| 60 | /// minted on behalf of a test that did not seal the environment. |
| 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(¤t) { |
| 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 | (scope.owner == Some(std::thread::current().id())).then_some(EnvScopeTicket { |
| 96 | generation: scope.generation, |
| 97 | }) |
| 98 | } |
| 99 | |
| 100 | /// Enroll the calling thread in the ticket's env scope for as long as the |
| 101 | /// returned guard lives. |
| 102 | /// |
| 103 | /// [`with_test_env_lock`] stops a foreign test from resolving another test's |
| 104 | /// temporary `HOME`. A helper thread doing work for the sealing test must see |
| 105 | /// that same environment without blocking on the mutex its owner holds. |
| 106 | pub(crate) fn join_env_scope(ticket: Option<EnvScopeTicket>) -> Option<EnvScopeMembership> { |
| 107 | let ticket = ticket?; |
| 108 | let mut scope = lock_env_scope(); |
| 109 | if scope.owner.is_none() || scope.generation != ticket.generation { |
| 110 | return None; |
| 111 | } |
| 112 | let thread = std::thread::current().id(); |
| 113 | if !scope.adopted.contains(&thread) { |
| 114 | scope.adopted.push(thread); |
| 115 | } |
| 116 | Some(EnvScopeMembership { |
| 117 | generation: ticket.generation, |
| 118 | thread, |
| 119 | }) |
| 120 | } |
| 121 | |
| 122 | pub(crate) struct EnvScopeMembership { |
| 123 | generation: u64, |
| 124 | thread: ThreadId, |
| 125 | } |
| 126 | |
| 127 | impl Drop for EnvScopeMembership { |
| 128 | fn drop(&mut self) { |
| 129 | let mut scope = lock_env_scope(); |
| 130 | if scope.generation == self.generation { |
| 131 | scope.adopted.retain(|thread| *thread != self.thread); |
| 132 | } |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | /// Owned process-wide test-environment lock. |
| 137 | /// |
| 138 | /// Clearing the owner before the underlying mutex unlocks keeps re-entrant |
| 139 | /// reader detection exact. Closing the scope also evicts adopted workers, so |
| 140 | /// enrollment cannot outlive the test that granted it. |
| 141 | pub(crate) struct TestEnvLock { |
| 142 | _guard: MutexGuard<'static, ()>, |
| 143 | } |
| 144 | |
| 145 | impl Drop for TestEnvLock { |
| 146 | fn drop(&mut self) { |
| 147 | let mut scope = lock_env_scope(); |
| 148 | if scope.owner == Some(std::thread::current().id()) { |
| 149 | scope.owner = None; |
| 150 | scope.adopted.clear(); |
| 151 | } |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | /// Acquire the process-wide env-var mutex. |
| 156 | /// |
| 157 | /// If a prior test panicked while holding the lock, recover the guard instead |
| 158 | /// of cascading failures across unrelated tests. |
| 159 | pub(crate) fn lock_test_env() -> TestEnvLock { |
| 160 | let guard = match env_lock().lock() { |
| 161 | Ok(guard) => guard, |
| 162 | Err(poisoned) => poisoned.into_inner(), |
| 163 | }; |
| 164 | open_env_scope(); |
| 165 | TestEnvLock { _guard: guard } |
| 166 | } |
| 167 | |
| 168 | /// Read process-global test environment while respecting [`lock_test_env`]. |
| 169 | /// |
| 170 | /// The owner check makes the barrier re-entrant for a test that reads its own |
| 171 | /// guarded override. |
| 172 | pub(crate) fn with_test_env_lock<T>(read: impl FnOnce() -> T) -> T { |
| 173 | if current_thread_owns_contended_env_lock() { |
| 174 | return read(); |
| 175 | } |
| 176 | |
| 177 | let _guard = lock_test_env(); |
| 178 | read() |
| 179 | } |
| 180 | |
| 181 | pub(crate) fn current_thread_holds_test_env_lock() -> bool { |
| 182 | match env_lock().try_lock() { |
| 183 | Ok(guard) => { |
| 184 | drop(guard); |
| 185 | false |
| 186 | } |
| 187 | Err(TryLockError::Poisoned(poisoned)) => { |
| 188 | drop(poisoned.into_inner()); |
| 189 | false |
| 190 | } |
| 191 | Err(TryLockError::WouldBlock) => current_thread_owns_contended_env_lock(), |
| 192 | } |
| 193 | } |
| 194 |