package httpapi import ( "context" "errors" "net" "time" "golang.org/x/xerrors" "cdr.dev/slog/v3" "github.com/coder/quartz" "github.com/coder/websocket" ) const HeartbeatInterval time.Duration = 15 * time.Second // HeartbeatClose loops to ping a WebSocket to keep it alive. // It calls `exit` on ping failure. func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn) { heartbeatCloseWith(ctx, logger, exit, conn, quartz.NewReal(), HeartbeatInterval) } // HeartbeatCloseWithClock is like HeartbeatClose, but uses the provided // clock so tests can drive heartbeat ticks deterministically. func HeartbeatCloseWithClock(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn, clk quartz.Clock) { heartbeatCloseWith(ctx, logger, exit, conn, clk, HeartbeatInterval) } func heartbeatCloseWith(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn, clk quartz.Clock, interval time.Duration) { ticker := clk.NewTicker(interval, "HeartbeatClose") defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: } err := pingWithTimeout(ctx, conn, interval) if err != nil { // These errors are all expected during normal connection // teardown and should not be logged at error level: // - context.DeadlineExceeded: client disconnected // without sending a close frame. // - context.Canceled: request context was canceled. // - net.ErrClosed: connection was already closed by // another goroutine (e.g. handler returned). // - websocket.CloseError: a close frame was // received or sent. if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || errors.Is(err, net.ErrClosed) || websocket.CloseStatus(err) != -1 { logger.Debug(ctx, "heartbeat ping stopped", slog.Error(err)) } else { logger.Error(ctx, "failed to heartbeat ping", slog.Error(err)) } _ = conn.Close(websocket.StatusGoingAway, "Ping failed") exit() return } } } func pingWithTimeout(ctx context.Context, conn *websocket.Conn, timeout time.Duration) error { ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() err := conn.Ping(ctx) if err != nil { return xerrors.Errorf("failed to ping: %w", err) } return nil }