Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 20 additions & 3 deletions pkg/websocket/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -207,9 +207,14 @@ func (c *Client) Run(ctx context.Context) {
"prefix": "websocket.client.Run",
}).Debug("Disconnected from Hookdeck")

c.ConnectionLost()
close(c.stopReadPump)
close(c.stopWritePump)
// Close the old socket before reconnecting. Otherwise it stays open
// until the server's idle check notices it is dead.
c.stateMu.Lock()
_ = c.conn.Close()
c.stateMu.Unlock()
c.ConnectionLost()
c.wg.Wait()
}
}
Expand Down Expand Up @@ -456,6 +461,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),
Expand Down Expand Up @@ -576,7 +587,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 {
Expand Down Expand Up @@ -610,7 +621,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

Expand Down
87 changes: 87 additions & 0 deletions pkg/websocket/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -292,3 +292,90 @@ 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 <-serverReadErr:
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)
}
}
Loading