From fe0ff032bd0d23742525f39a009e421e27b2c832 Mon Sep 17 00:00:00 2001 From: Kim Altintop Date: Thu, 17 Jul 2025 11:04:24 +0200 Subject: [PATCH] Put the state updates after close, like before --- crates/client-api/src/routes/subscribe.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 7aea07a1b6..107d99efd6 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -1048,11 +1048,6 @@ async fn ws_send_loop_inner( match msg { UnorderedWsMessage::Close(close_frame) => { log::trace!("intiating close"); - state.close(); - // We won't be polling `messages` anymore, - // so let senders know. - messages.close(); - // Send outstanding frames until one that has the FIN // bit set. Ensures the client won't receive partial // messages before we shut down. @@ -1074,6 +1069,20 @@ async fn ws_send_loop_inner( log::warn!("error sending close frame: {e:#}"); break; } + + // Lastly, update the state. + // + // NOTE: It's ok to not update the state if we fail to + // send the close frame, because we assume that the main + // loop with exit when this future terminates. + // We shouldn't set the state to closed before sending + // the close frame, however, as we would start dropping + // messages immediately (defeating the purpose of the + // close handshake). + state.close(); + // We won't be polling `messages` anymore, + // so let senders know. + messages.close(); }, UnorderedWsMessage::Ping(bytes) => { log::trace!("sending ping");