From eb71f44e9f953e0e383c1fa0c247d4e3b97dd97d Mon Sep 17 00:00:00 2001
From: IQuant <quant3234@gmail.com>
Date: Sun, 1 Dec 2024 13:03:03 +0300
Subject: [PATCH] Maybe fix that random issue with the connection?

---
 ewext/src/lib.rs            |  7 +++++++
 ewext/src/net.rs            | 21 ++++++++++++++++++---
 quant.ew/files/core/net.lua |  1 +
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/ewext/src/lib.rs b/ewext/src/lib.rs
index cfb62d19..a525238e 100644
--- a/ewext/src/lib.rs
+++ b/ewext/src/lib.rs
@@ -161,6 +161,12 @@ fn netmanager_send(lua: LuaState) -> eyre::Result<()> {
     Ok(())
 }
 
+fn netmanager_flush(_lua: LuaState) -> eyre::Result<()> {
+    let mut binding = NETMANAGER.lock().unwrap();
+    let netmanager = binding.as_mut().unwrap();
+    netmanager.flush()
+}
+
 impl LuaFnRet for InitKV {
     fn do_return(self, lua: LuaState) -> c_int {
         lua.create_table(2, 0);
@@ -306,6 +312,7 @@ pub unsafe extern "C" fn luaopen_ewext0(lua: *mut lua_State) -> c_int {
         add_lua_fn!(netmanager_connect);
         add_lua_fn!(netmanager_recv);
         add_lua_fn!(netmanager_send);
+        add_lua_fn!(netmanager_flush);
 
         add_lua_fn!(module_on_world_update);
     }
diff --git a/ewext/src/net.rs b/ewext/src/net.rs
index ac7cb7a8..5150347f 100644
--- a/ewext/src/net.rs
+++ b/ewext/src/net.rs
@@ -30,14 +30,17 @@ impl NetManager {
     }
 
     pub(crate) fn switch_to_non_blocking(&mut self) -> eyre::Result<()> {
-        self.ws.get_mut().set_read_timeout(None)?;
-        self.ws.get_mut().set_nonblocking(true)?;
+        let stream_ref = self.ws.get_mut();
+        stream_ref.set_nonblocking(true)?;
+        stream_ref.set_read_timeout(Some(Duration::from_millis(1)))?;
+        // Set write timeout to a somewhat high value just in case.
+        stream_ref.set_write_timeout(Some(Duration::from_secs(5)))?;
         Ok(())
     }
 
     pub(crate) fn send(&mut self, msg: &NoitaOutbound) -> eyre::Result<()> {
         self.ws
-            .send(tungstenite::Message::Binary(bitcode::encode(msg)))?;
+            .write(tungstenite::Message::Binary(bitcode::encode(msg)))?;
         Ok(())
     }
 
@@ -55,6 +58,18 @@ impl NetManager {
             }
         }
     }
+
+    pub(crate) fn flush(&mut self) -> eyre::Result<()> {
+        match self.ws.flush() {
+            Ok(()) => Ok(()),
+            Err(tungstenite::Error::Io(err))
+                if err.kind() == ErrorKind::WouldBlock || err.kind() == ErrorKind::TimedOut =>
+            {
+                Ok(())
+            }
+            Err(err) => Err(err.into()),
+        }
+    }
 }
 
 impl Drop for NetManager {
diff --git a/quant.ew/files/core/net.lua b/quant.ew/files/core/net.lua
index 16bb79db..f86cc5b7 100644
--- a/quant.ew/files/core/net.lua
+++ b/quant.ew/files/core/net.lua
@@ -163,6 +163,7 @@ function net.update()
     end
     handle_message(msg)
   end
+  ewext.netmanager_flush()
 end
 
 function net.init()