fix(semaphore): defer waiter waker drops

asyncband/src/internal/semaphore.rs

@@ -254,10 +254,13 @@ impl Drop for Acquire<'_> {
                 node.permits = 0;
                 true
             });
-            waiters.remove_unlinked_waiter(index);
+            let waiter = waiters.remove_unlinked_waiter(index);
             if acquired > 0 {
                 self.semaphore.insert_permits_with_lock(acquired, waiters);
+            } else {
+                drop(waiters);
             }
+            drop(waiter);
         }
     }
 }
@@ -275,6 +278,7 @@ impl Acquire<'_> {
             return Poll::Ready(());
         }
 
+        let mut old_waker = None;
         match index {
             Some(idx) => {
                 let mut waiters = semaphore.waiters.lock();
@@ -286,7 +290,7 @@ impl Acquire<'_> {
                             .as_ref()
                             .is_none_or(|current| !current.will_wake(waker));
                         if update_waker {
-                            node.waker = Some(waker.clone());
+                            old_waker = node.waker.replace(waker.clone());
                         }
                         false
                     } else {
@@ -311,6 +315,7 @@ impl Acquire<'_> {
             }
         };
 
+        drop(old_waker);
         Poll::Pending
     }
 }

tests-integration/tests/semaphore_test.rs

@@ -20,6 +20,7 @@ use std::pin::pin;
 use std::sync::Arc;
 use std::task::Context;
 use std::task::Poll;
+use std::task::Wake;
 use std::task::Waker;
 use std::vec::Vec;
 
@@ -186,6 +187,41 @@ fn wake_then_drop() {
     assert_eq!(s.available_permits(), 2);
 }
 
+#[test]
+fn cancellation_restores_permits_before_dropping_waker() {
+    struct AssertPermitsOnDrop {
+        semaphore: Arc<Semaphore>,
+    }
+
+    impl Wake for AssertPermitsOnDrop {
+        fn wake(self: Arc<Self>) {
+            panic!("cancellation must not wake the removed waiter");
+        }
+    }
+
+    impl Drop for AssertPermitsOnDrop {
+        fn drop(&mut self) {
+            assert_eq!(self.semaphore.available_permits(), 1);
+        }
+    }
+
+    let semaphore = Arc::new(Semaphore::new(1));
+    let waker = Waker::from(Arc::new(AssertPermitsOnDrop {
+        semaphore: semaphore.clone(),
+    }));
+    let mut acquire = Box::pin(semaphore.acquire(2));
+
+    assert!(
+        acquire
+            .as_mut()
+            .poll(&mut Context::from_waker(&waker))
+            .is_pending()
+    );
+    drop(waker);
+    drop(acquire);
+    assert_eq!(Arc::strong_count(&semaphore), 1);
+}
+
 #[test]
 fn reduce_permits_takes_priority_over_pending_acquires() {
     let s = Semaphore::new(0);