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

Code before:
}

cfg_rt! {
    /// Forcibly removes the budgeting constraints early.
    ///
    /// Returns the remaining budget
    pub(crate) fn stop() -> Budget {
        context::budget(|cell| {
            let prev = cell.get();
            cell.set(Budget::unconstrained());
            prev
        }).unwrap_or(Budget::unconstrained())
    }
}

cfg_coop! {
    use pin_project_lite::pin_project;
    use std::cell::Cell;
    use std::future::Future;
    use std::pin::Pin;
    use std::task::{ready, Context, Poll};

    /// Value returned by the [`poll_proceed`](poll_proceed) method.
    #[derive(Debug)]
    #[must_use]
    pub struct RestoreOnPending(Cell<Budget>);

    impl RestoreOnPending {
        /// Signals that the task that obtained this `RestoreOnPending` was able to make
        /// progress. This prevents the task budget from being restored to the value
        /// it had prior to obtaining this instance when it is dropped.
        pub fn made_progress(&self) {
            self.0.set(Budget::unconstrained());
        }
    }

    impl Drop for RestoreOnPending {
        fn drop(&mut self) {
            // Don't reset if budget was unconstrained or if we made progress.
            // They are both represented as the remembered budget being unconstrained.
            let budget = self.0.get();
            if !budget.is_unconstrained() {
                let _ = context::budget(|cell| {
                    cell.set(budget);
                });
            }
        }
    }

    /// Decrements the task budget and returns `Poll::Pending` if the budget is depleted.
    /// This indicates that the task should yield to the scheduler. Otherwise, returns

Review comment (suggestion):
New types should be added to `tests/async_send_sync.rs`.

Code after:
}

cfg_rt! {
    /// Forcibly removes the budgeting constraints early.
    ///
    /// Returns the remaining budget
    pub(crate) fn stop() -> Budget {
        context::budget(|cell| {
            let prev = cell.get();
            cell.set(Budget::unconstrained());
            prev
        }).unwrap_or(Budget::unconstrained())
    }
}

cfg_coop! {
    use pin_project_lite::pin_project;
    use std::cell::Cell;
    use std::future::Future;
    use std::marker::PhantomData;
    use std::pin::Pin;
    use std::task::{ready, Context, Poll};

    /// Value returned by the [`poll_proceed`] method.
    #[derive(Debug)]
    #[must_use]
    pub struct RestoreOnPending(Cell<Budget>, PhantomData<*mut ()>);

    impl RestoreOnPending {
        fn new(budget: Budget) -> Self {
            RestoreOnPending(
                Cell::new(budget),
                PhantomData,
            )
        }

        /// Signals that the task that obtained this `RestoreOnPending` was able to make
        /// progress. This prevents the task budget from being restored to the value
        /// it had prior to obtaining this instance when it is dropped.
        pub fn made_progress(&self) {
            self.0.set(Budget::unconstrained());
        }
    }

    impl Drop for RestoreOnPending {
        fn drop(&mut self) {
            // Don't reset if budget was unconstrained or if we made progress.
            // They are both represented as the remembered budget being unconstrained.
            let budget = self.0.get();
            if !budget.is_unconstrained() {
                let _ = context::budget(|cell| {