diff --git a/library/std/src/sync/once/tests.rs b/library/std/src/sync/once/tests.rs index 0c35597e11c51..484656709047a 100644 --- a/library/std/src/sync/once/tests.rs +++ b/library/std/src/sync/once/tests.rs @@ -1,7 +1,9 @@ use super::Once; use crate::panic; +use crate::sync::atomic::{AtomicBool, Ordering::Relaxed}; use crate::sync::mpsc::channel; use crate::thread; +use crate::time::Duration; #[test] fn smoke_once() { @@ -114,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/library/std/src/sys/sync/once/queue.rs b/library/std/src/sys/sync/once/queue.rs index 338c1242273f6..6bcae405c7aa2 100644 --- a/library/std/src/sys/sync/once/queue.rs +++ b/library/std/src/sys/sync/once/queue.rs @@ -157,7 +157,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); } } } @@ -220,14 +220,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), @@ -239,7 +243,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; }