From d38f5eb81c3ad132077c02cacc46c9a091de8a39 Mon Sep 17 00:00:00 2001 From: CodeShell <122738806+CodeShellDev@users.noreply.github.com> Date: Sun, 7 Dec 2025 17:59:04 +0100 Subject: [PATCH 1/3] fix --- internals/server/http.go | 46 ++++++++++++++++++---------------------- 1 file changed, 21 insertions(+), 25 deletions(-) diff --git a/internals/server/http.go b/internals/server/http.go index e133abb..61299b7 100644 --- a/internals/server/http.go +++ b/internals/server/http.go @@ -97,32 +97,28 @@ func httpHandler(w http.ResponseWriter, req *http.Request) { } if body.StartupTime != nil { - if err == nil { - time.Sleep(time.Duration(*body.StartupTime) * time.Second) - - reachable, err = tryPing(client, body.Addr, - func() (bool, error) { - sendToClient(client, map[string]any{ - "success": true, - "message": "Host is now reachable", - }) - return true, nil - }, - func() (bool, error) { - sendToClient(client, map[string]any{ - "success": false, - "error": true, - "message": "Host is still unreachable", - }) - return false, nil - }, - ) + time.Sleep(time.Duration(*body.StartupTime) * time.Second) + + reachable, err = tryPing(client, body.Addr, + func() (bool, error) { + sendToClient(client, map[string]any{ + "success": true, + "message": "Host is now reachable", + }) + return true, nil + }, + func() (bool, error) { + sendToClient(client, map[string]any{ + "success": false, + "error": true, + "message": "Host is still unreachable", + }) + return false, nil + }, + ) - closeClient(client) - return - } else { - logger.Error("Could not parse startupTime: ", err.Error()) - } + closeClient(client) + return } success, err := tryPingInterval(client, config.ENV.PING_INTERVAL, config.ENV.PING_RETRIES, body.Addr) From 69f278f996c1708041e19c0c1b0565420335bbf9 Mon Sep 17 00:00:00 2001 From: CodeShell <122738806+CodeShellDev@users.noreply.github.com> Date: Sun, 7 Dec 2025 18:08:40 +0100 Subject: [PATCH 2/3] move cleanup out of register due to waitfor --- internals/server/websocket.go | 43 +++++++++++++++++------------------ 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/internals/server/websocket.go b/internals/server/websocket.go index bc4a48a..022cd26 100644 --- a/internals/server/websocket.go +++ b/internals/server/websocket.go @@ -22,9 +22,18 @@ func websocketHandler(w http.ResponseWriter, req *http.Request) { if err != nil { return } - defer socket.Close() - register(req, socket) + clientID := register(req, socket) + if clientID == "" { + return + } + + defer func() { + clientsMutex.Lock() + delete(clients, clientID) + clientsMutex.Unlock() + socket.Close() + }() keepAlive(socket) @@ -37,29 +46,19 @@ func websocketHandler(w http.ResponseWriter, req *http.Request) { } } -func register(req *http.Request, socket *websocket.Conn) { - clientID := req.URL.Query().Get("client_id") - if clientID == "" { - socket.Close() - return - } - - clientsMutex.Lock() - clients[clientID] = socket - clientsMutex.Unlock() - defer func() { - clientsMutex.Lock() - - delete(clients, clientID) +func register(req *http.Request, socket *websocket.Conn) string { + clientID := req.URL.Query().Get("client_id") + if clientID == "" { + socket.Close() + return "" + } - clientsMutex.Unlock() - socket.Close() - }() -} + clientsMutex.Lock() + clients[clientID] = socket + clientsMutex.Unlock() -func getClient(id string) *websocket.Conn { - return clients[id] + return clientID } func WaitForClient(clientID string, timeout time.Duration) (*websocket.Conn, error) { From 00751f3223b35f281cf5098615daba7899f65f43 Mon Sep 17 00:00:00 2001 From: CodeShell <122738806+CodeShellDev@users.noreply.github.com> Date: Sun, 7 Dec 2025 18:15:24 +0100 Subject: [PATCH 3/3] add waiters queue --- internals/server/http.go | 2 +- internals/server/websocket.go | 45 ++++++++++++++++++++++------------- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/internals/server/http.go b/internals/server/http.go index 61299b7..e7b258a 100644 --- a/internals/server/http.go +++ b/internals/server/http.go @@ -43,7 +43,7 @@ func httpHandler(w http.ResponseWriter, req *http.Request) { w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(resp) - client, err := WaitForClient(clientID, time.Duration(15 * time.Second)) + client, err := waitForClient(clientID, time.Duration(15 * time.Second)) if err != nil { logger.Error("Could not get client: ", err.Error()) diff --git a/internals/server/websocket.go b/internals/server/websocket.go index 022cd26..89a4e77 100644 --- a/internals/server/websocket.go +++ b/internals/server/websocket.go @@ -14,6 +14,9 @@ var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, } +var waiters = make(map[string]chan *websocket.Conn) +var waitersMutex = &sync.Mutex{} + var clients = make(map[string]*websocket.Conn) var clientsMutex = &sync.Mutex{} @@ -46,7 +49,6 @@ func websocketHandler(w http.ResponseWriter, req *http.Request) { } } - func register(req *http.Request, socket *websocket.Conn) string { clientID := req.URL.Query().Get("client_id") if clientID == "" { @@ -57,28 +59,37 @@ func register(req *http.Request, socket *websocket.Conn) string { clientsMutex.Lock() clients[clientID] = socket clientsMutex.Unlock() + waitersMutex.Lock() + if ch, ok := waiters[clientID]; ok { + ch <- socket + + close(ch) + delete(waiters, clientID) + } + waitersMutex.Unlock() return clientID } -func WaitForClient(clientID string, timeout time.Duration) (*websocket.Conn, error) { - deadline := time.Now().Add(timeout) - - for time.Now().Before(deadline) { - clientsMutex.Lock() - conn, exists := clients[clientID] - clientsMutex.Unlock() - - if exists && conn != nil { - return conn, nil - } - - time.Sleep(100 * time.Millisecond) - } - - return nil, errors.New("Timed out waiting for client") +func waitForClient(clientID string, timeout time.Duration) (*websocket.Conn, error) { + waitCh := make(chan *websocket.Conn, 1) + + waitersMutex.Lock() + waiters[clientID] = waitCh + waitersMutex.Unlock() + + select { + case conn := <-waitCh: + return conn, nil + case <-time.After(timeout): + waitersMutex.Lock() + delete(waiters, clientID) + waitersMutex.Unlock() + return nil, errors.New("Timed out waiting for client") + } } + func sendToClient(client *websocket.Conn, data map[string]any) error { return client.WriteJSON(data) }