Wakers can be lost when simultaneously flushing and reading
This is similar to #92 but happens with a slightly different task poll order.
`poll_flush` doesn't check that the read side is already flushing, so it overwrites the read side's waker, resulting in `task2 woke` being printed instead of `task1 woke`.
Reproduction code:
```rust
use futures_util::{SinkExt, StreamExt};
use std::{
io::Result,
pin::Pin,
task::{Context, Poll, Waker},
};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio_websockets::{ClientBuilder, Message};
use waker_fn::waker_fn;
struct MockStream {
tx: Waker,
}
impl AsyncRead for MockStream {
fn poll_read(
self: Pin<&mut Self>,
_: &mut Context<'_>,
_: &mut ReadBuf<'_>,
) -> Poll<Result<()>> {
Poll::Ready(Ok(()))
}
}
impl AsyncWrite for MockStream {
fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, _: &[u8]) -> Poll<Result<usize>> {
self.tx = cx.waker().clone();
Poll::Pending
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
self.tx = cx.waker().clone();
Poll::Pending
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
self.tx = cx.waker().clone();
Poll::Pending
}
}
fn main() {
let mut ws = ClientBuilder::new().take_over(MockStream {
tx: Waker::noop().clone(),
});
// queue a message to trigger a flush on next read
let _ = ws.start_send_unpin(Message::text("message"));
let task1 = waker_fn(|| println!("task1 woke"));
let task2 = waker_fn(|| println!("task2 woke"));
// queued message triggers a flush and registers task1 with the stream
let _ = ws.poll_next_unpin(&mut Context::from_waker(&task1));
// assert!(ws.flush_in_progress);
// explicit flush overwrites task1's waker
let _ = ws.poll_flush_unpin(&mut Context::from_waker(&task2));
ws.get_ref().tx.wake_by_ref();
}
```
```toml
[package]
name = "tokio-ws-reproduce"
version = "0.1.0"
edition = "2024"
[dependencies]
futures-util = { version = "0.3.31", features = ["sink"] }
tokio = "1.44.2"
tokio-websockets = { version = "0.11.3", features = ["client", "sha1_smol", "rand"] }
waker-fn = "1.2.0"
```
4 条评论