Repository: tokio-rs/tokio
File: tokio/src/task/join_set.rs (Rust)

Code before:
    /// This method is cancel safe. If `join_next_with_id` is used as the event in a `tokio::select!`
    /// statement and some other branch completes first, it is guaranteed that no tasks were
    /// removed from this `JoinSet`.
    ///
    /// [task ID]: crate::task::Id
    /// [`JoinError::id`]: fn@crate::task::JoinError::id
    #[cfg(tokio_unstable)]
    #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
    pub async fn join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> {
        crate::future::poll_fn(|cx| self.poll_join_next_with_id(cx)).await
    }

    /// Tries to join one of the tasks in the set that has completed and return its output.
    ///
    /// Returns `None` if the set is empty.    
    pub fn try_join_next(&mut self) -> Option<Result<T, JoinError>> {
        let mut entry = self.inner.try_pop_notified()?;

        let res = entry.with_value_and_context(|jh, ctx| Pin::new(jh).poll(ctx));

        if let Poll::Ready(res) = res {
            let _entry = entry.remove();
            Some(res)
        } else {
            None
        }
    }

    /// Aborts all tasks and waits for them to finish shutting down.
    ///
    /// Calling this method is equivalent to calling [`abort_all`] and then calling [`join_next`] in
    /// a loop until it returns `None`.
    ///
    /// This method ignores any panics in the tasks shutting down. When this call returns, the
    /// `JoinSet` will be empty.
    ///
    /// [`abort_all`]: fn@Self::abort_all
    /// [`join_next`]: fn@Self::join_next
    pub async fn shutdown(&mut self) {
        self.abort_all();
        while self.join_next().await.is_some() {}
    }

    /// Aborts all tasks on this `JoinSet`.
    ///
    /// This does not remove the tasks from the `JoinSet`. To wait for the tasks to complete
    /// cancellation, you should call `join_next` in a loop until the `JoinSet` is empty.
    pub fn abort_all(&mut self) {
        self.inner.for_each(|jh| jh.abort());
    }


Review comment (suggestion):
This case happens if we successfully pop a `JoinHandle` from the list of notified tasks, but it turns out that the `JoinHandle` is not actually ready yet. In this case, we should not return `None` because there might be another `JoinHandle` in the list of notified task that is actually ready.

Instead, please wrap this function in a loop and go around the loop to try again in this case.

Code after:
    /// This method is cancel safe. If `join_next_with_id` is used as the event in a `tokio::select!`
    /// statement and some other branch completes first, it is guaranteed that no tasks were
    /// removed from this `JoinSet`.
    ///
    /// [task ID]: crate::task::Id
    /// [`JoinError::id`]: fn@crate::task::JoinError::id
    #[cfg(tokio_unstable)]
    #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
    pub async fn join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> {
        crate::future::poll_fn(|cx| self.poll_join_next_with_id(cx)).await
    }

    /// Tries to join one of the tasks in the set that has completed and return its output.
    ///
    /// Returns `None` if the set is empty.    
    pub fn try_join_next(&mut self) -> Option<Result<T, JoinError>> {
        // Loop over all notified `JoinHandle`s to find one that's ready, or until none are left.
        loop {
            let mut entry = self.inner.try_pop_notified()?;

            let res = entry.with_value_and_context(|jh, ctx| {
                // Since this function is not async and cannot be forced to yield, we should
                // disable budgeting when we want to check for the `JoinHandle` readiness.
                Pin::new(&mut unconstrained(jh)).poll(ctx)
            });

            if let Poll::Ready(res) = res {
                let _entry = entry.remove();

                return Some(res);
            }
        }
    }

    /// Tries to join one of the tasks in the set that has completed and return its output,
    /// along with the [task ID] of the completed task.
    ///
    /// Returns `None` if the set is empty.
    ///
    /// When this method returns an error, then the id of the task that failed can be accessed
    /// using the [`JoinError::id`] method.
    ///
    /// [task ID]: crate::task::Id
    /// [`JoinError::id`]: fn@crate::task::JoinError::id
    #[cfg(tokio_unstable)]
    #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
    pub fn try_join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> {
        // Loop over all notified `JoinHandle`s to find one that's ready, or until none are left.
        loop {
            let mut entry = self.inner.try_pop_notified()?;