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

Code before:
            items: Vec::with_capacity(max_size),
            cap: max_size,
        }
    }

    /// Consumes the [`ChunksTimeout`] and then returns all buffered items.
    ///
    /// ```
    /// use tokio::time::Duration;
    /// use tokio_stream::{self as stream, StreamExt};
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let chunked = stream::iter([1, 2, 3, 4])
    ///        .throttle(Duration::from_millis(2))
    ///        .chunks_timeout(4, Duration::from_millis(8));
    ///
    ///     tokio::pin!(chunked);
    ///
    ///     // race with another future
    ///     tokio::select! {
    ///         Some(_chunk) = chunked.next() => {}
    ///         _ = tokio::time::sleep(Duration::from_millis(3)) => {}
    ///     }
    ///
    ///     assert_eq!(chunked.into_remainder(), vec![1, 2]);
    /// }
    /// ```
    pub fn into_remainder(mut self: Pin<&mut Self>) -> Vec<S::Item> {
        let me = self.as_mut().project();
        std::mem::take(me.items)
    }
}

impl<S: Stream> Stream for ChunksTimeout<S> {
    type Item = Vec<S::Item>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut me = self.as_mut().project();
        loop {
            match me.stream.as_mut().poll_next(cx) {
                Poll::Pending => break,
                Poll::Ready(Some(item)) => {
                    if me.items.is_empty() {
                        me.deadline.set(Some(sleep(*me.duration)));
                        me.items.reserve_exact(*me.cap);
                    }
                    me.items.push(item);
                    if me.items.len() >= *me.cap {
                        return Poll::Ready(Some(std::mem::take(me.items)));
                    }

Review comment (refactor):
Please try to avoid introducing timers into tests, this can help us to make the test less flaky.  I suggest either improving this test or remove this test.

Code after:
            items: Vec::with_capacity(max_size),
            cap: max_size,
        }
    }

    /// Consumes the [`ChunksTimeout`] and then returns all buffered items.
    pub fn into_remainder(mut self: Pin<&mut Self>) -> Vec<S::Item> {
        let me = self.as_mut().project();
        std::mem::take(me.items)
    }
}

impl<S: Stream> Stream for ChunksTimeout<S> {
    type Item = Vec<S::Item>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut me = self.as_mut().project();
        loop {
            match me.stream.as_mut().poll_next(cx) {
                Poll::Pending => break,
                Poll::Ready(Some(item)) => {
                    if me.items.is_empty() {
                        me.deadline.set(Some(sleep(*me.duration)));
                        me.items.reserve_exact(*me.cap);
                    }
                    me.items.push(item);
                    if me.items.len() >= *me.cap {
                        return Poll::Ready(Some(std::mem::take(me.items)));
                    }
                }
                Poll::Ready(None) => {
                    // Returning Some here is only correct because we fuse the inner stream.
                    let last = if me.items.is_empty() {
                        None
                    } else {
                        Some(std::mem::take(me.items))
                    };

                    return Poll::Ready(last);
                }
            }
        }

        if !me.items.is_empty() {
            if let Some(deadline) = me.deadline.as_pin_mut() {
                ready!(deadline.poll(cx));
            }
            return Poll::Ready(Some(std::mem::take(me.items)));
        }

        Poll::Pending