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

Code before:
    /// as many items as it can up to the given limit. This method returns
    /// zero only if the `StreamMap` is empty (or if `limit` is zero).
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe. If `next_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 next_many(&mut self, buffer: &mut Vec<(K, V::Item)>, limit: usize) -> usize {
        poll_fn(|cx| self.poll_next_many(cx, buffer, limit)).await
    }

    /// Polls to receive multiple items on this `StreamMap`, extending the provided `buffer`.
    ///
    /// This method returns:
    /// * `Poll::Pending` if no items are available but the `StreamMap` is not empty.
    /// * `Poll::Ready(count)` where `count` is the number of items successfully received and
    ///   stored in `buffer`. This can be less than, or equal to, `limit`.
    /// * `Poll::Ready(0)` if `limit` is set to zero or when the `StreamMap` is empty.
    ///
    /// When the method returns `Poll::Pending`, the `Waker` in the provided
    /// `Context` is scheduled to receive a wakeup when an item is sent on any of the
    /// underlying stream. Note that on multiple calls to `poll_next_many`, only
    /// the `Waker` from the `Context` passed to the most recent call is scheduled
    /// to receive a wakeup.
    ///
    /// Note that this method does not guarantee that exactly `limit` items
    /// are received. Rather, if at least one item is available, it returns
    /// as many items as it can up to the given limit. This method returns
    /// zero only if the `StreamMap` is empty (or if `limit` is zero).
    pub fn poll_next_many(
        &mut self,
        cx: &mut Context<'_>,
        buffer: &mut Vec<(K, V::Item)>,
        limit: usize,
    ) -> Poll<usize> {
        if limit == 0 || self.entries.is_empty() {
            return Poll::Ready(0);
        }

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

        while remaining > 0 {
            match self.poll_one(cx, start) {
                Poll::Ready(Some((idx, val))) => {
                    remaining -= 1;
                    let key = self.entries[idx].0.clone();
                    buffer.push((key, val));


Review comment (suggestion):
I think this behavior depends on the underlying `Stream` (or Runtime) implementation. I guess `StreamMap` is expected to be used by various Stream implementations, so the description here may be unnecessary.

Code after:
    /// This method is cancel safe. If `next_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 next_many(&mut self, buffer: &mut Vec<(K, V::Item)>, limit: usize) -> usize {
        poll_fn(|cx| self.poll_next_many(cx, buffer, limit)).await
    }

    /// Polls to receive multiple items on this `StreamMap`, extending the provided `buffer`.
    ///
    /// This method returns:
    /// * `Poll::Pending` if no items are available but the `StreamMap` is not empty.
    /// * `Poll::Ready(count)` where `count` is the number of items successfully received and
    ///   stored in `buffer`. This can be less than, or equal to, `limit`.
    /// * `Poll::Ready(0)` if `limit` is set to zero or when the `StreamMap` is empty.
    ///
    /// Note that this method does not guarantee that exactly `limit` items
    /// are received. Rather, if at least one item is available, it returns
    /// as many items as it can up to the given limit. This method returns
    /// zero only if the `StreamMap` is empty (or if `limit` is zero).
    pub fn poll_next_many(
        &mut self,
        cx: &mut Context<'_>,
        buffer: &mut Vec<(K, V::Item)>,
        limit: usize,
    ) -> Poll<usize> {
        if limit == 0 || self.entries.is_empty() {
            return Poll::Ready(0);
        }

        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;