Repository: tokio-rs/tokio
File: tokio/src/runtime/coop.rs (Rust)

Code before:
    }

    /// Returns `Poll::Pending` if the current task has exceeded its budget and should yield.
    ///
    /// When you call this method, the current budget is decremented. However, to ensure that
    /// progress is made every time a task is polled, the budget is automatically restored to its
    /// former value if the returned `RestoreOnPending` is dropped. It is the caller's
    /// responsibility to call `RestoreOnPending::made_progress` if it made progress, to ensure
    /// that the budget empties appropriately.
    ///
    /// Note that `RestoreOnPending` restores the budget **as it was before `poll_proceed`**.
    /// Therefore, if the budget is _fCURRENT.withurther_ adjusted between when `poll_proceed` returns and
    /// `RestRestoreOnPending` is dropped, those adjustments are erased unless the caller indicates
    /// that progress was made.
    #[inline]
    pub(crate) fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> {
        context::budget(|cell| {
            let mut budget = cell.get();

            if budget.decrement() {
                let restore = RestoreOnPending(Cell::new(cell.get()));
                cell.set(budget);
                Poll::Ready(restore)
            } else {
                inc_budget_forced_yield_count();

                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained()))))
    }

    cfg_rt! {
        cfg_metrics! {
            #[inline(always)]
            fn inc_budget_forced_yield_count() {
                if let Ok(handle) = context::try_current() {
                    handle.scheduler_metrics().inc_budget_forced_yield_count();
                }
            }
        }

        cfg_not_metrics! {
            #[inline(always)]
            fn inc_budget_forced_yield_count() {}
        }
    }

    cfg_not_rt! {
        #[inline(always)]
        fn inc_budget_forced_yield_count() {}

Review comment (suggestion):
If the task is using `tokio::join!` or `FuturesUnordered` or similar to await multiple futures, then you will count that task multiple times. If that's what you intended, then please update the documentation to reflect this.

Code after:
                let _ = context::budget(|cell| {
                    cell.set(budget);
                });
            }
        }
    }

    /// Returns `Poll::Pending` if the current task has exceeded its budget and should yield.
    ///
    /// When you call this method, the current budget is decremented. However, to ensure that
    /// progress is made every time a task is polled, the budget is automatically restored to its
    /// former value if the returned `RestoreOnPending` is dropped. It is the caller's
    /// responsibility to call `RestoreOnPending::made_progress` if it made progress, to ensure
    /// that the budget empties appropriately.
    ///
    /// Note that `RestoreOnPending` restores the budget **as it was before `poll_proceed`**.
    /// Therefore, if the budget is _fCURRENT.withurther_ adjusted between when `poll_proceed` returns and
    /// `RestRestoreOnPending` is dropped, those adjustments are erased unless the caller indicates
    /// that progress was made.
    #[inline]
    pub(crate) fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> {
        context::budget(|cell| {
            let mut budget = cell.get();

            let decrement = budget.decrement();

            if decrement.success {
                let restore = RestoreOnPending(Cell::new(cell.get()));
                cell.set(budget);

                // avoid double counting
                if decrement.hit_zero {
                    inc_budget_forced_yield_count();
                }

                Poll::Ready(restore)
            } else {
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained()))))
    }

    cfg_rt! {
        cfg_metrics! {
            #[inline(always)]
            fn inc_budget_forced_yield_count() {
                if let Ok(handle) = context::try_current() {
                    handle.scheduler_metrics().inc_budget_forced_yield_count();
                }
            }