Add more logging and break out of loop on close frame

Author: joshua-spacetime

crates/client-api/src/routes/subscribe.rs

@@ -9,7 +9,7 @@ use axum_extra::TypedHeader;
 use bytes::Bytes;
 use bytestring::ByteString;
 use futures::future::MaybeDone;
-use futures::{Future, FutureExt, SinkExt, StreamExt};
+use futures::{Future, SinkExt, StreamExt};
 use http::{HeaderValue, StatusCode};
 use scopeguard::ScopeGuard;
 use serde::Deserialize;
@@ -182,6 +182,7 @@ where
 }
 
 const LIVELINESS_TIMEOUT: Duration = Duration::from_secs(60);
+const SEND_TIMEOUT: Duration = Duration::from_secs(5);
 
 async fn ws_client_actor(client: ClientConnection, ws: WebSocketStream, sendrx: MeteredReceiver<SerializableMessage>) {
     // ensure that even if this task gets cancelled, we always cleanup the connection
@@ -276,54 +277,72 @@ async fn ws_client_actor_inner(
 
             // If we have an outgoing message to send, send it off.
             // No incoming `message` to handle, so `continue`.
-            Some(n) = sendrx.recv_many(&mut rx_buf, 32).map(|n| (n != 0).then_some(n)) => {
+            n = sendrx.recv_many(&mut rx_buf, 32) => {
+                // The receiver has been closed and there are no pending messages.
+                // This is due to the client sending a close frame, so break from the loop.
+                if n == 0 {
+                    break;
+                }
                 if closed {
                     // TODO: this isn't great. when we receive a close request from the peer,
                     //       tungstenite doesn't let us send any new messages on the socket,
                     //       even though the websocket RFC allows it. should we fork tungstenite?
                     log::info!("dropping {n} messages due to ws already being closed");
                     log::debug!("dropped messages: {:?}", &rx_buf[..n]);
-                } else {
-                    let send_all = async {
-                        for msg in rx_buf.drain(..n) {
-                            let workload = msg.workload();
-                            let num_rows = msg.num_rows();
-
-                            // Serialize the message, report metrics,
-                            // and keep a handle to the buffer.
-                            let (msg_alloc, msg_data) = serialize(msg_buffer, msg, client.config);
-                            report_ws_sent_metrics(&addr, workload, num_rows, &msg_data);
-
-                            // Buffer the message without necessarily sending it.
-                            let res = ws.feed(datamsg_to_wsmsg(msg_data)).await;
-
-                            // At this point,
-                            // the underlying allocation of `msg_data` should have a single referent
-                            // and this should be `msg_alloc`.
-                            // We can put this back into our pool.
-                            msg_buffer = msg_alloc.try_reclaim()
-                                .expect("should have a unique referent to `msg_alloc`");
-
-                            if res.is_err() {
-                                return (res, msg_buffer);
-                            }
+                    break;
+                }
+                let send_all = async {
+                    for msg in rx_buf.drain(..n) {
+                        let workload = msg.workload();
+                        let num_rows = msg.num_rows();
+
+                        // Serialize the message, report metrics,
+                        // and keep a handle to the buffer.
+                        let (msg_alloc, msg_data) = serialize(msg_buffer, msg, client.config);
+                        report_ws_sent_metrics(&addr, workload, num_rows, &msg_data);
+
+                        // Buffer the message without necessarily sending it.
+                        let res = ws.feed(datamsg_to_wsmsg(msg_data)).await;
+
+                        // At this point,
+                        // the underlying allocation of `msg_data` should have a single referent
+                        // and this should be `msg_alloc`.
+                        // We can put this back into our pool.
+                        msg_buffer = msg_alloc.try_reclaim()
+                            .expect("should have a unique referent to `msg_alloc`");
+
+                        if res.is_err() {
+                            return (res, msg_buffer);
                         }
-                        // now we flush all the messages to the socket
-                        (ws.flush().await, msg_buffer)
-                    };
-                    // Flush the websocket while continuing to poll the `handle_queue`,
-                    // to avoid deadlocks or delays due to enqueued futures holding resources.
-                    let send_all = also_poll(send_all, make_progress(&mut current_message));
-                    let t1 = Instant::now();
-                    let (send_all_result, buf) = send_all.await;
-                    msg_buffer = buf;
-                    if let Err(error) = send_all_result {
-                        log::warn!("Websocket send error: {error}")
                     }
-                    let time = t1.elapsed();
-                    if time > Duration::from_millis(50) {
-                        tracing::warn!(?time, "send_all took a very long time");
+                    // now we flush all the messages to the socket
+                    (ws.flush().await, msg_buffer)
+                };
+                // build a future that both times out and drives the send
+                let send_all = async {
+                    tokio::time::timeout(SEND_TIMEOUT, send_all).await
+                };
+                // Flush the websocket while continuing to poll the `handle_queue`,
+                // to avoid deadlocks or delays due to enqueued futures holding resources.
+                let send_all = also_poll(send_all, make_progress(&mut current_message));
+                let t1 = Instant::now();
+                let (send_all_result, buf) = match send_all.await {
+                    Ok((send_all_result, buf)) => {
+                        (send_all_result, buf)
                     }
+                    Err(e) => {
+                        // Our send timed out; drop client without trying to send them a Close
+                        log::warn!("send_all timed out after: {e}");
+                        break;
+                    }
+                };
+                msg_buffer = buf;
+                if let Err(error) = send_all_result {
+                    log::warn!("Websocket send error: {error}")
+                }
+                let time = t1.elapsed();
+                if time > Duration::from_millis(50) {
+                    tracing::warn!(?time, "send_all took a very long time");
                 }
                 continue;
             }
@@ -349,13 +368,25 @@ async fn ws_client_actor_inner(
             }
 
             // If it's time to send a ping...
-            _ = liveness_check_interval.tick() => {
+            _ = liveness_check_interval.tick(), if got_pong => {
                 // If we received a pong at some point, send a fresh ping.
                 if mem::take(&mut got_pong) {
+                    // build a future that both times out and drives the send
+                    let ping_with_timeout = async {
+                        tokio::time::timeout(SEND_TIMEOUT, ws.send(WsMessage::Ping(Bytes::new()))).await
+                    };
                     // Send a ping message while continuing to poll the `handle_queue`,
                     // to avoid deadlocks or delays due to enqueued futures holding resources.
-                    if let Err(e) = also_poll(ws.send(WsMessage::Ping(Bytes::new())), make_progress(&mut current_message)).await {
-                        log::warn!("error sending ping: {e:#}");
+                    match also_poll(ping_with_timeout, make_progress(&mut current_message)).await {
+                        Ok(Err(e)) => {
+                            log::warn!("error sending ping: {e:#}");
+                        }
+                        Err(e) => {
+                            // Our ping timed out; drop them without trying to send them a Close
+                            log::warn!("ping timed out after: {e}");
+                            break;
+                        }
+                        _ => {}
                     }
                     continue;
                 } else {
@@ -380,7 +411,7 @@ async fn ws_client_actor_inner(
             Item::HandleResult(res) => {
                 if let Err(e) = res {
                     if let MessageHandleError::Execution(err) = e {
-                        log::error!("{err:#}");
+                        log::error!("reducer execution error: {err:#}");
                         // Serialize the message and keep a handle to the buffer.
                         let (msg_alloc, msg_data) = serialize(msg_buffer, err, client.config);
 
@@ -399,7 +430,7 @@ async fn ws_client_actor_inner(
 
                         continue;
                     }
-                    log::debug!("Client caused error on text message: {}", e);
+                    log::warn!("Client caused error on text message: {}", e);
                     if let Err(e) = ws
                         .close(Some(CloseFrame {
                             code: CloseCode::Error,
@@ -439,7 +470,7 @@ async fn ws_client_actor_inner(
                         .with_label_values(&addr)
                         .inc();
                 }
-                closed = true;
+                break;
             }
         }
     }