diff --git a/.gitignore b/.gitignore index 94ec8a5..8767c86 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +__debug* agent.json .vscode/ komari-agent.exe diff --git a/server/websocket.go b/server/websocket.go index 3263c38..f79411d 100644 --- a/server/websocket.go +++ b/server/websocket.go @@ -33,43 +33,57 @@ func EstablishWebSocketConnection() { interval = flags.Interval - 1 } - ticker := time.NewTicker(time.Duration(interval * float64(time.Second))) - defer ticker.Stop() + dataTicker := time.NewTicker(time.Duration(interval * float64(time.Second))) + defer dataTicker.Stop() - for range ticker.C { - // If no connection, attempt to connect - if conn == nil { - log.Println("Attempting to connect to WebSocket...") - retry := 0 - for retry <= flags.MaxRetries { - if retry > 0 { - log.Println("Retrying websocket connection, attempt:", retry) + heartbeatTicker := time.NewTicker(30 * time.Second) + defer heartbeatTicker.Stop() + + for { + select { + case <-dataTicker.C: + if conn == nil { + log.Println("Attempting to connect to WebSocket...") + retry := 0 + for retry <= flags.MaxRetries { + if retry > 0 { + log.Println("Retrying websocket connection, attempt:", retry) + } + conn, err = connectWebSocket(websocketEndpoint) + if err == nil { + log.Println("WebSocket connected") + go handleWebSocketMessages(conn, make(chan struct{})) + break + } else { + log.Println("Failed to connect to WebSocket:", err) + } + retry++ + time.Sleep(time.Duration(flags.ReconnectInterval) * time.Second) } - conn, err = connectWebSocket(websocketEndpoint) - if err == nil { - log.Println("WebSocket connected") - go handleWebSocketMessages(conn, make(chan struct{})) - break - } else { - log.Println("Failed to connect to WebSocket:", err) + + if retry > flags.MaxRetries { + log.Println("Max retries reached.") + return } - retry++ - time.Sleep(time.Duration(flags.ReconnectInterval) * time.Second) } - if retry > flags.MaxRetries { - log.Println("Max retries reached.") - return + data := monitoring.GenerateReport() + err = conn.WriteMessage(websocket.TextMessage, data) + if err != nil { + log.Println("Failed to send WebSocket message:", err) + conn.Close() + conn = nil // Mark connection as dead + continue + } + case <-heartbeatTicker.C: + if conn != nil { + err := conn.WriteMessage(websocket.PingMessage, nil) + if err != nil { + log.Println("Failed to send heartbeat:", err) + conn.Close() + conn = nil // Mark connection as dead + } } - } - - data := monitoring.GenerateReport() - err = conn.WriteMessage(websocket.TextMessage, data) - if err != nil { - log.Println("Failed to send WebSocket message:", err) - conn.Close() - conn = nil // Mark connection as dead - continue } } }