fix: harden websocket heartbeat under load

server/hdl_websock.go

@@ -28,14 +28,35 @@ import (
 const (
 	// Time allowed to write a message to the peer.
 	writeWait = 10 * time.Second
-
-	// Time allowed to read the next pong message from the peer.
-	pongWait = idleSessionTimeout
-
-	// Send pings to peer with this period. Must be less than pongWait.
-	pingPeriod = 15 * time.Second
+	// Heartbeat defaults tolerate transient mobile and public-network stalls while
+	// still detecting dead connections in a bounded amount of time.
+	defaultWSPongWait   = 90 * time.Second
+	defaultWSPingPeriod = 15 * time.Second
 )
 
+var (
+	pongWait   = defaultWSPongWait
+	pingPeriod = defaultWSPingPeriod
+)
+
+func configureWSHeartbeat(pingSeconds, pongSeconds int) error {
+	ping := defaultWSPingPeriod
+	pong := defaultWSPongWait
+	if pingSeconds > 0 {
+		ping = time.Duration(pingSeconds) * time.Second
+	}
+	if pongSeconds > 0 {
+		pong = time.Duration(pongSeconds) * time.Second
+	}
+	if ping >= pong {
+		return errors.New("websocket ping period must be shorter than pong wait")
+	}
+	pingPeriod = ping
+	pongWait = pong
+	immetrics.SetWSHeartbeatConfig(ping, pong)
+	return nil
+}
+
 func (sess *Session) closeWS() {
 	if sess.proto == WEBSOCK {
 		sess.ws.Close()
@@ -52,6 +73,7 @@ func (sess *Session) readLoop() {
 	sess.ws.SetReadDeadline(time.Now().Add(pongWait))
 	sess.ws.SetPongHandler(func(string) error {
 		atomic.StoreInt64(&sess.lastAction, types.TimeNow().UnixNano())
+		immetrics.IncWSHeartbeat("pong_received")
 		sess.refreshPresence("ws.pong")
 		sess.ws.SetReadDeadline(time.Now().Add(pongWait))
 		return nil

server/hdl_websock_config_test.go

@@ -0,0 +1,29 @@
+package main
+
+import (
+	"testing"
+	"time"
+)
+
+func TestConfigureWSHeartbeat(t *testing.T) {
+	originalPing, originalPong := pingPeriod, pongWait
+	t.Cleanup(func() {
+		pingPeriod, pongWait = originalPing, originalPong
+	})
+
+	if err := configureWSHeartbeat(20, 120); err != nil {
+		t.Fatal(err)
+	}
+	if pingPeriod != 20*time.Second || pongWait != 120*time.Second {
+		t.Fatalf("heartbeat ping=%s pong=%s", pingPeriod, pongWait)
+	}
+	if err := configureWSHeartbeat(90, 90); err == nil {
+		t.Fatal("expected invalid heartbeat configuration to fail")
+	}
+	if err := configureWSHeartbeat(0, 0); err != nil {
+		t.Fatal(err)
+	}
+	if pingPeriod != defaultWSPingPeriod || pongWait != defaultWSPongWait {
+		t.Fatalf("default heartbeat ping=%s pong=%s", pingPeriod, pongWait)
+	}
+}

server/main.go

@@ -322,6 +322,10 @@ type configType struct {
 	// If true, do not attempt to negotiate websocket per message compression (RFC 7692.4).
 	// It should be disabled (set to true) if you are using MSFT IIS as a reverse proxy.
 	WSCompressionDisabled bool `json:"ws_compression_disabled"`
+	// WebSocket heartbeat period and the maximum time to wait for a Pong, in seconds.
+	// Zero values use the built-in defaults.
+	WSPingPeriodSeconds int `json:"ws_ping_period_seconds"`
+	WSPongWaitSeconds   int `json:"ws_pong_wait_seconds"`
 	// Address:port to listen for gRPC clients. If blank gRPC support will not be initialized.
 	// Could be overridden from the command line with --grpc_listen.
 	GrpcListen string `json:"grpc_listen"`
@@ -716,6 +720,10 @@ func main() {
 
 	// Websocket compression.
 	globals.wsCompression = !config.WSCompressionDisabled
+	if err = configureWSHeartbeat(config.WSPingPeriodSeconds, config.WSPongWaitSeconds); err != nil {
+		logs.Err.Fatal("Invalid WebSocket heartbeat configuration: ", err)
+	}
+	logs.Info.Printf("WebSocket heartbeat ping=%s pong_wait=%s", pingPeriod, pongWait)
 
 	if config.Media != nil {
 		if config.Media.UseHandler == "" {

server/metrics/metrics.go

@@ -71,6 +71,14 @@ var (
 		Name: "im_ws_disconnect_total",
 		Help: "WebSocket read-loop disconnects by pod and bounded reason.",
 	}, []string{"pod", "reason"})
+	IMWSHeartbeatTotal = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Name: "im_ws_heartbeat_total",
+		Help: "WebSocket heartbeat events by pod and event.",
+	}, []string{"pod", "event"})
+	IMWSHeartbeatSeconds = prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Name: "im_ws_heartbeat_seconds",
+		Help: "Effective WebSocket heartbeat timing by pod and setting.",
+	}, []string{"pod", "setting"})
 	IMWSSendQueueCurrent = prometheus.NewGaugeVec(prometheus.GaugeOpts{
 		Name: "im_ws_send_queue_current",
 		Help: "Current number of queued WebSocket outbound channel entries across sessions on a pod.",
@@ -216,7 +224,8 @@ var (
 
 func init() {
 	Registry.MustRegister(WSConnectionsActive, IMWSConnectionsActive, IMWSWriteTotal, IMWSWriteDuration,
-		IMWSWriteBytesTotal, IMWSDisconnectTotal, IMWSSendQueueCurrent, IMWSSendQueueCapacity,
+		IMWSWriteBytesTotal, IMWSDisconnectTotal, IMWSHeartbeatTotal, IMWSHeartbeatSeconds,
+		IMWSSendQueueCurrent, IMWSSendQueueCapacity,
 		IMWSSendQueueEnqueueTotal, IMWSSendQueueDepth, IMWSSendQueueWaitDuration,
 		IMClusterLivePeers, IMClusterExpectedPeers,
 		HTTPRequestsTotal, HTTPRequestDuration, DBQueryDuration,
@@ -332,6 +341,16 @@ func IncWSDisconnect(reason string) {
 	IMWSDisconnectTotal.WithLabelValues(metricPodName(), normalizeLabel(reason)).Inc()
 }
 
+func IncWSHeartbeat(event string) {
+	IMWSHeartbeatTotal.WithLabelValues(metricPodName(), normalizeLabel(event)).Inc()
+}
+
+func SetWSHeartbeatConfig(pingPeriod, pongWait time.Duration) {
+	pod := metricPodName()
+	IMWSHeartbeatSeconds.WithLabelValues(pod, "ping_period").Set(pingPeriod.Seconds())
+	IMWSHeartbeatSeconds.WithLabelValues(pod, "pong_wait").Set(pongWait.Seconds())
+}
+
 func ObserveWSSendQueueEnqueue(depth, capacity int, accepted bool) {
 	result := "ok"
 	if !accepted {