spsc.rs

#![feature(test)]
extern crate test;

mod spsc {
    use std::cell::UnsafeCell;
    use std::mem;
    use std::ptr;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicPtr, Ordering};

    pub fn channel<T: Send>() -> (Sender<T>, Receiver<T>) {
        let a = Arc::new(Queue::new());
        (Sender { inner: a.clone() }, Receiver { inner: a })
    }

    pub struct Sender<T: Send> {
        inner: Arc<Queue<T>>,
    }

    impl<T: Send> Sender<T> {
        pub fn send(&self, t: T) -> Result<(), T> {
            self.inner.push(t);
            Ok(()) // disconnection check omitted
        }
    }

    pub struct Receiver<T: Send> {
        inner: Arc<Queue<T>>,
    }

    impl<T: Send> Receiver<T> {
        pub fn recv(&self) -> Result<T, ()> {
            match self.inner.pop() {
                Some(t) => Ok(t),
                None => Err(())
            }
        }
    }

    struct Queue<T: Send> {
        head: UnsafeCell<*mut Node<T>>,
        tail: AtomicPtr<Node<T>>,
    }

    impl<T: Send> Queue<T> {
        fn new() -> Queue<T> {
            let stub = Node::new(None);
            Queue {
                head: UnsafeCell::new(stub),
                tail: AtomicPtr::new(stub),
            }
        }

        fn push(&self, t: T) {
            unsafe {
                let n = Node::new(Some(t));
                (**self.head.get()).next.store(n, Ordering::Release);
                *self.head.get() = n;
            }
        }

        fn pop(&self) -> Option<T> {
            unsafe {
                let tail = self.tail.load(Ordering::Acquire);
                let next = (*tail).next.load(Ordering::Acquire);
                if next.is_null() {
                    return None;
                }
                let ret = (*next).value.take();
                self.tail.store(next, Ordering::Release);
                let _ = Box::from_raw(tail);
                ret
            }
        }
    }

    struct Node<T> {
        value: Option<T>,
        next: AtomicPtr<Node<T>>,
    }

    impl<T> Node<T> {
        fn new(v: Option<T>) -> *mut Node<T> {
            let mut b = Box::new(Node {
                value: v,
                next: AtomicPtr::new(ptr::null_mut())
            });
            let n = &mut *b as *mut _;
            mem::forget(b);
            n
        }
    }
}

mod spsc_plus_cas {
    use std::cell::UnsafeCell;
    use std::mem;
    use std::ptr;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicPtr, Ordering};

    pub fn channel<T: Send>() -> (Sender<T>, Receiver<T>) {
        let a = Arc::new(Queue::new());
        (Sender { inner: a.clone() }, Receiver { inner: a })
    }

    pub struct Sender<T: Send> {
        inner: Arc<Queue<T>>,
    }

    impl<T: Send> Sender<T> {
        pub fn send(&self, t: T) -> Result<(), T> {
            self.inner.push(t);
            Ok(()) // disconnection check omitted
        }
    }

    pub struct Receiver<T: Send> {
        inner: Arc<Queue<T>>,
    }

    impl<T: Send> Receiver<T> {
        pub fn recv(&self) -> Result<T, ()> {
            match self.inner.pop() {
                Some(t) => Ok(t),
                None => Err(())
            }
        }
    }

    struct Queue<T: Send> {
        head: UnsafeCell<*mut Node<T>>,
        tail: AtomicPtr<Node<T>>,
    }

    impl<T: Send> Queue<T> {
        fn new() -> Queue<T> {
            let stub = Node::new(None);
            Queue {
                head: UnsafeCell::new(stub),
                tail: AtomicPtr::new(stub),
            }
        }

        fn push(&self, t: T) {
            unsafe {
                let n = Node::new(Some(t));
                loop {
                    if ptr::null_mut() == (**self.head.get()).next.compare_and_swap(ptr::null_mut(), n, Ordering::AcqRel) {
                        *self.head.get() = n;
                        break;
                    }
                }
            }
        }

        fn pop(&self) -> Option<T> {
            unsafe {
                let tail = self.tail.load(Ordering::Acquire);
                let next = (*tail).next.load(Ordering::Acquire);
                if next.is_null() {
                    return None;
                }
                let ret = (*next).value.take();
                self.tail.store(next, Ordering::Release);
                let _ = Box::from_raw(tail);
                ret
            }
        }
    }

    struct Node<T> {
        value: Option<T>,
        next: AtomicPtr<Node<T>>,
    }

    impl<T> Node<T> {
        fn new(v: Option<T>) -> *mut Node<T> {
            let mut b = Box::new(Node {
                value: v,
                next: AtomicPtr::new(ptr::null_mut())
            });
            let n = &mut *b as *mut _;
            mem::forget(b);
            n
        }
    }
}

#[bench]
fn bench_spsc(b: &mut test::Bencher) {
    let (tx, rx) = spsc::channel();
    b.iter(move || {
        tx.send(5).unwrap();
        tx.send(6).unwrap();
        tx.send(7).unwrap();
        assert_eq!(rx.recv(), Ok(5));
        assert_eq!(rx.recv(), Ok(6));
        assert_eq!(rx.recv(), Ok(7));
    });
}


#[bench]
fn bench_spsc_plus_cas(b: &mut test::Bencher) {
    let (tx, rx) = spsc_plus_cas::channel();
    b.iter(move || {
        tx.send(5).unwrap();
        tx.send(6).unwrap();
        tx.send(7).unwrap();
        assert_eq!(rx.recv(), Ok(5));
        assert_eq!(rx.recv(), Ok(6));
        assert_eq!(rx.recv(), Ok(7));
    });
}