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

Code before:
    /// * `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_recv_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_recv_many(
        &mut self,
        cx: &mut Context<'_>,
        buffer: &mut Vec<(K, V::Item)>,
        limit: usize,
    ) -> Poll<usize> {
        if limit == 0 {
            return Poll::Ready(0);
        }

        let mut remaining = limit;

        while remaining > 0 {
            match self.poll_next_entry(cx) {
                Poll::Ready(Some((idx, val))) => {
                    remaining -= 1;
                    let key = self.entries[idx].0.clone();
                    buffer.push((key, val));
                }
                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

Review comment (suggestion):
If we're going to the effort of providing an implementation, then I think we should not go for a version that just calls `poll_next_entry` in a loop. At that point, it might as well be a generic piece of functionality that works for any `Stream`.

Each call to that method generates a random number to pick a stream to start at, but I think it would make more sense to pick a random number only once, and then keep going around the loop until we have `limit` items, or until we've polled `self.entries.len()` times in a row without getting an item.

Code after:
    ///
    /// # 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.
    ///
    /// 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));