testing

Author: CodeShellDev

internals/server/websocket.go

@@ -59,10 +59,13 @@ 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
-
+		select {
+		case ch <- socket:
+		default:
+		}
 		close(ch)
 		delete(waiters, clientID)
 	}
@@ -72,24 +75,30 @@ func register(req *http.Request, socket *websocket.Conn) string {
 }
 
 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")
-	}
+    clientsMutex.Lock()
+    if conn, ok := clients[clientID]; ok {
+        clientsMutex.Unlock()
+        return conn, nil
+    }
+    clientsMutex.Unlock()
+
+    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)
 }