Repository: tokio-rs/tokio
File: tokio-stream/src/stream_map.rs (Rust)

Code before:
                Poll::Ready(None) | Poll::Pending => break,
            }
        }

        let added = limit - remaining;

        if added > 0 {
            Poll::Ready(added)
        } else if self.entries.is_empty() {
            Poll::Ready(0)
        } else {
            Poll::Pending
        }
    }

    /// Receives multiple items on this [`StreamMap`], extending the provided `buffer`.
    ///
    /// This method returns the number of items that is appended to the `buffer`.
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe. If `recv_many` is used as the event in a
    /// [`tokio::select!`](tokio::select) statement and some other branch
    /// completes first, it is guaranteed that no items were received on any of
    /// the underlying streams.
    pub async fn recv_many(&mut self, buffer: &mut Vec<(K, V::Item)>, limit: usize) -> usize {
        poll_fn(|cx| self.poll_recv_many(cx, buffer, limit)).await
    }
}

impl<K, V> Stream for StreamMap<K, V>
where
    K: Clone + Unpin,
    V: Stream + Unpin,
{
    type Item = (K, V::Item);

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if let Some((idx, val)) = ready!(self.poll_next_entry(cx)) {
            let key = self.entries[idx].0.clone();
            Poll::Ready(Some((key, val)))
        } else {
            Poll::Ready(None)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let mut ret = (0, Some(0));

        for (_, stream) in &self.entries {
            let hint = stream.size_hint();

Review comment (suggestion):
The behavior of this method can be a bit confusing, so the documentation needs to be really clear about what the function does. Please clarify that it can return less than `limit` items even if the streams are not closed.

Code after:
        }

        let mut added = 0;

        let start = self::rand::thread_rng_n(self.entries.len() as u32) as usize;
        let mut idx = start;

        while added < limit {
            // Indicates whether at least one stream returned a value when polled or not
            let mut should_loop = false;

            for _ in 0..self.entries.len() {
                let (_, stream) = &mut self.entries[idx];

                match Pin::new(stream).poll_next(cx) {
                    Poll::Ready(Some(val)) => {
                        added += 1;

                        let key = self.entries[idx].0.clone();
                        buffer.push((key, val));

                        should_loop = true;

                        idx = idx.wrapping_add(1) % self.entries.len();
                    }
                    Poll::Ready(None) => {
                        // Remove the entry
                        self.entries.swap_remove(idx);

                        // Check if this was the last entry, if so the cursor needs
                        // to wrap
                        if idx == self.entries.len() {
                            idx = 0;
                        } else if idx < start && start <= self.entries.len() {
                            // The stream being swapped into the current index has
                            // already been polled, so skip it.
                            idx = idx.wrapping_add(1) % self.entries.len();
                        }
                    }
                    Poll::Pending => {
                        idx = idx.wrapping_add(1) % self.entries.len();
                    }
                }
            }

            if !should_loop {
                break;
            }
        }

        if added > 0 {