util: Do not unregister waker on wake in AtomicWaker.

3 files changed, 15 insertions, 21 deletions

diff --git a/embassy-nrf/src/gpiote.rs b/embassy-nrf/src/gpiote.rs
index 9a7642c51..920d69236 100644
--- a/embassy-nrf/src/gpiote.rs
+++ b/embassy-nrf/src/gpiote.rs
@@ -9,7 +9,7 @@ use core::ptr;
 use core::task::{Context, Poll};
 use embassy::interrupt::InterruptExt;
 use embassy::traits::gpio::{WaitForHigh, WaitForLow};
-use embassy::util::{AtomicWakerRegistration, PeripheralBorrow, Signal};
+use embassy::util::{AtomicWaker, PeripheralBorrow, Signal};
 use embedded_hal::digital::v2::{InputPin, OutputPin, StatefulOutputPin};
 
 use crate::gpio::sealed::Pin as _;
@@ -68,9 +68,9 @@ impl ChannelID for ChAny {
     }
 }
 
-const NEW_AWR: AtomicWakerRegistration = AtomicWakerRegistration::new();
-static CHANNEL_WAKERS: [AtomicWakerRegistration; CHANNEL_COUNT] = [NEW_AWR; CHANNEL_COUNT];
-static PORT_WAKERS: [AtomicWakerRegistration; PIN_COUNT] = [NEW_AWR; PIN_COUNT];
+const NEW_AWR: AtomicWaker = AtomicWaker::new();
+static CHANNEL_WAKERS: [AtomicWaker; CHANNEL_COUNT] = [NEW_AWR; CHANNEL_COUNT];
+static PORT_WAKERS: [AtomicWaker; PIN_COUNT] = [NEW_AWR; PIN_COUNT];
 
 pub enum InputChannelPolarity {
     None,
diff --git a/embassy/src/util/waker.rs b/embassy/src/util/waker.rs
index 2b72fd560..cd53cca61 100644
--- a/embassy/src/util/waker.rs
+++ b/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) };
         }
     }
 }
diff --git a/embassy/src/util/waker_agnostic.rs b/embassy/src/util/waker_agnostic.rs
index b4234c0f5..3f6ad373a 100644
--- a/embassy/src/util/waker_agnostic.rs
+++ b/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));
             }
         })
     }