diff --git a/std/src/sync/once/tests.rs b/std/src/sync/once/tests.rs index d43dabc1cf137..ce96468aeb6e1 100644 --- a/std/src/sync/once/tests.rs +++ b/std/src/sync/once/tests.rs @@ -1,5 +1,8 @@ use super::Once; +use crate::sync::atomic::AtomicBool; +use crate::sync::atomic::Ordering::Relaxed; use crate::sync::mpsc::channel; +use crate::time::Duration; use crate::{panic, thread}; #[test] @@ -113,3 +116,47 @@ fn wait_for_force_to_finish() { assert!(t1.join().is_ok()); assert!(t2.join().is_ok()); } + +#[test] +fn wait() { + for _ in 0..50 { + let val = AtomicBool::new(false); + let once = Once::new(); + + thread::scope(|s| { + for _ in 0..4 { + s.spawn(|| { + once.wait(); + assert!(val.load(Relaxed)); + }); + } + + once.call_once(|| val.store(true, Relaxed)); + }); + } +} + +#[test] +fn wait_on_poisoned() { + let once = Once::new(); + + panic::catch_unwind(|| once.call_once(|| panic!())).unwrap_err(); + panic::catch_unwind(|| once.wait()).unwrap_err(); +} + +#[test] +fn wait_force_on_poisoned() { + let once = Once::new(); + + thread::scope(|s| { + panic::catch_unwind(|| once.call_once(|| panic!())).unwrap_err(); + + s.spawn(|| { + thread::sleep(Duration::from_millis(100)); + + once.call_once_force(|_| {}); + }); + + once.wait_force(); + }) +} diff --git a/std/src/sys/sync/once/queue.rs b/std/src/sys/sync/once/queue.rs index 7a020c94080bf..86f72c82008bc 100644 --- a/std/src/sys/sync/once/queue.rs +++ b/std/src/sys/sync/once/queue.rs @@ -153,7 +153,7 @@ impl Once { panic!("Once instance has previously been poisoned"); } _ => { - current = wait(&self.state_and_queue, current); + current = wait(&self.state_and_queue, current, !ignore_poisoning); } } } @@ -216,14 +216,18 @@ impl Once { // All other values must be RUNNING with possibly a // pointer to the waiter queue in the more significant bits. assert!(state == RUNNING); - current = wait(&self.state_and_queue, current); + current = wait(&self.state_and_queue, current, true); } } } } } -fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAndQueue { +fn wait( + state_and_queue: &AtomicPtr<()>, + mut current: StateAndQueue, + return_on_poisoned: bool, +) -> StateAndQueue { let node = &Waiter { thread: Cell::new(Some(thread::current())), signaled: AtomicBool::new(false), @@ -235,7 +239,7 @@ fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAnd let queue = to_queue(current); // If initialization has finished, return. - if matches!(state, POISONED | COMPLETE) { + if state == COMPLETE || (return_on_poisoned && state == POISONED) { return current; }