AtomicWaker: add take method

embassy-executor/tests/test.rs

@@ -156,8 +156,7 @@ fn waking_after_completion_does_not_poll() {
     let (executor, trace) = setup();
     executor.spawner().spawn(task1(trace.clone(), waker)).unwrap();
 
-    unsafe { executor.poll() };
-    waker.wake();
+    // Task registers waker, then exits
     unsafe { executor.poll() };
 
     // Exited task may be waken but is not polled
@@ -176,7 +175,6 @@ fn waking_after_completion_does_not_poll() {
             "pend",       // spawning a task pends the executor
             "poll task1", //
             "pend",       // manual wake, gets cleared by poll
-            "pend",       // manual wake, single pend for two wakes
             "pend",       // respawning a task pends the executor
             "poll task1", //
         ]

embassy-sync/CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- AtomicWaker now drops the Waker after waking it.
+- Added `AtomicWaker::take`.
+
 ## 0.6.2 - 2025-01-15
 
 - Add dynamic dispatch variant of `Pipe`.

embassy-sync/src/waitqueue/atomic_waker.rs

@@ -30,14 +30,16 @@ impl<M: RawMutex> GenericAtomicWaker<M> {
         })
     }
 
+    /// Extracts the previously registered waker, leaving `None` in its place.
+    pub fn take(&self) -> Option<Waker> {
+        self.waker.lock(|cell| cell.take())
+    }
+
     /// Wake the registered waker, if any.
     pub fn wake(&self) {
-        self.waker.lock(|cell| {
-            if let Some(w) = cell.replace(None) {
-                w.wake_by_ref();
-                cell.set(Some(w));
-            }
-        })
+        if let Some(waker) = self.take() {
+            waker.wake();
+        }
     }
 }
 
@@ -59,6 +61,11 @@ impl AtomicWaker {
         self.waker.register(w);
     }
 
+    /// Extracts the previously registered waker, leaving `None` in its place.
+    pub fn take(&self) -> Option<Waker> {
+        self.waker.take()
+    }
+
     /// Wake the registered waker, if any.
     pub fn wake(&self) {
         self.waker.wake();