Skip to content

Commit c259165

Browse files
committed
test(subscriber): test with custom self_wake() function
Part of the testing performed in the `console-subscriber` integration tests is detecting self wakes. This relied upon the `yield_now()` from Tokio. However, the behavior of this function was changed in tokio-rs/tokio#5223 and since Tokio 1.23 the wake doesn't occur in the task that `yield_now()` is called from. This breaks the test when using a newer version of Tokio. This change replaces the use of `yield_now()` with a custom `self_wake()` function that returns a future which does perform a self wake (wakes the task from within itself before returning `Poll::Pending`). The same custom `self_wake()` is also included in the `app` example so that it shows self wakes correctly. Tokio has been updated to 1.28.2 in the lock file (the last with compatible MSRV) so that this fix is tested. Ref #512
1 parent 28a27fc commit c259165

File tree

5 files changed

+82
-16
lines changed

5 files changed

+82
-16
lines changed

Cargo.lock

+6-8
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

console-subscriber/examples/app.rs

+28-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::time::Duration;
1+
use std::{future::Future, task::Poll, time::Duration};
22

33
static HELP: &str = r#"
44
Example console-instrumented app
@@ -121,7 +121,7 @@ async fn burn(min: u64, max: u64) {
121121
loop {
122122
for i in min..max {
123123
for _ in 0..i {
124-
tokio::task::yield_now().await;
124+
self_wake().await;
125125
}
126126
tokio::time::sleep(Duration::from_secs(i - min)).await;
127127
}
@@ -152,3 +152,29 @@ async fn spawn_blocking(seconds: u64) {
152152
.await;
153153
}
154154
}
155+
156+
fn self_wake() -> impl Future<Output = ()> {
157+
struct SelfWake {
158+
yielded: bool,
159+
}
160+
161+
impl Future for SelfWake {
162+
type Output = ();
163+
164+
fn poll(
165+
mut self: std::pin::Pin<&mut Self>,
166+
cx: &mut std::task::Context<'_>,
167+
) -> Poll<Self::Output> {
168+
if self.yielded == true {
169+
return Poll::Ready(());
170+
}
171+
172+
self.yielded = true;
173+
cx.waker().wake_by_ref();
174+
175+
Poll::Pending
176+
}
177+
}
178+
179+
SelfWake { yielded: false }
180+
}

console-subscriber/tests/framework.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ fn self_wakes() {
6363
.match_default_name()
6464
.expect_self_wakes(1);
6565

66-
let future = async { task::yield_now().await };
66+
let future = async { support::self_wake().await };
6767

6868
assert_task(expected_task, future);
6969
}

console-subscriber/tests/support/mod.rs

+45-3
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
1-
use futures::Future;
1+
use std::{future::Future, task::Poll};
2+
3+
use tokio::task::JoinHandle;
24

35
mod state;
46
mod subscriber;
57
mod task;
68

79
use subscriber::run_test;
8-
910
pub(crate) use subscriber::MAIN_TASK_NAME;
1011
pub(crate) use task::ExpectedTask;
11-
use tokio::task::JoinHandle;
1212

1313
/// Assert that an `expected_task` is recorded by a console-subscriber
1414
/// when driving the provided `future` to completion.
@@ -62,3 +62,45 @@ where
6262
.spawn(f)
6363
.expect(&format!("spawning task '{name}' failed"))
6464
}
65+
66+
/// Wakes itself from within this task.
67+
///
68+
/// This function returns a future which will wake itself and then
69+
/// return `Poll::Pending` the first time it is called. The next time
70+
/// it will return `Poll::Ready`.
71+
///
72+
/// This is the old behavior of Tokio's [`yield_now()`] function, before it
73+
/// was improved in [tokio-rs/tokio#5223] to avoid starving the resource
74+
/// drivers.
75+
///
76+
/// Awaiting the future returned from this function will result in a
77+
/// self-wake being recorded.
78+
///
79+
/// [`yield_now()`]: fn@tokio::task::yield_now
80+
/// [tokio-rs/tokio#5223]: https://github.com/tokio-rs/tokio/pull/5223
81+
#[allow(dead_code)]
82+
pub(crate) fn self_wake() -> impl Future<Output = ()> {
83+
struct SelfWake {
84+
yielded: bool,
85+
}
86+
87+
impl Future for SelfWake {
88+
type Output = ();
89+
90+
fn poll(
91+
mut self: std::pin::Pin<&mut Self>,
92+
cx: &mut std::task::Context<'_>,
93+
) -> Poll<Self::Output> {
94+
if self.yielded == true {
95+
return Poll::Ready(());
96+
}
97+
98+
self.yielded = true;
99+
cx.waker().wake_by_ref();
100+
101+
Poll::Pending
102+
}
103+
}
104+
105+
SelfWake { yielded: false }
106+
}

console-subscriber/tests/wake.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ mod support;
22
use std::time::Duration;
33

44
use support::{assert_task, ExpectedTask};
5-
use tokio::{task, time::sleep};
5+
use tokio::time::sleep;
66

77
#[test]
88
fn sleep_wakes() {
@@ -41,7 +41,7 @@ fn self_wake() {
4141
.expect_self_wakes(1);
4242

4343
let future = async {
44-
task::yield_now().await;
44+
support::self_wake().await;
4545
};
4646

4747
assert_task(expected_task, future);

0 commit comments

Comments
 (0)