feat: implement configurable ping timeout for nhooyr.io/websocket #88

- added `pingTimeout   time.Duration` to `internal/nhooyr.io/websocket/conn.go`'s `Conn struct`
- implemented `SetPingTimeout`, following how `nhooyr.io/websocket` exposes configurable conn values
- if `conn`'s `pingTimeout` is nonpositive, then use `handleControl`'s 5 second timeout https://github.com/xmidt-org/xmidt-agent/blob/78dffe0cad394ab82f581940e3a8117f04941077/internal/nhooyr.io/websocket/read.go#L301

internal/nhooyr.io/websocket/conn.go

@@ -17,6 +17,7 @@ import (
 	"strconv"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
 // MessageType represents the type of a WebSocket message.
@@ -79,6 +80,7 @@ type Conn struct {
 	closeErr   error
 	wroteClose bool
 
+	pingTimeout   time.Duration
 	pingCounter   int32
 	activePingsMu sync.Mutex
 	activePings   map[string]chan<- struct{}

internal/nhooyr.io/websocket/read.go

@@ -95,6 +95,13 @@ func (c *Conn) SetPongListener(f func(context.Context, []byte)) {
 	c.pongListener = f
 }
 
+// SetPingTimeout sets the maximum time allowed between PINGs for the connection
+// before the connection is closed.
+// Nonpositive PingTimeout will default to handleControl's 5 second timeout.
+func (c *Conn) SetPingTimeout(d time.Duration) {
+	c.pingTimeout = d
+}
+
 // SetReadLimit sets the max number of bytes to read for a single message.
 // It applies to the Reader and Read methods.
 //
@@ -313,6 +320,11 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
 
 	switch h.opcode {
 	case opPing:
+		if c.pingTimeout > 0 {
+			ctx, cancel = context.WithTimeout(ctx, c.pingTimeout)
+			defer cancel()
+		}
+
 		c.pingListener(ctx, b)
 		return c.writeControl(ctx, opPong, b)
 	case opPong:

internal/websocket/ws.go

@@ -233,7 +233,11 @@ func (ws *Websocket) run(ctx context.Context) {
 			// Store the connection so writing can take place.
 			ws.m.Lock()
 			ws.conn = conn
-			ws.conn.SetPingListener((func(context.Context, []byte) {
+			ws.conn.SetPingListener((func(ctx context.Context, b []byte) {
+				if ctx.Err() != nil {
+					return
+				}
+
 				ws.heartbeatListeners.Visit(func(l event.HeartbeatListener) {
 					l.OnHeartbeat(event.Heartbeat{
 						At:   ws.nowFunc(),
@@ -242,6 +246,10 @@ func (ws *Websocket) run(ctx context.Context) {
 				})
 			}))
 			ws.conn.SetPongListener(func(ctx context.Context, b []byte) {
+				if ctx.Err() != nil {
+					return
+				}
+
 				ws.heartbeatListeners.Visit(func(l event.HeartbeatListener) {
 					l.OnHeartbeat(event.Heartbeat{
 						At:   ws.nowFunc(),
@@ -334,6 +342,7 @@ func (ws *Websocket) dial(ctx context.Context, mode ipMode) (*nhws.Conn, *http.R
 	}
 
 	conn.SetReadLimit(ws.maxMessageBytes)
+	conn.SetPingTimeout(ws.pingTimeout)
 	return conn, resp, nil
 }