feat(websocket): write-progress timeout on the WS write pump (WS-18)

- DEFAULT_WS_WRITE_TIMEOUT (60 s, WS-01 knob family): one outbound WS
  send that stays unsent past the window (peer stopped reading)
  evicts the connection — the send path was slot-bounded but
  time-unbounded
- bound is per send call: a slow-but-draining peer resets it with
  every message emitted; only a fully stalled sink trips it
- on timeout the pump signals the stream error (InvalidData naming
  the write stall) and ends WITHOUT a close frame — a clogged socket
  cannot receive one and the close send would park on it
- test-support split helper with both knobs explicit backs the
  scaled eviction test (clogged duplex, eviction well inside the
  5 s slack)

Verification: scripts/verify.sh OK (343 passed), clippy -D warnings
clean, fmt clean
This commit is contained in:
2026-08-30 21:36:53 +00:00
parent d6d5509c6d
commit da65e5b6db
2 changed files with 186 additions and 8 deletions
+184 -6
View File
@@ -147,6 +147,14 @@ pub const WS_INTERNAL_ERROR: u16 = 1011;
/// failed), not a protocol error.
pub const WS_GOING_AWAY: u16 = 1001;
/// Default write-progress timeout for the WS write pump (WS-18): one
/// outbound WS send that stays unsent past this window (the peer
/// stopped reading) evicts the connection — the write-side analog of
/// the idle-read knob. The bound is per `send` call, not per
/// connection: a slow-but-draining peer resets it with every message
/// that gets out.
pub const DEFAULT_WS_WRITE_TIMEOUT: Duration = Duration::from_secs(60);
/// Default idle-read timeout for the WS pumps (WS-01, WS-13): a
/// connection whose inbound stream completes **no chunk** within this
/// window is evicted — the deadline resets on demux progress (complete
@@ -431,10 +439,30 @@ async fn fail_write_pump<M, S>(
/// frame with the requested code and its per-cause reason (WS-15) —
/// the code alone is ambiguous across causes (idle 1001 vs protocol
/// 1002 vs internal 1011 arms), so the pump never invents one.
/// The write pump: drains the queued byte spans, parses the pending
/// bytes for complete chunks (8-byte header, length validated against
/// `MAX_CHUNK_LEN` — WS-04/HY-09), byte-caps the accumulator
/// (`PENDING_BUFFER_CAP` — WS-05), and emits one WS binary message per
/// `WS_MESSAGE_CAP` piece. Ends with a WS Close after the queue
/// drains; `CloseWith` requests (`WriteMsg::CloseWith`) map to a close
/// frame with the requested code and its per-cause reason (WS-15) —
/// the code alone is ambiguous across causes (idle 1001 vs protocol
/// 1002 vs internal 1011 arms), so the pump never invents one.
///
/// Write-progress timeout (WS-18): a peer that stops reading parks
/// the pump inside a WS send — the slot-bounded queue bounds memory,
/// but the stall was time-unbounded. When `write_timeout` is
/// `Some(window)`, every WS send must complete within the window:
/// exceeding it signals the stream error (naming the stall) and ends
/// the pump **without** a close frame — the peer cannot receive one
/// through a clogged socket, and the close send would park on it too.
/// The window is per send, so a slow-but-moving peer survives; only a
/// fully stalled sink trips it.
async fn run_write_pump<M, S>(
mut ws_sink: futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
mut write_rx: futures_mpsc::Receiver<WriteMsg>,
write_error_slot: WriteErrorSlot,
write_timeout: Option<Duration>,
) where
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
M: WsFraming,
@@ -443,9 +471,15 @@ async fn run_write_pump<M, S>(
while let Some(msg) = write_rx.next().await {
match msg {
WriteMsg::CloseWith(code, reason) => {
let _ = ws_sink
.send(<M as WsFraming>::close_message(code, reason))
.await;
let close = <M as WsFraming>::close_message(code, reason);
match write_timeout {
None => {
let _ = ws_sink.send(close).await;
}
Some(window) => {
let _ = tokio::time::timeout(window, ws_sink.send(close)).await;
}
}
break;
}
WriteMsg::Bytes(b) => {
@@ -474,10 +508,8 @@ async fn run_write_pump<M, S>(
}
let chunk: Vec<u8> = pending.drain(..total).collect();
for piece in chunk.chunks(WS_MESSAGE_CAP) {
if ws_sink
.send(<M as WsFraming>::binary_message(piece.to_vec()))
if !send_bounded::<M, _>(&mut ws_sink, piece, write_timeout, &write_error_slot)
.await
.is_err()
{
return;
}
@@ -496,6 +528,45 @@ async fn run_write_pump<M, S>(
let _ = ws_sink.close().await;
}
/// One WS send under the WS-18 write-progress bound: a send that
/// outlasts the window means the peer stopped reading, so the pump
/// signals the error slot (`InvalidData` on the `AsyncWrite` half) and
/// ends without a close frame — a clogged socket cannot receive one,
/// and the close send would park on the same stall. On send error, end
/// silently (the sink is already broken). Returns whether the pump
/// should continue.
async fn send_bounded<M, S>(
ws_sink: &mut futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
piece: &[u8],
write_timeout: Option<Duration>,
slot: &WriteErrorSlot,
) -> bool
where
M: WsFraming,
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
{
match write_timeout {
None => ws_sink
.send(<M as WsFraming>::binary_message(piece.to_vec()))
.await
.is_ok(),
Some(window) => {
let send = ws_sink.send(<M as WsFraming>::binary_message(piece.to_vec()));
match tokio::time::timeout(window, send).await {
Ok(Ok(())) => true,
Ok(Err(_)) => false,
Err(_elapsed) => {
send_write_error(
slot,
"connection made no write progress past the write timeout",
);
false
}
}
}
}
}
/// The `AsyncRead + AsyncWrite` view of a `WebSocket` handed to
/// alkcall's channels machinery (demux/mux).
pub struct WsByteStream {
@@ -619,6 +690,7 @@ pub fn split_ws_to_bytes_idle(
ws_sink,
write_rx,
write_error_slot,
Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT),
));
(
@@ -887,6 +959,7 @@ where
ws_sink,
write_rx,
write_error_slot,
Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT),
));
(
@@ -909,11 +982,116 @@ where
)
}
/// Test-only split with both knobs explicit (WS-18 acceptance: the
/// stall test scales the write window down; other tests keep the
/// default).
#[cfg(test)]
pub(crate) fn split_tungstenite_to_bytes_idle_with_write<S>(
socket: tokio_tungstenite::WebSocketStream<S>,
idle_timeout: Option<Duration>,
write_timeout: Option<Duration>,
) -> (WsByteStream, WsPumps)
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
let (ws_sink, ws_stream) = socket.split();
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_error_slot, write_error_rx) = make_write_error_slot();
let read_eof = tokio::sync::watch::channel(false).0;
let write_tx_for_read = write_tx.clone();
let write_error_slot_for_read = Arc::clone(&write_error_slot);
let read_eof_for_task = read_eof.clone();
let read_task = tokio::spawn(run_read_pump::<TungsteniteFraming, _, _>(
ws_stream,
read_tx,
write_tx_for_read,
write_error_slot_for_read,
idle_timeout,
move || {
let _ = read_eof_for_task.send(true);
},
));
let write_task = tokio::spawn(run_write_pump::<TungsteniteFraming, _>(
ws_sink,
write_rx,
write_error_slot,
write_timeout,
));
(
WsByteStream {
read_rx,
read_buf: Vec::new(),
read_pos: 0,
eof: false,
write_tx: Some(write_tx),
write_open: true,
write_error: write_error_rx,
write_failed: false,
},
WsPumps {
read_task,
write_task,
read_eof,
},
)
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
/// WS-18 acceptance: a peer that stops reading (the sink clogs —
/// nothing is drained from the duplex) parks the write pump inside
/// a WS send; the pump must end within the knob (+ slack), surface
/// the stream error naming the write stall, and fail the
/// `AsyncWrite` half with `InvalidData`. Scaled test: both knobs
/// idle-read `None` (so only the write knob can evict) and the
/// write window at 150 ms.
#[tokio::test]
async fn tungstenite_write_stall_is_evicted_within_the_write_timeout() {
let (client_io, _server_io) = tokio::io::duplex(64);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let knob = std::time::Duration::from_millis(150);
let (mut stream, _pumps) = split_tungstenite_to_bytes_idle_with_write(ws, None, Some(knob));
let mut chunk = vec![0u8; 8];
chunk[4..8].copy_from_slice(&64u32.to_be_bytes());
chunk.extend_from_slice(&[0u8; 64]);
stream.write_all(&chunk).await.expect("chunk written");
stream.flush().await.expect("flush");
let started = tokio::time::Instant::now();
let err = loop {
match stream.write_all(&[0u8; 8]).await {
Err(e) => break e,
Ok(_) => assert!(
started.elapsed() < std::time::Duration::from_secs(5),
"stream never failed while the sink stayed clogged"
),
}
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(
err.to_string().contains("write timeout"),
"error names the violation: {err}"
);
assert!(
started.elapsed() < std::time::Duration::from_secs(5),
"eviction must happen within the knob (+ slack), not hang"
);
}
/// An 8-byte bogus chunk header claiming `MAX_CHUNK_LEN + 1`
/// payload bytes (WS-04/HY-09 acceptance: the parser must fail the
/// stream loudly instead of silently waiting for ~4 GiB).