Repository: alexcrichton/futures-rs
Pull request #404 (2017-03-06 23:28:57)

Add a first-class `Sink::close` method
This commit adds a new method to the `Sink` trait, `close`. The intention of
this method is to model closing a sink with an indication that no more messages
will be pushed into it. This indication can then be acted upon to do blocking
and/or other asynchronous work necessary to perform a clean shutdown.

For many protocols this is then followed by some subsequent work to shut down
the sink. For example TLS has a specific shutdown handshake which is applicable
for any TLS-backed `Sink` types.

The upcoming [tokio-io] crate recently [added] a `shutdown` method to the
`AsyncWrite` trait which this method is intended to connect up with. A call to
`Sink::close` will translate to an `AsyncWrite::shutdown` call.

Ideally this method would be a required method, but in the interest of not
breaking API compatibility just yet the method is introduced as a default method
on the `Sink` trait which calls `poll_complete`. It is strongly recommended,
however, that all implementations of `Sink` will implement this method.

More documentation about the method itself can be found in the commit with the
associated documentation.


Patch:
diff --git a/src/sink/buffer.rs b/src/sink/buffer.rs
index 3a8ae0f9..3d0db7ed 100644
--- a/src/sink/buffer.rs
+++ b/src/sink/buffer.rs
@@ -80,4 +80,12 @@ impl<S: Sink> Sink for Buffer<S> {
         debug_assert!(self.buf.is_empty());
         self.sink.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), Self::SinkError> {
+        if self.buf.len() > 0 {
+            try_ready!(self.try_empty_buffer());
+        }
+        assert_eq!(self.buf.len(), 0);
+        self.sink.close()
+    }
 }
diff --git a/src/sink/map_err.rs b/src/sink/map_err.rs
index 4ae7fa9c..6a429788 100644
--- a/src/sink/map_err.rs
+++ b/src/sink/map_err.rs
@@ -28,4 +28,8 @@ impl<S, F, E> Sink for SinkMapErr<S, F>
     fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
         self.sink.poll_complete().map_err(|e| self.f.take().expect("cannot use MapErr after an error")(e))
     }
+
+    fn close(&mut self) -> Poll<(), Self::SinkError> {
+        self.sink.close().map_err(|e| self.f.take().expect("cannot use MapErr after an error")(e))
+    }
 }
diff --git a/src/sink/mod.rs b/src/sink/mod.rs
index b0c82448..4d828ed5 100644
--- a/src/sink/mod.rs
+++ b/src/sink/mod.rs
@@ -42,6 +42,10 @@ if_std! {
         fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
             Ok(::Async::Ready(()))
         }
+
+        fn close(&mut self) -> Poll<(), Self::SinkError> {
+            Ok(::Async::Ready(()))
+        }
     }
 
     /// A type alias for `Box<Stream + Send>`
@@ -60,6 +64,10 @@ if_std! {
         fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
             (**self).poll_complete()
         }
+
+        fn close(&mut self) -> Poll<(), Self::SinkError> {
+            (**self).close()
+        }
     }
 }
 
