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

Code before:
    /// Returns `None` if the queue is empty.
    ///
    /// When this method returns an error, then the id of the task that failed can be accessed
    /// using the [`JoinError::id`] method.
    ///
    /// # Cancel Safety
    ///
    /// 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 [`JoinQueue`].
    ///
    /// [task ID]: tokio::task::Id
    /// [`JoinError::id`]: fn@tokio::task::JoinError::id
    pub async fn join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> {
        std::future::poll_fn(|cx| self.poll_join_next_with_id(cx)).await
    }

    /// Tries to join the next task in FIFO order if it has completed.
    ///
    /// Returns `None` if the queue is empty or if the next task is not yet ready.
    pub fn try_join_next(&mut self) -> Option<Result<T, JoinError>> {
        let jh = self.0.front_mut()?;
        // 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.
        let jh = tokio::task::coop::unconstrained(jh);
        let res = jh.now_or_never()?;
        // Use `detach` to avoid calling `abort` on a task that has already completed.
        // Dropping `AbortOnDropHandle` would abort the task, but since it is finished,
        // we only need to drop the `JoinHandle` for cleanup.
        drop(self.0.pop_front().unwrap().detach());
        Some(res)
    }

    /// Tries to join the next task in FIFO order if it has completed and return its output,
    /// along with its [task ID].
    ///
    /// Returns `None` if the queue is empty or if the next task is not yet ready.
    ///
    /// When this method returns an error, then the id of the task that failed can be accessed
    /// using the [`JoinError::id`] method.
    ///
    /// [task ID]: tokio::task::Id
    /// [`JoinError::id`]: fn@tokio::task::JoinError::id
    pub fn try_join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> {
        let jh = self.0.front_mut()?;
        // 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.
        let jh = tokio::task::coop::unconstrained(jh);
        let res = jh.now_or_never()?;
        // Use `detach` to avoid calling `abort` on a task that has already completed.
        // Dropping `AbortOnDropHandle` would abort the task, but since it is finished,

Review comment (refactor):
nit, the semantic of `now_or_never` is a bit confusing here, could we write an helper function and it internally uses the noop_waker?

Code after:
    ///
    /// When this method returns an error, then the id of the task that failed can be accessed
    /// using the [`JoinError::id`] method.
    ///
    /// # Cancel Safety
    ///
    /// 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 [`JoinQueue`].
    ///
    /// [task ID]: tokio::task::Id
    /// [`JoinError::id`]: fn@tokio::task::JoinError::id
    pub async fn join_next_with_id(&mut self) -> Option<Result<(Id, T), JoinError>> {
        std::future::poll_fn(|cx| self.poll_join_next_with_id(cx)).await
    }

    /// Tries to poll an `AbortOnDropHandle` without blocking or yielding.
    ///
    /// Note that on success the handle will panic on subsequent polls
    /// since it becomes consumed.
    fn try_poll_handle(jh: &mut AbortOnDropHandle<T>) -> Option<Result<T, JoinError>> {
        let waker = futures_util::task::noop_waker();
        let mut cx = Context::from_waker(&waker);

        // 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.
        let jh = std::pin::pin!(tokio::task::coop::unconstrained(jh));
        if let Poll::Ready(res) = jh.poll(&mut cx) {
            Some(res)
        } else {
            None
        }
    }

    /// Tries to join the next task in FIFO order if it has completed.
    ///
    /// Returns `None` if the queue is empty or if the next task is not yet ready.
    pub fn try_join_next(&mut self) -> Option<Result<T, JoinError>> {
        let jh = self.0.front_mut()?;
        let res = Self::try_poll_handle(jh)?;
        // Use `detach` to avoid calling `abort` on a task that has already completed.
        // Dropping `AbortOnDropHandle` would abort the task, but since it is finished,
        // we only need to drop the `JoinHandle` for cleanup.
        drop(self.0.pop_front().unwrap().detach());
        Some(res)
    }

    /// Tries to join the next task in FIFO order if it has completed and return its output,
    /// along with its [task ID].
    ///
    /// Returns `None` if the queue is empty or if the next task is not yet ready.