stream: honor StreamMap::next_many limit (#8215)

tokio-stream/src/stream_map.rs

@@ -632,6 +632,10 @@ where
                         should_loop = true;
 
                         idx = idx.wrapping_add(1) % self.entries.len();
+
+                        if added == limit {
+                            break;
+                        }
                     }
                     Poll::Ready(None) => {
                         // Remove the entry

tokio-stream/tests/stream_stream_map.rs

@@ -429,6 +429,27 @@ async fn poll_next_many_enough() {
     assert!(buffer.contains(&(1, 1)));
 }
 
+#[tokio::test]
+async fn poll_next_many_does_not_exceed_limit() {
+    let mut stream_map: StreamMap<usize, UsizeStream> = StreamMap::new();
+
+    stream_map.insert(0, Box::pin(iter([0usize].into_iter())) as UsizeStream);
+    stream_map.insert(1, Box::pin(iter([1usize].into_iter())) as UsizeStream);
+
+    let mut buffer = vec![];
+    let n = poll_fn(|cx| stream_map.poll_next_many(cx, &mut buffer, 1)).await;
+
+    assert_eq!(n, 1);
+    assert_eq!(buffer.len(), 1);
+
+    let n = poll_fn(|cx| stream_map.poll_next_many(cx, &mut buffer, 1)).await;
+
+    assert_eq!(n, 1);
+    assert_eq!(buffer.len(), 2);
+    assert!(buffer.contains(&(0, 0)));
+    assert!(buffer.contains(&(1, 1)));
+}
+
 #[tokio::test]
 async fn poll_next_many_correctly_loops_around() {
     for _ in 0..10 {
@@ -543,6 +564,27 @@ async fn next_many_enough() {
     assert!(buffer.contains(&(1, 1)));
 }
 
+#[tokio::test]
+async fn next_many_does_not_exceed_limit() {
+    let mut stream_map: StreamMap<usize, UsizeStream> = StreamMap::new();
+
+    stream_map.insert(0, Box::pin(iter([0usize].into_iter())) as UsizeStream);
+    stream_map.insert(1, Box::pin(iter([1usize].into_iter())) as UsizeStream);
+
+    let mut buffer = vec![];
+    let n = poll_fn(|cx| pin!(stream_map.next_many(&mut buffer, 1)).poll(cx)).await;
+
+    assert_eq!(n, 1);
+    assert_eq!(buffer.len(), 1);
+
+    let n = poll_fn(|cx| pin!(stream_map.next_many(&mut buffer, 1)).poll(cx)).await;
+
+    assert_eq!(n, 1);
+    assert_eq!(buffer.len(), 2);
+    assert!(buffer.contains(&(0, 0)));
+    assert!(buffer.contains(&(1, 1)));
+}
+
 #[tokio::test]
 async fn next_many_correctly_loops_around() {
     for _ in 0..10 {