From 0927830744db6e9bf4cb0fbef202da8d5beb4823 Mon Sep 17 00:00:00 2001
From: Mauve <mauve@google.com>
Date: Sun, 1 Oct 2023 21:14:43 -0400
Subject: [PATCH] Upgrade tokio-tungstenite, add explicit flushes

---
 Cargo.lock                                      |  8 ++++----
 src/libs/ml2_mods/Cargo.toml                    |  2 +-
 src/libs/ml2_mods/src/spelunkyfyi/web_socket.rs | 12 ++++--------
 3 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f8e94cdc3..58946df35 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4428,9 +4428,9 @@ dependencies = [
 
 [[package]]
 name = "tokio-tungstenite"
-version = "0.19.0"
+version = "0.20.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ec509ac96e9a0c43427c74f003127d953a265737636129424288d27cb5c4b12c"
+checksum = "212d5dcb2a1ce06d81107c3d0ffa3121fe974b73f068c8282cb1c32328113b6c"
 dependencies = [
  "futures-util",
  "log",
@@ -4625,9 +4625,9 @@ checksum = "3528ecfd12c466c6f163363caf2d02a71161dd5e1cc6ae7b34207ea2d42d81ed"
 
 [[package]]
 name = "tungstenite"
-version = "0.19.0"
+version = "0.20.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "15fba1a6d6bb030745759a9a2a588bfe8490fc8b4751a277db3a0be1c9ebbf67"
+checksum = "9e3dac10fd62eaf6617d3a904ae222845979aec67c615d1c842b4002c7666fb9"
 dependencies = [
  "byteorder",
  "bytes",
diff --git a/src/libs/ml2_mods/Cargo.toml b/src/libs/ml2_mods/Cargo.toml
index 858abb9db..66a8b129c 100644
--- a/src/libs/ml2_mods/Cargo.toml
+++ b/src/libs/ml2_mods/Cargo.toml
@@ -35,7 +35,7 @@ tokio = { version = "1.32", features = [
     "tracing",
 ] }
 tokio-graceful-shutdown = "0.13"
-tokio-tungstenite = { version = "0.19", features = ["native-tls"] }
+tokio-tungstenite = { version = "0.20", features = ["native-tls"] }
 tower = { version = "0.4", features = ["util"] }
 tracing = "0.1"
 zip = "0.6"
diff --git a/src/libs/ml2_mods/src/spelunkyfyi/web_socket.rs b/src/libs/ml2_mods/src/spelunkyfyi/web_socket.rs
index b2d3e23c1..769de587c 100644
--- a/src/libs/ml2_mods/src/spelunkyfyi/web_socket.rs
+++ b/src/libs/ml2_mods/src/spelunkyfyi/web_socket.rs
@@ -120,14 +120,8 @@ impl WebSocketClient {
             .header(SEC_WEBSOCKET_KEY, generate_key())
             .body(())?;
 
-        // This is the maximum number of messages that will be queued in Tungstenite. This doesn't
-        // include pong or close messages. All messages are buffered before they're written. So,
-        // the smallest functional value is 1. Notably, the size of the messages and the underlying
-        // TCP send buffer aren't considered.
-        let config = WebSocketConfig {
-            max_send_queue: Some(2),
-            ..Default::default()
-        };
+        let mut config = WebSocketConfig::default();
+        config.max_write_buffer_size = 2 * config.write_buffer_size;
 
         let mut stream = tokio_tungstenite::connect_async_with_config(request, Some(config), false)
             .await
@@ -163,6 +157,7 @@ impl WebSocketClient {
                         stream
                             .send(Message::Ping(payload))
                             .await?;
+                        stream.flush().await?;
                         check_state = Check::Pong();
                         check_sleep = Box::pin(time::sleep(self.pong_timeout));
                     }
@@ -312,6 +307,7 @@ async fn send_message(
 ) -> Result<(), ConnectionError> {
     let reply = serde_json::to_string(&msg)?;
     stream.send(Message::Text(reply)).await?;
+    stream.flush().await?;
     Ok(())
 }