stream: stop polling the underlying stream once map_while yields None (#8233)

tokio-stream/src/stream_ext/map_while.rs

@@ -3,6 +3,7 @@ use crate::Stream;
 use core::fmt;
 use core::pin::Pin;
 use core::task::{Context, Poll};
+use futures_core::FusedStream;
 use pin_project_lite::pin_project;
 
 pin_project! {
@@ -12,6 +13,7 @@ pin_project! {
         #[pin]
         stream: St,
         f: F,
+        done: bool,
     }
 }
 
@@ -22,13 +24,18 @@ where
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("MapWhile")
             .field("stream", &self.stream)
+            .field("done", &self.done)
             .finish()
     }
 }
 
 impl<St, F> MapWhile<St, F> {
     pub(super) fn new(stream: St, f: F) -> Self {
-        MapWhile { stream, f }
+        MapWhile {
+            stream,
+            f,
+            done: false,
+        }
     }
 }
 
@@ -41,12 +48,36 @@ where
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
         let me = self.project();
+        if *me.done {
+            return Poll::Ready(None);
+        }
+
         let f = me.f;
-        me.stream.poll_next(cx).map(|opt| opt.and_then(f))
+        let done = me.done;
+        me.stream.poll_next(cx).map(|opt| {
+            let mapped = opt.and_then(f);
+            if mapped.is_none() {
+                *done = true;
+            }
+            mapped
+        })
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.done {
+            return (0, Some(0));
+        }
+
         let (_, upper) = self.stream.size_hint();
         (0, upper)
     }
 }
+
+impl<St, F> FusedStream for MapWhile<St, F>
+where
+    Self: Stream,
+{
+    fn is_terminated(&self) -> bool {
+        self.done
+    }
+}