Add BufferOne sink combinator

src/sink/buffer_one.rs

@@ -0,0 +1,128 @@
+use {Poll, Async};
+use {StartSend, AsyncSink};
+use sink::Sink;
+use stream::Stream;
+
+/// Adds buffering for a single value to the underlying sink.
+///
+/// Unlike `Sink::buffer`, this combinator is able to only buffer a single
+/// value. As such, the buffer does not require any allocation. `BufferOne` also
+/// provides `poll_ready`, which allows checking if sending a value will be
+/// accepted.
+#[derive(Debug)]
+#[must_use = "sinks do nothing unless polled"]
+pub struct BufferOne<S: Sink> {
+    sink: S,
+    buf: Option<Result<S::SinkItem, S::SinkError>>,
+}
+
+impl<S: Sink> BufferOne<S> {
+    /// Decorate `sink` with `BufferOne`.
+    ///
+    /// The returned `BufferOne` will have an empty buffer and will be
+    /// guaranteed to be able to accept a single value.
+    pub fn new(sink: S) -> BufferOne<S> {
+        BufferOne {
+            sink: sink,
+            buf: None,
+        }
+    }
+
+    /// Get a shared reference to the inner sink.
+    pub fn get_ref(&self) -> &S {
+        &self.sink
+    }
+
+    /// Get a mutable reference to the inner sink.
+    pub fn get_mut(&mut self) -> &mut S {
+        &mut self.sink
+    }
+
+    /// Consumes the `BufferOne`, returning its underlying sink.
+    ///
+    /// If a value is currently buffered, it will be lost.
+    pub fn into_inner(self) -> S {
+        self.sink
+    }
+
+    /// Polls the sink to detect whether a call to `start_send` will be accepted
+    /// or not.
+    ///
+    /// If this function returns `Async::Ready`, then a call to `start_send` is
+    /// guaranted to return either `Ok(AsyncSink::Ready)` or `Err(_)`.
+    ///
+    /// If this function returns `Async::NotReady`, then the current task will
+    /// become notified once the sink becomes ready to accept a new value.
+    pub fn poll_ready(&mut self) -> Async<()> {
+        match self.try_empty_buffer() {
+            Err(e) => {
+                self.buf = Some(Err(e));
+                Async::Ready(())
+            }
+            Ok(ready) => ready,
+        }
+    }
+
+    fn try_empty_buffer(&mut self) -> Poll<(), S::SinkError> {
+        match self.buf.take() {
+            Some(Ok(item)) => {
+                match try!(self.sink.start_send(item)) {
+                    AsyncSink::Ready => Ok(Async::Ready(())),
+                    AsyncSink::NotReady(item) => {
+                        self.buf = Some(Ok(item));
+
+                        try!(self.sink.poll_complete());
+
+                        Ok(Async::NotReady)
+                    }
+                }
+            }
+            Some(Err(e)) => Err(e),
+            None => {
+                Ok(Async::Ready(()))
+            }
+        }
+    }
+}
+
+impl<S> Stream for BufferOne<S> where S: Sink + Stream {
+    type Item = S::Item;
+    type Error = S::Error;
+
+    fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
+        self.sink.poll()
+    }
+}
+
+impl<S: Sink> Sink for BufferOne<S> {
+    type SinkItem = S::SinkItem;
+    type SinkError = S::SinkError;
+
+    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
+        try!(self.try_empty_buffer());
+
+        if self.buf.is_some() {
+            return Ok(AsyncSink::NotReady(item));
+        }
+
+        self.buf = Some(Ok(item));
+
+        Ok(AsyncSink::Ready)
+    }
+
+    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
+        try_ready!(self.try_empty_buffer());
+        debug_assert!(self.buf.is_none());
+        self.sink.poll_complete()
+    }
+
+    fn close(&mut self) -> Poll<(), Self::SinkError> {
+        if self.buf.is_some() {
+            try_ready!(self.try_empty_buffer());
+        }
+
+        assert!(self.buf.is_none());
+
+        self.sink.close()
+    }
+}

src/sink/mod.rs

