coder/coderd/httpapi/websocket.go

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

// ProbeResult classifies the outcome of a single WebSocket liveness
// probe so that callers (typically a Prometheus recorder) can track
// successes and the various failure modes independently.
type ProbeResult string

const (
	ProbeOK         ProbeResult = "ok"
	ProbeTimeout    ProbeResult = "timeout"
	ProbePeerClosed ProbeResult = "peer_closed"
	ProbeCanceled   ProbeResult = "canceled"
	ProbeError      ProbeResult = "error"
)

// ProbeRecorder is called once per liveness probe with its outcome.
// It may be nil, in which case probes are still run but not recorded.
type ProbeRecorder func(ctx context.Context, result ProbeResult)

// PingCloser is the minimal interface for WebSocket liveness probing.
// *websocket.Conn satisfies this interface.
type PingCloser interface {
	Ping(ctx context.Context) error
	Close(code websocket.StatusCode, reason string) error
}

// WSWatcher supervises WebSocket connections for liveness by
// periodically sending ping frames. On probe failure, the watcher
// closes the connection with StatusGoingAway and cancels the
// returned context; the caller owns closing the connection on
// normal teardown.
type WSWatcher struct {
	rec      ProbeRecorder
	clk      quartz.Clock
	interval time.Duration
}

// NewWSWatcher creates a WSWatcher. Pass nil for rec when no
// recording is needed (e.g. agent-side code without a Prometheus
// registry).
func NewWSWatcher(clk quartz.Clock, rec ProbeRecorder) *WSWatcher {
	return &WSWatcher{
		rec:      rec,
		clk:      clk,
		interval: HeartbeatInterval,
	}
}

// Watch supervises conn for liveness. The returned context is
// canceled when parent is canceled or when conn fails a probe.
// Watch closes conn on probe failure with StatusGoingAway; the
// caller owns close on normal teardown.
func (w *WSWatcher) Watch(parent context.Context, log slog.Logger, conn PingCloser) context.Context {
	if w == nil {
		panic("developer error: WSWatcher is nil")
	}
	ctx, cancel := context.WithCancel(parent)
	go func() {
		defer cancel()
		w.supervise(ctx, log, conn)
	}()
	return ctx
}

func (w *WSWatcher) supervise(ctx context.Context, log slog.Logger, conn PingCloser) {
	ticker := w.clk.NewTicker(w.interval, "WSWatcher")
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}

		result, err := probe(ctx, conn, w.interval)
		if w.rec != nil {
			w.rec(ctx, result)
		}
		if result == ProbeOK {
			continue
		}
		if result == ProbeError {
			log.Error(ctx, "websocket probe failed", slog.Error(err))
		} else {
			log.Debug(ctx, "websocket probe stopped",
				slog.F("result", string(result)), slog.Error(err))
		}
		_ = conn.Close(websocket.StatusGoingAway, "liveness probe failed")
		return
	}
}

func probe(ctx context.Context, conn PingCloser, timeout time.Duration) (ProbeResult, error) {
	pingCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	err := conn.Ping(pingCtx)
	switch {
	case err == nil:
		return ProbeOK, nil
	case errors.Is(err, context.Canceled):
		return ProbeCanceled, err
	case errors.Is(err, context.DeadlineExceeded):
		return ProbeTimeout, err
	case errors.Is(err, net.ErrClosed) || websocket.CloseStatus(err) != -1:
		return ProbePeerClosed, err
	default:
		return ProbeError, xerrors.Errorf("ping: %w", err)
	}
}

// HeartbeatClose is a legacy helper that pings conn in a loop and
// calls exit on failure. Callers that need metric recording should
// use WSWatcher directly.
func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn) {
	w := NewWSWatcher(quartz.NewReal(), nil)
	watchCtx := w.Watch(ctx, logger, conn)
	<-watchCtx.Done()
	// Only call exit when the probe failed; if the parent context was
	// canceled the caller is already shutting down.
	if ctx.Err() == nil {
		exit()
	}
}