From c8b87575945aae4d92d835f0da5c76e94e7420bc Mon Sep 17 00:00:00 2001
From: LGUG2Z <LGUG2Z@fastmail.com>
Date: Fri, 19 Jan 2024 15:53:45 -0800
Subject: [PATCH] fix(wm): add retries to socket connections

---
 komorebi/src/process_command.rs | 75 ++++++++++++++++++++++++++-------
 komorebic/src/main.rs           | 12 +++++-
 2 files changed, 69 insertions(+), 18 deletions(-)

diff --git a/komorebi/src/process_command.rs b/komorebi/src/process_command.rs
index 197115f1..e4b3daf5 100644
--- a/komorebi/src/process_command.rs
+++ b/komorebi/src/process_command.rs
@@ -743,8 +743,14 @@ impl WindowManager {
                 };
 
                 let socket = DATA_DIR.join("komorebic.sock");
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(state.as_bytes())?;
+
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(state.as_bytes())?;
+                    }
+                }
             }
             SocketMessage::VisibleWindows => {
                 let mut monitor_visible_windows = HashMap::new();
@@ -768,8 +774,14 @@ impl WindowManager {
                     };
 
                 let socket = DATA_DIR.join("komorebic.sock");
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(visible_windows_state.as_bytes())?;
+
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(visible_windows_state.as_bytes())?;
+                    }
+                }
             }
 
             SocketMessage::Query(query) => {
@@ -789,8 +801,14 @@ impl WindowManager {
                 .to_string();
 
                 let socket = DATA_DIR.join("komorebic.sock");
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(response.as_bytes())?;
+
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(response.as_bytes())?;
+                    }
+                }
             }
             SocketMessage::ResizeWindowEdge(direction, sizing) => {
                 self.resize_window(direction, sizing, self.resize_delta, true)?;
@@ -1258,39 +1276,64 @@ impl WindowManager {
                 let schema = serde_json::to_string_pretty(&asc)?;
                 let socket = DATA_DIR.join("komorebic.sock");
 
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(schema.as_bytes())?;
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(schema.as_bytes())?;
+                    }
+                }
             }
             SocketMessage::NotificationSchema => {
                 let notification = schema_for!(Notification);
                 let schema = serde_json::to_string_pretty(&notification)?;
                 let socket = DATA_DIR.join("komorebic.sock");
 
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(schema.as_bytes())?;
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(schema.as_bytes())?;
+                    }
+                }
             }
             SocketMessage::SocketSchema => {
                 let socket_message = schema_for!(SocketMessage);
                 let schema = serde_json::to_string_pretty(&socket_message)?;
                 let socket = DATA_DIR.join("komorebic.sock");
 
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(schema.as_bytes())?;
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(schema.as_bytes())?;
+                    }
+                }
             }
             SocketMessage::StaticConfigSchema => {
                 let socket_message = schema_for!(StaticConfig);
                 let schema = serde_json::to_string_pretty(&socket_message)?;
                 let socket = DATA_DIR.join("komorebic.sock");
 
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(schema.as_bytes())?;
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(schema.as_bytes())?;
+                    }
+                }
             }
             SocketMessage::GenerateStaticConfig => {
                 let config = serde_json::to_string_pretty(&StaticConfig::from(&*self))?;
                 let socket = DATA_DIR.join("komorebic.sock");
 
-                let mut stream = UnixStream::connect(socket)?;
-                stream.write_all(config.as_bytes())?;
+                let mut connected = false;
+                while !connected {
+                    if let Ok(mut stream) = UnixStream::connect(&socket) {
+                        connected = true;
+                        stream.write_all(config.as_bytes())?;
+                    }
+                }
             }
             SocketMessage::RemoveTitleBar(_, ref id) => {
                 let mut identifiers = NO_TITLEBAR.lock();
diff --git a/komorebic/src/main.rs b/komorebic/src/main.rs
index c1a4bb0c..8dc8218a 100644
--- a/komorebic/src/main.rs
+++ b/komorebic/src/main.rs
@@ -1136,8 +1136,16 @@ enum SubCommand {
 
 pub fn send_message(bytes: &[u8]) -> Result<()> {
     let socket = DATA_DIR.join("komorebi.sock");
-    let mut stream = UnixStream::connect(socket)?;
-    Ok(stream.write_all(bytes)?)
+
+    let mut connected = false;
+    while !connected {
+        if let Ok(mut stream) = UnixStream::connect(&socket) {
+            connected = true;
+            stream.write_all(bytes)?;
+        }
+    }
+
+    Ok(())
 }
 
 fn with_komorebic_socket<F: Fn() -> Result<()>>(f: F) -> Result<()> {