diff --git a/tokio-test/src/task.rs b/tokio-test/src/task.rs index 2e646d44b..8e88a890b 100644 --- a/tokio-test/src/task.rs +++ b/tokio-test/src/task.rs @@ -70,6 +70,9 @@ const IDLE: usize = 0; const WAKE: usize = 1; const SLEEP: usize = 2; +/// Default maximum number of poll iterations in [`Spawn::poll_until_idle`]. +const POLL_UNTIL_IDLE_MAX_ITERATIONS: usize = 150; + impl Spawn { /// Consumes `self` returning the inner value pub fn into_inner(self) -> T @@ -123,6 +126,43 @@ impl Spawn { let fut = self.future.as_mut(); self.task.enter(|cx| fut.poll(cx)) } + + /// Polls the future until it is idle. + /// + /// A future is considered idle when it either completes, or returns + /// [`Poll::Pending`] without a pending wake notification. + /// + /// Unlike [`poll`](Self::poll), this method keeps polling while the future + /// returns [`Poll::Pending`] but has received a wake notification, advancing + /// the future as far as possible without waiting for external events. + /// + /// Polling is bounded to avoid infinite loops when a future wakes without + /// making progress. + /// + /// # Panics + /// + /// Panics if the iteration limit is exceeded. + /// + /// # Example + /// + /// ``` + /// use tokio_test::task; + /// + /// let mut task = task::spawn(async { 42 }); + /// + /// assert!(task.poll_until_idle().is_ready()); + /// ``` + pub fn poll_until_idle(&mut self) -> Poll { + for _ in 0..POLL_UNTIL_IDLE_MAX_ITERATIONS { + let result = self.poll(); + if result.is_ready() || !self.is_woken() { + return result; + } + } + panic!( + "poll_until_idle exceeded {POLL_UNTIL_IDLE_MAX_ITERATIONS} iterations; future may be waking without making progress" + ); + } } impl Spawn { diff --git a/tokio-test/tests/task.rs b/tokio-test/tests/task.rs index a464bf9b7..5df0b9ffb 100644 --- a/tokio-test/tests/task.rs +++ b/tokio-test/tests/task.rs @@ -1,3 +1,4 @@ +use std::future::{pending, Future}; use std::pin::Pin; use std::task::{Context, Poll}; use tokio_stream::Stream; @@ -23,3 +24,81 @@ fn test_spawn_stream_size_hint() { let spawn = task::spawn(SizedStream); assert_eq!(spawn.size_hint(), (100, Some(200))); } + +#[test] +fn poll_until_idle_ready() { + let mut task = task::spawn(async { 42 }); + assert_eq!(task.poll_until_idle(), Poll::Ready(42)); +} + +#[test] +fn poll_until_idle_pending_not_woken() { + let mut task = task::spawn(pending::<()>()); + assert!(task.poll_until_idle().is_pending()); +} + +struct WakeThenReady { + step: u8, +} + +impl Future for WakeThenReady { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + match self.step { + 0 => { + self.step = 1; + cx.waker().wake_by_ref(); + Poll::Pending + } + _ => Poll::Ready(()), + } + } +} + +#[test] +fn poll_until_idle_advances_on_wake() { + let mut task = task::spawn(WakeThenReady { step: 0 }); + assert!(task.poll_until_idle().is_ready()); +} + +struct WakeNTimes { + remaining: u8, +} + +impl Future for WakeNTimes { + type Output = u8; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + if self.remaining == 0 { + return Poll::Ready(0); + } + self.remaining -= 1; + cx.waker().wake_by_ref(); + Poll::Pending + } +} + +#[test] +fn poll_until_idle_multiple_wakes() { + let mut task = task::spawn(WakeNTimes { remaining: 3 }); + assert_eq!(task.poll_until_idle(), Poll::Ready(0)); +} + +struct WakeForever; + +impl Future for WakeForever { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + cx.waker().wake_by_ref(); + Poll::Pending + } +} + +#[test] +#[should_panic(expected = "poll_until_idle exceeded 150 iterations")] +fn poll_until_idle_panics_on_infinite_wake() { + let mut task = task::spawn(WakeForever); + let _ = task.poll_until_idle(); +}