#![deny(missing_docs)]
#![deny(missing_debug_implementations)]
extern crate futures;
extern crate slab;
use std::cell::UnsafeCell;
use std::fmt::{self, Debug};
use std::iter::FromIterator;
use std::marker::PhantomData;
use std::mem;
use std::ops::{Index, IndexMut};
use std::ptr;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release, SeqCst};
use std::sync::atomic::{AtomicBool, AtomicPtr};
use std::sync::{Arc, Weak};
use std::usize;
use futures::executor::{self, Notify, NotifyHandle, UnsafeNotify};
use futures::task::AtomicTask;
use futures::{task, Async, Poll, Stream};
#[must_use = "streams do nothing unless polled"]
pub struct StreamUnordered<S> {
inner: Arc<Inner>,
streams: slab::Slab<S>,
head_all: *const Node,
}
unsafe impl<S: Send> Send for StreamUnordered<S> {}
unsafe impl<S: Sync> Sync for StreamUnordered<S> {}
#[allow(missing_debug_implementations)]
struct Inner {
parent: AtomicTask,
head_readiness: AtomicPtr<Node>,
tail_readiness: UnsafeCell<*const Node>,
stub: Arc<Node>,
}
struct Node {
stream: UnsafeCell<Option<usize>>,
next_all: UnsafeCell<*const Node>,
prev_all: UnsafeCell<*const Node>,
next_readiness: AtomicPtr<Node>,
queue: Weak<Inner>,
queued: AtomicBool,
}
enum Dequeue {
Data(*const Node),
Empty,
Inconsistent,
}
impl<S> Default for StreamUnordered<S> {
fn default() -> Self {
StreamUnordered::new()
}
}
#[derive(Debug)]
pub struct StreamSlot<'a, S: 'a> {
entry: slab::VacantEntry<'a, S>,
backref: *mut StreamUnordered<S>,
}
impl<'a, S: 'a> StreamSlot<'a, S> {
pub fn insert(self, stream: S) -> &'a mut S {
let token = self.entry.key();
{
self.entry.insert(stream);
}
let this = unsafe { &mut *self.backref };
let node = Arc::new(Node {
stream: UnsafeCell::new(Some(token)),
next_all: UnsafeCell::new(ptr::null_mut()),
prev_all: UnsafeCell::new(ptr::null_mut()),
next_readiness: AtomicPtr::new(ptr::null_mut()),
queued: AtomicBool::new(true),
queue: Arc::downgrade(&this.inner),
});
let ptr = this.link(node);
this.inner.enqueue(ptr);
&mut this[token]
}
pub fn token(&self) -> usize {
self.entry.key()
}
}
impl<S> StreamUnordered<S> {
pub fn new() -> StreamUnordered<S> {
let stub = Arc::new(Node {
stream: UnsafeCell::new(None),
next_all: UnsafeCell::new(ptr::null()),
prev_all: UnsafeCell::new(ptr::null()),
next_readiness: AtomicPtr::new(ptr::null_mut()),
queued: AtomicBool::new(true),
queue: Weak::new(),
});
let stub_ptr = &*stub as *const Node;
let inner = Arc::new(Inner {
parent: AtomicTask::new(),
head_readiness: AtomicPtr::new(stub_ptr as *mut _),
tail_readiness: UnsafeCell::new(stub_ptr),
stub: stub,
});
StreamUnordered {
streams: slab::Slab::new(),
head_all: ptr::null_mut(),
inner: inner,
}
}
pub fn len(&self) -> usize {
self.streams.len()
}
pub fn is_empty(&self) -> bool {
self.streams.is_empty()
}
pub fn stream_slot(&mut self) -> StreamSlot<S> {
let this = self as *mut _;
StreamSlot {
entry: self.streams.vacant_entry(),
backref: this,
}
}
pub fn push(&mut self, stream: S) -> usize {
let s = self.stream_slot();
let token = s.token();
s.insert(stream);
token
}
pub fn get(&self, stream: usize) -> Option<&S> {
self.streams.get(stream)
}
pub fn get_mut(&mut self, stream: usize) -> Option<&mut S> {
self.streams.get_mut(stream)
}
pub fn iter_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut S> {
self.streams.iter_mut().map(|(_, s)| s)
}
fn release_node(&mut self, node: Arc<Node>) {
let prev = node.queued.swap(true, SeqCst);
if let Some(idx) = unsafe { (*node.stream.get()).take() } {
drop(self.streams.remove(idx));
}
if prev {
mem::forget(node);
}
}
fn link(&mut self, node: Arc<Node>) -> *const Node {
let ptr = arc2ptr(node);
unsafe {
*(*ptr).next_all.get() = self.head_all;
if !self.head_all.is_null() {
*(*self.head_all).prev_all.get() = ptr;
}
}
self.head_all = ptr;
return ptr;
}
unsafe fn unlink(&mut self, node: *const Node) -> Arc<Node> {
let node = ptr2arc(node);
let next = *node.next_all.get();
let prev = *node.prev_all.get();
*node.next_all.get() = ptr::null_mut();
*node.prev_all.get() = ptr::null_mut();
if !next.is_null() {
*(*next).prev_all.get() = prev;
}
if !prev.is_null() {
*(*prev).next_all.get() = next;
} else {
self.head_all = next;
}
return node;
}
}
impl<S> Index<usize> for StreamUnordered<S> {
type Output = S;
fn index(&self, stream: usize) -> &Self::Output {
&self.streams[stream]
}
}
impl<S> IndexMut<usize> for StreamUnordered<S> {
fn index_mut(&mut self, stream: usize) -> &mut Self::Output {
&mut self.streams[stream]
}
}
#[derive(Debug)]
pub enum StreamYield<S>
where
S: Stream,
{
Item(S::Item),
Finished(S),
}
impl<S> Stream for StreamUnordered<S>
where
S: Stream,
{
type Item = (StreamYield<S>, usize);
type Error = S::Error;
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
self.inner.parent.register();
loop {
let node = match unsafe { self.inner.dequeue() } {
Dequeue::Empty => {
if self.is_empty() {
return Ok(Async::Ready(None));
} else {
return Ok(Async::NotReady);
}
}
Dequeue::Inconsistent => {
task::current().notify();
return Ok(Async::NotReady);
}
Dequeue::Data(node) => node,
};
debug_assert!(node != self.inner.stub());
unsafe {
let stream = match (*(*node).stream.get()).take() {
Some(stream) => stream,
None => {
let node = ptr2arc(node);
assert!((*node.next_all.get()).is_null());
assert!((*node.prev_all.get()).is_null());
continue;
}
};
let prev = (*node).queued.swap(false, SeqCst);
assert!(prev);
struct Bomb<'a, T: 'a> {
queue: &'a mut StreamUnordered<T>,
node: Option<Arc<Node>>,
}
impl<'a, T> Drop for Bomb<'a, T> {
fn drop(&mut self) {
if let Some(node) = self.node.take() {
self.queue.release_node(node);
}
}
}
let mut bomb = Bomb {
node: Some(self.unlink(node)),
queue: self,
};
let res = {
let notify = NodeToHandle(bomb.node.as_ref().unwrap());
let mut stream = bomb.queue.streams.get_mut(stream).unwrap();
executor::with_notify(¬ify, 0, || stream.poll())
};
break match res {
Ok(Async::NotReady) => {
let node = bomb.node.take().unwrap();
*node.stream.get() = Some(stream);
bomb.queue.link(node);
continue;
}
Ok(Async::Ready(Some(e))) => {
NotifyHandle::from(NodeToHandle(bomb.node.as_ref().unwrap())).notify(0);
let node = bomb.node.take().unwrap();
*node.stream.get() = Some(stream);
bomb.queue.link(node);
Ok(Async::Ready(Some((StreamYield::Item(e), stream))))
}
Ok(Async::Ready(None)) => {
let s = bomb.queue.streams.remove(stream);
Ok(Async::Ready(Some((StreamYield::Finished(s), stream))))
}
Err(e) => Err(e),
};
}
}
}
}
impl<S: Debug> Debug for StreamUnordered<S> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
write!(fmt, "StreamUnordered {{ ... }}")
}
}
impl<S> Drop for StreamUnordered<S> {
fn drop(&mut self) {
unsafe {
while !self.head_all.is_null() {
let head = self.head_all;
let node = self.unlink(head);
self.release_node(node);
}
}
}
}
impl<S> FromIterator<S> for StreamUnordered<S> {
fn from_iter<I>(iter: I) -> Self
where
I: IntoIterator<Item = S>,
{
let mut new = StreamUnordered::new();
for stream in iter.into_iter() {
new.push(stream);
}
new
}
}
impl Inner {
fn enqueue(&self, node: *const Node) {
unsafe {
debug_assert!((*node).queued.load(Relaxed));
(*node).next_readiness.store(ptr::null_mut(), Relaxed);
let node = node as *mut _;
let prev = self.head_readiness.swap(node, AcqRel);
(*prev).next_readiness.store(node, Release);
}
}
unsafe fn dequeue(&self) -> Dequeue {
let mut tail = *self.tail_readiness.get();
let mut next = (*tail).next_readiness.load(Acquire);
if tail == self.stub() {
if next.is_null() {
return Dequeue::Empty;
}
*self.tail_readiness.get() = next;
tail = next;
next = (*next).next_readiness.load(Acquire);
}
if !next.is_null() {
*self.tail_readiness.get() = next;
debug_assert!(tail != self.stub());
return Dequeue::Data(tail);
}
if self.head_readiness.load(Acquire) as *const _ != tail {
return Dequeue::Inconsistent;
}
self.enqueue(self.stub());
next = (*tail).next_readiness.load(Acquire);
if !next.is_null() {
*self.tail_readiness.get() = next;
return Dequeue::Data(tail);
}
Dequeue::Inconsistent
}
fn stub(&self) -> *const Node {
&*self.stub
}
}
impl Drop for Inner {
fn drop(&mut self) {
unsafe {
loop {
match self.dequeue() {
Dequeue::Empty => break,
Dequeue::Inconsistent => abort("inconsistent in drop"),
Dequeue::Data(ptr) => drop(ptr2arc(ptr)),
}
}
}
}
}
#[allow(missing_debug_implementations)]
struct NodeToHandle<'a>(&'a Arc<Node>);
impl<'a> Clone for NodeToHandle<'a> {
fn clone(&self) -> Self {
NodeToHandle(self.0)
}
}
impl<'a> From<NodeToHandle<'a>> for NotifyHandle {
fn from(handle: NodeToHandle<'a>) -> NotifyHandle {
unsafe {
let ptr = handle.0.clone();
let ptr = mem::transmute::<Arc<Node>, *mut ArcNode>(ptr);
NotifyHandle::new(hide_lt(ptr))
}
}
}
struct ArcNode(PhantomData<()>);
unsafe impl Send for ArcNode {}
unsafe impl Sync for ArcNode {}
impl Notify for ArcNode {
fn notify(&self, _id: usize) {
unsafe {
let me: *const ArcNode = self;
let me: *const *const ArcNode = &me;
let me = me as *const Arc<Node>;
Node::notify(&*me)
}
}
}
unsafe impl UnsafeNotify for ArcNode {
unsafe fn clone_raw(&self) -> NotifyHandle {
let me: *const ArcNode = self;
let me: *const *const ArcNode = &me;
let me = &*(me as *const Arc<Node>);
NodeToHandle(me).into()
}
unsafe fn drop_raw(&self) {
let mut me: *const ArcNode = self;
let me = &mut me as *mut *const ArcNode as *mut Arc<Node>;
ptr::drop_in_place(me);
}
}
unsafe fn hide_lt(p: *mut ArcNode) -> *mut UnsafeNotify {
mem::transmute(p as *mut UnsafeNotify)
}
impl Node {
fn notify(me: &Arc<Node>) {
let inner = match me.queue.upgrade() {
Some(inner) => inner,
None => return,
};
let prev = me.queued.swap(true, SeqCst);
if !prev {
inner.enqueue(&**me);
inner.parent.notify();
}
}
}
impl Drop for Node {
fn drop(&mut self) {
unsafe {
if (*self.stream.get()).is_some() {
abort("stream still here when dropping");
}
}
}
}
fn arc2ptr<T>(ptr: Arc<T>) -> *const T {
let addr = &*ptr as *const T;
mem::forget(ptr);
return addr;
}
unsafe fn ptr2arc<T>(ptr: *const T) -> Arc<T> {
let anchor = mem::transmute::<usize, Arc<T>>(0x10);
let addr = &*anchor as *const T;
mem::forget(anchor);
let offset = addr as isize - 0x10;
mem::transmute::<isize, Arc<T>>(ptr as isize - offset)
}
fn abort(s: &str) -> ! {
struct DoublePanic;
impl Drop for DoublePanic {
fn drop(&mut self) {
panic!("panicking twice to abort the program");
}
}
let _bomb = DoublePanic;
panic!("{}", s);
}