fix(listen): stop reconnecting when a websocket pong is a little late

pkg/websocket/client.go

@@ -207,13 +207,27 @@ func (c *Client) Run(ctx context.Context) {
 			"prefix": "websocket.client.Run",
 		}).Debug("Disconnected from Hookdeck")
 
-		c.ConnectionLost()
 		close(c.stopReadPump)
 		close(c.stopWritePump)
+		c.closeForReconnect()
+		c.ConnectionLost()
 		c.wg.Wait()
 	}
 }
 
+// closeForReconnect drops the socket without a close frame so the CLI can
+// reconnect. Otherwise the socket stays open until the server's idle check
+// notices it is dead.
+//
+// It must not send a 1000 close like Stop does. The server treats a 1000 close
+// as an intentional shutdown and deletes the session immediately, whereas an
+// abnormal close (1006) keeps the session for the reconnect grace window.
+func (c *Client) closeForReconnect() {
+	c.stateMu.Lock()
+	defer c.stateMu.Unlock()
+	_ = c.conn.Close()
+}
+
 // ConnectionLost sends NotifyExpired
 func (c *Client) ConnectionLost() {
 	c.NotifyExpired <- struct{}{}
@@ -456,6 +470,12 @@ func (c *Client) readPump() {
 			return
 		}
 
+		// Any frame proves the connection is alive, so a busy connection
+		// doesn't drop just because a pong arrived late.
+		if err := c.conn.SetReadDeadline(time.Now().Add(c.cfg.PongWait)); err != nil {
+			c.cfg.Log.Debug("SetReadDeadline error: ", err)
+		}
+
 		c.cfg.Log.WithFields(log.Fields{
 			"prefix":  "websocket.Client.readPump",
 			"message": string(data),
@@ -576,7 +596,7 @@ func NewClient(url string, webSocketID string, CLIKey string, teamID string, con
 	}
 
 	if cfg.PingPeriod == 0 {
-		cfg.PingPeriod = (cfg.PongWait * 9) / 10
+		cfg.PingPeriod = min(defaultPingPeriod, (cfg.PongWait*9)/10)
 	}
 
 	if cfg.WriteWait == 0 {
@@ -610,7 +630,13 @@ func NewClient(url string, webSocketID string, CLIKey string, teamID string, con
 const (
 	defaultConnectAttemptWait = 10 * time.Second
 
-	defaultPongWait = 10 * time.Second
+	// defaultPongWait leaves room for a few late pongs (slow network, busy
+	// server) before the connection is treated as dead.
+	defaultPongWait = 30 * time.Second
+
+	// defaultPingPeriod stays well under the server's idle check so the server
+	// keeps seeing the connection as alive.
+	defaultPingPeriod = 10 * time.Second
 
 	defaultWriteWait = 10 * time.Second
 

pkg/websocket/client_test.go

@@ -292,3 +292,95 @@ func TestLastConnectErr(t *testing.T) {
 		}
 	})
 }
+
+// serverWithoutPongs upgrades the connection and reads from it without
+// answering pings, as a server would when its pongs are delayed.
+func serverWithoutPongs(t *testing.T, onConn func(conn *ws.Conn)) *httptest.Server {
+	t.Helper()
+	var captured http.Header
+	return upgradeTestServer(t, &captured, func(conn *ws.Conn) {
+		conn.SetPingHandler(func(string) error { return nil })
+		onConn(conn)
+	})
+}
+
+// Data frames prove the connection is alive, so a connection that keeps
+// receiving messages must not drop when pongs are late.
+func TestIncomingMessagesKeepConnectionAlive(t *testing.T) {
+	server := serverWithoutPongs(t, func(conn *ws.Conn) {
+		go func() {
+			for {
+				if _, _, err := conn.ReadMessage(); err != nil {
+					return
+				}
+			}
+		}()
+		for i := 0; i < 10; i++ {
+			if err := conn.WriteMessage(ws.TextMessage, []byte(`{}`)); err != nil {
+				return
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+	})
+	defer server.Close()
+
+	client := NewClient(wsURL(server), "cses_test", "cli-key", "tm_test", nil, "", &Config{
+		PongWait:   300 * time.Millisecond,
+		PingPeriod: 100 * time.Millisecond,
+	})
+	startClient(t, client)
+
+	select {
+	case <-client.NotifyExpired:
+		t.Fatal("client dropped a connection that was receiving messages")
+	case <-time.After(900 * time.Millisecond):
+	}
+}
+
+// When the read deadline expires, the client must close its socket before
+// reconnecting rather than leave it open for the server to reap.
+func TestReadTimeoutClosesSocket(t *testing.T) {
+	serverReadErr := make(chan error, 1)
+	server := serverWithoutPongs(t, func(conn *ws.Conn) {
+		for {
+			if _, _, err := conn.ReadMessage(); err != nil {
+				serverReadErr <- err
+				return
+			}
+		}
+	})
+	defer server.Close()
+
+	client := NewClient(wsURL(server), "cses_test", "cli-key", "tm_test", nil, "", &Config{
+		PongWait:   300 * time.Millisecond,
+		PingPeriod: 100 * time.Millisecond,
+	})
+	startClient(t, client)
+
+	select {
+	case <-client.NotifyExpired:
+	case <-time.After(5 * time.Second):
+		t.Fatal("client did not report connection loss after the read deadline")
+	}
+
+	select {
+	case err := <-serverReadErr:
+		// A 1000 close would make the server delete the session instead of
+		// keeping it for the reconnect.
+		if !ws.IsCloseError(err, ws.CloseAbnormalClosure) {
+			t.Fatalf("server saw %v, want an abnormal close (1006) so the session survives the reconnect", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("client left the old socket open after the read deadline")
+	}
+}
+
+func TestDefaultKeepalive(t *testing.T) {
+	client := NewClient("ws://127.0.0.1:1", "cses_test", "cli-key", "tm_test", nil, "", nil)
+	if client.cfg.PongWait != 30*time.Second {
+		t.Errorf("PongWait = %s, want 30s", client.cfg.PongWait)
+	}
+	if client.cfg.PingPeriod != 10*time.Second {
+		t.Errorf("PingPeriod = %s, want 10s", client.cfg.PingPeriod)
+	}
+}