[MM-64244] Add websocket disconnect reason metric (#31032)

We've recently spent some effort improving websocket reconnection logic. With this commit, I've augmented the websocket reconnect metric to include a disconnect reason. This will help us measure the impact of these changes in production.

server/channels/api4/websocket.go

@@ -5,6 +5,7 @@ package api4
 
 import (
 	"net/http"
+	"strconv"
 
 	"github.com/gorilla/websocket"
 
@@ -15,11 +16,39 @@ import (
 )
 
 const (
-	connectionIDParam   = "connection_id"
+	connectionIDParam      = "connection_id"
-	sequenceNumberParam = "sequence_number"
+	sequenceNumberParam    = "sequence_number"
-	postedAckParam      = "posted_ack"
+	postedAckParam         = "posted_ack"
+	disconnectErrCodeParam = "disconnect_err_code"
+
+	clientPingTimeoutErrCode      = 4000
+	clientSequenceMismatchErrCode = 4001
 )
 
+// validateDisconnectErrCode ensures the specified disconnect error code
+// is a valid websocket close code
+func validateDisconnectErrCode(errCode string) bool {
+	if errCode == "" {
+		return false
+	}
+
+	// Ensure the disconnect code is a standard close code
+	code, err := strconv.Atoi(errCode)
+	if err != nil {
+		return false
+	}
+
+	// We only support the standard close codes between
+	// 1000 and 1016, and a few custom application codes
+	if (code < 1000 || code > 1016) &&
+		code != clientPingTimeoutErrCode &&
+		code != clientSequenceMismatchErrCode {
+		return false
+	}
+
+	return true
+}
+
 func (api *API) InitWebSocket() {
 	// Optionally supports a trailing slash
 	api.BaseRoutes.APIRoot.Handle("/{websocket:websocket(?:\\/)?}", api.APIHandlerTrustRequester(connectWebSocket)).Methods(http.MethodGet)
@@ -53,6 +82,12 @@ func connectWebSocket(c *Context, w http.ResponseWriter, r *http.Request) {
 		RemoteAddress: c.AppContext.IPAddress(),
 		XForwardedFor: c.AppContext.XForwardedFor(),
 	}
+
+	disconnectErrCode := r.URL.Query().Get(disconnectErrCodeParam)
+	if codeValid := validateDisconnectErrCode(disconnectErrCode); codeValid {
+		cfg.DisconnectErrCode = disconnectErrCode
+	}
+
 	// The WebSocket upgrade request coming from mobile is missing the
 	// user agent so we need to fallback on the session's metadata.
 	if c.AppContext.Session().IsMobileApp() {

server/channels/api4/websocket_test.go

@@ -465,3 +465,74 @@ func TestWebSocketUpgrade(t *testing.T) {
 	require.NoError(t, th.TestLogger.Flush())
 	testlib.AssertLog(t, buffer, mlog.LvlDebug.Name, "URL Blocked because of CORS. Url: ")
 }
+
+func TestValidateDisconnectErrCode(t *testing.T) {
+	testCases := []struct {
+		name    string
+		errCode string
+		valid   bool
+	}{
+		{
+			name:    "empty string",
+			errCode: "",
+			valid:   false,
+		},
+		{
+			name:    "non-numeric string",
+			errCode: "not-a-number",
+			valid:   false,
+		},
+		{
+			name:    "valid standard close code - 1000",
+			errCode: "1000",
+			valid:   true,
+		},
+		{
+			name:    "valid standard close code - 1001",
+			errCode: "1001",
+			valid:   true,
+		},
+		{
+			name:    "valid standard close code - 1015",
+			errCode: "1015",
+			valid:   true,
+		},
+		{
+			name:    "valid standard close code - 1016",
+			errCode: "1016",
+			valid:   true,
+		},
+		{
+			name:    "out of range (too low)",
+			errCode: "999",
+			valid:   false,
+		},
+		{
+			name:    "out of range (too high)",
+			errCode: "1017",
+			valid:   false,
+		},
+		{
+			name:    "valid custom code - client ping timeout",
+			errCode: "4000",
+			valid:   true,
+		},
+		{
+			name:    "valid custom code - client sequence mismatch",
+			errCode: "4001",
+			valid:   true,
+		},
+		{
+			name:    "invalid custom code",
+			errCode: "5000",
+			valid:   false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := validateDisconnectErrCode(tc.errCode)
+			require.Equal(t, tc.valid, result)
+		})
+	}
+}

server/channels/app/platform/web_conn.go

@@ -62,17 +62,18 @@ type pluginWSPostedHook struct {
 }
 
 type WebConnConfig struct {
-	WebSocket     *websocket.Conn
+	WebSocket         *websocket.Conn
-	Session       model.Session
+	Session           model.Session
-	TFunc         i18n.TranslateFunc
+	TFunc             i18n.TranslateFunc
-	Locale        string
+	Locale            string
-	ConnectionID  string
+	ConnectionID      string
-	Active        bool
+	Active            bool
-	ReuseCount    int
+	ReuseCount        int
-	OriginClient  string
+	OriginClient      string
-	PostedAck     bool
+	PostedAck         bool
-	RemoteAddress string
+	RemoteAddress     string
-	XForwardedFor string
+	XForwardedFor     string
+	DisconnectErrCode string
 
 	// These aren't necessary to be exported to api layer.
 	sequence         int64
@@ -85,16 +86,17 @@ type WebConnConfig struct {
 // It contains all the necessary state to manage sending/receiving data to/from
 // a websocket.
 type WebConn struct {
-	sessionExpiresAt int64 // This should stay at the top for 64-bit alignment of 64-bit words accessed atomically
+	sessionExpiresAt  int64 // This should stay at the top for 64-bit alignment of 64-bit words accessed atomically
-	Platform         *PlatformService
+	Platform          *PlatformService
-	Suite            SuiteIFace
+	Suite             SuiteIFace
-	HookRunner       HookRunner
+	HookRunner        HookRunner
-	WebSocket        *websocket.Conn
+	WebSocket         *websocket.Conn
-	T                i18n.TranslateFunc
+	T                 i18n.TranslateFunc
-	Locale           string
+	Locale            string
-	Sequence         int64
+	Sequence          int64
-	UserId           string
+	UserId            string
-	PostedAck        bool
+	PostedAck         bool
+	DisconnectErrCode string
 
 	allChannelMembers         map[string]string
 	lastAllChannelMembersTime int64
@@ -246,6 +248,7 @@ func (ps *PlatformService) NewWebConn(cfg *WebConnConfig, suite SuiteIFace, runn
 		T:                  cfg.TFunc,
 		Locale:             cfg.Locale,
 		PostedAck:          cfg.PostedAck,
+		DisconnectErrCode:  cfg.DisconnectErrCode,
 		reuseCount:         cfg.ReuseCount,
 		endWritePump:       make(chan struct{}),
 		pumpFinished:       make(chan struct{}),
@@ -523,7 +526,7 @@ func (wc *WebConn) writePump() {
 				return
 			}
 			if m := wc.Platform.metricsIFace; m != nil {
-				m.IncrementWebsocketReconnectEvent(reconnectFound)
+				m.IncrementWebsocketReconnectEventWithDisconnectErrCode(reconnectFound, wc.DisconnectErrCode)
 			}
 		} else if wc.hasMsgLoss() {
 			// If the seq number is not in dead queue, but it was supposed to be,
@@ -541,11 +544,11 @@ func (wc *WebConn) writePump() {
 				return
 			}
 			if m := wc.Platform.metricsIFace; m != nil {
-				m.IncrementWebsocketReconnectEvent(reconnectNotFound)
+				m.IncrementWebsocketReconnectEventWithDisconnectErrCode(reconnectNotFound, wc.DisconnectErrCode)
 			}
 		} else {
 			if m := wc.Platform.metricsIFace; m != nil {
-				m.IncrementWebsocketReconnectEvent(reconnectLossless)
+				m.IncrementWebsocketReconnectEventWithDisconnectErrCode(reconnectLossless, wc.DisconnectErrCode)
 			}
 		}
 	}

server/einterfaces/metrics.go

@@ -48,7 +48,7 @@ type MetricsInterface interface {
 	DecrementWebSocketBroadcastBufferSize(hub string, amount float64)
 	IncrementWebSocketBroadcastUsersRegistered(hub string, amount float64)
 	DecrementWebSocketBroadcastUsersRegistered(hub string, amount float64)
-	IncrementWebsocketReconnectEvent(eventType string)
+	IncrementWebsocketReconnectEventWithDisconnectErrCode(eventType string, disconnectErrCode string)
 
 	IncrementHTTPWebSockets(originClient string)
 	DecrementHTTPWebSockets(originClient string)

server/einterfaces/mocks/MetricsInterface.go

@@ -298,9 +298,9 @@ func (_m *MetricsInterface) IncrementWebsocketEvent(eventType model.WebsocketEve
 	_m.Called(eventType)
 }
 
-// IncrementWebsocketReconnectEvent provides a mock function with given fields: eventType
+// IncrementWebsocketReconnectEventWithDisconnectErrCode provides a mock function with given fields: eventType, disconnectErrCode
-func (_m *MetricsInterface) IncrementWebsocketReconnectEvent(eventType string) {
+func (_m *MetricsInterface) IncrementWebsocketReconnectEventWithDisconnectErrCode(eventType string, disconnectErrCode string) {
-	_m.Called(eventType)
+	_m.Called(eventType, disconnectErrCode)
 }
 
 // ObserveAPIEndpointDuration provides a mock function with given fields: endpoint, method, statusCode, originClient, pageLoadContext, elapsed

server/enterprise/metrics/metrics.go

@@ -726,7 +726,7 @@ func New(ps *platform.PlatformService, driver, dataSource string) *MetricsInterf
 			Help:        "Total number of websocket reconnect attempts",
 			ConstLabels: additionalLabels,
 		},
-		[]string{"type"},
+		[]string{"type", "disconnect_err_code"},
 	)
 	m.Registry.MustRegister(m.WebSocketReconnectCounter)
 
@@ -1822,8 +1822,14 @@ func (mi *MetricsInterfaceImpl) IncrementWebsocketEvent(eventType model.Websocke
 	mi.WebsocketEventCounters.With(prometheus.Labels{"type": string(eventType)}).Inc()
 }
 
-func (mi *MetricsInterfaceImpl) IncrementWebsocketReconnectEvent(eventType string) {
+func (mi *MetricsInterfaceImpl) IncrementWebsocketReconnectEventWithDisconnectErrCode(eventType string, disconnectErrCode string) {
-	mi.WebSocketReconnectCounter.With(prometheus.Labels{"type": eventType}).Inc()
+	if disconnectErrCode == "" {
+		disconnectErrCode = "unknown"
+	}
+	mi.WebSocketReconnectCounter.With(prometheus.Labels{
+		"type":                eventType,
+		"disconnect_err_code": disconnectErrCode,
+	}).Inc()
 }
 
 func (mi *MetricsInterfaceImpl) IncrementWebSocketBroadcastBufferSize(hub string, amount float64) {

webapp/channels/src/components/post_view/post_body_additional_content/__snapshots__/post_body_additional_content.test.tsx.snap

@@ -115,6 +115,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "eventCallback": null,
         "firstConnectCallback": null,
         "firstConnectListeners": Set {},
+        "lastErrCode": null,
         "messageListeners": Set {},
         "missedEventCallback": null,
         "missedMessageListeners": Set {},
@@ -175,6 +176,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "eventCallback": null,
         "firstConnectCallback": null,
         "firstConnectListeners": Set {},
+        "lastErrCode": null,
         "messageListeners": Set {},
         "missedEventCallback": null,
         "missedMessageListeners": Set {},

webapp/platform/client/src/websocket.test.ts

@@ -49,6 +49,22 @@ if (typeof Event === 'undefined') {
     };
 }
 
+// Mock CloseEvent class if it's not defined
+if (typeof CloseEvent === 'undefined') {
+    (global as any).CloseEvent = class MockCloseEvent extends (global as any).Event {
+        code: number;
+        reason: string;
+        wasClean: boolean;
+
+        constructor(type: string, options?: {code?: number; reason?: string; wasClean?: boolean}) {
+            super(type);
+            this.code = options?.code || 0;
+            this.reason = options?.reason || '';
+            this.wasClean = options?.wasClean || false;
+        }
+    };
+}
+
 class MockWebSocket {
     readonly binaryType: BinaryType = 'blob';
     readonly bufferedAmount: number = 0;

webapp/platform/client/src/websocket.ts

@@ -19,6 +19,10 @@ export type WebSocketClientConfig = {
     clientPingInterval: number;
 }
 
+// Custom close error codes must be in the range of 4000-4999
+const clientPingTimeoutErrCode = 4000;
+const clientSequenceMismatchErrCode = 4001;
+
 const defaultWebSocketClientConfig: WebSocketClientConfig = {
     maxWebSocketFails: 7,
     minWebSocketRetryTime: 3000, // 3 seconds
@@ -45,6 +49,7 @@ export default class WebSocketClient {
     private serverSequence: number;
     private connectFailCount: number;
     private responseCallbacks: {[x: number]: ((msg: any) => void)};
+    private lastErrCode: string | null;
 
     /**
      * @deprecated Use messageListeners instead
@@ -110,6 +115,7 @@ export default class WebSocketClient {
         this.config = {...defaultWebSocketClientConfig, ...config};
         this.pingInterval = null;
         this.waitingForPong = false;
+        this.lastErrCode = null;
     }
 
     // on connect, only send auth cookie and blank state.
@@ -198,19 +204,32 @@ export default class WebSocketClient {
         // Add connection id, and last_sequence_number to the query param.
         // We cannot use a cookie because it will bleed across tabs.
         // We cannot also send it as part of the auth_challenge, because the session cookie is already sent with the request.
-        const websocketUrl = `${connectionUrl}?connection_id=${this.connectionId}&sequence_number=${this.serverSequence}${this.postedAck ? '&posted_ack=true' : ''}`;
+        let websocketUrl = `${connectionUrl}?connection_id=${this.connectionId}&sequence_number=${this.serverSequence}`;
+
+        if (this.postedAck) {
+            websocketUrl += '&posted_ack=true';
+        }
+
+        if (this.lastErrCode) {
+            websocketUrl += `&disconnect_err_code=${encodeURIComponent(this.lastErrCode)}`;
+        }
+
         if (this.config.newWebSocketFn) {
             this.conn = this.config.newWebSocketFn(websocketUrl);
         } else {
             this.conn = new WebSocket(websocketUrl);
         }
 
-        const onclose = () => {
+        const onclose = (event: CloseEvent) => {
             this.conn = null;
             this.responseSequence = 1;
 
+            if (!this.lastErrCode && event && event.code) {
+                this.lastErrCode = `${event.code}`;
+            }
+
             if (this.connectFailCount === 0) {
-                console.log('websocket closed'); //eslint-disable-line no-console
+                console.log(`websocket closed: ${this.lastErrCode}`); //eslint-disable-line no-console
             }
 
             this.connectFailCount++;
@@ -254,6 +273,8 @@ export default class WebSocketClient {
                 this.sendMessage('authentication_challenge', {token});
             }
 
+            this.lastErrCode = null;
+
             if (this.connectFailCount > 0) {
                 console.log('websocket re-established connection'); //eslint-disable-line no-console
 
@@ -294,6 +315,11 @@ export default class WebSocketClient {
 
                     console.log('ping received no response within time limit: re-establishing websocket'); //eslint-disable-line no-console
 
+                    const closeEvent = new CloseEvent('close', {
+                        code: clientPingTimeoutErrCode,
+                        wasClean: false,
+                    });
+
                     // Calling conn.close() will trigger the onclose callback,
                     // but sometimes with a significant delay. So instead, we
                     // call the onclose callback ourselves immediately. We also
@@ -303,7 +329,7 @@ export default class WebSocketClient {
                     this.responseSequence = 1;
                     this.conn.onclose = () => {};
                     this.conn.close();
-                    onclose();
+                    onclose(closeEvent);
                 },
                 this.config.clientPingInterval);
 
@@ -369,10 +395,24 @@ export default class WebSocketClient {
                 // we just disconnect and reconnect.
                 if (msg.seq !== this.serverSequence) {
                     console.log('missed websocket event, act_seq=' + msg.seq + ' exp_seq=' + this.serverSequence); //eslint-disable-line no-console
-                    // We are not calling this.close() because we need to auto-restart.
+
+                    const closeEvent = new CloseEvent('close', {
+                        code: clientSequenceMismatchErrCode,
+                        wasClean: false,
+                    });
+
+                    // Calling conn.close() will trigger the onclose callback,
+                    // but sometimes with a significant delay. So instead, we
+                    // call the onclose callback ourselves immediately. We also
+                    // unset the callback on the old connection to ensure it
+                    // is only called once.
                     this.connectFailCount = 0;
                     this.responseSequence = 1;
-                    this.conn?.close(); // Will auto-reconnect after MIN_WEBSOCKET_RETRY_TIME.
+                    if (this.conn) {
+                        this.conn.onclose = () => {};
+                        this.conn.close();
+                        onclose(closeEvent);
+                    }
                     return;
                 }
                 this.serverSequence = msg.seq + 1;
@@ -507,13 +547,14 @@ export default class WebSocketClient {
         this.connectFailCount = 0;
         this.responseSequence = 1;
         this.clearReconnectTimeout();
+        this.lastErrCode = null;
         this.stopPingInterval();
 
         if (this.conn && this.conn.readyState === WebSocket.OPEN) {
             this.conn.onclose = () => {};
             this.conn.close();
             this.conn = null;
-            console.log('websocket closed'); //eslint-disable-line no-console
+            console.log('websocket closed manually'); //eslint-disable-line no-console
         }
 
         if (this.onlineHandler) {