ITADN

Wakers can be lost when simultaneously flushing and reading

#106Openr58Playz 创建于 2025-04-11
R
r58Playzcommented
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 条评论