Fix SplitSink performance problem caused by flushing inner sink every…

… time
rust-lang · Nov 13, 2019 · 5b63dfd · 5b63dfd
1 parent f3d262e
commit 5b63dfd
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/futures-util/src/stream/stream/split.rs b/futures-util/src/stream/stream/split.rs
@@ -61,6 +61,19 @@ impl<S: Sink<Item> + Unpin, Item> SplitSink<S, Item> {
     }
 }
 
+impl<S: Sink<Item>, Item> SplitSink<S, Item> {
+    fn poll_flush_slot(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
+        let this = &mut *self;
+        let mut inner = ready!(this.lock.poll_lock(cx));
+        if this.slot.is_some() {
+            ready!(inner.as_pin_mut().poll_ready(cx))?;
+            Poll::Ready(inner.as_pin_mut().start_send(this.slot.take().unwrap()))
+        } else {
+            Poll::Ready(Ok(()))
+        }
+    }
+}
+
 impl<S: Sink<Item>, Item> Sink<Item> for SplitSink<S, Item> {
     type Error = S::Error;
 
@@ -69,7 +82,7 @@ impl<S: Sink<Item>, Item> Sink<Item> for SplitSink<S, Item> {
             if self.slot.is_none() {
                 return Poll::Ready(Ok(()));
             }
-            ready!(self.as_mut().poll_flush(cx))?;
+            ready!(self.as_mut().poll_flush_slot(cx))?;
         }
     }