actix/actix-net
1
0
Fork 0
mirror of https://github.com/fafhrd91/actix-net synced 2025-06-26 20:57:43 +02:00
Code Issues Releases Wiki Activity

Simplify oneshot and mpsc implementations

Browse Source
This commit is contained in:
Nikolay Kim
2019-12-11 11:28:09 +06:00
parent e315cf2893
commit 2e5e69c9ba
12 changed files with 129 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
actix-utils/src/cell.rs
View File

@ -2,16 +2,12 @@
use std::cell::UnsafeCell;
use std::fmt;
use std::rc::{Rc, Weak};
use std::rc::Rc;
pub(crate) struct Cell<T> {
pub(crate) inner: Rc<UnsafeCell<T>>,
}
pub(crate) struct WeakCell<T> {
inner: Weak<UnsafeCell<T>>,
}
impl<T> Clone for Cell<T> {
fn clone(&self) -> Self {
Self {
@ -27,39 +23,25 @@ impl<T: fmt::Debug> fmt::Debug for Cell<T> {
}
impl<T> Cell<T> {
pub fn new(inner: T) -> Self {
pub(crate) fn new(inner: T) -> Self {
Self {
inner: Rc::new(UnsafeCell::new(inner)),
}
}
pub fn downgrade(&self) -> WeakCell<T> {
WeakCell {
inner: Rc::downgrade(&self.inner),
}
pub(crate) fn strong_count(&self) -> usize {
Rc::strong_count(&self.inner)
}
pub fn get_ref(&self) -> &T {
pub(crate) fn get_ref(&self) -> &T {
unsafe { &*self.inner.as_ref().get() }
}
pub fn get_mut(&mut self) -> &mut T {
pub(crate) fn get_mut(&mut self) -> &mut T {
unsafe { &mut *self.inner.as_ref().get() }
}
}
impl<T> WeakCell<T> {
pub fn upgrade(&self) -> Option<Cell<T>> {
if let Some(inner) = self.inner.upgrade() {
Some(Cell { inner })
} else {
None
}
}
}
impl<T: fmt::Debug> fmt::Debug for WeakCell<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.inner.fmt(f)
pub(crate) unsafe fn get_mut_unsafe(&self) -> &mut T {
&mut *self.inner.as_ref().get()
}
}

134
actix-utils/src/mpsc.rs
View File

@ -1,34 +1,27 @@
//! A multi-producer, single-consumer, futures-aware, FIFO queue with back
//! pressure, for use communicating between tasks on the same thread.
//!
//! These queues are the same as those in `futures::sync`, except they're not
//! intended to be sent across threads.
//! A multi-producer, single-consumer, futures-aware, FIFO queue.
use std::any::Any;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::pin::Pin;
use std::rc::{Rc, Weak};
use std::task::{Context, Poll};
use std::{fmt, mem};
use futures::{Sink, Stream};
use crate::cell::Cell;
use crate::task::LocalWaker;
/// Creates a unbounded in-memory channel with buffered storage.
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
let shared = Rc::new(RefCell::new(Shared {
let shared = Cell::new(Shared {
has_receiver: true,
buffer: VecDeque::new(),
blocked_recv: LocalWaker::new(),
}));
});
let sender = Sender {
shared: Rc::downgrade(&shared),
};
let receiver = Receiver {
state: State::Open(shared),
shared: shared.clone(),
};
let receiver = Receiver { shared };
(sender, receiver)
}
@ -36,6 +29,7 @@ pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
struct Shared<T> {
buffer: VecDeque<T>,
blocked_recv: LocalWaker,
has_receiver: bool,
}
/// The transmission end of a channel.
@ -43,18 +37,18 @@ struct Shared<T> {
/// This is created by the `channel` function.
#[derive(Debug)]
pub struct Sender<T> {
shared: Weak<RefCell<Shared<T>>>,
shared: Cell<Shared<T>>,
}
impl<T> Unpin for Sender<T> {}
impl<T> Sender<T> {
/// Sends the provided message along this channel.
pub fn send(&self, item: T) -> Result<(), SendError<T>> {
let shared = match self.shared.upgrade() {
Some(shared) => shared,
None => return Err(SendError(item)), // receiver was dropped
let shared = unsafe { self.shared.get_mut_unsafe() };
if !shared.has_receiver {
return Err(SendError(item)); // receiver was dropped
};
let mut shared = shared.borrow_mut();
shared.buffer.push_back(item);
shared.blocked_recv.wake();
Ok(())
@ -91,17 +85,13 @@ impl<T> Sink<T> for Sender<T> {
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
let shared = match self.shared.upgrade() {
Some(shared) => shared,
None => return,
};
// The number of existing `Weak` indicates if we are possibly the last
// `Sender`. If we are the last, we possibly must notify a blocked
// `Receiver`. `self.shared` is always one of the `Weak` to this shared
// data. Therefore the smallest possible Rc::weak_count(&shared) is 1.
if Rc::weak_count(&shared) == 1 {
let count = self.shared.strong_count();
let shared = self.shared.get_mut();
// check is last sender is about to drop
if shared.has_receiver && count == 2 {
// Wake up receiver as its stream has ended
shared.borrow_mut().blocked_recv.wake();
shared.blocked_recv.wake();
}
}
}
@ -111,56 +101,23 @@ impl<T> Drop for Sender<T> {
/// This is created by the `channel` function.
#[derive(Debug)]
pub struct Receiver<T> {
state: State<T>,
shared: Cell<Shared<T>>,
}
impl<T> Unpin for Receiver<T> {}
/// Possible states of a receiver. We're either Open (can receive more messages)
/// or we're closed with a list of messages we have left to receive.
#[derive(Debug)]
enum State<T> {
Open(Rc<RefCell<Shared<T>>>),
Closed(VecDeque<T>),
}
impl<T> Receiver<T> {
/// Closes the receiving half
///
/// This prevents any further messages from being sent on the channel while
/// still enabling the receiver to drain messages that are buffered.
pub fn close(&mut self) {
let items = match self.state {
State::Open(ref state) => {
let mut state = state.borrow_mut();
mem::replace(&mut state.buffer, VecDeque::new())
}
State::Closed(_) => return,
};
self.state = State::Closed(items);
}
}
impl<T> Stream for Receiver<T> {
type Item = T;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let me = match self.state {
State::Open(ref mut me) => me,
State::Closed(ref mut items) => return Poll::Ready(items.pop_front()),
};
if let Some(shared) = Rc::get_mut(me) {
if self.shared.strong_count() == 1 {
// All senders have been dropped, so drain the buffer and end the
// stream.
return Poll::Ready(shared.borrow_mut().buffer.pop_front());
}
let mut shared = me.borrow_mut();
if let Some(msg) = shared.buffer.pop_front() {
Poll::Ready(self.shared.get_mut().buffer.pop_front())
} else if let Some(msg) = self.shared.get_mut().buffer.pop_front() {
Poll::Ready(Some(msg))
} else {
shared.blocked_recv.register(cx.waker());
self.shared.get_mut().blocked_recv.register(cx.waker());
Poll::Pending
}
}
@ -168,7 +125,9 @@ impl<T> Stream for Receiver<T> {
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
self.close();
let shared = self.shared.get_mut();
shared.buffer.clear();
shared.has_receiver = false;
}
}
@ -200,3 +159,38 @@ impl<T> SendError<T> {
self.0
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::future::lazy;
use futures::{Stream, StreamExt};
#[actix_rt::test]
async fn test_mpsc() {
let (tx, mut rx) = channel();
tx.send("test").unwrap();
assert_eq!(rx.next().await.unwrap(), "test");
let tx2 = tx.clone();
tx2.send("test2").unwrap();
assert_eq!(rx.next().await.unwrap(), "test2");
assert_eq!(
lazy(|cx| Pin::new(&mut rx).poll_next(cx)).await,
Poll::Pending
);
drop(tx2);
assert_eq!(
lazy(|cx| Pin::new(&mut rx).poll_next(cx)).await,
Poll::Pending
);
drop(tx);
assert_eq!(rx.next().await, None);
let (tx, rx) = channel();
tx.send("test").unwrap();
drop(rx);
assert!(tx.send("test").is_err());
}
}

132
actix-utils/src/oneshot.rs
View File

@ -1,69 +1,45 @@
//! A one-shot, futures-aware channel
//!
//! This channel is similar to that in `sync::oneshot` but cannot be sent across
//! threads.
//! A one-shot, futures-aware channel.
use std::future::Future;
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Context, Poll};
pub use futures::channel::oneshot::Canceled;
use crate::cell::{Cell, WeakCell};
use crate::cell::Cell;
use crate::task::LocalWaker;
/// Creates a new futures-aware, one-shot channel.
///
/// This function is the same as `sync::oneshot::channel` except that the
/// returned values cannot be sent across threads.
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
let inner = Cell::new(Inner {
value: None,
rx_task: LocalWaker::new(),
});
let tx = Sender {
inner: inner.downgrade(),
};
let rx = Receiver {
state: State::Open(inner),
inner: inner.clone(),
};
let rx = Receiver { inner };
(tx, rx)
}
/// Represents the completion half of a oneshot through which the result of a
/// computation is signaled.
///
/// This is created by the `unsync::oneshot::channel` function and is equivalent
/// in functionality to `sync::oneshot::Sender` except that it cannot be sent
/// across threads.
#[derive(Debug)]
pub struct Sender<T> {
inner: WeakCell<Inner<T>>,
inner: Cell<Inner<T>>,
}
/// A future representing the completion of a computation happening elsewhere in
/// memory.
///
/// This is created by the `unsync::oneshot::channel` function and is equivalent
/// in functionality to `sync::oneshot::Receiver` except that it cannot be sent
/// across threads.
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Receiver<T> {
state: State<T>,
inner: Cell<Inner<T>>,
}
// The channels do not ever project Pin to the inner T
impl<T> Unpin for Receiver<T> {}
impl<T> Unpin for Sender<T> {}
#[derive(Debug)]
enum State<T> {
Open(Cell<Inner<T>>),
Closed(Option<T>),
}
#[derive(Debug)]
struct Inner<T> {
value: Option<T>,
@ -78,12 +54,12 @@ impl<T> Sender<T> {
/// represents.
///
/// If the value is successfully enqueued for the remote end to receive,
/// then `Ok(())` is returned. If the receiving end was deallocated before
/// then `Ok(())` is returned. If the receiving end was dropped before
/// this function was called, however, then `Err` is returned with the value
/// provided.
pub fn send(self, val: T) -> Result<(), T> {
if let Some(mut inner) = self.inner.upgrade() {
let inner = inner.get_mut();
pub fn send(mut self, val: T) -> Result<(), T> {
if self.inner.strong_count() == 2 {
let inner = self.inner.get_mut();
inner.value = Some(val);
inner.rx_task.wake();
Ok(())
@ -91,47 +67,12 @@ impl<T> Sender<T> {
Err(val)
}
}
/// Tests to see whether this `Sender`'s corresponding `Receiver`
/// has gone away.
///
/// This function can be used to learn about when the `Receiver` (consumer)
/// half has gone away and nothing will be able to receive a message sent
/// from `send`.
///
/// Note that this function is intended to *not* be used in the context of a
/// future. If you're implementing a future you probably want to call the
/// `poll_cancel` function which will block the current task if the
/// cancellation hasn't happened yet. This can be useful when working on a
/// non-futures related thread, though, which would otherwise panic if
/// `poll_cancel` were called.
pub fn is_canceled(&self) -> bool {
self.inner.upgrade().is_none()
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
if let Some(inner) = self.inner.upgrade() {
inner.get_ref().rx_task.wake();
};
}
}
impl<T> Receiver<T> {
/// Gracefully close this receiver, preventing sending any future messages.
///
/// Any `send` operation which happens after this method returns is
/// guaranteed to fail. Once this method is called the normal `poll` method
/// can be used to determine whether a message was actually sent or not. If
/// `Canceled` is returned from `poll` then no message was sent.
pub fn close(&mut self) {
match self.state {
State::Open(ref mut inner) => {
let value = inner.get_mut().value.take();
self.state = State::Closed(value);
}
State::Closed(_) => {}
if self.inner.strong_count() == 2 {
self.inner.get_ref().rx_task.wake();
};
}
}
@ -142,33 +83,48 @@ impl<T> Future for Receiver<T> {
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let inner = match this.state {
State::Open(ref mut inner) => inner,
State::Closed(ref mut item) => match item.take() {
Some(item) => return Poll::Ready(Ok(item)),
None => return Poll::Ready(Err(Canceled)),
},
};
// If we've got a value, then skip the logic below as we're done.
if let Some(val) = inner.get_mut().value.take() {
if let Some(val) = this.inner.get_mut().value.take() {
return Poll::Ready(Ok(val));
}
// If we can get mutable access, then the sender has gone away. We
// didn't see a value above, so we're canceled. Otherwise we park
// our task and wait for a value to come in.
if Rc::get_mut(&mut inner.inner).is_some() {
// Check if sender is dropped and return error if it is.
if this.inner.strong_count() == 1 {
Poll::Ready(Err(Canceled))
} else {
inner.get_ref().rx_task.register(cx.waker());
this.inner.get_ref().rx_task.register(cx.waker());
Poll::Pending
}
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
self.close();
#[cfg(test)]
mod tests {
use super::*;
use futures::future::lazy;
#[actix_rt::test]
async fn test_oneshot() {
let (tx, rx) = channel();
tx.send("test").unwrap();
assert_eq!(rx.await.unwrap(), "test");
let (tx, rx) = channel();
drop(rx);
assert!(tx.send("test").is_err());
let (tx, rx) = channel::<&'static str>();
drop(tx);
assert!(rx.await.is_err());
let (tx, mut rx) = channel::<&'static str>();
assert_eq!(lazy(|cx| Pin::new(&mut rx).poll(cx)).await, Poll::Pending);
tx.send("test").unwrap();
assert_eq!(rx.await.unwrap(), "test");
let (tx, mut rx) = channel::<&'static str>();
assert_eq!(lazy(|cx| Pin::new(&mut rx).poll(cx)).await, Poll::Pending);
drop(tx);
assert!(rx.await.is_err());
}
}