waker registration somehow lost, back again

rtic-channel/src/waker_registration.rs

@@ -0,0 +1,64 @@
+use core::cell::UnsafeCell;
+use core::task::Waker;
+
+/// A critical section based waker handler.
+pub struct CriticalSectionWakerRegistration {
+    waker: UnsafeCell<Option<Waker>>,
+}
+
+unsafe impl Send for CriticalSectionWakerRegistration {}
+unsafe impl Sync for CriticalSectionWakerRegistration {}
+
+impl CriticalSectionWakerRegistration {
+    /// Create a new waker registration.
+    pub const fn new() -> Self {
+        Self {
+            waker: UnsafeCell::new(None),
+        }
+    }
+
+    /// Register a waker.
+    /// This will overwrite the previous waker if there was one.
+    pub fn register(&self, new_waker: &Waker) {
+        critical_section::with(|_| {
+            // SAFETY: This access is protected by the critical section.
+            let self_waker = unsafe { &mut *self.waker.get() };
+
+            // From embassy
+            // https://github.com/embassy-rs/embassy/blob/b99533607ceed225dd12ae73aaa9a0d969a7365e/embassy-sync/src/waitqueue/waker.rs#L59-L61
+            match self_waker {
+                // Optimization: If both the old and new Wakers wake the same task, we can simply
+                // keep the old waker, skipping the clone. (In most executor implementations,
+                // cloning a waker is somewhat expensive, comparable to cloning an Arc).
+                Some(ref w2) if (w2.will_wake(new_waker)) => {}
+                _ => {
+                    // clone the new waker and store it
+                    if let Some(old_waker) = core::mem::replace(self_waker, Some(new_waker.clone()))
+                    {
+                        // We had a waker registered for another task. Wake it, so the other task can
+                        // reregister itself if it's still interested.
+                        //
+                        // If two tasks are waiting on the same thing concurrently, this will cause them
+                        // to wake each other in a loop fighting over this WakerRegistration. This wastes
+                        // CPU but things will still work.
+                        //
+                        // If the user wants to have two tasks waiting on the same thing they should use
+                        // a more appropriate primitive that can store multiple wakers.
+                        old_waker.wake()
+                    }
+                }
+            }
+        });
+    }
+
+    /// Wake the waker.
+    pub fn wake(&self) {
+        critical_section::with(|_| {
+            // SAFETY: This access is protected by the critical section.
+            let self_waker = unsafe { &mut *self.waker.get() };
+            if let Some(waker) = self_waker.take() {
+                waker.wake()
+            }
+        });
+    }
+}