@@ -22,9 +22,11 @@ mod map_err;
 
 if_std! {
     mod buffer;
+    mod buffer_one;
     mod wait;
 
     pub use self::buffer::Buffer;
+    pub use self::buffer_one::BufferOne;
     pub use self::wait::Wait;
 
     // TODO: consider expanding this via e.g. FromIterator

tests/buffer_one.rs

@@ -0,0 +1,169 @@
+extern crate futures;
+
+use futures::{Sink, StartSend, Poll, Async, AsyncSink};
+use futures::sink::BufferOne;
+use std::collections::VecDeque;
+
+#[test]
+fn always_ready() {
+    let mut sink = BufferOne::new(
+        Mock::default()
+            .expect_start_send(1, Ok(AsyncSink::Ready))
+            .expect_poll_complete(Ok(Async::Ready(())))
+            .expect_start_send(2, Ok(AsyncSink::Ready))
+            .expect_start_send(3, Ok(AsyncSink::Ready))
+            .expect_poll_complete(Ok(Async::Ready(())))
+            );
+
+    assert!(sink.start_send(1).unwrap().is_ready());
+    assert_eq!(5, sink.get_ref().len());
+
+    assert!(sink.poll_complete().unwrap().is_ready());
+    assert_eq!(3, sink.get_ref().len());
+
+    assert!(sink.start_send(2).unwrap().is_ready());
+    assert!(sink.start_send(3).unwrap().is_ready());
+    assert_eq!(2, sink.get_ref().len());
+
+    assert!(sink.poll_complete().unwrap().is_ready());
+    assert_eq!(0, sink.get_ref().len());
+}
+
+#[test]
+fn start_send_push_back() {
+    let mut sink = BufferOne::new(
+        Mock::default()
+            .expect_start_send(1, Ok(AsyncSink::NotReady(1)))
+            .expect_poll_complete(Ok(Async::Ready(())))
+            .expect_start_send(1, Ok(AsyncSink::Ready))
+            .expect_poll_complete(Ok(Async::Ready(())))
+            );
+
+    assert!(sink.start_send(1).unwrap().is_ready());
+    assert_eq!(4, sink.get_ref().len());
+
+    assert!(!sink.poll_complete().unwrap().is_ready());
+    assert_eq!(2, sink.get_ref().len());
+
+    assert!(sink.poll_complete().unwrap().is_ready());
+    assert_eq!(0, sink.get_ref().len());
+}
+
+#[test]
+fn poll_complete_push_back() {
+    let mut sink = BufferOne::new(
+        Mock::default()
+            .expect_start_send(1, Ok(AsyncSink::Ready))
+            .expect_poll_complete(Ok(Async::NotReady))
+            .expect_start_send(2, Ok(AsyncSink::Ready))
+            .expect_poll_complete(Ok(Async::Ready(())))
+            );
+
+    assert!(sink.start_send(1).unwrap().is_ready());
+    assert_eq!(4, sink.get_ref().len());
+
+    assert!(!sink.poll_complete().unwrap().is_ready());
+    assert_eq!(2, sink.get_ref().len());
+
+    assert!(sink.start_send(2).unwrap().is_ready());
+    assert_eq!(2, sink.get_ref().len());
+
+    assert!(sink.poll_complete().unwrap().is_ready());
+    assert_eq!(0, sink.get_ref().len());
+}
+
+#[test]
+fn poll_ready() {
+    let mut sink = BufferOne::new(
+        Mock::default()
+            .expect_start_send(1, Ok(AsyncSink::Ready))
+            .expect_start_send(2, Ok(AsyncSink::NotReady(2)))
+            .expect_poll_complete(Ok(Async::NotReady))
+            .expect_start_send(2, Ok(AsyncSink::Ready))
+            );
+
+    assert!(sink.poll_ready().is_ready());
+    assert_eq!(4, sink.get_ref().len());
+
+    assert!(sink.start_send(1).unwrap().is_ready());
+    assert_eq!(4, sink.get_ref().len());
+
+    assert!(sink.poll_ready().is_ready());
+    assert_eq!(3, sink.get_ref().len());
+
+    assert!(sink.start_send(2).unwrap().is_ready());
+    assert!(!sink.poll_ready().is_ready());
+    assert_eq!(1, sink.get_ref().len());
+
+    assert!(sink.poll_ready().is_ready());
+    assert_eq!(0, sink.get_ref().len());
+}
+
+#[test]
+fn poll_ready_err() {
+    let mut sink = BufferOne::new(
+        Mock::default()
+            .expect_start_send(1, Err(()))
+            );
+
+    assert!(sink.start_send(1).unwrap().is_ready());
+
+    assert!(sink.poll_ready().is_ready());
+    assert_eq!(0, sink.get_ref().len());
+
+    assert!(sink.start_send(2).is_err());
+}
+
+#[derive(Default)]
+struct Mock {
+    expect: VecDeque<Call>,
+}
+
+enum Call {
+    StartSend(u32, StartSend<u32, ()>),
+    PollComplete(Poll<(), ()>),
+}
+
+impl Mock {
+    fn expect_start_send(mut self, val: u32, ret: StartSend<u32, ()>) -> Self {
+        self.expect.push_back(Call::StartSend(val, ret));
+        self
+    }
+
+    fn expect_poll_complete(mut self, ret: Poll<(), ()>) -> Self {
+        self.expect.push_back(Call::PollComplete(ret));
+        self
+    }
+
+    fn len(&self) -> usize {
+        self.expect.len()
+    }
+}
+
+impl Sink for Mock {
+    type SinkItem = u32;
+    type SinkError = ();
+
+    fn start_send(&mut self, item: u32) -> StartSend<u32, ()> {
+        match self.expect.pop_front() {
+            Some(Call::StartSend(v, ret)) => {
+                assert_eq!(item, v);
+                ret
+            }
+            Some(Call::PollComplete(..)) => panic!("expected poll_complete but got start_send"),
+            None => panic!("expected no further calls"),
+        }
+    }
+
+    fn poll_complete(&mut self) -> Poll<(), ()> {
+        match self.expect.pop_front() {
+            Some(Call::PollComplete(ret)) => ret,
+            Some(Call::StartSend(..)) => panic!("expected start_send but got poll_complete"),
+            None => panic!("expected no further calls"),
+        }
+    }
+
+    fn close(&mut self) -> Poll<(), ()> {
+        self.poll_complete()
+    }
+}