Removed UB code around the send future

This commit is contained in:
huntc 2021-07-14 11:31:15 +10:00
parent baab52d40c
commit 7c723d2bfd

View file

@ -210,7 +210,7 @@ where
pub async fn send(&self, message: T) -> Result<(), SendError<T>> { pub async fn send(&self, message: T) -> Result<(), SendError<T>> {
SendFuture { SendFuture {
sender: self.clone(), sender: self.clone(),
message: UnsafeCell::new(message), message: Some(message),
} }
.await .await
} }
@ -266,7 +266,7 @@ where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,
{ {
sender: Sender<'ch, M, T, N>, sender: Sender<'ch, M, T, N>,
message: UnsafeCell<T>, message: Option<T>,
} }
impl<'ch, M, T, const N: usize> Future for SendFuture<'ch, M, T, N> impl<'ch, M, T, const N: usize> Future for SendFuture<'ch, M, T, N>
@ -275,25 +275,23 @@ where
{ {
type Output = Result<(), SendError<T>>; type Output = Result<(), SendError<T>>;
fn poll(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 match self.message.take() {
.sender Some(m) => match self.sender.channel.get().try_send_with_context(m, Some(cx)) {
.channel Ok(..) => Poll::Ready(Ok(())),
.get() Err(TrySendError::Closed(m)) => Poll::Ready(Err(SendError(m))),
.try_send_with_context(unsafe { self.message.get().read() }, Some(cx)) Err(TrySendError::Full(m)) => {
{ self.message.insert(m);
Ok(..) => Poll::Ready(Ok(())), Poll::Pending
Err(TrySendError::Closed(m)) => Poll::Ready(Err(SendError(m))), }
Err(TrySendError::Full(..)) => { },
Poll::Pending None => panic!("Message cannot be None"),
// Note we leave the existing UnsafeCell contents - they still
// contain the original message. We could create another UnsafeCell
// with the message of Full, but there's no real need.
}
} }
} }
} }
impl<'ch, M, T, const N: usize> Unpin for SendFuture<'ch, M, T, N> where M: Mutex<Data = ()> {}
struct CloseFuture<'ch, M, T, const N: usize> struct CloseFuture<'ch, M, T, const N: usize>
where where
M: Mutex<Data = ()>, M: Mutex<Data = ()>,