From 52c33ec48a1e7d3acde5e6c2f81dd52253385fae Mon Sep 17 00:00:00 2001
From: lixiangwuxian <mail@lxtend.com>
Date: Thu, 26 Dec 2024 10:41:46 +0800
Subject: [PATCH] feat: implement WebSocket client reconnection logic and
 enhance client structure

---
 main.go             | 23 ++++++++++++++++++-----
 ws_client/client.go | 17 ++++++++++-------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/main.go b/main.go
index 5672244..658819f 100644
--- a/main.go
+++ b/main.go
@@ -20,11 +20,24 @@ func main() {
 	defer exec.DockerContainer.RemoveContainer()
 	gin.SetMode(gin.ReleaseMode)
 	sqlite3.InitDB()
-	client, err := wsclient.NewWebSocketClient("ws", "localhost:3001", "")
-	if err != nil {
-		log.Print("Error creating WebSocket client:", err)
-	}
-	defer client.Close()
+	go func() {
+		const reconnectDelay = 5 * time.Second
+		for {
+			client, err := wsclient.NewWebSocketClient("ws", "localhost:3001", "")
+			if err != nil {
+				log.Printf("WebSocket连接失败: %v, %v 后重试", err, reconnectDelay)
+				time.Sleep(reconnectDelay)
+				continue
+			}
+
+			// 等待连接关闭
+			<-client.Done
+			client.Close()
+
+			log.Printf("WebSocket连接断开, %v 后重连", reconnectDelay)
+			time.Sleep(reconnectDelay)
+		}
+	}()
 	startRouter()
 	for {
 		time.Sleep(1000 * time.Second)
diff --git a/ws_client/client.go b/ws_client/client.go
index 876428d..dc560f1 100644
--- a/ws_client/client.go
+++ b/ws_client/client.go
@@ -13,8 +13,9 @@ import (
 )
 
 type WebSocketClient struct {
-	conn *websocket.Conn
-	done chan struct{}
+	conn   *websocket.Conn
+	closed bool
+	Done   chan struct{}
 }
 
 func NewWebSocketClient(scheme, host, path string) (*WebSocketClient, error) {
@@ -27,8 +28,9 @@ func NewWebSocketClient(scheme, host, path string) (*WebSocketClient, error) {
 	}
 
 	client := &WebSocketClient{
-		conn: conn,
-		done: make(chan struct{}),
+		conn:   conn,
+		closed: false,
+		Done:   make(chan struct{}),
 	}
 	defer func() {
 		if r := recover(); r != nil {
@@ -45,8 +47,8 @@ func NewWebSocketClient(scheme, host, path string) (*WebSocketClient, error) {
 }
 
 func (c *WebSocketClient) receiveMessages() {
-	defer close(c.done)
-	for {
+	defer close(c.Done)
+	for !c.closed {
 		_, message, err := c.conn.ReadMessage()
 		if err != nil {
 			log.Println("Error reading message:", err)
@@ -90,10 +92,11 @@ func (c *WebSocketClient) SendMessage(messageType int, message []byte) error {
 }
 
 func (c *WebSocketClient) Close() error {
+	c.closed = true
 	err := c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
 	if err != nil {
 		return err
 	}
-	<-c.done
+	<-c.Done
 	return c.conn.Close()
 }