From b71c1c4738979998b8487a7a1cc7dc7027ffbd56 Mon Sep 17 00:00:00 2001
From: Siyuan Miao <i@xswan.net>
Date: Thu, 10 Apr 2025 15:39:03 +0200
Subject: [PATCH] chore(websocket): only show warning if websocket is closed
 abnormally

---
 cloud.go | 10 ++++++++--
 web.go   | 10 ++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/cloud.go b/cloud.go
index 89666a1..76b281b 100644
--- a/cloud.go
+++ b/cloud.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -280,17 +281,22 @@ func runWebsocketClient() error {
 			return true
 		},
 	})
+	// if the context is canceled, we don't want to return an error
 	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			cloudLogger.Infof("websocket connection canceled")
+			return nil
+		}
 		return err
 	}
 	defer c.CloseNow() //nolint:errcheck
 	cloudLogger.Infof("websocket connected to %s", wsURL)
 
 	// set the metrics when we successfully connect to the cloud.
-	wsResetMetrics(true, "cloud", "")
+	wsResetMetrics(true, "cloud", wsURL.Host)
 
 	// we don't have a source for the cloud connection
-	return handleWebRTCSignalWsMessages(c, true, "")
+	return handleWebRTCSignalWsMessages(c, true, wsURL.Host)
 }
 
 func authenticateSession(ctx context.Context, c *websocket.Conn, req WebRTCSessionRequest) error {
diff --git a/web.go b/web.go
index 0258dc6..1b1f8af 100644
--- a/web.go
+++ b/web.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"embed"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/fs"
 	"net/http"
@@ -251,6 +252,15 @@ func handleWebRTCSignalWsMessages(wsCon *websocket.Conn, isCloudConnection bool,
 		for {
 			time.Sleep(WebsocketPingInterval)
 
+			if ctxErr := runCtx.Err(); ctxErr != nil {
+				if !errors.Is(ctxErr, context.Canceled) {
+					logWarnf("websocket connection closed: %v", ctxErr)
+				} else {
+					logTracef("websocket connection closed as the context was canceled: %v")
+				}
+				return
+			}
+
 			// set the timer for the ping duration
 			timer := prometheus.NewTimer(prometheus.ObserverFunc(func(v float64) {
 				metricConnectionLastPingDuration.WithLabelValues(sourceType, source).Set(v)