Repository: MystenLabs/sui
File: crates/sui-futures/src/stream.rs (Rust)

Code before:
        let mut error = None;

        let mut self_ = pin!(self);

        loop {
            // Eager inner loop: spawn tasks while permits allow and items are ready,
            // avoiding select! overhead when items are immediately available.
            while !draining && permits > 0 {
                match poll_fn(|cx| self_.as_mut().poll_next(cx)).now_or_never() {
                    Some(Some(item)) => {
                        permits -= 1;
                        join_set.spawn(f(item));
                    }
                    Some(None) => {
                        // If the stream is empty, signal that the worker pool is going to
                        // start draining now, so that once we get all our permits back, we
                        // know we can wind down the pool.
                        draining = true;
                    }
                    None => break,
                }
            }

            if draining && permits == limit {
                break;
            }

            tokio::select! {
                biased;

                Some(res) = join_set.join_next() => {
                    match res {
                        Ok(Err(e)) if error.is_none() => {
                            error = Some(e);
                            permits += 1;
                            draining = true;
                        }

                        Ok(_) => permits += 1,

                        // Worker panicked, propagate the panic.
                        Err(e) if e.is_panic() => {
                            panic::resume_unwind(e.into_panic())
                        }

                        // Worker was cancelled -- this can only happen if its join handle was
                        // cancelled (not possible because that was created in this function),
                        // or the runtime it was running in was wound down, in which case,
                        // prepare the worker pool to drain.
                        Err(e) => {
                            assert!(e.is_cancelled());

Review comment (suggestion):
Do you need this check here and below? If this condition is true then:

- `join_set` is empty, so `join_set.join_next()` will return `None`
- the guard on the `next` pattern is false (because it's guarded by `!draining`)

So we will evaluate the `else` condition straight away and exit.

Code after:
        let mut error = None;

        let mut self_ = pin!(self);

        loop {
            // Eager inner loop: spawn tasks while permits allow and items are ready,
            // avoiding select! overhead when items are immediately available.
            while !draining && permits > 0 {
                match poll_fn(|cx| self_.as_mut().poll_next(cx)).now_or_never() {
                    Some(Some(item)) => {
                        permits -= 1;
                        join_set.spawn(f(item));
                    }
                    Some(None) => {
                        // If the stream is empty, signal that the worker pool is going to
                        // start draining now, so that once we get all our permits back, we
                        // know we can wind down the pool.
                        draining = true;
                    }
                    None => break,
                }
            }

            tokio::select! {
                biased;

                Some(res) = join_set.join_next() => {
                    match res {
                        Ok(Err(e)) if error.is_none() => {
                            error = Some(e);
                            permits += 1;
                            draining = true;
                        }

                        Ok(_) => permits += 1,

                        // Worker panicked, propagate the panic.
                        Err(e) if e.is_panic() => {
                            panic::resume_unwind(e.into_panic())
                        }

                        // Worker was cancelled -- this can only happen if its join handle was
                        // cancelled (not possible because that was created in this function),
                        // or the runtime it was running in was wound down, in which case,
                        // prepare the worker pool to drain.
                        Err(e) => {
                            assert!(e.is_cancelled());
                            permits += 1;
                            draining = true;
                        }
                    }