From 79799d241c78510ebcbae801b519bd2422d12b6b Mon Sep 17 00:00:00 2001 From: Emil Fresk Date: Sat, 28 Jan 2023 13:21:44 +0100 Subject: [PATCH] Use `Pin` in the linked lists --- rtic-channel/src/lib.rs | 18 +++++++++++++----- rtic-channel/src/wait_queue.rs | 16 ++++++++++------ rtic-time/src/lib.rs | 18 +++++++++++++++--- rtic-time/src/linked_list.rs | 5 ++++- 4 files changed, 42 insertions(+), 15 deletions(-) diff --git a/rtic-channel/src/lib.rs b/rtic-channel/src/lib.rs index 5ee2c710a9..20086ac7fb 100644 --- a/rtic-channel/src/lib.rs +++ b/rtic-channel/src/lib.rs @@ -8,6 +8,7 @@ use core::{ cell::UnsafeCell, future::poll_fn, mem::MaybeUninit, + pin::Pin, ptr, task::{Poll, Waker}, }; @@ -177,10 +178,11 @@ impl<'a, T, const N: usize> Sender<'a, T, N> { pub async fn send(&mut self, val: T) -> Result<(), NoReceiver> { if self.is_closed() {} - let mut __hidden_link: Option> = None; + let mut link_ptr: Option> = None; + + // Make this future `Drop`-safe, also shadow the original definition so we can't abuse it. + let link_ptr = &mut link_ptr as *mut Option>; - // Make this future `Drop`-safe - let link_ptr = &mut __hidden_link as *mut Option>; let dropper = OnDrop::new(|| { // SAFETY: We only run this closure and dereference the pointer if we have // exited the `poll_fn` below in the `drop(dropper)` call. The other dereference @@ -198,12 +200,18 @@ impl<'a, T, const N: usize> Sender<'a, T, N> { // Do all this in one critical section, else there can be race conditions let queue_idx = critical_section::with(|cs| { if !self.0.wait_queue.is_empty() || self.0.access(cs).freeq.is_empty() { - // SAFETY: This pointer is only dereferenced here and on drop of the future. + // SAFETY: This pointer is only dereferenced here and on drop of the future + // which happens outside this `poll_fn`'s stack frame. let link = unsafe { &mut *link_ptr }; if link.is_none() { // Place the link in the wait queue on first run. let link_ref = link.insert(wait_queue::Link::new(cx.waker().clone())); - self.0.wait_queue.push(link_ref); + + // SAFETY: The address to the link is stable as it is hidden behind + // `link_ptr`, and `link_ptr` shadows the original making it unmovable. + self.0 + .wait_queue + .push(unsafe { Pin::new_unchecked(link_ref) }); } return None; diff --git a/rtic-channel/src/wait_queue.rs b/rtic-channel/src/wait_queue.rs index 5b59983261..ba05e6bb75 100644 --- a/rtic-channel/src/wait_queue.rs +++ b/rtic-channel/src/wait_queue.rs @@ -1,6 +1,7 @@ //! ... use core::marker::PhantomPinned; +use core::pin::Pin; use core::ptr::null_mut; use core::sync::atomic::{AtomicPtr, Ordering}; use core::task::Waker; @@ -65,13 +66,16 @@ impl LinkedList { } /// Put an element at the back of the queue. - pub fn push(&self, link: &mut Link) { + pub fn push(&self, link: Pin<&mut Link>) { cs::with(|_| { // Make sure all previous writes are visible core::sync::atomic::fence(Ordering::SeqCst); let tail = self.tail.load(Self::R); + // SAFETY: This datastructure does not move the underlying value. + let link = unsafe { link.get_unchecked_mut() }; + if let Some(tail_ref) = unsafe { tail.as_ref() } { // Queue is not empty link.prev.store(tail, Self::R); @@ -221,11 +225,11 @@ mod tests { let mut i4 = Link::new(13); let mut i5 = Link::new(14); - wq.push(&mut i1); - wq.push(&mut i2); - wq.push(&mut i3); - wq.push(&mut i4); - wq.push(&mut i5); + wq.push(unsafe { Pin::new_unchecked(&mut i1) }); + wq.push(unsafe { Pin::new_unchecked(&mut i2) }); + wq.push(unsafe { Pin::new_unchecked(&mut i3) }); + wq.push(unsafe { Pin::new_unchecked(&mut i4) }); + wq.push(unsafe { Pin::new_unchecked(&mut i5) }); wq.print(); diff --git a/rtic-time/src/lib.rs b/rtic-time/src/lib.rs index eeecd86c2d..6b23f7653d 100644 --- a/rtic-time/src/lib.rs +++ b/rtic-time/src/lib.rs @@ -7,6 +7,7 @@ #![feature(async_fn_in_trait)] use core::future::{poll_fn, Future}; +use core::pin::Pin; use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use core::task::{Poll, Waker}; use futures_util::{ @@ -185,7 +186,10 @@ impl TimerQueue { ); } - let mut link = None; + let mut link_ptr: Option>> = None; + + // Make this future `Drop`-safe, also shadow the original definition so we can't abuse it. + let link_ptr = &mut link_ptr as *mut Option>>; let queue = &self.queue; let marker = &AtomicUsize::new(0); @@ -199,6 +203,9 @@ impl TimerQueue { return Poll::Ready(()); } + // SAFETY: This pointer is only dereferenced here and on drop of the future + // which happens outside this `poll_fn`'s stack frame. + let link = unsafe { &mut *link_ptr }; if link.is_none() { let mut link_ref = link.insert(Link::new(WaitingWaker { waker: cx.waker().clone(), @@ -206,7 +213,9 @@ impl TimerQueue { was_poped: AtomicBool::new(false), })); - let (was_empty, addr) = queue.insert(&mut link_ref); + // SAFETY: The address to the link is stable as it is defined outside this stack + // frame. + let (was_empty, addr) = queue.insert(unsafe { Pin::new_unchecked(&mut link_ref) }); marker.store(addr, Ordering::Relaxed); if was_empty { @@ -219,7 +228,10 @@ impl TimerQueue { }) .await; - if let Some(link) = link { + // SAFETY: We only run this and dereference the pointer if we have + // exited the `poll_fn` below in the `drop(dropper)` call. The other dereference + // of this pointer is in the `poll_fn`. + if let Some(link) = unsafe { &mut *link_ptr } { if link.val.was_poped.load(Ordering::Relaxed) { // If it was poped from the queue there is no need to run delete dropper.defuse(); diff --git a/rtic-time/src/linked_list.rs b/rtic-time/src/linked_list.rs index 52a955bd46..de5ea2a48a 100644 --- a/rtic-time/src/linked_list.rs +++ b/rtic-time/src/linked_list.rs @@ -1,6 +1,7 @@ //! ... use core::marker::PhantomPinned; +use core::pin::Pin; use core::sync::atomic::{AtomicPtr, Ordering}; use critical_section as cs; @@ -92,8 +93,10 @@ impl LinkedList { /// Insert a new link into the linked list. /// The return is (was_empty, address), where the address of the link is for use with `delete`. - pub fn insert(&self, val: &mut Link) -> (bool, usize) { + pub fn insert(&self, val: Pin<&mut Link>) -> (bool, usize) { cs::with(|_| { + // SAFETY: This datastructure does not move the underlying value. + let val = unsafe { val.get_unchecked_mut() }; let addr = val as *const _ as usize; // Make sure all previous writes are visible