fix(websocket): lossless EOF signal + pending sweep (WS-02, CON-02)
Replace the axum/tungstenite pump paths' Notify-based read-EOF signal with a retained tokio watch channel so a late subscriber observes EOF regardless of when it fired. Extend the from_wss drop monitor to sweep the pending map (1 s interval) once EOF is observed, so calls registered after the initial fail_all also resolve retryable instead of hanging. cargo test; cargo clippy --all-targets -- -D warnings (default + all-features); cargo fmt --check
This commit is contained in:
@@ -36,9 +36,6 @@ use std::{
|
||||
task::{Context, Poll},
|
||||
};
|
||||
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::ws::{CloseFrame, Message as AxumMessage, WebSocket};
|
||||
use futures::channel::mpsc as futures_mpsc;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
@@ -84,7 +81,7 @@ pub struct WsPumps {
|
||||
read_task: tokio::task::JoinHandle<()>,
|
||||
write_task: tokio::task::JoinHandle<()>,
|
||||
#[cfg(feature = "wss")]
|
||||
read_eof: Arc<tokio::sync::Notify>,
|
||||
read_eof: tokio::sync::watch::Sender<bool>,
|
||||
}
|
||||
|
||||
impl WsPumps {
|
||||
@@ -93,12 +90,35 @@ impl WsPumps {
|
||||
self.write_task.abort();
|
||||
}
|
||||
|
||||
/// Fires when the WS read side reaches EOF (socket close from either
|
||||
/// side) — used by `from_wss`'s connection-drop monitor to await
|
||||
/// socket EOF (ADR-070).
|
||||
/// A lossless receiver for the WS read-EOF signal (socket close from
|
||||
/// either side): the watch channel retains the latest value, so an
|
||||
/// EOF signaled at any point — including before the receiver is
|
||||
/// taken or the observer starts awaiting — is still observed, and
|
||||
/// may be observed repeatedly. Used by `from_wss`'s
|
||||
/// connection-drop monitor (ADR-070); await it with
|
||||
/// [`wait_for_eof`].
|
||||
#[cfg(feature = "wss")]
|
||||
pub(crate) async fn read_eof(&self) {
|
||||
self.read_eof.notified().await;
|
||||
pub(crate) fn read_eof(&self) -> tokio::sync::watch::Receiver<bool> {
|
||||
self.read_eof.subscribe()
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolves once the given WS read-EOF watch receiver observes `true`.
|
||||
///
|
||||
/// Lossless by construction: a `watch::Receiver` retains the latest
|
||||
/// value, so an EOF signaled before this call (or before the receiver
|
||||
/// existed) is still observed — the property a one-shot `Notify` lacked
|
||||
/// (WS-02). Returns early (treated as EOF) if the sender half is
|
||||
/// dropped, e.g. the read pump was aborted.
|
||||
#[cfg(feature = "wss")]
|
||||
pub(crate) async fn wait_for_eof(rx: &mut tokio::sync::watch::Receiver<bool>) {
|
||||
loop {
|
||||
if *rx.borrow_and_update() {
|
||||
return;
|
||||
}
|
||||
if rx.changed().await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -111,11 +131,11 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
|
||||
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||||
|
||||
#[cfg(feature = "wss")]
|
||||
let read_eof = Arc::new(tokio::sync::Notify::new());
|
||||
let read_eof = tokio::sync::watch::channel(false).0;
|
||||
|
||||
let write_tx_for_read = write_tx.clone();
|
||||
#[cfg(feature = "wss")]
|
||||
let read_eof_for_task = Arc::clone(&read_eof);
|
||||
let read_eof_for_task = read_eof.clone();
|
||||
let read_task = tokio::spawn(async move {
|
||||
while let Some(msg) = ws_stream.next().await {
|
||||
match msg {
|
||||
@@ -136,7 +156,9 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
|
||||
}
|
||||
}
|
||||
#[cfg(feature = "wss")]
|
||||
read_eof_for_task.notify_waiters();
|
||||
{
|
||||
let _ = read_eof_for_task.send(true);
|
||||
}
|
||||
});
|
||||
|
||||
let write_task = tokio::spawn(async move {
|
||||
@@ -293,10 +315,10 @@ where
|
||||
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
|
||||
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||||
|
||||
let read_eof = Arc::new(tokio::sync::Notify::new());
|
||||
let read_eof = tokio::sync::watch::channel(false).0;
|
||||
|
||||
let write_tx_for_read = write_tx.clone();
|
||||
let read_eof_for_task = Arc::clone(&read_eof);
|
||||
let read_eof_for_task = read_eof.clone();
|
||||
let read_task = tokio::spawn(async move {
|
||||
while let Some(msg) = ws_stream.next().await {
|
||||
match msg {
|
||||
@@ -316,7 +338,7 @@ where
|
||||
Ok(_) => {}
|
||||
}
|
||||
}
|
||||
read_eof_for_task.notify_waiters();
|
||||
let _ = read_eof_for_task.send(true);
|
||||
});
|
||||
|
||||
let write_task = tokio::spawn(async move {
|
||||
|
||||
Reference in New Issue
Block a user