tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-09-06 00:00:10 +02:00
Code Issues Packages Projects Releases Wiki Activity

Introduce tokio-sync crate containing synchronization primitives. (#839)

Browse Source 
Introduce a tokio-sync crate containing useful synchronization primitives for programs
written using Tokio.

The initial release contains:

* An mpsc channel
* A oneshot channel
* A semaphore implementation
* An `AtomicTask` primitive.

The `oneshot` and `mpsc` channels are new implementations providing improved
performance characteristics. In some benchmarks, the new mpsc channel shows
up to 7x improvement over the version provided by the `futures` crate. Unfortunately,
the `oneshot` implementation only provides a slight performance improvement as it
is mostly limited by the `futures` 0.1 task system. Once updated to the `std` version
of `Future` (currently nightly only), much greater performance improvements should
be achievable by `oneshot`.

Additionally, he implementations provided here are checked using
[Loom](http://github.com/carllerche/loom/), which provides greater confidence of
correctness.
This commit is contained in:
Carl Lerche
2019-01-22 11:37:26 -08:00
committed by GitHub
parent 91f20e33a4
commit 13083153aa
43 changed files with 5145 additions and 2384 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio-sync/src/lib.rs
+29
View File
@@ -0,0 +1,29 @@
#![doc(html_root_url = "https://docs.rs/tokio-sync/0.1.0")]
#![deny(missing_debug_implementations, missing_docs, unreachable_pub)]
#![cfg_attr(test, deny(warnings))]
//! Asynchronous synchronization primitives.
//!
//! This crate provides primitives for synchronizing asynchronous tasks.
extern crate futures;
macro_rules! debug {
($($t:tt)*) => {
if false {
println!($($t)*);
}
}
}
macro_rules! if_fuzz {
($($t:tt)*) => {{
if false { $($t)* }
}}
}
mod loom;
pub mod oneshot;
pub mod mpsc;
pub mod semaphore;
pub mod task;
tokio-sync/src/loom.rs
+36
View File
@@ -0,0 +1,36 @@
pub(crate) mod futures {
pub(crate) use futures::task;
pub(crate) use ::task::AtomicTask;
}
pub(crate) mod sync {
pub(crate) use std::sync::atomic;
use std::cell::UnsafeCell;
pub(crate) struct CausalCell<T>(UnsafeCell<T>);
impl<T> CausalCell<T> {
pub(crate) fn new(data: T) -> CausalCell<T> {
CausalCell(UnsafeCell::new(data))
}
pub(crate) fn with<F, R>(&self, f: F) -> R
where
F: FnOnce(*const T) -> R,
{
f(self.0.get())
}
pub(crate) fn with_mut<F, R>(&self, f: F) -> R
where
F: FnOnce(*mut T) -> R,
{
f(self.0.get())
}
}
}
pub(crate) fn yield_now() {
::std::sync::atomic::spin_loop_hint();
}
tokio-sync/src/mpsc/block.rs
+395
View File
@@ -0,0 +1,395 @@
use loom::{
self,
sync::CausalCell,
sync::atomic::{
AtomicPtr,
AtomicUsize,
},
};
use std::mem::{self, ManuallyDrop};
use std::ops;
use std::ptr::{self, NonNull};
use std::sync::atomic::Ordering::{self, Acquire, Release, AcqRel};
/// A block in a linked list.
///
/// Each block in the list can hold up to `BLOCK_CAP` messages.
pub(crate) struct Block<T> {
/// The start index of this block.
///
/// Slots in this block have indices in `start_index .. start_index + BLOCK_CAP`.
start_index: usize,
/// The next block in the linked list.
next: AtomicPtr<Block<T>>,
/// Bitfield tracking slots that are ready to have their values consumed.
ready_slots: AtomicUsize,
/// The observed `tail_position` value *after* the block has been passed by
/// `block_tail`.
observed_tail_position: CausalCell<usize>,
/// Array containing values pushed into the block. Values are stored in a
/// continuous array in order to improve cache line behavior when reading.
/// The values must be manually dropped.
values: Values<T>,
}
pub(crate) enum Read<T> {
Value(T),
Closed,
}
struct Values<T>([CausalCell<ManuallyDrop<T>>; BLOCK_CAP]);
use super::BLOCK_CAP;
/// Masks an index to get the block identifier
const BLOCK_MASK: usize = !(BLOCK_CAP - 1);
/// Masks an index to get the value offset in a block.
const SLOT_MASK: usize = BLOCK_CAP - 1;
/// Flag tracking that a block has gone through the sender's release routine.
///
/// When this is set, the receiver may consider freeing the block.
const RELEASED: usize = 1 << BLOCK_CAP;
/// Flag tracking all senders dropped.
///
/// When this flag is set, the send half of the channel has closed.
const TX_CLOSED: usize = RELEASED << 1;
/// Mask covering all bits used to track slot readiness.
const READY_MASK: usize = RELEASED - 1;
/// Returns the index of the first slot in the block referenced by `slot_index`.
#[inline(always)]
pub(crate) fn start_index(slot_index: usize) -> usize {
BLOCK_MASK & slot_index
}
/// Returns the offset into the block referenced by `slot_index`.
#[inline(always)]
pub(crate) fn offset(slot_index: usize) -> usize {
SLOT_MASK & slot_index
}
impl<T> Block<T> {
pub(crate) fn new(start_index: usize) -> Block<T> {
Block {
// The absolute index in the channel of the first slot in the block.
start_index,
// Pointer to the next block in the linked list.
next: AtomicPtr::new(ptr::null_mut()),
ready_slots: AtomicUsize::new(0),
observed_tail_position: CausalCell::new(0),
// Value storage
values: unsafe { Values::uninitialized() },
}
}
/// Returns `true` if the block matches the given index
pub(crate) fn is_at_index(&self, index: usize) -> bool {
debug_assert!(offset(index) == 0);
self.start_index == index
}
/// Returns the number of blocks between `self` and the block at the
/// specified index.
///
/// `start_index` must represent a block *after* `self`.
pub(crate) fn distance(&self, other_index: usize) -> usize {
debug_assert!(offset(other_index) == 0);
other_index.wrapping_sub(self.start_index) / BLOCK_CAP
}
/// Read the value at the given offset.
///
/// Returns `None` if the slot is empty.
///
/// # Safety
///
/// To maintain safety, the caller must ensure:
///
/// * No concurrent access to the slot.
pub(crate) unsafe fn read(&self, slot_index: usize) -> Option<Read<T>> {
let offset = offset(slot_index);
let ready_bits = self.ready_slots.load(Acquire);
if !is_ready(ready_bits, offset) {
if is_tx_closed(ready_bits) {
return Some(Read::Closed);
}
return None;
}
// Get the value
let value = self.values[offset].with(|ptr| {
ptr::read(ptr)
});
Some(Read::Value(ManuallyDrop::into_inner(value)))
}
/// Write a value to the block at the given offset.
///
/// # Safety
///
/// To maintain safety, the caller must ensure:
///
/// * The slot is empty.
/// * No concurrent access to the slot.
pub(crate) unsafe fn write(&self, slot_index: usize, value: T) {
// Get the offset into the block
let slot_offset = offset(slot_index);
self.values[slot_offset].with_mut(|ptr| {
ptr::write(ptr, ManuallyDrop::new(value));
});
// Release the value. After this point, the slot ref may no longer
// be used. It is possible for the receiver to free the memory at
// any point.
self.set_ready(slot_offset);
}
/// Signal to the receiver that the sender half of the list is closed.
pub(crate) unsafe fn tx_close(&self) {
self.ready_slots.fetch_or(TX_CLOSED, Release);
}
/// Reset the block to a blank state. This enables reusing blocks in the
/// channel.
///
/// # Safety
///
/// To maintain safety, the caller must ensure:
///
/// * All slots are empty.
/// * The caller holds a unique pointer to the block.
pub(crate) unsafe fn reclaim(&mut self) {
self.start_index = 0;
self.next = AtomicPtr::new(ptr::null_mut());
self.ready_slots = AtomicUsize::new(0);
}
/// Release the block to the rx half for freeing.
///
/// This function is called by the tx half once it can be guaranteed that no
/// more senders will attempt to access the block.
///
/// # Safety
///
/// To maintain safety, the caller must ensure:
///
/// * The block will no longer be accessed by any sender.
pub(crate) unsafe fn tx_release(&self, tail_position: usize) {
// Track the observed tail_position. Any sender targetting a greater
// tail_position is guaranteed to not access this block.
self.observed_tail_position.with_mut(|ptr| *ptr = tail_position);
// Set the released bit, signalling to the receiver that it is safe to
// free the block's memory as soon as all slots **prior** to
// `observed_tail_position` have been filled.
self.ready_slots.fetch_or(RELEASED, Release);
}
/// Mark a slot as ready
fn set_ready(&self, slot: usize) {
let mask = 1 << slot;
self.ready_slots.fetch_or(mask, Release);
}
/// Returns `true` when all slots have their `ready` bits set.
///
/// This indicates that the block is in its final state and will no longer
/// be mutated.
///
/// # Implementation
///
/// The implementation walks each slot checking the `ready` flag. It might
/// be that it would make more sense to coalesce ready flags as bits in a
/// single atomic cell. However, this could have negative impact on cache
/// behavior as there would be many more mutations to a single slot.
pub(crate) fn is_final(&self) -> bool {
self.ready_slots.load(Acquire) & READY_MASK == READY_MASK
}
/// Returns the `observed_tail_position` value, if set
pub(crate) fn observed_tail_position(&self) -> Option<usize> {
if 0 == RELEASED & self.ready_slots.load(Acquire) {
None
} else {
Some(self.observed_tail_position.with(|ptr| unsafe { *ptr }))
}
}
/// Load the next block
pub(crate) fn load_next(&self, ordering: Ordering) -> Option<NonNull<Block<T>>> {
let ret = NonNull::new(self.next.load(ordering));
debug_assert!(unsafe {
ret.map(|block| {
block.as_ref().start_index == self.start_index.wrapping_add(BLOCK_CAP)
}).unwrap_or(true)
});
ret
}
/// Push `block` as the next block in the link.
///
/// Returns Ok if successful, otherwise, a pointer to the next block in
/// the list is returned.
///
/// This requires that the next pointer is null.
///
/// # Ordering
///
/// This performs a compare-and-swap on `next` using AcqRel ordering.
///
/// # Safety
///
/// To maintain safety, the caller must ensure:
///
/// * `block` is not freed until it has been removed from the list.
pub(crate) unsafe fn try_push(&self, block: &mut NonNull<Block<T>>, ordering: Ordering)
-> Result<(), NonNull<Block<T>>>
{
block.as_mut().start_index =
self.start_index.wrapping_add(BLOCK_CAP);
let next_ptr = self.next.compare_and_swap(
ptr::null_mut(), block.as_ptr(), ordering);
match NonNull::new(next_ptr) {
Some(next_ptr) => Err(next_ptr),
None => Ok(()),
}
}
/// Grow the `Block` linked list by allocating and appending a new block.
///
/// The next block in the linked list is returned. This may or may not be
/// the one allocated by the function call.
///
/// # Implementation
///
/// It is assumed that `self.next` is null. A new block is allocated with
/// `start_index` set to be the next block. A compare-and-swap is performed
/// with AcqRel memory ordering. If the compare-and-swap is successful, the
/// newly allocated block is released to other threads walking the block
/// linked list. If the compare-and-swap fails, the current thread acquires
/// the next block in the linked list, allowing the current thread to access
/// the slots.
pub(crate) fn grow(&self) -> NonNull<Block<T>> {
// Create the new block. It is assumed that the block will become the
// next one after `&self`. If this turns out to not be the case,
// `start_index` is updated accordingly.
let new_block = Box::new(
Block::new(self.start_index + BLOCK_CAP));
let mut new_block = unsafe {
NonNull::new_unchecked(Box::into_raw(new_block))
};
// Attempt to store the block. The first compare-and-swap attempt is
// "unrolled" due to minor differences in logic
//
// `AcqRel` is used as the ordering **only** when attempting the
// compare-and-swap on self.next.
//
// If the compare-and-swap fails, then the actual value of the cell is
// returned from this function and accessed by the caller. Given this,
// the memory must be acquired.
//
// `Release` ensures that the newly allocated block is available to
// other threads acquiring the next pointer.
let next = NonNull::new(
self.next.compare_and_swap(
ptr::null_mut(), new_block.as_ptr(), AcqRel));
let next = match next {
Some(next) => next,
None => {
// The compare-and-swap succeeded and the newly allocated block
// is successfully pushed.
return new_block;
}
};
// There already is a next block in the linked list. The newly allocated
// block could be dropped and the discovered next block returned;
// however, that would be wasteful. Instead, the linked list is walked
// by repeatedly attempting to compare-and-swap the pointer into the
// `next` register until the compare-and-swap succeed.
//
// Care is taken to update new_block's start_index field as appropriate.
let mut curr = next;
// TODO: Should this iteration be capped?
loop {
let actual = unsafe {
curr.as_ref().try_push(&mut new_block, AcqRel)
};
curr = match actual {
Ok(_) => {
return next;
}
Err(curr) => curr,
};
// When running outside of loom, this calls `spin_loop_hint`.
loom::yield_now();
}
}
}
/// Returns `true` if the specificed slot has a value ready to be consumed.
fn is_ready(bits: usize, slot: usize) -> bool {
let mask = 1 << slot;
mask == mask & bits
}
/// Returns `true` if the closed flag has been set.
fn is_tx_closed(bits: usize) -> bool {
TX_CLOSED == bits & TX_CLOSED
}
impl<T> Values<T> {
unsafe fn uninitialized() -> Values<T> {
let mut vals = mem::uninitialized();
// When fuzzing, `CausalCell` needs to be initialized.
if_fuzz! {
use std::ptr;
for v in &mut vals {
ptr::write(
v as *mut _,
CausalCell::new(mem::zeroed()));
}
}
Values(vals)
}
}
impl<T> ops::Index<usize> for Values<T> {
type Output = CausalCell<ManuallyDrop<T>>;
fn index(&self, index: usize) -> &Self::Output {
self.0.index(index)
}
}
tokio-sync/src/mpsc/bounded.rs
+243
View File
@@ -0,0 +1,243 @@
use super::chan;
use futures::{Poll, Sink, StartSend, Stream};
use std::fmt;
/// Send values to the associated `Receiver`.
///
/// Instances are created by the [`channel`](fn.channel.html) function.
#[derive(Debug, Clone)]
pub struct Sender<T> {
chan: chan::Tx<T, Semaphore>,
}
/// Receive values from the associated `Sender`.
///
/// Instances are created by the [`channel`](fn.channel.html) function.
#[derive(Debug)]
pub struct Receiver<T> {
/// The channel receiver
chan: chan::Rx<T, Semaphore>,
}
/// Error returned by the `Sender`.
#[derive(Debug)]
pub struct SendError(());
/// Error returned by `Sender::try_send`.
#[derive(Debug)]
pub struct TrySendError<T> {
kind: ErrorKind,
value: T,
}
#[derive(Debug)]
enum ErrorKind {
Closed,
NoCapacity,
}
/// Error returned by `Receiver`.
#[derive(Debug)]
pub struct RecvError(());
/// Create a bounded mpsc channel for communicating between asynchronous tasks,
/// returning the sender/receiver halves.
///
/// All data sent on `Sender` will become available on `Receiver` in the same
/// order as it was sent.
///
/// The `Sender` can be cloned to `send` to the same channel from multiple code
/// locations. Only one `Receiver` is supported.
///
/// If the `Receiver` is disconnected while trying to `send`, the `send` method
/// will return a `SendError`. Similarly, if `Sender` is disconnected while
/// trying to `recv`, the `recv` method will return a `RecvError`.
///
/// # Examples
///
/// ```rust
/// extern crate futures;
/// extern crate tokio;
///
/// use tokio::sync::mpsc::channel;
/// use tokio::prelude::*;
/// use futures::future::lazy;
///
/// # fn some_computation() -> impl Future<Item = (), Error = ()> + Send {
/// # futures::future::ok::<(), ()>(())
/// # }
///
/// tokio::run(lazy(|| {
/// let (tx, rx) = channel(100);
///
/// tokio::spawn({
/// some_computation()
/// .and_then(|value| {
/// tx.send(value)
/// .map_err(|_| ())
/// })
/// .map(|_| ())
/// .map_err(|_| ())
/// });
///
/// rx.for_each(|value| {
/// println!("got value = {:?}", value);
/// Ok(())
/// })
/// .map(|_| ())
/// .map_err(|_| ())
/// }));
/// ```
pub fn channel<T>(buffer: usize) -> (Sender<T>, Receiver<T>) {
let semaphore = (::semaphore::Semaphore::new(buffer), buffer);
let (tx, rx) = chan::channel(semaphore);
let tx = Sender::new(tx);
let rx = Receiver::new(rx);
(tx, rx)
}
/// Channel semaphore is a tuple of the semaphore implementation and a `usize`
/// representing the channel bound.
type Semaphore = (::semaphore::Semaphore, usize);
impl<T> Receiver<T> {
pub(crate) fn new(chan: chan::Rx<T, Semaphore>) -> Receiver<T> {
Receiver { chan }
}
/// Closes the receiving half of a channel, without dropping it.
///
/// 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) {
self.chan.close();
}
}
impl<T> Stream for Receiver<T> {
type Item = T;
type Error = RecvError;
fn poll(&mut self) -> Poll<Option<T>, Self::Error> {
self.chan.recv()
.map_err(|_| RecvError(()))
}
}
impl<T> Sender<T> {
pub(crate) fn new(chan: chan::Tx<T, Semaphore>) -> Sender<T> {
Sender { chan }
}
/// Check if the `Sender` is ready to handle a value.
///
/// Polls the channel to determine if there is guaranteed capacity to send
/// at least one item without waiting.
///
/// When `poll_ready` returns `Ready`, the channel reserves capacity for one
/// message for this `Sender` instance. The capacity is held until a message
/// is send or the `Sender` instance is dropped. Callers should ensure a
/// message is sent in a timely fashion in order to not starve other
/// `Sender` instances.
///
/// # Return value
///
/// This method returns:
///
/// - `Ok(Async::Ready(_))` if capacity is reserved for a single message.
/// - `Ok(Async::NotReady)` if the channel may not have capacity, in which
/// case the current task is queued to be notified once
/// capacity is available;
/// - `Err(SendError)` if the receiver has been dropped.
pub fn poll_ready(&mut self) -> Poll<(), SendError> {
self.chan.poll_ready()
.map_err(|_| SendError(()))
}
/// Attempts to send a message on this `Sender`, returning the message
/// if there was an error.
pub fn try_send(&mut self, message: T) -> Result<(), TrySendError<T>> {
self.chan.try_send(message)?;
Ok(())
}
}
impl<T> Sink for Sender<T> {
type SinkItem = T;
type SinkError = SendError;
fn start_send(&mut self, msg: T) -> StartSend<T, Self::SinkError> {
use futures::AsyncSink;
use futures::Async::*;
match self.poll_ready()? {
Ready(_) => {
self.try_send(msg).map_err(|_| SendError(()))?;
Ok(AsyncSink::Ready)
}
NotReady => {
Ok(AsyncSink::NotReady(msg))
}
}
}
fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
use futures::Async::Ready;
Ok(Ready(()))
}
fn close(&mut self) -> Poll<(), Self::SinkError> {
use futures::Async::Ready;
Ok(Ready(()))
}
}
// ===== impl SendError =====
impl fmt::Display for SendError {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::error::Error;
write!(fmt, "{}", self.description())
}
}
impl ::std::error::Error for SendError {
fn description(&self) -> &str {
"channel closed"
}
}
// ===== impl TrySendError =====
impl<T: fmt::Debug> fmt::Display for TrySendError<T> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::error::Error;
write!(fmt, "{}", self.description())
}
}
impl<T: fmt::Debug> ::std::error::Error for TrySendError<T> {
fn description(&self) -> &str {
match self.kind {
ErrorKind::Closed => "channel closed",
ErrorKind::NoCapacity => "no available capacity",
}
}
}
impl<T> From<(T, chan::TrySendError)> for TrySendError<T> {
fn from((value, err): (T, chan::TrySendError)) -> TrySendError<T> {
TrySendError {
value,
kind: match err {
chan::TrySendError::Closed => ErrorKind::Closed,
chan::TrySendError::NoPermits => ErrorKind::NoCapacity,
}
}
}
}
tokio-sync/src/mpsc/chan.rs
+402
View File
@@ -0,0 +1,402 @@
use super::list;
use futures::Poll;
use futures::task::AtomicTask;
use std::cell::UnsafeCell;
use std::process;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{AcqRel, Relaxed};
/// Channel sender
#[derive(Debug)]
pub(crate) struct Tx<T, S: Semaphore> {
inner: Arc<Chan<T, S>>,
permit: S::Permit,
}
/// Channel receiver
#[derive(Debug)]
pub(crate) struct Rx<T, S: Semaphore> {
inner: Arc<Chan<T, S>>,
}
#[derive(Debug, Eq, PartialEq)]
pub(crate) enum TrySendError {
Closed,
NoPermits,
}
pub(crate) trait Semaphore: Sync {
type Permit;
fn new_permit() -> Self::Permit;
/// The permit is dropped without a value being sent. In this case, the
/// permit must be returned to the semaphore.
fn drop_permit(&self, permit: &mut Self::Permit);
fn is_idle(&self) -> bool;
fn add_permit(&self);
fn poll_acquire(&self, permit: &mut Self::Permit) -> Poll<(), ()>;
fn try_acquire(&self, permit: &mut Self::Permit) -> Result<(), TrySendError>;
/// A value was sent into the channel and the permit held by `tx` is
/// dropped. In this case, the permit should not immeditely be returned to
/// the semaphore. Instead, the permit is returnred to the semaphore once
/// the sent value is read by the rx handle.
fn forget(&self, permit: &mut Self::Permit);
fn close(&self);
}
#[derive(Debug)]
struct Chan<T, S> {
/// Handle to the push half of the lock-free list.
tx: list::Tx<T>,
/// Coordinates access to channel's capacity.
semaphore: S,
/// Receiver task. Notified when a value is pushed into the channel.
rx_task: AtomicTask,
/// Tracks the number of outstanding sender handles.
///
/// When this drops to zero, the send half of the channel is closed.
tx_count: AtomicUsize,
/// Only accessed by `Rx` handle.
rx_fields: UnsafeCell<RxFields<T>>,
}
/// Fields only accessed by `Rx` handle.
#[derive(Debug)]
struct RxFields<T> {
/// Channel receiver. This field is only accessed by the `Receiver` type.
list: list::Rx<T>,
/// `true` if `Rx::close` is called.
rx_closed: bool,
}
unsafe impl<T: Send, S: Send> Send for Chan<T, S> {}
unsafe impl<T: Send, S: Sync> Sync for Chan<T, S> {}
pub(crate) fn channel<T, S>(semaphore: S) -> (Tx<T, S>, Rx<T, S>)
where
S: Semaphore,
{
let (tx, rx) = list::channel();
let chan = Arc::new(Chan {
tx,
semaphore,
rx_task: AtomicTask::new(),
tx_count: AtomicUsize::new(1),
rx_fields: UnsafeCell::new(RxFields {
list: rx,
rx_closed: false,
}),
});
(Tx::new(chan.clone()), Rx::new(chan))
}
// ===== impl Tx =====
impl<T, S> Tx<T, S>
where
S: Semaphore,
{
fn new(chan: Arc<Chan<T, S>>) -> Tx<T, S> {
Tx {
inner: chan,
permit: S::new_permit(),
}
}
/// TODO: Docs
pub(crate) fn poll_ready(&mut self) -> Poll<(), ()> {
self.inner.semaphore.poll_acquire(&mut self.permit)
}
/// Send a message and notify the receiver.
pub(crate) fn try_send(&mut self, value: T) -> Result<(), (T, TrySendError)> {
if let Err(e) = self.inner.semaphore.try_acquire(&mut self.permit) {
return Err((value, e));
}
// Push the value
self.inner.tx.push(value);
// Notify the rx task
self.inner.rx_task.notify();
// Release the permit
self.inner.semaphore.forget(&mut self.permit);
Ok(())
}
}
impl<T, S> Clone for Tx<T, S>
where
S: Semaphore,
{
fn clone(&self) -> Tx<T, S> {
// Using a Relaxed ordering here is sufficient as the caller holds a
// strong ref to `self`, preventing a concurrent decrement to zero.
self.inner.tx_count.fetch_add(1, Relaxed);
Tx {
inner: self.inner.clone(),
permit: S::new_permit(),
}
}
}
impl<T, S> Drop for Tx<T, S>
where
S: Semaphore,
{
fn drop(&mut self) {
self.inner.semaphore.drop_permit(&mut self.permit);
if self.inner.tx_count.fetch_sub(1, AcqRel) != 1 {
return;
}
// Close the list, which sends a `Close` message
self.inner.tx.close();
// Notify the receiver
self.inner.rx_task.notify();
}
}
// ===== impl Rx =====
impl<T, S> Rx<T, S>
where
S: Semaphore,
{
fn new(chan: Arc<Chan<T, S>>) -> Rx<T, S> {
Rx { inner: chan }
}
pub(crate) fn close(&mut self) {
let rx_fields = unsafe { &mut *self.inner.rx_fields.get() };
if rx_fields.rx_closed {
return;
}
rx_fields.rx_closed = true;
self.inner.semaphore.close();
}
/// Receive the next value
pub(crate) fn recv(&mut self) -> Poll<Option<T>, ()> {
use super::block::Read::*;
use futures::Async::*;
let rx_fields = unsafe { &mut *self.inner.rx_fields.get() };
macro_rules! try_recv {
() => {
match rx_fields.list.pop(&self.inner.tx) {
Some(Value(value)) => {
self.inner.semaphore.add_permit();
return Ok(Ready(Some(value)));
}
Some(Closed) => {
// TODO: This check may not be required as it most
// likely can only return `true` at this point. A
// channel is closed when all tx handles are dropped.
// Dropping a tx handle releases memory, which ensures
// that if dropping the tx handle is visible, then all
// messages sent are also visible.
assert!(self.inner.semaphore.is_idle());
return Ok(Ready(None));
}
None => {} // fall through
}
}
}
try_recv!();
self.inner.rx_task.register();
// It is possible that a value was pushed between attempting to read and
// registering the task, so we have to check the channel a second time
// here.
try_recv!();
debug!("recv; rx_closed = {:?}; is_idle = {:?}",
rx_fields.rx_closed, self.inner.semaphore.is_idle());
if rx_fields.rx_closed && self.inner.semaphore.is_idle() {
Ok(Ready(None))
} else {
Ok(NotReady)
}
}
}
impl<T, S> Drop for Rx<T, S>
where
S: Semaphore,
{
fn drop(&mut self) {
use super::block::Read::Value;
self.close();
let rx_fields = unsafe { &mut *self.inner.rx_fields.get() };
while let Some(Value(_)) = rx_fields.list.pop(&self.inner.tx) {
self.inner.semaphore.add_permit();
}
}
}
// ===== impl Chan =====
impl<T, S> Drop for Chan<T, S> {
fn drop(&mut self) {
use super::block::Read::Value;
let rx_fields = unsafe { &mut *self.rx_fields.get() };
while let Some(Value(_)) = rx_fields.list.pop(&self.tx) {
}
}
}
use semaphore::TryAcquireError;
impl From<TryAcquireError> for TrySendError {
fn from(src: TryAcquireError) -> TrySendError {
if src.is_closed() {
TrySendError::Closed
} else if src.is_no_permits() {
TrySendError::NoPermits
} else {
unreachable!();
}
}
}
// ===== impl Semaphore for (::Semaphore, capacity) =====
use semaphore::Permit;
impl Semaphore for (::semaphore::Semaphore, usize) {
type Permit = Permit;
fn new_permit() -> Permit {
Permit::new()
}
fn drop_permit(&self, permit: &mut Permit) {
if permit.is_acquired() {
permit.release(&self.0);
}
}
fn add_permit(&self) {
self.0.add_permits(1)
}
fn is_idle(&self) -> bool {
self.0.available_permits() == self.1
}
fn poll_acquire(&self, permit: &mut Permit) -> Poll<(), ()> {
permit.poll_acquire(&self.0)
.map_err(|_| ())
}
fn try_acquire(&self, permit: &mut Permit) -> Result<(), TrySendError> {
permit.try_acquire(&self.0)?;
Ok(())
}
fn forget(&self, permit: &mut Self::Permit) {
permit.forget()
}
fn close(&self) {
self.0.close();
}
}
// ===== impl Semaphore for AtomicUsize =====
use std::sync::atomic::Ordering::{Acquire, Release};
use std::usize;
impl Semaphore for AtomicUsize {
type Permit = ();
fn new_permit() {
}
fn drop_permit(&self, _permit: &mut ()) {
}
fn add_permit(&self) {
let prev = self.fetch_sub(2, Release);
if prev >> 1 == 0 {
// Something went wrong
process::abort();
}
}
fn is_idle(&self) -> bool {
self.load(Acquire) >> 1 == 0
}
fn poll_acquire(&self, permit: &mut ()) -> Poll<(), ()> {
use futures::Async::Ready;
self.try_acquire(permit)
.map(Ready)
.map_err(|_| ())
}
fn try_acquire(&self, _permit: &mut ()) -> Result<(), TrySendError> {
let mut curr = self.load(Acquire);
loop {
if curr & 1 == 1 {
return Err(TrySendError::Closed);
}
if curr == usize::MAX ^ 1 {
// Overflowed the ref count. There is no safe way to recover, so
// abort the process. In practice, this should never happen.
process::abort()
}
match self.compare_exchange(curr, curr + 2, AcqRel, Acquire) {
Ok(_) => return Ok(()),
Err(actual) => {
curr = actual;
}
}
}
}
fn forget(&self, _permit: &mut ()) {
}
fn close(&self) {
self.fetch_or(1, Release);
}
}
tokio-sync/src/mpsc/list.rs
+329
View File
@@ -0,0 +1,329 @@
//! A concurrent, lock-free, FIFO list.
use super::block::{self, Block};
use loom::{
self,
sync::atomic::{AtomicUsize, AtomicPtr},
};
use std::fmt;
use std::ptr::NonNull;
use std::sync::atomic::Ordering::{Acquire, Release, AcqRel, Relaxed};
/// List queue transmit handle
pub(crate) struct Tx<T> {
/// Tail in the `Block` mpmc list.
block_tail: AtomicPtr<Block<T>>,
/// Position to push the next message. This reference a block and offset
/// into the block.
tail_position: AtomicUsize,
}
/// List queue receive handle
pub(crate) struct Rx<T> {
/// Pointer to the block being processed
head: NonNull<Block<T>>,
/// Next slot index to process
index: usize,
/// Pointer to the next block pending release
free_head: NonNull<Block<T>>,
}
pub(crate) fn channel<T>() -> (Tx<T>, Rx<T>) {
// Create the initial block shared between the tx and rx halves.
let initial_block = Box::new(Block::new(0));
let initial_block_ptr = Box::into_raw(initial_block);
let tx = Tx {
block_tail: AtomicPtr::new(initial_block_ptr),
tail_position: AtomicUsize::new(0),
};
let head = NonNull::new(initial_block_ptr).unwrap();
let rx = Rx {
head,
index: 0,
free_head: head,
};
(tx, rx)
}
impl<T> Tx<T> {
/// Push a value into the list.
pub(crate) fn push(&self, value: T) {
// First, claim a slot for the value. `Acquire` is used here to
// synchronize with the `fetch_add` in `free_blocks`.
let slot_index = self.tail_position
.fetch_add(1, Acquire);
// Load the current block and write the value
let block = self.find_block(slot_index);
unsafe {
// Write the value to the block
block.as_ref().write(slot_index, value);
}
}
/// Close the send half of the list
///
/// Similar process as pushing a value, but instead of writing the value &
/// setting the ready flag, the TX_CLOSED flag is set on the block.
pub(crate) fn close(&self) {
// First, claim a slot for the value. This is the last slot that will be
// claimed.
let slot_index = self.tail_position
.fetch_add(1, Acquire);
let block = self.find_block(slot_index);
unsafe {
block.as_ref().tx_close()
}
}
fn find_block(&self, slot_index: usize) -> NonNull<Block<T>> {
// The start index of the block that contains `index`.
let start_index = block::start_index(slot_index);
// The index offset into the block
let offset = block::offset(slot_index);
// Load the current head of the block
let mut block_ptr = self.block_tail.load(Acquire);
let block = unsafe { &*block_ptr };
// Calculate the distance between the tail ptr and the target block
let distance = block.distance(start_index);
// Decide if this call to `find_block` should attempt to update the
// `block_tail` pointer.
//
// Updating `block_tail` is not always performed in order to reduce
// contention.
//
// When set, as the routine walks the linked list, it attempts to update
// `block_tail`. If the update cannot be performed, `try_updating_tail`
// is unset.
let mut try_updating_tail = distance > offset;
// Walk the linked list of blocks until the block with `start_index` is
// found.
loop {
let block = unsafe { &(*block_ptr) };
if block.is_at_index(start_index) {
return unsafe { NonNull::new_unchecked(block_ptr) };
}
let next_block = block.load_next(Acquire)
// There is no allocated next block, grow the linked list.
.unwrap_or_else(|| block.grow());
// If the block is **not** final, then the tail pointer cannot be
// advanced any more.
try_updating_tail &= block.is_final();
if try_updating_tail {
// Advancing `block_tail` must happen when walking the linked
// list. `block_tail` may not advance passed any blocks that are
// not "final". At the point a block is finalized, it is unknown
// if there are any prior blocks that are unfinalized, which
// makes it impossible to advance `block_tail`.
//
// While walking the linked list, `block_tail` can be advanced
// as long as finalized blocks are traversed.
//
// Release ordering is used to ensure that any subsequent reads
// are able to see the memory pointed to by `block_tail`.
//
// Acquire is not needed as any "actual" value is not accessed.
// At this point, the linked list is walked to acquire blocks.
let actual = self.block_tail.compare_and_swap(
block_ptr, next_block.as_ptr(), Release);
if actual == block_ptr {
// Synchronize with any senders
let tail_position =
self.tail_position.fetch_add(0, Release);
unsafe { block.tx_release(tail_position); }
} else {
// A concurrent sender is also working on advancing
// `block_tail` and this thread is falling behind.
//
// Stop trying to advance the tail pointer
try_updating_tail = false;
}
}
block_ptr = next_block.as_ptr();
loom::yield_now();
}
}
pub(crate) unsafe fn reclaim_block(&self, mut block: NonNull<Block<T>>) {
// The block has been removed from the linked list and ownership
// is reclaimed.
//
// Before dropping the block, see if it can be reused by
// inserting it back at the end of the linked list.
//
// First, reset the data
block.as_mut().reclaim();
let mut reused = false;
// Attempt to insert the block at the end
//
// Walk at most three times
//
let curr_ptr = self.block_tail.load(Acquire);
// The pointer can never be null
debug_assert!(!curr_ptr.is_null());
let mut curr = NonNull::new_unchecked(curr_ptr);
// TODO: Unify this logic with Block::grow
for _ in 0..3 {
match curr.as_ref().try_push(&mut block, AcqRel) {
Ok(_) => {
reused = true;
break;
}
Err(next) => {
curr = next;
}
}
}
if !reused {
let _ = Box::from_raw(block.as_ptr());
}
}
}
impl<T: fmt::Debug> fmt::Debug for Tx<T> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::sync::atomic::Ordering::Relaxed;
fmt.debug_struct("Tx")
.field("block_tail", &self.block_tail.load(Relaxed))
.field("tail_position", &self.tail_position.load(Relaxed))
.finish()
}
}
impl<T> Rx<T> {
/// Pop the next value off the queue
pub(crate) fn pop(&mut self, tx: &Tx<T>) -> Option<block::Read<T>> {
// Advance `head`, if needed
if !self.try_advancing_head() {
debug!("+ !self.try_advancing_head() -> false");
return None;
}
self.free_blocks(tx);
unsafe {
let block = self.head.as_ref();
let ret = block.read(self.index);
if let Some(block::Read::Value(..)) = ret {
self.index = self.index.wrapping_add(1);
}
ret
}
}
/// Try advancing the block pointer to the block referenced by `self.index`.
///
/// Returns `true` if successful, `false` if there is no next block to load.
fn try_advancing_head(&mut self) -> bool {
let block_index = block::start_index(self.index);
loop {
let next_block = {
let block = unsafe { self.head.as_ref() };
if block.is_at_index(block_index) {
return true;
}
block.load_next(Acquire)
};
let next_block = match next_block {
Some(next_block) => next_block,
None => {
return false;
}
};
self.head = next_block;
loom::yield_now();
}
}
fn free_blocks(&mut self, tx: &Tx<T>) {
debug!("+ free_blocks()");
while self.free_head != self.head {
unsafe {
// Get a handle to the block that will be freed and update
// `free_head` to point to the next block.
let block = self.free_head;
let observed_tail_position =
block.as_ref().observed_tail_position();
let required_index = match observed_tail_position {
Some(i) => i,
None => return,
};
if required_index > self.index {
return;
}
// We may read the next pointer with `Relaxed` ordering as it is
// guaranteed that the `free_blocks` routine trails the `recv`
// routine. Any memory accessed by `free_blocks` has already
// been acquired by `recv`.
let next_block =
block.as_ref().load_next(Relaxed);
// Update the free list head
self.free_head = next_block.unwrap();
// Push the emptied block onto the back of the queue, making it
// available to senders.
tx.reclaim_block(block);
}
loom::yield_now();
}
}
}
impl<T: fmt::Debug> fmt::Debug for Rx<T> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("Rx")
.field("head", &self.head)
.field("index", &self.index)
.field("free_head", &self.free_head)
.finish()
}
}
tokio-sync/src/mpsc/mod.rs
+79
View File
@@ -0,0 +1,79 @@
//! A multi-producer, single-consumer queue for sending values across
//! asynchronous tasks.
//!
//! Similar to `std`, channel creation provides [`Receiver`] and [`Sender`]
//! handles. [`Receiver`] implements `Stream` and allows a task to read values
//! out of the channel. If there is no message to read, the current task will be
//! notified when a new value is sent. [`Sender`] implements the `Sink` trait
//! and allows sending messages into the channel. If the channel is at capacity,
//! the send is rejected and the task will be notified when additional capacity
//! is available. In other words, the channel provides backpressure.
//!
//! Unbounded channels are also available using the `unbounded_channel`
//! constructor.
//!
//! # Disconnection
//!
//! When all [`Sender`] handles have been dropped, it is no longer
//! possible to send values into the channel. This is considered the termination
//! event of the stream. As such, `Receiver::poll` returns `Ok(Ready(None))`.
//!
//! If the [`Receiver`] handle is dropped, then messages can no longer
//! be read out of the channel. In this case, all further attempts to send will
//! result in an error.
//!
//! # Clean Shutdown
//!
//! When the [`Receiver`] is dropped, it is possible for unprocessed messages to
//! remain in the channel. Instead, it is usually desirable to perform a "clean"
//! shutdown. To do this, the receiver first calls `close`, which will prevent
//! any further messages to be sent into the channel. Then, the receiver
//! consumes the channel to completion, at which point the receiver can be
//! dropped.
//!
//! [`Sender`]: struct.Sender.html
//! [`Receiver`]: struct.Receiver.html
mod block;
mod bounded;
mod chan;
mod list;
mod unbounded;
pub use self::bounded::{
channel,
Receiver,
Sender
};
pub use self::unbounded::{
unbounded_channel,
UnboundedReceiver,
UnboundedSender,
};
pub mod error {
//! Channel error types
pub use super::bounded::{
SendError,
TrySendError,
RecvError,
};
pub use super::unbounded::{
UnboundedSendError,
UnboundedTrySendError,
UnboundedRecvError,
};
}
/// The number of values a block can contain.
///
/// This value must be a power of 2. It also must be smaller than the number of
/// bits in `usize`.
#[cfg(target_pointer_width = "64")]
const BLOCK_CAP: usize = 32;
#[cfg(not(target_pointer_width = "64"))]
const BLOCK_CAP: usize = 16;
tokio-sync/src/mpsc/rx.rs
+1
View File
@@ -0,0 +1 @@
tokio-sync/src/mpsc/tx.rs
View File
tokio-sync/src/mpsc/unbounded.rs
+157
View File
@@ -0,0 +1,157 @@
use super::chan;
use loom::sync::atomic::AtomicUsize;
use futures::{Poll, Sink, StartSend, Stream};
use std::fmt;
/// Send values to the associated `UnboundedReceiver`.
///
/// Instances are created by the
/// [`unbounded_channel`](fn.unbounded_channel.html) function.
#[derive(Debug, Clone)]
pub struct UnboundedSender<T> {
chan: chan::Tx<T, Semaphore>,
}
/// Receive values from the associated `UnboundedSender`.
///
/// Instances are created by the
/// [`unbounded_channel`](fn.unbounded_channel.html) function.
#[derive(Debug)]
pub struct UnboundedReceiver<T> {
/// The channel receiver
chan: chan::Rx<T, Semaphore>,
}
/// Error returned by the `UnboundedSender`.
#[derive(Debug)]
pub struct UnboundedSendError(());
/// Error returned by `UnboundedSender::try_send`.
#[derive(Debug)]
pub struct UnboundedTrySendError<T>(T);
/// Error returned by `UnboundedReceiver`.
#[derive(Debug)]
pub struct UnboundedRecvError(());
/// Create an unbounded mpsc channel for communicating between asynchronous
/// tasks.
///
/// A `send` on this channel will always succeed as long as the receive half has
/// not been closed. If the receiver falls behind, messages will be arbitrarily
/// buffered.
///
/// **Note** that the amount of available system memory is an implicit bound to
/// the channel. Using an `unbounded` channel has the ability of causing the
/// process to run out of memory. In this case, the process will be aborted.
pub fn unbounded_channel<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
let (tx, rx) = chan::channel(AtomicUsize::new(0));
let tx = UnboundedSender::new(tx);
let rx = UnboundedReceiver::new(rx);
(tx, rx)
}
/// No capacity
type Semaphore = AtomicUsize;
impl<T> UnboundedReceiver<T> {
pub(crate) fn new(chan: chan::Rx<T, Semaphore>) -> UnboundedReceiver<T> {
UnboundedReceiver { chan }
}
/// Closes the receiving half of a channel, without dropping it.
///
/// 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) {
self.chan.close();
}
}
impl<T> Stream for UnboundedReceiver<T> {
type Item = T;
type Error = UnboundedRecvError;
fn poll(&mut self) -> Poll<Option<T>, Self::Error> {
self.chan.recv()
.map_err(|_| UnboundedRecvError(()))
}
}
impl<T> UnboundedSender<T> {
pub(crate) fn new(chan: chan::Tx<T, Semaphore>) -> UnboundedSender<T> {
UnboundedSender { chan }
}
/// Attempts to send a message on this `UnboundedSender` without blocking.
pub fn try_send(&mut self, message: T)
-> Result<(), UnboundedTrySendError<T>>
{
self.chan.try_send(message)?;
Ok(())
}
}
impl<T> Sink for UnboundedSender<T> {
type SinkItem = T;
type SinkError = UnboundedSendError;
fn start_send(&mut self, msg: T) -> StartSend<T, Self::SinkError> {
use futures::AsyncSink;
self.try_send(msg).map_err(|_| UnboundedSendError(()))?;
Ok(AsyncSink::Ready)
}
fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
use futures::Async::Ready;
Ok(Ready(()))
}
fn close(&mut self) -> Poll<(), Self::SinkError> {
use futures::Async::Ready;
Ok(Ready(()))
}
}
// ===== impl UnboundedSendError =====
impl fmt::Display for UnboundedSendError {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::error::Error;
write!(fmt, "{}", self.description())
}
}
impl ::std::error::Error for UnboundedSendError {
fn description(&self) -> &str {
"channel closed"
}
}
// ===== impl TrySendError =====
impl<T: fmt::Debug> fmt::Display for UnboundedTrySendError<T> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::error::Error;
write!(fmt, "{}", self.description())
}
}
impl<T: fmt::Debug> ::std::error::Error for UnboundedTrySendError<T> {
fn description(&self) -> &str {
"channel closed"
}
}
impl<T> From<(T, chan::TrySendError)> for UnboundedTrySendError<T> {
fn from((value, err): (T, chan::TrySendError)) -> UnboundedTrySendError<T> {
assert_eq!(chan::TrySendError::Closed, err);
UnboundedTrySendError(value)
}
}
tokio-sync/src/oneshot.rs
+509
View File
@@ -0,0 +1,509 @@
//! A channel for sending a single message between asynchronous tasks.
use loom::{
futures::task::{self, Task},
sync::CausalCell,
sync::atomic::AtomicUsize,
};
use futures::{Async, Future, Poll};
use std::fmt;
use std::mem::{self, ManuallyDrop};
use std::sync::Arc;
use std::sync::atomic::Ordering::{self, Acquire, AcqRel};
/// Sends a value to the associated `Receiver`.
///
/// Instances are created by the [`channel`](fn.channel.html) function.
#[derive(Debug)]
pub struct Sender<T> {
inner: Option<Arc<Inner<T>>>,
}
/// Receive a value from the associated `Sender`.
///
/// Instances are created by the [`channel`](fn.channel.html) function.
#[derive(Debug)]
pub struct Receiver<T> {
inner: Option<Arc<Inner<T>>>,
}
pub mod error {
//! Oneshot error types
/// Error returned by the `Future` implementation for `Receiver`.
#[derive(Debug)]
pub struct RecvError(pub(super) ());
/// Error returned by the `try_recv` function on `Receiver`.
#[derive(Debug)]
pub struct TryRecvError(pub(super) ());
}
use self::error::*;
struct Inner<T> {
/// Manages the state of the inner cell
state: AtomicUsize,
/// The value. This is set by `Sender` and read by `Receiver`. The state of
/// the cell is tracked by `state`.
value: CausalCell<Option<T>>,
/// The task to notify when the receiver drops without consuming the value.
tx_task: CausalCell<ManuallyDrop<Task>>,
/// The task to notify when the value is sent.
rx_task: CausalCell<ManuallyDrop<Task>>,
}
#[derive(Clone, Copy)]
struct State(usize);
/// Create a new one-shot channel for sending single values across asynchronous
/// tasks.
///
/// The function returns separate "send" and "receive" handles. The `Sender`
/// handle is used by the producer to send the value. The `Receiver` handle is
/// used by the consumer to receive the value.
///
/// Each handle can be used on separate tasks.
///
/// # Examples
///
/// ```
/// extern crate futures;
/// extern crate tokio;
///
/// use tokio::sync::oneshot;
/// use futures::Future;
/// use std::thread;
///
/// let (sender, receiver) = oneshot::channel::<i32>();
///
/// # let t =
/// thread::spawn(|| {
/// let future = receiver.map(|i| {
/// println!("got: {:?}", i);
/// });
/// // ...
/// # return future;
/// });
///
/// sender.send(3).unwrap();
/// # t.join().unwrap().wait().unwrap();
/// ```
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
let inner = Arc::new(Inner {
state: AtomicUsize::new(State::new().as_usize()),
value: CausalCell::new(None),
tx_task: CausalCell::new(ManuallyDrop::new(unsafe { mem::uninitialized() })),
rx_task: CausalCell::new(ManuallyDrop::new(unsafe { mem::uninitialized() })),
});
let tx = Sender { inner: Some(inner.clone()) };
let rx = Receiver { inner: Some(inner) };
(tx, rx)
}
impl<T> Sender<T> {
/// Completes this oneshot with a successful result.
///
/// The function consumes `self` and notifies the `Receiver` handle that a
/// value is ready to be received.
///
/// If the value is successfully enqueued for the remote end to receive,
/// 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(mut self, t: T) -> Result<(), T> {
let inner = self.inner.take().unwrap();
inner.value.with_mut(|ptr| {
unsafe { *ptr = Some(t); }
});
if !inner.complete() {
return Err(inner.value.with_mut(|ptr| {
unsafe { (*ptr).take() }.unwrap()
}));
}
Ok(())
}
/// Check if the associated [`Receiver`] handle has been dropped.
///
/// # Return values
///
/// If `Ok(Ready)` is returned then the associated `Receiver` has been
/// dropped, which means any work required for sending should be canceled.
///
/// If `Ok(NotReady)` is returned then the associated `Receiver` is still
/// alive and may be able to receive a message if sent. The current task is
/// registered to receive a notification if the `Receiver` handle goes away.
///
/// [`Receiver`]: struct.Receiver.html
pub fn poll_close(&mut self) -> Poll<(), ()> {
let inner = self.inner.as_ref().unwrap();
let mut state = State::load(&inner.state, Acquire);
if state.is_closed() {
return Ok(Async::Ready(()));
}
if state.is_tx_task_set() {
let tx_task = unsafe { inner.tx_task() };
if !tx_task.will_notify_current() {
state = State::unset_tx_task(&inner.state);
if state.is_closed() {
return Ok(Async::Ready(()));
}
}
}
if !state.is_tx_task_set() {
// Attempt to set the task
unsafe { inner.set_tx_task(); }
// Update the state
state = State::set_tx_task(&inner.state);
if state.is_closed() {
return Ok(Async::Ready(()));
}
}
Ok(Async::NotReady)
}
/// Check if the associated [`Receiver`] handle has been dropped.
///
/// Unlike [`poll_close`], this function does not register a task for
/// wakeup upon close.
///
/// [`Receiver`]: struct.Receiver.html
/// [`poll_close`]: struct.Sender.html#method.poll_close
pub fn is_closed(&self) -> bool {
let inner = self.inner.as_ref().unwrap();
let state = State::load(&inner.state, Acquire);
state.is_closed()
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
if let Some(inner) = self.inner.as_ref() {
inner.complete();
}
}
}
impl<T> Receiver<T> {
/// Prevent the associated [`Sender`] handle from sending a value.
///
/// Any `send` operation which happens after calling `close` is guaranteed
/// to fail. After calling `close`, `Receiver::poll`] should be called to
/// receive a value if one was sent **before** the call to `close`
/// completed.
///
/// [`Sender`]: struct.Sender.html
pub fn close(&mut self) {
let inner = self.inner.as_ref().unwrap();
inner.close();
}
/// Attempts to receive a value outside of the context of a task.
///
/// Does not register a task if no value has been sent.
///
/// A return value of `None` must be considered immediately stale (out of
/// date) unless [`close`] has been called first.
///
/// Returns an error if the sender was dropped.
///
/// [`close`]: #method.close
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
let result = if let Some(inner) = self.inner.as_ref() {
let state = State::load(&inner.state, Acquire);
if state.is_complete() {
match unsafe { inner.consume_value() } {
Some(value) => Ok(value),
None => Err(TryRecvError(())),
}
} else if state.is_closed() {
Err(TryRecvError(()))
} else {
// Not ready, this does not clear `inner`
return Err(TryRecvError(()));
}
} else {
panic!("called after complete");
};
self.inner = None;
result
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
if let Some(inner) = self.inner.as_ref() {
inner.close();
}
}
}
impl<T> Future for Receiver<T> {
type Item = T;
type Error = RecvError;
fn poll(&mut self) -> Poll<T, RecvError> {
use futures::Async::{Ready, NotReady};
// If `inner` is `None`, then `poll()` has already completed.
let ret = if let Some(inner) = self.inner.as_ref() {
match inner.poll_recv() {
Ok(Ready(v)) => Ok(Ready(v)),
Ok(NotReady) => return Ok(NotReady),
Err(e) => Err(e),
}
} else {
panic!("called after complete");
};
self.inner = None;
ret
}
}
impl<T> Inner<T> {
fn complete(&self) -> bool {
let prev = State::set_complete(&self.state);
if prev.is_closed() {
return false;
}
if prev.is_rx_task_set() {
let rx_task = unsafe { self.rx_task() };
rx_task.notify();
}
true
}
fn poll_recv(&self) -> Poll<T, RecvError> {
use futures::Async::{Ready, NotReady};
// Load the state
let mut state = State::load(&self.state, Acquire);
if state.is_complete() {
match unsafe { self.consume_value() } {
Some(value) => Ok(Ready(value)),
None => Err(RecvError(())),
}
} else if state.is_closed() {
Err(RecvError(()))
} else {
if state.is_rx_task_set() {
let rx_task = unsafe { self.rx_task() };
// Check if the task is still the same
if !rx_task.will_notify_current() {
// Unset the task
state = State::unset_rx_task(&self.state);
if state.is_complete() {
return match unsafe { self.consume_value() } {
Some(value) => Ok(Ready(value)),
None => Err(RecvError(())),
};
}
}
}
if !state.is_rx_task_set() {
// Attempt to set the task
unsafe { self.set_rx_task(); }
// Update the state
state = State::set_rx_task(&self.state);
if state.is_complete() {
match unsafe { self.consume_value() } {
Some(value) => Ok(Ready(value)),
None => Err(RecvError(())),
}
} else {
return Ok(NotReady);
}
} else {
return Ok(NotReady);
}
}
}
/// Called by `Receiver` to indicate that the value will never be received.
fn close(&self) {
let prev = State::set_closed(&self.state);
if prev.is_tx_task_set() && !prev.is_complete() {
let tx_task = unsafe { self.tx_task() };
tx_task.notify();
}
}
/// Consume the value. This function does not check `state`.
unsafe fn consume_value(&self) -> Option<T> {
self.value.with_mut(|ptr| {
(*ptr).take()
})
}
unsafe fn rx_task(&self) -> &Task {
&*self.rx_task.with(|ptr| ptr)
}
unsafe fn set_rx_task(&self) {
self.rx_task.with_mut(|ptr| {
*ptr = ManuallyDrop::new(task::current())
});
}
unsafe fn tx_task(&self) -> &Task {
&*self.tx_task.with(|ptr| ptr)
}
unsafe fn set_tx_task(&self) {
self.tx_task.with_mut(|ptr| {
*ptr = ManuallyDrop::new(task::current())
});
}
}
unsafe impl<T: Send> Send for Inner<T> {}
unsafe impl<T: Send> Sync for Inner<T> {}
impl<T> Drop for Inner<T> {
fn drop(&mut self) {
let state = State(*self.state.get_mut());
if state.is_rx_task_set() {
self.rx_task.with_mut(|ptr| {
unsafe {
ManuallyDrop::drop(&mut *ptr);
}
});
}
if state.is_tx_task_set() {
self.tx_task.with_mut(|ptr| {
unsafe {
ManuallyDrop::drop(&mut *ptr);
}
});
}
}
}
impl<T: fmt::Debug> fmt::Debug for Inner<T> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::sync::atomic::Ordering::Relaxed;
fmt.debug_struct("Inner")
.field("state", &State::load(&self.state, Relaxed))
.finish()
}
}
const RX_TASK_SET: usize = 0b00001;
const VALUE_SENT: usize = 0b00010;
const CLOSED: usize = 0b00100;
const TX_TASK_SET: usize = 0b01000;
impl State {
fn new() -> State {
State(0)
}
fn is_complete(&self) -> bool {
self.0 & VALUE_SENT == VALUE_SENT
}
fn set_complete(cell: &AtomicUsize) -> State {
// TODO: This could be `Release`, followed by an `Acquire` fence *if*
// the `RX_TASK_SET` flag is set. However, `loom` does not support
// fences yet.
let val = cell.fetch_or(VALUE_SENT, AcqRel);
State(val)
}
fn is_rx_task_set(&self) -> bool {
self.0 & RX_TASK_SET == RX_TASK_SET
}
fn set_rx_task(cell: &AtomicUsize) -> State {
let val = cell.fetch_or(RX_TASK_SET, AcqRel);
State(val | RX_TASK_SET)
}
fn unset_rx_task(cell: &AtomicUsize) -> State {
let val = cell.fetch_and(!RX_TASK_SET, AcqRel);
State(val & !RX_TASK_SET)
}
fn is_closed(&self) -> bool {
self.0 & CLOSED == CLOSED
}
fn set_closed(cell: &AtomicUsize) -> State {
// Acquire because we want all later writes (attempting to poll) to be
// ordered after this.
let val = cell.fetch_or(CLOSED, Acquire);
State(val)
}
fn set_tx_task(cell: &AtomicUsize) -> State {
let val = cell.fetch_or(TX_TASK_SET, AcqRel);
State(val | TX_TASK_SET)
}
fn unset_tx_task(cell: &AtomicUsize) -> State {
let val = cell.fetch_and(!TX_TASK_SET, AcqRel);
State(val & !TX_TASK_SET)
}
fn is_tx_task_set(&self) -> bool {
self.0 & TX_TASK_SET == TX_TASK_SET
}
fn as_usize(self) -> usize {
self.0
}
fn load(cell: &AtomicUsize, order: Ordering) -> State {
let val = cell.load(order);
State(val)
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("State")
.field("is_complete", &self.is_complete())
.field("is_closed", &self.is_closed())
.field("is_rx_task_set", &self.is_rx_task_set())
.field("is_tx_task_set", &self.is_tx_task_set())
.finish()
}
}
tokio-sync/src/semaphore.rs
+1116
View File
File diff suppressed because it is too large Load Diff
tokio-sync/src/task/atomic_task.rs
+286
View File
@@ -0,0 +1,286 @@
use ::loom::{
futures::task::{self, Task},
sync::CausalCell,
sync::atomic::AtomicUsize,
};
use std::fmt;
use std::sync::atomic::Ordering::{Acquire, Release, AcqRel};
/// A synchronization primitive for task notification.
///
/// `AtomicTask` will coordinate concurrent notifications with the consumer
/// potentially "updating" the underlying task to notify. This is useful in
/// scenarios where a computation completes in another thread and wants to
/// notify the consumer, but the consumer is in the process of being migrated to
/// a new logical task.
///
/// Consumers should call `register` before checking the result of a computation
/// and producers should call `notify` after producing the computation (this
/// differs from the usual `thread::park` pattern). It is also permitted for
/// `notify` to be called **before** `register`. This results in a no-op.
///
/// A single `AtomicTask` may be reused for any number of calls to `register` or
/// `notify`.
///
/// `AtomicTask` does not provide any memory ordering guarantees, as such the
/// user should use caution and use other synchronization primitives to guard
/// the result of the underlying computation.
pub struct AtomicTask {
state: AtomicUsize,
task: CausalCell<Option<Task>>,
}
// `AtomicTask` is a multi-consumer, single-producer transfer cell. The cell
// stores a `Task` value produced by calls to `register` and many threads can
// race to take the task (to notify it) by calling `notify.
//
// If a new `Task` instance is produced by calling `register` before an existing
// one is consumed, then the existing one is overwritten.
//
// While `AtomicTask` is single-producer, the implementation ensures memory
// safety. In the event of concurrent calls to `register`, there will be a
// single winner whose task will get stored in the cell. The losers will not
// have their tasks notified. As such, callers should ensure to add
// synchronization to calls to `register`.
//
// The implementation uses a single `AtomicUsize` value to coordinate access to
// the `Task` cell. There are two bits that are operated on independently. These
// are represented by `REGISTERING` and `NOTIFYING`.
//
// The `REGISTERING` bit is set when a producer enters the critical section. The
// `NOTIFYING` bit is set when a consumer enters the critical section. Neither
// bit being set is represented by `WAITING`.
//
// A thread obtains an exclusive lock on the task cell by transitioning the
// state from `WAITING` to `REGISTERING` or `NOTIFYING`, depending on the
// operation the thread wishes to perform. When this transition is made, it is
// guaranteed that no other thread will access the task cell.
//
// # Registering
//
// On a call to `register`, an attempt to transition the state from WAITING to
// REGISTERING is made. On success, the caller obtains a lock on the task cell.
//
// If the lock is obtained, then the thread sets the task cell to the task
// provided as an argument. Then it attempts to transition the state back from
// `REGISTERING` -> `WAITING`.
//
// If this transition is successful, then the registering process is complete
// and the next call to `notify` will observe the task.
//
// If the transition fails, then there was a concurrent call to `notify` that
// was unable to access the task cell (due to the registering thread holding the
// lock). To handle this, the registering thread removes the task it just set
// from the cell and calls `notify` on it. This call to notify represents the
// attempt to notify by the other thread (that set the `NOTIFYING` bit). The
// state is then transitioned from `REGISTERING | NOTIFYING` back to `WAITING`.
// This transition must succeed because, at this point, the state cannot be
// transitioned by another thread.
//
// # Notifying
//
// On a call to `notify`, an attempt to transition the state from `WAITING` to
// `NOTIFYING` is made. On success, the caller obtains a lock on the task cell.
//
// If the lock is obtained, then the thread takes ownership of the current value
// in teh task cell, and calls `notify` on it. The state is then transitioned
// back to `WAITING`. This transition must succeed as, at this point, the state
// cannot be transitioned by another thread.
//
// If the thread is unable to obtain the lock, the `NOTIFYING` bit is still.
// This is because it has either been set by the current thread but the previous
// value included the `REGISTERING` bit **or** a concurrent thread is in the
// `NOTIFYING` critical section. Either way, no action must be taken.
//
// If the current thread is the only concurrent call to `notify` and another
// thread is in the `register` critical section, when the other thread **exits**
// the `register` critical section, it will observe the `NOTIFYING` bit and
// handle the notify itself.
//
// If another thread is in the `notify` critical section, then it will handle
// notifying the task.
//
// # A potential race (is safely handled).
//
// Imagine the following situation:
//
// * Thread A obtains the `notify` lock and notifies a task.
//
// * Before thread A releases the `notify` lock, the notified task is scheduled.
//
// * Thread B attempts to notify the task. In theory this should result in the
// task being notified, but it cannot because thread A still holds the notify
// lock.
//
// This case is handled by requiring users of `AtomicTask` to call `register`
// **before** attempting to observe the application state change that resulted
// in the task being notified. The notifiers also change the application state
// before calling notify.
//
// Because of this, the task will do one of two things.
//
// 1) Observe the application state change that Thread B is notifying on. In
// this case, it is OK for Thread B's notification to be lost.
//
// 2) Call register before attempting to observe the application state. Since
// Thread A still holds the `notify` lock, the call to `register` will result
// in the task notifying itself and get scheduled again.
/// Idle state
const WAITING: usize = 0;
/// A new task value is being registered with the `AtomicTask` cell.
const REGISTERING: usize = 0b01;
/// The task currently registered with the `AtomicTask` cell is being notified.
const NOTIFYING: usize = 0b10;
impl AtomicTask {
/// Create an `AtomicTask` initialized with the given `Task`
pub fn new() -> AtomicTask {
AtomicTask {
state: AtomicUsize::new(WAITING),
task: CausalCell::new(None),
}
}
/// Registers the current task to be notified on calls to `notify`.
///
/// This is the same as calling `register_task` with `task::current()`.
pub fn register(&self) {
self.register_task(task::current());
}
/// Registers the provided task to be notified on calls to `notify`.
///
/// The new task will take place of any previous tasks that were registered
/// by previous calls to `register`. Any calls to `notify` that happen after
/// a call to `register` (as defined by the memory ordering rules), will
/// notify the `register` caller's task.
///
/// It is safe to call `register` with multiple other threads concurrently
/// calling `notify`. This will result in the `register` caller's current
/// task being notified once.
///
/// This function is safe to call concurrently, but this is generally a bad
/// idea. Concurrent calls to `register` will attempt to register different
/// tasks to be notified. One of the callers will win and have its task set,
/// but there is no guarantee as to which caller will succeed.
pub fn register_task(&self, task: Task) {
debug!(" + register_task");
match self.state.compare_and_swap(WAITING, REGISTERING, Acquire) {
WAITING => {
unsafe {
// Locked acquired, update the waker cell
self.task.with_mut(|t| *t = Some(task));
// Release the lock. If the state transitioned to include
// the `NOTIFYING` bit, this means that a notify has been
// called concurrently, so we have to remove the task and
// notify it.`
//
// Start by assuming that the state is `REGISTERING` as this
// is what we jut set it to.
let res = self.state.compare_exchange(
REGISTERING, WAITING, AcqRel, Acquire);
match res {
Ok(_) => {}
Err(actual) => {
// This branch can only be reached if a
// concurrent thread called `notify`. In this
// case, `actual` **must** be `REGISTERING |
// `NOTIFYING`.
debug_assert_eq!(actual, REGISTERING | NOTIFYING);
// Take the task to notify once the atomic operation has
// completed.
let notify = self.task.with_mut(|t| (*t).take()).unwrap();
// Just swap, because no one could change state
// while state == `Registering | `Waking`
self.state.swap(WAITING, AcqRel);
// The atomic swap was complete, now
// notify the task and return.
notify.notify();
}
}
}
}
NOTIFYING => {
// Currently in the process of notifying the task, i.e.,
// `notify` is currently being called on the old task handle.
// So, we call notify on the new task handle
task.notify();
}
state => {
// In this case, a concurrent thread is holding the
// "registering" lock. This probably indicates a bug in the
// caller's code as racing to call `register` doesn't make much
// sense.
//
// We just want to maintain memory safety. It is ok to drop the
// call to `register`.
debug_assert!(
state == REGISTERING ||
state == REGISTERING | NOTIFYING);
}
}
}
/// Notifies the task that last called `register`.
///
/// If `register` has not been called yet, then this does nothing.
pub fn notify(&self) {
debug!(" + notify");
// AcqRel ordering is used in order to acquire the value of the `task`
// cell as well as to establish a `release` ordering with whatever
// memory the `AtomicTask` is associated with.
match self.state.fetch_or(NOTIFYING, AcqRel) {
WAITING => {
debug!(" + WAITING");
// The notifying lock has been acquired.
let task = unsafe { self.task.with_mut(|t| (*t).take()) };
// Release the lock
self.state.fetch_and(!NOTIFYING, Release);
debug!(" + Done notifying");
if let Some(task) = task {
task.notify();
}
}
state => {
debug!(" + state = {:?}", state);
// There is a concurrent thread currently updating the
// associated task.
//
// Nothing more to do as the `NOTIFYING` bit has been set. It
// doesn't matter if there are concurrent registering threads or
// not.
//
debug_assert!(
state == REGISTERING ||
state == REGISTERING | NOTIFYING ||
state == NOTIFYING);
}
}
}
}
impl Default for AtomicTask {
fn default() -> Self {
AtomicTask::new()
}
}
impl fmt::Debug for AtomicTask {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
write!(fmt, "AtomicTask")
}
}
unsafe impl Send for AtomicTask {}
unsafe impl Sync for AtomicTask {}
tokio-sync/src/task/mod.rs
+5
View File
@@ -0,0 +1,5 @@
//! Thread-safe task notification primitives.
mod atomic_task;
pub use self::atomic_task::AtomicTask;