use fnv::FnvHashMap;
use futures::task::AtomicTask;
use futures::{Async, AsyncSink, Poll, Sink, StartSend, Stream};
use std::ops;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::{Arc, Mutex, RwLock, RwLockReadGuard, Weak};
#[derive(Debug)]
pub struct Receiver<T> {
shared: Arc<Shared<T>>,
inner: Arc<WatchInner>,
id: u64,
ver: usize,
}
#[derive(Debug)]
pub struct Sender<T> {
shared: Weak<Shared<T>>,
}
#[derive(Debug)]
pub struct Ref<'a, T: 'a> {
inner: RwLockReadGuard<'a, T>,
}
pub mod error {
use std::fmt;
#[derive(Debug)]
pub struct RecvError {
pub(crate) _p: (),
}
#[derive(Debug)]
pub struct SendError<T> {
pub(crate) inner: T,
}
impl fmt::Display for RecvError {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
use std::error::Error;
write!(fmt, "{}", self.description())
}
}
impl ::std::error::Error for RecvError {
fn description(&self) -> &str {
"channel closed"
}
}
impl<T: fmt::Debug> fmt::Display for SendError<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 SendError<T> {
fn description(&self) -> &str {
"channel closed"
}
}
}
#[derive(Debug)]
struct Shared<T> {
value: RwLock<T>,
version: AtomicUsize,
watchers: Mutex<Watchers>,
cancel: AtomicTask,
}
#[derive(Debug)]
struct Watchers {
next_id: u64,
watchers: FnvHashMap<u64, Arc<WatchInner>>,
}
#[derive(Debug)]
struct WatchInner {
task: AtomicTask,
}
const CLOSED: usize = 1;
pub fn channel<T>(init: T) -> (Sender<T>, Receiver<T>) {
const INIT_ID: u64 = 0;
let inner = Arc::new(WatchInner::new());
let mut watchers = FnvHashMap::with_capacity_and_hasher(0, Default::default());
watchers.insert(INIT_ID, inner.clone());
let shared = Arc::new(Shared {
value: RwLock::new(init),
version: AtomicUsize::new(2),
watchers: Mutex::new(Watchers {
next_id: INIT_ID + 1,
watchers,
}),
cancel: AtomicTask::new(),
});
let tx = Sender {
shared: Arc::downgrade(&shared),
};
let rx = Receiver {
shared,
inner,
id: INIT_ID,
ver: 0,
};
(tx, rx)
}
impl<T> Receiver<T> {
pub fn get_ref(&self) -> Ref<T> {
let inner = self.shared.value.read().unwrap();
Ref { inner }
}
pub fn poll_ref(&mut self) -> Poll<Option<Ref<T>>, error::RecvError> {
self.inner.task.register();
let state = self.shared.version.load(SeqCst);
let version = state & !CLOSED;
if version != self.ver {
self.ver = version;
let inner = self.shared.value.read().unwrap();
return Ok(Some(Ref { inner }).into());
}
if CLOSED == state & CLOSED {
return Ok(None.into());
}
Ok(Async::NotReady)
}
}
impl<T: Clone> Stream for Receiver<T> {
type Item = T;
type Error = error::RecvError;
fn poll(&mut self) -> Poll<Option<T>, error::RecvError> {
let item = try_ready!(self.poll_ref());
Ok(Async::Ready(item.map(|v_ref| v_ref.clone())))
}
}
impl<T> Clone for Receiver<T> {
fn clone(&self) -> Self {
let inner = Arc::new(WatchInner::new());
let shared = self.shared.clone();
let id = {
let mut watchers = shared.watchers.lock().unwrap();
let id = watchers.next_id;
watchers.next_id += 1;
watchers.watchers.insert(id, inner.clone());
id
};
let ver = self.ver;
Receiver {
shared: shared,
inner,
id,
ver,
}
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
let mut watchers = self.shared.watchers.lock().unwrap();
watchers.watchers.remove(&self.id);
}
}
impl WatchInner {
fn new() -> Self {
WatchInner {
task: AtomicTask::new(),
}
}
}
impl<T> Sender<T> {
pub fn broadcast(&mut self, value: T) -> Result<(), error::SendError<T>> {
let shared = match self.shared.upgrade() {
Some(shared) => shared,
None => return Err(error::SendError { inner: value }),
};
{
let mut lock = shared.value.write().unwrap();
*lock = value;
}
shared.version.fetch_add(2, SeqCst);
notify_all(&*shared);
Ok(())
}
pub fn poll_close(&mut self) -> Poll<(), ()> {
match self.shared.upgrade() {
Some(shared) => {
shared.cancel.register();
Ok(Async::NotReady)
}
None => Ok(Async::Ready(())),
}
}
}
impl<T> Sink for Sender<T> {
type SinkItem = T;
type SinkError = error::SendError<T>;
fn start_send(&mut self, item: T) -> StartSend<T, error::SendError<T>> {
let _ = self.broadcast(item)?;
Ok(AsyncSink::Ready)
}
fn poll_complete(&mut self) -> Poll<(), error::SendError<T>> {
Ok(().into())
}
}
fn notify_all<T>(shared: &Shared<T>) {
let watchers = shared.watchers.lock().unwrap();
for watcher in watchers.watchers.values() {
watcher.task.notify();
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
if let Some(shared) = self.shared.upgrade() {
shared.version.fetch_or(CLOSED, SeqCst);
notify_all(&*shared);
}
}
}
impl<'a, T: 'a> ops::Deref for Ref<'a, T> {
type Target = T;
fn deref(&self) -> &T {
self.inner.deref()
}
}
impl<T> Drop for Shared<T> {
fn drop(&mut self) {
self.cancel.notify();
}
}