@@ -198,6 +206,80 @@ pub trait Sink {
     /// yet.
     fn poll_complete(&mut self) -> Poll<(), Self::SinkError>;
 
+    /// A method to indicate that no more values will ever be pushed into this
+    /// sink.
+    ///
+    /// This method is used to indicate that a sink will no longer even be given
+    /// another value by the caller. That is, the `start_send` method above will
+    /// be called no longer (nor `poll_complete`). This method is intended to
+    /// model "graceful shutdown" in various protocols where the intent to shut
+    /// down is followed by a little more blocking work.
+    ///
+    /// Callers of this function should work it it in a similar fashion to
+    /// `poll_complete`. Once called it may return `NotReady` which indicates
+    /// that more external work needs to happen to make progress. The current
+    /// task will be scheduled to receive a notification in such an event,
+    /// however.
+    ///
+    /// Note that this function will imply `poll_complete` above. That is, if a
+    /// sink has buffered data, then it'll be flushed out during a `close`
+    /// operation. It is not necessary to have `poll_complete` return `Ready`
+    /// before a `close` is called. Once a `close` is called, though,
+    /// `poll_complete` cannot be called.
+    ///
+    /// # Return value
+    ///
+    /// This function, like `poll_complete`, returns a `Poll`. The value is
+    /// `Ready` once the close operation has completed. At that point it should
+    /// be safe to drop the sink and deallocate associated resources.
+    ///
+    /// If the value returned is `NotReady` then the sink is not yet closed and
+    /// work needs to be done to close it. The work has been scheduled and the
+    /// current task will recieve a notification when it's next ready to call
+    /// this method again.
+    ///
+    /// Finally, this function may also return an error.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an `Err` if any operation along the way during
+    /// the close operation fails. An error typically is fatal for a sink and is
+    /// unable to be recovered from, but in specific situations this may not
+    /// always be true.
+    ///
+    /// Note that it's also typically an error to call `start_send` or
+    /// `poll_complete` after the `close` function is called. This method will
+    /// *initiate* a close, and continuing to send values after that (or attempt
+    /// to flush) may result in strange behavior, panics, errors, etc. Once this
+    /// method is called, it must be the only method called on this `Sink`.
+    ///
+    /// # Panics
+    ///
+    /// This method may panic or cause panics if:
+    ///
+    /// * It is called outside the context of a future's task
+    /// * It is called and then `start_send` or `poll_complete` is called
+    ///
+    /// # Compatibility notes
+    ///
+    /// Note that this function is currently by default a provided function,
+    /// defaulted to calling `poll_complete` above. This function was added
+    /// in the 0.1 series of the crate as a backwards-compatible addition. It
+    /// is intended that in the 0.2 series the method will no longer be a
+    /// default method.
+    ///
+    /// It is highly recommended to consider this method a required method and
+    /// to implement it whenever you implement `Sink` locally. It is especially
+    /// crucial to be sure to close inner sinks, if applicable.
+    #[cfg(feature = "with-deprecated")]
+    fn close(&mut self) -> Poll<(), Self::SinkError> {
+        self.poll_complete()
+    }
+
+    /// dox (you should see the above, not this)
+    #[cfg(not(feature = "with-deprecated"))]
+    fn close(&mut self) -> Poll<(), Self::SinkError>;
+
     /// Creates a new object which will produce a synchronous sink.
     ///
     /// The sink returned does **not** implement the `Sink` trait, and instead
@@ -332,4 +414,8 @@ impl<'a, S: ?Sized + Sink> Sink for &'a mut S {
     fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
         (**self).poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), Self::SinkError> {
+        (**self).close()
+    }
 }
diff --git a/src/sink/send_all.rs b/src/sink/send_all.rs
index a5552f26..7c7df7f5 100644
--- a/src/sink/send_all.rs
+++ b/src/sink/send_all.rs
@@ -75,7 +75,7 @@ impl<T, U> Future for SendAll<T, U>
             match try!(self.stream_mut().poll()) {
                 Async::Ready(Some(item)) => try_ready!(self.try_start_send(item)),
                 Async::Ready(None) => {
-                    try_ready!(self.sink_mut().poll_complete());
+                    try_ready!(self.sink_mut().close());
                     return Ok(Async::Ready(self.take_result()))
                 }
                 Async::NotReady => {
diff --git a/src/sink/with.rs b/src/sink/with.rs
index 237bf161..5f510a2d 100644
--- a/src/sink/with.rs
+++ b/src/sink/with.rs
@@ -137,4 +137,9 @@ impl<S, U, F, Fut> Sink for With<S, U, F, Fut>
         try_ready!(self.sink.poll_complete());
         Ok(me_ready)
     }
+
+    fn close(&mut self) -> Poll<(), Fut::Error> {
+        try_ready!(self.poll());
+        Ok(try!(self.sink.close()))
+    }
 }
