From 2cecb924402de4a01701c7082ffda870a24b4fbb Mon Sep 17 00:00:00 2001
From: Akizon77 <i@akz.moe>
Date: Tue, 8 Jul 2025 15:08:17 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20https://github.com/komari-monitor/komar?=
 =?UTF-8?q?i/issues/108=20=20=E5=AE=9E=E7=8E=B0=E5=BA=94=E7=94=A8=E5=B1=82?=
 =?UTF-8?q?=E7=9A=84=E5=BF=83=E8=B7=B3=E6=9C=BA=E5=88=B6=20git=E5=BF=BD?=
 =?UTF-8?q?=E7=95=A5=5F=5Fdebug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore          |  1 +
 server/websocket.go | 76 +++++++++++++++++++++++++++------------------
 2 files changed, 46 insertions(+), 31 deletions(-)

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