Skip to content

fix(envoy-client): bound websocket writer memory - #5754

Open
RitwijParmar wants to merge 2 commits into
rivet-dev:mainfrom
RitwijParmar:codex/bounded-envoy-writer
Open

RitwijParmar wants to merge 2 commits into
rivet-dev:mainfrom
RitwijParmar:codex/bounded-envoy-writer

Conversation

@RitwijParmar

Copy link
Copy Markdown

Closes #5468

  • Replace the unbounded legacy WebSocket writer with bounded data and control lanes.
  • Enforce byte and message limits while reserving capacity for acknowledgements, pongs, metadata, and shutdown.
  • Record admission failures and keep KV and SQLite requests retryable when the writer is saturated.
  • Apply the same writer behavior to native and WASM transports.
  • Add saturation coverage and keep the existing envoy-client test suite passing.

Checked locally:

  • cargo test -p rivet-envoy-client --lib
  • cargo check -p rivet-envoy-client
  • cargo check -p rivet-envoy-client --no-default-features --features wasm-transport
  • git diff --check

The repository Clippy command is currently blocked by existing warnings in envoy-protocol/build.rs and util-serde/src/lib.rs.

@the-company-company the-company-company Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 1 high · 🟠 2 medium

Reviewed commit 7517c9b.

.serialize(protocol::PROTOCOL_VERSION)
.expect("failed to encode message");
if tx.send(WsTxMessage::Send(encoded)).is_err() {
if tx.try_send(to_ws_lane(message_kind), encoded).is_err() {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 Medium · Queue saturation strands requests until they time out

A full byte/message budget now makes ws_send_for_session return Unavailable while the WebSocket remains connected. KV and ordinary SQLite callers leave such requests marked unsent, but process_unsent_* is only invoked on ToEnvoyInit, so after the writer drains there is no retry and the request sits until the 30s cleanup timeout. Buffered tunnel messages and unacknowledged events have the same reconnect-only retry trigger. Admission needs to wait for capacity, schedule a retry when capacity becomes available, or force a reconnect so these retained messages are actually replayed.

@@ -165,13 +163,17 @@ async fn single_connection(

loop {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 High · Priority close discards already accepted data messages

The control-first biased selection lets WsTxMessage::Close overtake every message already admitted to the data and HTTP queues. During cleanup try_close() is called after the envoy loop has admitted its final actor events, but the writer can select the close immediately, break, and drop those queued events; the context is then cleared, so they are never replayed. The close path needs to act as a drain barrier across all lanes (and the wasm writer needs the same ordering) rather than sharing the priority queue with ordinary control messages.

@@ -165,13 +163,17 @@ async fn single_connection(

loop {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 Medium · Biased data selection can starve HTTP writes

With biased, a continuously nonempty control or data receiver always wins before http_ws_rx. Under sustained legacy traffic the HTTP queue can therefore remain unread indefinitely even though it has its own reserved capacity; ws_send_http_for_session waits for the writer acknowledgement and streaming responses stall. Use a fair/weighted scheduler that reserves control progress without making the data lane permanently higher priority than HTTP, and mirror it in the wasm transport.

@RitwijParmar

Copy link
Copy Markdown
Author

I pushed a follow-up for the three review points.

Close now blocks new admissions and drains already accepted legacy and HTTP messages before the close frame. Saturated KV, SQLite, tunnel and event sends are retried on a short loop. The writer select is fair in native and WASM so HTTP cannot be starved.

Local checks pass:

  • cargo test -p rivet-envoy-client --lib
  • cargo check -p rivet-envoy-client
  • cargo check -p rivet-envoy-client --no-default-features --features wasm-transport
  • git diff --check

The strict Clippy command still stops on existing warnings in envoy-protocol/build.rs and util-serde.

@the-company-company the-company-company Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 2 medium-severity findings

Reviewed commit 74356a0.

return Ok(());
}
self.control_tx
.try_send(WsTxMessage::Close)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 Medium · A saturated control queue prevents the connection from ever closing

closing is set before the nonblocking close enqueue, so if all 32 control slots are occupied this returns an error while permanently rejecting new admissions. Cleanup ignores that error, and any later try_close would return Ok(()) without enqueueing a close because the flag is already true. Once the writer drains the existing messages it therefore waits on the still-open channels indefinitely instead of sending the close frame. Reserve close capacity or use a close signal that cannot be rejected; at minimum, do not latch closing when enqueueing the close fails.

} else {
connection
.tx
.try_send(HttpWsTxMessage {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 Medium · A full HTTP message queue now aborts healthy streams

This changed HTTP admission from send().await to try_send, even though the byte semaphore has already bounded the queued and waiting payload memory. When 256 small frames occupy the queue, the next frame is rejected solely on message count while the socket is healthy; HttpResponseSender::send then queues a transport abort, truncating the response instead of applying backpressure until the fair writer frees a slot. Keep the closing check atomic with admission, but wait for channel capacity (or reserve a sender permit before entering the gate) rather than treating transient fullness as a transport failure.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(envoy-client): bound the shared WebSocket writer queue

1 participant