From 98104c36f1377cb12a24911f12fe51a7899e7d7a Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Wed, 1 Jul 2026 17:55:20 +0200
Subject: [PATCH] stream: honor `StreamMap::next_many` limit (#8215)

---
 tokio-stream/src/stream_map.rs          |  4 +++
 tokio-stream/tests/stream_stream_map.rs | 42 +++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/tokio-stream/src/stream_map.rs b/tokio-stream/src/stream_map.rs
index fef1c082b..bea730106 100644
--- a/tokio-stream/src/stream_map.rs
+++ b/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
diff --git a/tokio-stream/tests/stream_stream_map.rs b/tokio-stream/tests/stream_stream_map.rs
index 6290f85cd..d600db1ea 100644
--- a/tokio-stream/tests/stream_stream_map.rs
+++ b/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 {