max/embassy
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Partial borrow for receiver to enforce compile-time mpssc

Browse Source
This commit is contained in:
huntc 2021-07-14 22:10:14 +10:00
parent 076198a3b9
commit 6f78527aeb
1 changed files with 61 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

109
embassy/src/util/mpsc.rs
View File

@ -39,6 +39,7 @@
use core::cell::UnsafeCell; use core::cell::UnsafeCell;
use core::fmt; use core::fmt;
use core::marker::PhantomData;
use core::mem::MaybeUninit; use core::mem::MaybeUninit;
use core::pin::Pin; use core::pin::Pin;
use core::ptr; use core::ptr;
@ -61,7 +62,7 @@ pub struct Sender<'ch, M, T, const N: usize>
where where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
channel: &'ch Channel<M, T, N>, channel_cell: &'ch UnsafeCell<ChannelCell<M, T, N>>,
} }
// Safe to pass the sender around // Safe to pass the sender around
@ -77,7 +78,8 @@ pub struct Receiver<'ch, M, T, const N: usize>
where where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
channel: &'ch Channel<M, T, N>, channel_cell: &'ch UnsafeCell<ChannelCell<M, T, N>>,
_receiver_consumed: &'ch mut PhantomData<()>,
} }
// Safe to pass the receiver around // Safe to pass the receiver around
@ -111,18 +113,23 @@ unsafe impl<'ch, M, T, const N: usize> Sync for Receiver<'ch, M, T, N> where
/// ///
/// let (sender, receiver) = { /// let (sender, receiver) = {
/// let mut channel = Channel::<WithThreadModeOnly, u32, 3>::with_thread_mode_only(); /// let mut channel = Channel::<WithThreadModeOnly, u32, 3>::with_thread_mode_only();
/// mpsc::split(&channel) /// mpsc::split(&mut channel)
/// }; /// };
/// ``` /// ```
pub fn split<M, T, const N: usize>( pub fn split<M, T, const N: usize>(
channel: &Channel<M, T, N>, channel: &mut Channel<M, T, N>,
) -> (Sender<M, T, N>, Receiver<M, T, N>) ) -> (Sender<M, T, N>, Receiver<M, T, N>)
where where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
let sender = Sender { channel: &channel }; let sender = Sender {
let receiver = Receiver { channel: &channel }; channel_cell: &channel.channel_cell,
channel.lock(|c| { };
let receiver = Receiver {
channel_cell: &channel.channel_cell,
_receiver_consumed: &mut channel.receiver_consumed,
};
Channel::lock(&channel.channel_cell, |c| {
c.register_receiver(); c.register_receiver();
c.register_sender(); c.register_sender();
}); });
@ -154,11 +161,12 @@ where
} }
fn recv_poll(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> { fn recv_poll(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
self.channel Channel::lock(self.channel_cell, |c| {
.lock(|c| match c.try_recv_with_context(Some(cx)) { match c.try_recv_with_context(Some(cx)) {
Ok(v) => Poll::Ready(Some(v)), Ok(v) => Poll::Ready(Some(v)),
Err(TryRecvError::Closed) => Poll::Ready(None), Err(TryRecvError::Closed) => Poll::Ready(None),
Err(TryRecvError::Empty) => Poll::Pending, Err(TryRecvError::Empty) => Poll::Pending,
}
}) })
} }
@ -167,7 +175,7 @@ where
/// This method will either receive a message from the channel immediately or return an error /// This method will either receive a message from the channel immediately or return an error
/// if the channel is empty. /// if the channel is empty.
pub fn try_recv(&self) -> Result<T, TryRecvError> { pub fn try_recv(&self) -> Result<T, TryRecvError> {
self.channel.lock(|c| c.try_recv()) Channel::lock(self.channel_cell, |c| c.try_recv())
} }
/// Closes the receiving half of a channel without dropping it. /// Closes the receiving half of a channel without dropping it.
@ -181,7 +189,7 @@ where
/// until those are released. /// until those are released.
/// ///
pub fn close(&mut self) { pub fn close(&mut self) {
self.channel.lock(|c| c.close()) Channel::lock(self.channel_cell, |c| c.close())
} }
} }
@ -190,7 +198,7 @@ where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
fn drop(&mut self) { fn drop(&mut self) {
self.channel.lock(|c| c.deregister_receiver()) Channel::lock(self.channel_cell, |c| c.deregister_receiver())
} }
} }
@ -245,7 +253,7 @@ where
/// [`channel`]: channel /// [`channel`]: channel
/// [`close`]: Receiver::close /// [`close`]: Receiver::close
pub fn try_send(&self, message: T) -> Result<(), TrySendError<T>> { pub fn try_send(&self, message: T) -> Result<(), TrySendError<T>> {
self.channel.lock(|c| c.try_send(message)) Channel::lock(self.channel_cell, |c| c.try_send(message))
} }
/// Completes when the receiver has dropped. /// Completes when the receiver has dropped.
@ -266,7 +274,7 @@ where
/// [`Receiver`]: crate::sync::mpsc::Receiver /// [`Receiver`]: crate::sync::mpsc::Receiver
/// [`Receiver::close`]: crate::sync::mpsc::Receiver::close /// [`Receiver::close`]: crate::sync::mpsc::Receiver::close
pub fn is_closed(&self) -> bool { pub fn is_closed(&self) -> bool {
self.channel.lock(|c| c.is_closed()) Channel::lock(self.channel_cell, |c| c.is_closed())
} }
} }
@ -286,11 +294,9 @@ where
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.message.take() { match self.message.take() {
Some(m) => match self Some(m) => match Channel::lock(self.sender.channel_cell, |c| {
.sender c.try_send_with_context(m, Some(cx))
.channel }) {
.lock(|c| c.try_send_with_context(m, Some(cx)))
{
Ok(..) => Poll::Ready(Ok(())), Ok(..) => Poll::Ready(Ok(())),
Err(TrySendError::Closed(m)) => Poll::Ready(Err(SendError(m))), Err(TrySendError::Closed(m)) => Poll::Ready(Err(SendError(m))),
Err(TrySendError::Full(m)) => { Err(TrySendError::Full(m)) => {
@ -319,11 +325,9 @@ where
type Output = (); type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if self if Channel::lock(self.sender.channel_cell, |c| {
.sender c.is_closed_with_context(Some(cx))
.channel }) {
.lock(|c| c.is_closed_with_context(Some(cx)))
{
Poll::Ready(()) Poll::Ready(())
} else { } else {
Poll::Pending Poll::Pending
@ -336,7 +340,7 @@ where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
fn drop(&mut self) { fn drop(&mut self) {
self.channel.lock(|c| c.deregister_sender()) Channel::lock(self.channel_cell, |c| c.deregister_sender())
} }
} }
@ -346,9 +350,9 @@ where
{ {
#[allow(clippy::clone_double_ref)] #[allow(clippy::clone_double_ref)]
fn clone(&self) -> Self { fn clone(&self) -> Self {
self.channel.lock(|c| c.register_sender()); Channel::lock(self.channel_cell, |c| c.register_sender());
Sender { Sender {
channel: self.channel.clone(), channel_cell: self.channel_cell.clone(),
} }
} }
} }
@ -564,6 +568,7 @@ where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
channel_cell: UnsafeCell<ChannelCell<M, T, N>>, channel_cell: UnsafeCell<ChannelCell<M, T, N>>,
receiver_consumed: PhantomData<()>,
} }
struct ChannelCell<M, T, const N: usize> struct ChannelCell<M, T, const N: usize>
@ -588,7 +593,7 @@ impl<T, const N: usize> Channel<WithCriticalSections, T, N> {
/// // Declare a bounded channel of 3 u32s. /// // Declare a bounded channel of 3 u32s.
/// let mut channel = Channel::<WithCriticalSections, u32, 3>::with_critical_sections(); /// let mut channel = Channel::<WithCriticalSections, u32, 3>::with_critical_sections();
/// // once we have a channel, obtain its sender and receiver /// // once we have a channel, obtain its sender and receiver
/// let (sender, receiver) = mpsc::split(&channel); /// let (sender, receiver) = mpsc::split(&mut channel);
/// ``` /// ```
pub const fn with_critical_sections() -> Self { pub const fn with_critical_sections() -> Self {
let mutex = CriticalSectionMutex::new(()); let mutex = CriticalSectionMutex::new(());
@ -596,6 +601,7 @@ impl<T, const N: usize> Channel<WithCriticalSections, T, N> {
let channel_cell = ChannelCell { mutex, state }; let channel_cell = ChannelCell { mutex, state };
Channel { Channel {
channel_cell: UnsafeCell::new(channel_cell), channel_cell: UnsafeCell::new(channel_cell),
receiver_consumed: PhantomData,
} }
} }
} }
@ -615,7 +621,7 @@ impl<T, const N: usize> Channel<WithThreadModeOnly, T, N> {
/// // Declare a bounded channel of 3 u32s. /// // Declare a bounded channel of 3 u32s.
/// let mut channel = Channel::<WithThreadModeOnly, u32, 3>::with_thread_mode_only(); /// let mut channel = Channel::<WithThreadModeOnly, u32, 3>::with_thread_mode_only();
/// // once we have a channel, obtain its sender and receiver /// // once we have a channel, obtain its sender and receiver
/// let (sender, receiver) = mpsc::split(&channel); /// let (sender, receiver) = mpsc::split(&mut channel);
/// ``` /// ```
pub const fn with_thread_mode_only() -> Self { pub const fn with_thread_mode_only() -> Self {
let mutex = ThreadModeMutex::new(()); let mutex = ThreadModeMutex::new(());
@ -623,6 +629,7 @@ impl<T, const N: usize> Channel<WithThreadModeOnly, T, N> {
let channel_cell = ChannelCell { mutex, state }; let channel_cell = ChannelCell { mutex, state };
Channel { Channel {
channel_cell: UnsafeCell::new(channel_cell), channel_cell: UnsafeCell::new(channel_cell),
receiver_consumed: PhantomData,
} }
} }
} }
@ -639,7 +646,7 @@ impl<T, const N: usize> Channel<WithNoThreads, T, N> {
/// // Declare a bounded channel of 3 u32s. /// // Declare a bounded channel of 3 u32s.
/// let mut channel = Channel::<WithNoThreads, u32, 3>::with_no_threads(); /// let mut channel = Channel::<WithNoThreads, u32, 3>::with_no_threads();
/// // once we have a channel, obtain its sender and receiver /// // once we have a channel, obtain its sender and receiver
/// let (sender, receiver) = mpsc::split(&channel); /// let (sender, receiver) = mpsc::split(&mut channel);
/// ``` /// ```
pub const fn with_no_threads() -> Self { pub const fn with_no_threads() -> Self {
let mutex = NoopMutex::new(()); let mutex = NoopMutex::new(());
@ -647,6 +654,7 @@ impl<T, const N: usize> Channel<WithNoThreads, T, N> {
let channel_cell = ChannelCell { mutex, state }; let channel_cell = ChannelCell { mutex, state };
Channel { Channel {
channel_cell: UnsafeCell::new(channel_cell), channel_cell: UnsafeCell::new(channel_cell),
receiver_consumed: PhantomData,
} }
} }
} }
@ -655,9 +663,12 @@ impl<M, T, const N: usize> Channel<M, T, N>
where where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
fn lock<R>(&self, f: impl FnOnce(&mut ChannelState<T, N>) -> R) -> R { fn lock<R>(
channel_cell: &UnsafeCell<ChannelCell<M, T, N>>,
f: impl FnOnce(&mut ChannelState<T, N>) -> R,
) -> R {
unsafe { unsafe {
let channel_cell = &mut *(self.channel_cell.get()); let channel_cell = &mut *(channel_cell.get());
let mutex = &mut channel_cell.mutex; let mutex = &mut channel_cell.mutex;
let mut state = &mut channel_cell.state; let mut state = &mut channel_cell.state;
mutex.lock(|_| f(&mut state)) mutex.lock(|_| f(&mut state))
@ -747,16 +758,16 @@ mod tests {
#[test] #[test]
fn simple_send_and_receive() { fn simple_send_and_receive() {
let c = Channel::<WithNoThreads, u32, 3>::with_no_threads(); let mut c = Channel::<WithNoThreads, u32, 3>::with_no_threads();
let (s, r) = split(&c); let (s, r) = split(&mut c);
assert!(s.clone().try_send(1).is_ok()); assert!(s.clone().try_send(1).is_ok());
assert_eq!(r.try_recv().unwrap(), 1); assert_eq!(r.try_recv().unwrap(), 1);
} }
#[test] #[test]
fn should_close_without_sender() { fn should_close_without_sender() {
let c = Channel::<WithNoThreads, u32, 3>::with_no_threads(); let mut c = Channel::<WithNoThreads, u32, 3>::with_no_threads();
let (s, r) = split(&c); let (s, r) = split(&mut c);
drop(s); drop(s);
match r.try_recv() { match r.try_recv() {
Err(TryRecvError::Closed) => assert!(true), Err(TryRecvError::Closed) => assert!(true),
@ -766,8 +777,8 @@ mod tests {
#[test] #[test]
fn should_close_once_drained() { fn should_close_once_drained() {
let c = Channel::<WithNoThreads, u32, 3>::with_no_threads(); let mut c = Channel::<WithNoThreads, u32, 3>::with_no_threads();
let (s, r) = split(&c); let (s, r) = split(&mut c);
assert!(s.try_send(1).is_ok()); assert!(s.try_send(1).is_ok());
drop(s); drop(s);
assert_eq!(r.try_recv().unwrap(), 1); assert_eq!(r.try_recv().unwrap(), 1);
@ -779,8 +790,8 @@ mod tests {
#[test] #[test]
fn should_reject_send_when_receiver_dropped() { fn should_reject_send_when_receiver_dropped() {
let c = Channel::<WithNoThreads, u32, 3>::with_no_threads(); let mut c = Channel::<WithNoThreads, u32, 3>::with_no_threads();
let (s, r) = split(&c); let (s, r) = split(&mut c);
drop(r); drop(r);
match s.try_send(1) { match s.try_send(1) {
Err(TrySendError::Closed(1)) => assert!(true), Err(TrySendError::Closed(1)) => assert!(true),
@ -790,8 +801,8 @@ mod tests {
#[test] #[test]
fn should_reject_send_when_channel_closed() { fn should_reject_send_when_channel_closed() {
let c = Channel::<WithNoThreads, u32, 3>::with_no_threads(); let mut c = Channel::<WithNoThreads, u32, 3>::with_no_threads();
let (s, mut r) = split(&c); let (s, mut r) = split(&mut c);
assert!(s.try_send(1).is_ok()); assert!(s.try_send(1).is_ok());
r.close(); r.close();
assert_eq!(r.try_recv().unwrap(), 1); assert_eq!(r.try_recv().unwrap(), 1);
@ -808,7 +819,7 @@ mod tests {
static mut CHANNEL: Channel<WithCriticalSections, u32, 3> = static mut CHANNEL: Channel<WithCriticalSections, u32, 3> =
Channel::with_critical_sections(); Channel::with_critical_sections();
let (s, mut r) = split(unsafe { &CHANNEL }); let (s, mut r) = split(unsafe { &mut CHANNEL });
assert!(executor assert!(executor
.spawn(async move { .spawn(async move {
drop(s); drop(s);
@ -823,7 +834,7 @@ mod tests {
static mut CHANNEL: Channel<WithCriticalSections, u32, 3> = static mut CHANNEL: Channel<WithCriticalSections, u32, 3> =
Channel::with_critical_sections(); Channel::with_critical_sections();
let (s, mut r) = split(unsafe { &CHANNEL }); let (s, mut r) = split(unsafe { &mut CHANNEL });
assert!(executor assert!(executor
.spawn(async move { .spawn(async move {
assert!(s.try_send(1).is_ok()); assert!(s.try_send(1).is_ok());
@ -836,7 +847,7 @@ mod tests {
async fn sender_send_completes_if_capacity() { async fn sender_send_completes_if_capacity() {
static mut CHANNEL: Channel<WithCriticalSections, u32, 1> = static mut CHANNEL: Channel<WithCriticalSections, u32, 1> =
Channel::with_critical_sections(); Channel::with_critical_sections();
let (s, mut r) = split(unsafe { &CHANNEL }); let (s, mut r) = split(unsafe { &mut CHANNEL });
assert!(s.send(1).await.is_ok()); assert!(s.send(1).await.is_ok());
assert_eq!(r.recv().await, Some(1)); assert_eq!(r.recv().await, Some(1));
} }
@ -845,7 +856,7 @@ mod tests {
async fn sender_send_completes_if_closed() { async fn sender_send_completes_if_closed() {
static mut CHANNEL: Channel<WithCriticalSections, u32, 1> = static mut CHANNEL: Channel<WithCriticalSections, u32, 1> =
Channel::with_critical_sections(); Channel::with_critical_sections();
let (s, r) = split(unsafe { &CHANNEL }); let (s, r) = split(unsafe { &mut CHANNEL });
drop(r); drop(r);
match s.send(1).await { match s.send(1).await {
Err(SendError(1)) => assert!(true), Err(SendError(1)) => assert!(true),
@ -859,7 +870,7 @@ mod tests {
static mut CHANNEL: Channel<WithCriticalSections, u32, 1> = static mut CHANNEL: Channel<WithCriticalSections, u32, 1> =
Channel::with_critical_sections(); Channel::with_critical_sections();
let (s0, mut r) = split(unsafe { &CHANNEL }); let (s0, mut r) = split(unsafe { &mut CHANNEL });
assert!(s0.try_send(1).is_ok()); assert!(s0.try_send(1).is_ok());
let s1 = s0.clone(); let s1 = s0.clone();
let send_task_1 = executor.spawn_with_handle(async move { s0.send(2).await }); let send_task_1 = executor.spawn_with_handle(async move { s0.send(2).await });
@ -879,7 +890,7 @@ mod tests {
async fn sender_close_completes_if_closing() { async fn sender_close_completes_if_closing() {
static mut CHANNEL: Channel<WithCriticalSections, u32, 1> = static mut CHANNEL: Channel<WithCriticalSections, u32, 1> =
Channel::with_critical_sections(); Channel::with_critical_sections();
let (s, mut r) = split(unsafe { &CHANNEL }); let (s, mut r) = split(unsafe { &mut CHANNEL });
r.close(); r.close();
s.closed().await; s.closed().await;
} }
@ -888,7 +899,7 @@ mod tests {
async fn sender_close_completes_if_closed() { async fn sender_close_completes_if_closed() {
static mut CHANNEL: Channel<WithCriticalSections, u32, 1> = static mut CHANNEL: Channel<WithCriticalSections, u32, 1> =
Channel::with_critical_sections(); Channel::with_critical_sections();
let (s, r) = split(unsafe { &CHANNEL }); let (s, r) = split(unsafe { &mut CHANNEL });
drop(r); drop(r);
s.closed().await; s.closed().await;
} }