fix waitfor

Author: CodeShellDev

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)
 }