diff --git a/src/stream/and_then.rs b/src/stream/and_then.rs
index 4dc3cac5..dc28a5c4 100644
--- a/src/stream/and_then.rs
+++ b/src/stream/and_then.rs
@@ -41,6 +41,10 @@ impl<S, F, U: IntoFuture> ::sink::Sink for AndThen<S, F, U>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, F, U> Stream for AndThen<S, F, U>
diff --git a/src/stream/buffer_unordered.rs b/src/stream/buffer_unordered.rs
index a0f3aa8e..f339a8db 100644
--- a/src/stream/buffer_unordered.rs
+++ b/src/stream/buffer_unordered.rs
@@ -173,4 +173,8 @@ impl<S> ::sink::Sink for BufferUnordered<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
diff --git a/src/stream/buffered.rs b/src/stream/buffered.rs
index 75bd9892..3c0bc871 100644
--- a/src/stream/buffered.rs
+++ b/src/stream/buffered.rs
@@ -71,6 +71,10 @@ impl<S> ::sink::Sink for Buffered<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S> Stream for Buffered<S>
diff --git a/src/stream/chunks.rs b/src/stream/chunks.rs
index dccfc400..aafccac4 100644
--- a/src/stream/chunks.rs
+++ b/src/stream/chunks.rs
@@ -45,6 +45,10 @@ impl<S> ::sink::Sink for Chunks<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 
diff --git a/src/stream/filter.rs b/src/stream/filter.rs
index c876df91..5a564c43 100644
--- a/src/stream/filter.rs
+++ b/src/stream/filter.rs
@@ -36,6 +36,10 @@ impl<S, F> ::sink::Sink for Filter<S, F>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, F> Stream for Filter<S, F>
diff --git a/src/stream/filter_map.rs b/src/stream/filter_map.rs
index 9bcbf69c..9a11bb2a 100644
--- a/src/stream/filter_map.rs
+++ b/src/stream/filter_map.rs
@@ -36,6 +36,10 @@ impl<S, F> ::sink::Sink for FilterMap<S, F>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, F, B> Stream for FilterMap<S, F>
diff --git a/src/stream/flatten.rs b/src/stream/flatten.rs
index d335e0a0..b9b0f29b 100644
--- a/src/stream/flatten.rs
+++ b/src/stream/flatten.rs
@@ -39,6 +39,10 @@ impl<S> ::sink::Sink for Flatten<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S> Stream for Flatten<S>
diff --git a/src/stream/forward.rs b/src/stream/forward.rs
index 1c9f1a37..959a4397 100644
--- a/src/stream/forward.rs
+++ b/src/stream/forward.rs
@@ -77,7 +77,7 @@ impl<T, U> Future for Forward<T, U>
             match try!(self.stream_mut().poll()) {
                 Async::Ready(Some(item)) => try_ready!(self.try_start_send(item)),
                 Async::Ready(None) => {
-                    try_ready!(self.sink_mut().poll_complete());
+                    try_ready!(self.sink_mut().close());
                     return Ok(Async::Ready(self.take_result()))
                 }
                 Async::NotReady => {
diff --git a/src/stream/fuse.rs b/src/stream/fuse.rs
index 2871c55b..2ded2b3f 100644
--- a/src/stream/fuse.rs
+++ b/src/stream/fuse.rs
@@ -27,6 +27,10 @@ impl<S> ::sink::Sink for Fuse<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 pub fn new<S: Stream>(s: S) -> Fuse<S> {
diff --git a/src/stream/map.rs b/src/stream/map.rs
index f5b8eeab..0c1e9d72 100644
--- a/src/stream/map.rs
+++ b/src/stream/map.rs
@@ -36,6 +36,10 @@ impl<S, F> ::sink::Sink for Map<S, F>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, F, U> Stream for Map<S, F>
diff --git a/src/stream/map_err.rs b/src/stream/map_err.rs
index b85a79ce..54addf30 100644
--- a/src/stream/map_err.rs
+++ b/src/stream/map_err.rs
@@ -36,6 +36,10 @@ impl<S, F> ::sink::Sink for MapErr<S, F>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, F, U> Stream for MapErr<S, F>
diff --git a/src/stream/or_else.rs b/src/stream/or_else.rs
index 123dea43..2d15fa2b 100644
--- a/src/stream/or_else.rs
+++ b/src/stream/or_else.rs
@@ -41,6 +41,10 @@ impl<S, F, U> ::sink::Sink for OrElse<S, F, U>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, F, U> Stream for OrElse<S, F, U>
diff --git a/src/stream/peek.rs b/src/stream/peek.rs
index 75c19fc2..96e65766 100644
--- a/src/stream/peek.rs
+++ b/src/stream/peek.rs
@@ -35,6 +35,10 @@ impl<S> ::sink::Sink for Peekable<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S: Stream> Stream for Peekable<S> {
diff --git a/src/stream/skip.rs b/src/stream/skip.rs
index f1fe4ea2..bdfac97a 100644
--- a/src/stream/skip.rs
+++ b/src/stream/skip.rs
@@ -34,6 +34,10 @@ impl<S> ::sink::Sink for Skip<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S> Stream for Skip<S>
diff --git a/src/stream/skip_while.rs b/src/stream/skip_while.rs
index aa665295..3541613f 100644
--- a/src/stream/skip_while.rs
+++ b/src/stream/skip_while.rs
@@ -41,6 +41,10 @@ impl<S, P, R> ::sink::Sink for SkipWhile<S, P, R>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, P, R> Stream for SkipWhile<S, P, R>
diff --git a/src/stream/split.rs b/src/stream/split.rs
index 611474c7..6e53b978 100644
--- a/src/stream/split.rs
+++ b/src/stream/split.rs
@@ -40,6 +40,13 @@ impl<S: Sink> Sink for SplitSink<S> {
             Async::NotReady => Ok(Async::NotReady),
         }
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        match self.0.poll_lock() {
+            Async::Ready(mut inner) => inner.close(),
+            Async::NotReady => Ok(Async::NotReady),
+        }
+    }
 }
 
 pub fn split<S: Stream + Sink>(s: S) -> (SplitSink<S>, SplitStream<S>) {
diff --git a/src/stream/take.rs b/src/stream/take.rs
index 8b010c25..12eb601b 100644
--- a/src/stream/take.rs
+++ b/src/stream/take.rs
@@ -34,6 +34,10 @@ impl<S> ::sink::Sink for Take<S>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S> Stream for Take<S>
diff --git a/src/stream/take_while.rs b/src/stream/take_while.rs
index 3e2147b8..0b0bf249 100644
--- a/src/stream/take_while.rs
+++ b/src/stream/take_while.rs
@@ -41,6 +41,10 @@ impl<S, P, R> ::sink::Sink for TakeWhile<S, P, R>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, P, R> Stream for TakeWhile<S, P, R>
diff --git a/src/stream/then.rs b/src/stream/then.rs
index 6f64676f..cab338e9 100644
--- a/src/stream/then.rs
+++ b/src/stream/then.rs
@@ -41,6 +41,10 @@ impl<S, F, U> ::sink::Sink for Then<S, F, U>
     fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
         self.stream.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        self.stream.close()
+    }
 }
 
 impl<S, F, U> Stream for Then<S, F, U>
diff --git a/src/sync/mpsc/mod.rs b/src/sync/mpsc/mod.rs
index a27bce4b..bd014997 100644
--- a/src/sync/mpsc/mod.rs
+++ b/src/sync/mpsc/mod.rs
@@ -506,6 +506,10 @@ impl<T> Sink for Sender<T> {
     fn poll_complete(&mut self) -> Poll<(), SendError<T>> {
         Ok(Async::Ready(()))
     }
+
+    fn close(&mut self) -> Poll<(), SendError<T>> {
+        Ok(Async::Ready(()))
+    }
 }
 
 impl<T> UnboundedSender<T> {
@@ -530,6 +534,10 @@ impl<T> Sink for UnboundedSender<T> {
     fn poll_complete(&mut self) -> Poll<(), SendError<T>> {
         self.0.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), SendError<T>> {
+        Ok(Async::Ready(()))
+    }
 }
 
 impl<'a, T> Sink for &'a UnboundedSender<T> {
@@ -544,6 +552,10 @@ impl<'a, T> Sink for &'a UnboundedSender<T> {
     fn poll_complete(&mut self) -> Poll<(), SendError<T>> {
         Ok(Async::Ready(()))
     }
+
+    fn close(&mut self) -> Poll<(), SendError<T>> {
+        Ok(Async::Ready(()))
+    }
 }
 
 impl<T> Clone for UnboundedSender<T> {
diff --git a/src/unsync/mpsc.rs b/src/unsync/mpsc.rs
index 19ec8891..2a474438 100644
--- a/src/unsync/mpsc.rs
+++ b/src/unsync/mpsc.rs
@@ -102,6 +102,10 @@ impl<T> Sink for Sender<T> {
     fn poll_complete(&mut self) -> Poll<(), SendError<T>> {
         Ok(Async::Ready(()))
     }
+
+    fn close(&mut self) -> Poll<(), SendError<T>> {
+        Ok(Async::Ready(()))
+    }
 }
 
 impl<T> Drop for Sender<T> {
@@ -220,6 +224,9 @@ impl<T> Sink for UnboundedSender<T> {
     fn poll_complete(&mut self) -> Poll<(), SendError<T>> {
         Ok(Async::Ready(()))
     }
+    fn close(&mut self) -> Poll<(), SendError<T>> {
+        Ok(Async::Ready(()))
+    }
 }
 
 impl<'a, T> Sink for &'a UnboundedSender<T> {
@@ -233,6 +240,10 @@ impl<'a, T> Sink for &'a UnboundedSender<T> {
     fn poll_complete(&mut self) -> Poll<(), SendError<T>> {
         Ok(Async::Ready(()))
     }
+
+    fn close(&mut self) -> Poll<(), SendError<T>> {
+        Ok(Async::Ready(()))
+    }
 }
 
 impl<T> UnboundedSender<T> {


Test patch:
diff --git a/tests/sink.rs b/tests/sink.rs
index 97283a79..351cfca7 100644
--- a/tests/sink.rs
+++ b/tests/sink.rs
@@ -195,6 +195,10 @@ impl<T> Sink for ManualFlush<T> {
             Ok(Async::NotReady)
         }
     }
+
+    fn close(&mut self) -> Poll<(), ()> {
+        Ok(().into())
+    }
 }
 
 impl<T> ManualFlush<T> {
@@ -296,6 +300,10 @@ impl<T> Sink for ManualAllow<T> {
     fn poll_complete(&mut self) -> Poll<(), ()> {
         Ok(Async::Ready(()))
     }
+
+    fn close(&mut self) -> Poll<(), ()> {
+        Ok(().into())
+    }
 }
 
 fn manual_allow<T>() -> (ManualAllow<T>, Rc<Allow>) {
diff --git a/tests/split.rs b/tests/split.rs
index 06e6e0f3..b9309c7f 100644
--- a/tests/split.rs
+++ b/tests/split.rs
@@ -27,6 +27,10 @@ impl<T, U: Sink> Sink for Join<T, U> {
     fn poll_complete(&mut self) -> Poll<(), U::SinkError> {
         self.1.poll_complete()
     }
+
+    fn close(&mut self) -> Poll<(), U::SinkError> {
+        self.1.close()
+    }
 }
 
 #[test]