断线重连

cmd/client/client.go

@@ -2,13 +2,16 @@ package client
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"github.com/gorilla/websocket"
 	"github.com/spf13/viper"
 	"gopkg.in/toast.v1"
 	"gotify-client/internal/client"
 	"log"
+	"net/http"
 	"net/url"
+	"time"
 
 	"gotify-client/constants"
 	"gotify-client/pkg/config"
@@ -43,11 +46,10 @@ func Main() {
 	}
 
 	if err := viper.ReadInConfig(); err != nil {
-		if _, ok := err.(viper.ConfigFileNotFoundError); ok {
+		var configFileNotFoundError viper.ConfigFileNotFoundError
+		if errors.As(err, &configFileNotFoundError) {
 			_ = toml.GenerateConfig()
 			logger.Log().Fatalf("未找到配置文件, 已生成示例配置文件于运行路径下")
-		} else {
-			logger.Log().Fatalf("配置解析失败 %s", err)
 		}
 	}
 
@@ -66,23 +68,65 @@ func Main() {
 	u.RawQuery = params.Encode()
 
 	logger.Log().Infof("url %s", u.String())
-	ws, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	dialer := &websocket.Dialer{
+		Proxy:            http.ProxyFromEnvironment,
+		HandshakeTimeout: 10 * time.Second,
+	}
+
+	wsBuilder := func() (*websocket.Conn, error) {
+		ws, _, err := dialer.Dial(u.String(), nil)
+		return ws, err
+	}
+	ws, err := wsBuilder()
 	if err != nil {
 		logger.Log().Fatalf("ws 连接失败: %s, 请检查配置是否有误", err)
 	}
+
+	defer func() {
+		err := recover()
+		if err != nil {
+			log.Println("panic:", err)
+			ws, _ = wsBuilder()
+		}
+	}()
+
+	go func() {
+		hasError := false
+
+		preHandler := func() {
+			if hasError {
+				for {
+					if ws != nil {
+						_ = ws.Close()
+					}
+
+					log.Println("尝试断线重连...")
+					ws, err = wsBuilder()
+					if err != nil {
+						log.Println("ws 重连异常:", err)
+						hasError = true
+						time.Sleep(time.Second * 3)
+					} else {
+						hasError = false
+						break
+					}
+				}
+			}
+		}
+
+		handler := func() {
 			defer func() {
 				_ = ws.Close()
 			}()
 
-	done := make(chan struct{})
+			preHandler()
-
-	go func() {
-		defer close(done)
 			for {
 				_, rawMessage, err := ws.ReadMessage()
 				if err != nil {
-				log.Println("read:", err)
+					log.Println("ws 消息接收异常:", err)
-				return
+					hasError = true
+					preHandler()
+					continue
 				}
 
 				log.Printf("recv: %s", rawMessage)
@@ -98,6 +142,9 @@ func Main() {
 				notification.Message = retMessage
 				_ = notification.Push()
 			}
+		}
+
+		handler()
 	}()
 
 	quit := make(chan os.Signal)