From 0e8bb5dc0b59a490f679f82c3efc6c2994c2d1d9 Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Fri, 26 Mar 2021 23:20:53 +0100
Subject: [PATCH] util: Do not unregister waker on wake in AtomicWaker.

---
 embassy-nrf/src/gpiote.rs          |  8 ++++----
 embassy/src/util/waker.rs          | 15 ++++++---------
 embassy/src/util/waker_agnostic.rs | 13 +++++--------
 3 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/embassy-nrf/src/gpiote.rs b/embassy-nrf/src/gpiote.rs
index 9a7642c5..920d6923 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 2b72fd56..cd53cca6 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 b4234c0f..3f6ad373 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));
             }
         })
     }