feat: implement WebSocket client reconnection logic and enhance client structure

This commit is contained in:
lixiangwuxian 2024-12-26 10:41:46 +08:00
parent d2fc755df5
commit 52c33ec48a
2 changed files with 28 additions and 12 deletions

17
main.go
View File

@ -20,11 +20,24 @@ func main() {
defer exec.DockerContainer.RemoveContainer() defer exec.DockerContainer.RemoveContainer()
gin.SetMode(gin.ReleaseMode) gin.SetMode(gin.ReleaseMode)
sqlite3.InitDB() sqlite3.InitDB()
go func() {
const reconnectDelay = 5 * time.Second
for {
client, err := wsclient.NewWebSocketClient("ws", "localhost:3001", "") client, err := wsclient.NewWebSocketClient("ws", "localhost:3001", "")
if err != nil { if err != nil {
log.Print("Error creating WebSocket client:", err) log.Printf("WebSocket连接失败: %v, %v 后重试", err, reconnectDelay)
time.Sleep(reconnectDelay)
continue
} }
defer client.Close()
// 等待连接关闭
<-client.Done
client.Close()
log.Printf("WebSocket连接断开, %v 后重连", reconnectDelay)
time.Sleep(reconnectDelay)
}
}()
startRouter() startRouter()
for { for {
time.Sleep(1000 * time.Second) time.Sleep(1000 * time.Second)

View File

@ -14,7 +14,8 @@ import (
type WebSocketClient struct { type WebSocketClient struct {
conn *websocket.Conn conn *websocket.Conn
done chan struct{} closed bool
Done chan struct{}
} }
func NewWebSocketClient(scheme, host, path string) (*WebSocketClient, error) { func NewWebSocketClient(scheme, host, path string) (*WebSocketClient, error) {
@ -28,7 +29,8 @@ func NewWebSocketClient(scheme, host, path string) (*WebSocketClient, error) {
client := &WebSocketClient{ client := &WebSocketClient{
conn: conn, conn: conn,
done: make(chan struct{}), closed: false,
Done: make(chan struct{}),
} }
defer func() { defer func() {
if r := recover(); r != nil { if r := recover(); r != nil {
@ -45,8 +47,8 @@ func NewWebSocketClient(scheme, host, path string) (*WebSocketClient, error) {
} }
func (c *WebSocketClient) receiveMessages() { func (c *WebSocketClient) receiveMessages() {
defer close(c.done) defer close(c.Done)
for { for !c.closed {
_, message, err := c.conn.ReadMessage() _, message, err := c.conn.ReadMessage()
if err != nil { if err != nil {
log.Println("Error reading message:", err) log.Println("Error reading message:", err)
@ -90,10 +92,11 @@ func (c *WebSocketClient) SendMessage(messageType int, message []byte) error {
} }
func (c *WebSocketClient) Close() error { func (c *WebSocketClient) Close() error {
c.closed = true
err := c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) err := c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
if err != nil { if err != nil {
return err return err
} }
<-c.done <-c.Done
return c.conn.Close() return c.conn.Close()
} }