wait_list: make sure the node is removed before dropping `Prepare`

crates/eonix_sync/src/wait_list.rs

@@ -1,7 +1,7 @@
 mod prepare;
 mod wait_object;
 
-use crate::Spin;
+use crate::{LazyLock, Spin};
 use core::{fmt, sync::atomic::Ordering};
 use intrusive_collections::LinkedList;
 use wait_object::WaitObjectAdapter;
@@ -9,13 +9,13 @@ use wait_object::WaitObjectAdapter;
 pub use prepare::Prepare;
 
 pub struct WaitList {
-    waiters: Spin<LinkedList<WaitObjectAdapter>>,
+    waiters: LazyLock<Spin<LinkedList<WaitObjectAdapter>>>,
 }
 
 impl WaitList {
-    pub fn new() -> Self {
+    pub const fn new() -> Self {
         Self {
-            waiters: Spin::new(LinkedList::new(WaitObjectAdapter::new())),
+            waiters: LazyLock::new(|| Spin::new(LinkedList::new(WaitObjectAdapter::new()))),
         }
     }
 
@@ -43,7 +43,7 @@ impl WaitList {
     }
 
     pub fn notify_all(&self) -> usize {
-        let mut waiters = self.waiters.lock().take();
+        let mut waiters = self.waiters.lock();
         let mut waiter = waiters.front_mut();
         let mut count = 0;
 
@@ -76,6 +76,12 @@ impl WaitList {
     }
 }
 
+impl Default for WaitList {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl fmt::Debug for WaitList {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("WaitList").finish()

crates/eonix_sync/src/wait_list/prepare.rs

@@ -167,10 +167,29 @@ impl Future for Prepare<'_> {
 
 impl Drop for Prepare<'_> {
     fn drop(&mut self) {
-        assert_eq!(
-            self.state,
-            State::WokenUp,
-            "Prepare dropped before woken up."
-        );
+        match self.state {
+            State::Init | State::WokenUp => {}
+            State::OnList | State::WakerSet => {
+                let wait_object = self.wait_object();
+                if wait_object.woken_up.load(Ordering::Acquire) {
+                    // We've woken up by someone. It won't be long before they
+                    // remove us from the list. So spin until we are off the list.
+                    // And we're done.
+                    while wait_object.on_list() {}
+                } else {
+                    // Lock the list and try again.
+                    let mut waiters = self.wait_list.waiters.lock();
+
+                    if wait_object.on_list() {
+                        let mut cursor = unsafe {
+                            // SAFETY: The list is locked so no one could be polling nodes
+                            //         off while we are trying to remove it.
+                            waiters.cursor_mut_from_ptr(wait_object)
+                        };
+                        assert!(cursor.remove().is_some());
+                    }
+                }
+            }
+        }
     }
 }