example.go

@@ -16,13 +16,21 @@ const (
 func main() {
 	ws := websocket.NewWebSocketClient(
 		// onOpen
-		func(ws *websocket.WebSocketClient) {
+		func(ws *websocket.WebSocketClient, isReconnecting bool) {
+			if isReconnecting {
+				log.Println("Reconnected")
+			} else {
 				log.Println("Connected")
+			}
 		},
 
 		// onClose
-		func(ws *websocket.WebSocketClient) {
+		func(ws *websocket.WebSocketClient, isReconnecting bool) {
+			if isReconnecting {
+				log.Println("Reconnecting...")
+			} else {
 				log.Println("Disconnected")
+			}
 		},
 
 		// onMessage
@@ -32,7 +40,7 @@ func main() {
 	)
 
 	// Connect to server
-	if err := ws.Connect(WEBSOCKET_URL, ATTEMPTS, INTERVAL); err != nil {
+	if err := ws.Connect(WEBSOCKET_URL, ATTEMPTS, INTERVAL, false); err != nil {
 		log.Println("Failed to connect:", err)
 	}
 
@@ -48,11 +56,11 @@ func main() {
 	time.Sleep(2 * time.Second)
 
 	// Reconnecting
-	ws.Connect(WEBSOCKET_URL, ATTEMPTS, INTERVAL)
+	ws.Reconnect(WEBSOCKET_URL, ATTEMPTS, INTERVAL)
 	sendMessage()
 	time.Sleep(2 * time.Second)
 
-	ws.Disconnect()
+	ws.Disconnect(false)
 
 	log.Println("Done")
 }

websocket/client.go

@@ -17,14 +17,14 @@ type WebSocketClient struct {
 	wg        sync.WaitGroup
 	ctx       context.Context
 	cancel    context.CancelFunc
-	onOpen    func(ws *WebSocketClient)
+	onOpen    func(ws *WebSocketClient, isReconnecting bool)
-	onClose   func(ws *WebSocketClient)
+	onClose   func(ws *WebSocketClient, isReconnecting bool)
 	onMessage func(ws *WebSocketClient, payload []byte)
 }
 
 func NewWebSocketClient(
-	onOpen func(ws *WebSocketClient),
+	onOpen func(ws *WebSocketClient, isReconnecting bool),
-	onClose func(ws *WebSocketClient),
+	onClose func(ws *WebSocketClient, isReconnecting bool),
 	onMessage func(ws *WebSocketClient, payload []byte),
 ) *WebSocketClient {
 	ctx, cancel := context.WithCancel(context.Background())
@@ -41,10 +41,8 @@ func (ws *WebSocketClient) Connect(
 	webSocketUrl string,
 	attempts int,
 	interval time.Duration,
+	isReconnecting bool,
 ) (err error) {
-	ws.Disconnect()
-	ws.ctx, ws.cancel = context.WithCancel(context.Background())
-
 	ws.URL, err = url.Parse(webSocketUrl)
 	if err != nil {
 		return err
@@ -64,7 +62,7 @@ func (ws *WebSocketClient) Connect(
 	}
 
 	if ws.onOpen != nil {
-		ws.onOpen(ws)
+		ws.onOpen(ws, isReconnecting)
 	}
 
 	// Message Handler
@@ -74,12 +72,14 @@ func (ws *WebSocketClient) Connect(
 		for {
 			select {
 			case <-ws.ctx.Done():
+				// log.Println("websocket receive cancel")
 				return
 			default:
 				var payload []byte
 				if err := websocket.Message.Receive(ws.conn, &payload); err != nil {
+					// log.Println("receive error", err)
 					if ws.onClose != nil {
-						ws.onClose(ws)
+						ws.onClose(ws, isReconnecting)
 					}
 					return
 				}
@@ -91,7 +91,7 @@ func (ws *WebSocketClient) Connect(
 	return nil
 }
 
-func (ws *WebSocketClient) Disconnect() {
+func (ws *WebSocketClient) Disconnect(isReconnecting bool) {
 	if ws.conn != nil {
 		ws.conn.Close()
 	}
@@ -100,10 +100,16 @@ func (ws *WebSocketClient) Disconnect() {
 	ws.wg.Wait()
 
 	if ws.onClose != nil {
-		ws.onClose(ws)
+		ws.onClose(ws, isReconnecting)
 	}
 }
 
+func (ws *WebSocketClient) Reconnect(url string, attempts int, interval time.Duration) error {
+	ws.Disconnect(true)
+	ws.ctx, ws.cancel = context.WithCancel(context.Background())
+	return ws.Connect(url, attempts, interval, true)
+}
+
 func (ws *WebSocketClient) SendJSON(v any) error {
 	bytes, err := json.Marshal(v)
 	if err != nil {