std: fix busy-waiting in Once::wait_force, add more tests

joboet · joboet · commit 1d49aad84458 · 2024-07-31T17:44:10.000+02:00
diff --git a/library/std/src/sync/once/tests.rs b/library/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/library/std/src/sys/sync/once/queue.rs b/library/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;
         }
 

Original file line number	Diff line number	Diff line change
`@@ -153,7 +153,7 @@ impl Once {`
`153`	`153`	`panic!("Once instance has previously been poisoned");`
`154`	`154`	`}`
`155`	`155`	`_ => {`
`156`		`- current = wait(&self.state_and_queue, current);`
	`156`	`+ current = wait(&self.state_and_queue, current, !ignore_poisoning);`
`157`	`157`	`}`
`158`	`158`	`}`
`159`	`159`	`}`
`@@ -216,14 +216,18 @@ impl Once {`
`216`	`216`	`// All other values must be RUNNING with possibly a`
`217`	`217`	`// pointer to the waiter queue in the more significant bits.`
`218`	`218`	`assert!(state == RUNNING);`
`219`		`- current = wait(&self.state_and_queue, current);`
	`219`	`+ current = wait(&self.state_and_queue, current, true);`
`220`	`220`	`}`
`221`	`221`	`}`
`222`	`222`	`}`
`223`	`223`	`}`
`224`	`224`	`}`
`225`	`225`
`226`		`-fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAndQueue {`
	`226`	`+fn wait(`
	`227`	`+ state_and_queue: &AtomicPtr<()>,`
	`228`	`+ mut current: StateAndQueue,`
	`229`	`+ return_on_poisoned: bool,`
	`230`	`+) -> StateAndQueue {`
`227`	`231`	`let node = &Waiter {`
`228`	`232`	`thread: Cell::new(Some(thread::current())),`
`229`	`233`	`signaled: AtomicBool::new(false),`
`@@ -235,7 +239,7 @@ fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAnd`
`235`	`239`	`let queue = to_queue(current);`
`236`	`240`
`237`	`241`	`// If initialization has finished, return.`
`238`		`- if matches!(state, POISONED \| COMPLETE) {`
	`242`	`+ if state == COMPLETE \|\| (return_on_poisoned && state == POISONED) {`
`239`	`243`	`return current;`
`240`	`244`	`}`
`241`	`245`