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

Code before:
    ///
    /// 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_next_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;
        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));
                }
                Poll::Ready(None) | Poll::Pending => break,
            }

            // Some streams may have been removed during the call to `poll_one`,
            // so we may need to regenerate the start index.
            if start >= self.entries.len() {
                start = self::rand::thread_rng_n(self.entries.len() as u32) as usize;
            }
        }

        let added = limit - remaining;

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

Review comment (suggestion):
Hmm. I imagine that when we get many items, we would want to get them from different streams (probably in the order they're stored). But with this strategy, whichever stream is at `start` will be preferred every time.

Code after:
    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;

                        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();
                    }
                }