util: Do not unregister waker on wake in AtomicWaker.

embassy/src/util/waker.rs

@@ -48,11 +48,11 @@ impl WakerRegistration {
     }
 }
 
-pub struct AtomicWakerRegistration {
+pub struct AtomicWaker {
     waker: AtomicPtr<TaskHeader>,
 }
 
-impl AtomicWakerRegistration {
+impl AtomicWaker {
     pub const fn new() -> Self {
         Self {
             waker: AtomicPtr::new(ptr::null_mut()),
@@ -62,17 +62,14 @@ impl AtomicWakerRegistration {
     /// Register a waker. Overwrites the previous waker, if any.
     pub fn register(&self, w: &Waker) {
         let w = unsafe { task_from_waker(w) };
-        let w2 = self.waker.swap(w.as_ptr(), Ordering::Relaxed);
-        if !w2.is_null() && w2 != w.as_ptr() {
-            unsafe { wake_task(NonNull::new_unchecked(w2)) };
-        }
+        self.waker.store(w.as_ptr(), Ordering::Relaxed);
     }
 
     /// Wake the registered waker, if any.
     pub fn wake(&self) {
-        let w2 = self.waker.swap(ptr::null_mut(), Ordering::Relaxed);
-        if !w2.is_null() {
-            unsafe { wake_task(NonNull::new_unchecked(w2)) };
+        let w2 = self.waker.load(Ordering::Relaxed);
+        if let Some(w2) = NonNull::new(w2) {
+            unsafe { wake_task(w2) };
         }
     }
 }

embassy/src/util/waker_agnostic.rs

@@ -49,11 +49,11 @@ impl WakerRegistration {
 }
 
 /// Utility struct to register and wake a waker.
-pub struct AtomicWakerRegistration {
+pub struct AtomicWaker {
     waker: Mutex<Cell<Option<Waker>>>,
 }
 
-impl AtomicWakerRegistration {
+impl AtomicWaker {
     pub const fn new() -> Self {
         Self {
             waker: Mutex::new(Cell::new(None)),
@@ -66,11 +66,7 @@ impl AtomicWakerRegistration {
             let cell = self.waker.borrow(cs);
             cell.set(match cell.replace(None) {
                 Some(w2) if (w2.will_wake(w)) => Some(w2),
-                Some(w2) => {
-                    w2.wake();
-                    Some(w.clone())
-                }
-                None => Some(w.clone()),
+                _ => Some(w.clone()),
             })
         })
     }
@@ -80,7 +76,8 @@ impl AtomicWakerRegistration {
         cortex_m::interrupt::free(|cs| {
             let cell = self.waker.borrow(cs);
             if let Some(w) = cell.replace(None) {
-                w.wake()
+                w.wake_by_ref();
+                cell.set(Some(w));
             }
         })
     }