test(subscriber): test with custom self_wake() function

console-subscriber/examples/app.rs

@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{future::Future, task::Poll, time::Duration};
 
 static HELP: &str = r#"
 Example console-instrumented app
@@ -121,7 +121,7 @@ async fn burn(min: u64, max: u64) {
     loop {
         for i in min..max {
             for _ in 0..i {
-                tokio::task::yield_now().await;
+                self_wake().await;
             }
             tokio::time::sleep(Duration::from_secs(i - min)).await;
         }
@@ -152,3 +152,29 @@ async fn spawn_blocking(seconds: u64) {
         .await;
     }
 }
+
+fn self_wake() -> impl Future<Output = ()> {
+    struct SelfWake {
+        yielded: bool,
+    }
+
+    impl Future for SelfWake {
+        type Output = ();
+
+        fn poll(
+            mut self: std::pin::Pin<&mut Self>,
+            cx: &mut std::task::Context<'_>,
+        ) -> Poll<Self::Output> {
+            if self.yielded == true {
+                return Poll::Ready(());
+            }
+
+            self.yielded = true;
+            cx.waker().wake_by_ref();
+
+            Poll::Pending
+        }
+    }
+
+    SelfWake { yielded: false }
+}

console-subscriber/tests/framework.rs

@@ -63,7 +63,7 @@ fn self_wakes() {
         .match_default_name()
         .expect_self_wakes(1);
 
-    let future = async { task::yield_now().await };
+    let future = async { support::self_wake().await };
 
     assert_task(expected_task, future);
 }

console-subscriber/tests/support/mod.rs

@@ -1,14 +1,14 @@
-use futures::Future;
+use std::{future::Future, task::Poll};
+
+use tokio::task::JoinHandle;
 
 mod state;
 mod subscriber;
 mod task;
 
 use subscriber::run_test;
-
 pub(crate) use subscriber::MAIN_TASK_NAME;
 pub(crate) use task::ExpectedTask;
-use tokio::task::JoinHandle;
 
 /// Assert that an `expected_task` is recorded by a console-subscriber
 /// when driving the provided `future` to completion.
@@ -62,3 +62,45 @@ where
         .spawn(f)
         .expect(&format!("spawning task '{name}' failed"))
 }
+
+/// Wakes itself from within this task.
+///
+/// This function returns a future which will wake itself and then
+/// return `Poll::Pending` the first time it is called. The next time
+/// it will return `Poll::Ready`.
+///
+/// This is the old behavior of Tokio's [`yield_now()`] function, before it
+/// was improved in [tokio-rs/tokio#5223] to avoid starving the resource
+/// drivers.
+///
+/// Awaiting the future returned from this function will result in a
+/// self-wake being recorded.
+///
+/// [`yield_now()`]: fn@tokio::task::yield_now
+/// [tokio-rs/tokio#5223]: https://github.com/tokio-rs/tokio/pull/5223
+#[allow(dead_code)]
+pub(crate) fn self_wake() -> impl Future<Output = ()> {
+    struct SelfWake {
+        yielded: bool,
+    }
+
+    impl Future for SelfWake {
+        type Output = ();
+
+        fn poll(
+            mut self: std::pin::Pin<&mut Self>,
+            cx: &mut std::task::Context<'_>,
+        ) -> Poll<Self::Output> {
+            if self.yielded == true {
+                return Poll::Ready(());
+            }
+
+            self.yielded = true;
+            cx.waker().wake_by_ref();
+
+            Poll::Pending
+        }
+    }
+
+    SelfWake { yielded: false }
+}

console-subscriber/tests/wake.rs

@@ -2,7 +2,7 @@ mod support;
 use std::time::Duration;
 
 use support::{assert_task, ExpectedTask};
-use tokio::{task, time::sleep};
+use tokio::time::sleep;
 
 #[test]
 fn sleep_wakes() {
@@ -41,7 +41,7 @@ fn self_wake() {
         .expect_self_wakes(1);
 
     let future = async {
-        task::yield_now().await;
+        support::self_wake().await;
     };
 
     assert_task(expected_task, future);