websocket

package
v0.1.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Jun 7, 2026 License: MIT Imports: 16 Imported by: 0

Documentation

Overview

Package websocket provides a Hub-based WebSocket connection manager with signal-first broadcasting, ping/pong health checks, and graceful shutdown.

Hub lifecycle: NewHub → Start (blocks) → Stop (terminal, single-use).

This package is protocol-agnostic: it operates on the Conn interface. Use adapters/websocket for the github.com/coder/websocket binding.

Example wiring at composition root (managed mode — recommended):

hub := websocket.NewHub(clock.Real(), websocket.DefaultHubConfig(), msgHandler)
bootstrap.New(..., bootstrap.WithManagedResource(hub))
// bootstrap auto-starts the hub via Hub.Worker() (kernel/worker.Worker)
// and tears it down via Hub.Close() during phase10 LIFO shutdown.
// Do NOT also run `go hub.Start(ctx)` — duplicate Start returns
// ErrWSAlreadyStarted and breaks the bootstrap worker pipeline.

Manual mode (legacy / unit tests outside bootstrap):

hub := websocket.NewHub(clock.Real(), websocket.DefaultHubConfig(), msgHandler)
go func() { _ = hub.Start(ctx) }()
defer hub.Stop(ctx)

Index

Constants

View Source
const ProbeReady healthz.ProbeName = "websocket_hub_ready"

ProbeReady is the ops-contract name for the websocket hub readiness probe. healthz.ProbeName-typed, funneled by PROBENAME-SEALED-FUNNEL-01.

Variables

This section is empty.

Functions

func MustValidateHubConfig

func MustValidateHubConfig(cfg HubConfig)

MustValidateHubConfig panics if cfg contains values that would cause runtime panics or zero-budget shutdown later. Called by NewHub. Every panic in this function is wrapped with panicregister.Approved per PANIC-REGISTERED-01.

Validates:

  • ShutdownTimeout < 0 → would collapse external-cancel context to already-expired
  • ConcurrentCloseLimit < 0 → would panic at make(chan struct{}, limit) during shutdown

Zero values are accepted (NewHub falls back to package defaults).

Types

type Conn

type Conn interface {
	// ID returns the unique connection identifier.
	ID() string
	// Ping checks liveness. Returns an error if the peer is unresponsive.
	Ping(ctx context.Context) error
	// Read blocks until a message arrives or the context is canceled.
	Read(ctx context.Context) ([]byte, error)
	// Write sends a text message.
	Write(ctx context.Context, data []byte) error
	// Close closes the connection. The contract:
	//   - Idempotent: multiple calls return nil.
	//   - Does NOT require a graceful WebSocket close handshake.
	//   - Must cause any in-progress Read to return promptly.
	//   - Must not block longer than necessary (no long mutex waits).
	// If a graceful close is needed in the future, add CloseGracefully(ctx).
	Close() error

	// RemoteAddr returns the remote network address as a "host:port" string.
	// Used by Hub for diagnostic logging (eviction reasons, shutdown drain).
	//
	// Implementations MUST return a sanitized address (e.g. via
	// pkg/logutil.SafeAddr) — the Hub logs this value verbatim. Empty string
	// is not a valid return; use "unknown" if address is unavailable. The
	// address must be immutable after construction; the Hub may invoke
	// RemoteAddr() outside connMu.
	RemoteAddr() string

	// Principal returns the authenticated principal bound at handshake time.
	// It may return nil if the adapter did not bind a principal (e.g. test
	// fakes); the Hub treats nil as "no subject indexing, no expiry tracked"
	// and behaves as if the connection has Kind=PrincipalUnknown.
	//
	// The returned pointer is owned by the Conn implementation; consumers
	// must not mutate the Principal struct.
	//
	// Both Principal().Subject and RemoteAddr() must be immutable after
	// construction; the Hub may invoke them outside connMu.
	//
	// Principal data (Subject in particular) appears in slog at Warn/Info —
	// operators should treat WS structured logs as containing user identifiers
	// and apply log-pipeline redaction if PII handling requires it.
	Principal() *auth.Principal
}

Conn abstracts a WebSocket connection. Implementations live in adapters/ (e.g., adapters/websocket for github.com/coder/websocket).

Implementations must be safe for concurrent use:

  • Read is called from a single goroutine per connection.
  • Write, Ping, and Close may be called concurrently with Read.
  • Close must cause any in-progress Read to return an error.

type Hub

type Hub struct {
	// contains filtered or unexported fields
}

Hub manages WebSocket connections and provides signal-first broadcasting.

Lifecycle: NewHub → Start (blocks) → Stop (terminal, single-use). A stopped Hub cannot be restarted; create a new one instead.

Both Stop(ctx) and external cancellation of Start(ctx) converge on the same internal shutdown path. There is exactly one code path that drains connections and transitions to the terminal state.

ref: centrifugal/centrifuge hub.go — sharded hub, semaphore-bounded

concurrent close, snapshot-then-release drain pattern.

ref: olahol/melody hub.go — pointer-keyed session map, channel-based

exit signal (we use atomic CAS instead).

ref: coder/websocket chat example — CloseNow + defer pattern for

connection teardown; CloseRead for context-based lifecycle.

func NewHub

func NewHub(clk clock.Clock, cfg HubConfig, handler MessageHandler) *Hub

NewHub creates a Hub. No background goroutines are started until Start.

Fail-fast wiring: negative ShutdownTimeout / ConcurrentCloseLimit panic at construction (see MustValidateHubConfig). Zero values fall back to package defaults (10s / 64). This matches clock.MustHaveClock's panic-on-misconfig style — wiring bugs surface before any goroutine runs, never as a shutdown-time make-chan panic or an already-expired context.

func (*Hub) BroadcastFilter

func (h *Hub) BroadcastFilter(ctx context.Context, data []byte, filter func(Conn) bool) error

BroadcastFilter sends data to every connection for which filter returns true. filter must be non-nil; passing nil returns ErrWebsocketBroadcastFilterMissing (fail-closed: full-broadcast must be expressed as `func(Conn) bool { return true }`).

filter is invoked under no lock; implementations must be O(1) (a closure that reads Conn.Principal() is fine; database queries are an anti-pattern). filter may be invoked from a single goroutine (the caller's goroutine).

A connection whose send buffer is full is evicted (slow client; the broadcast continues to other connections). Returns nil on success even if some connections were evicted; the eviction is logged and counted via slog.

ref: olahol/melody melody.go BroadcastFilter

func (*Hub) BroadcastToSubject

func (h *Hub) BroadcastToSubject(ctx context.Context, subject string, data []byte) error

BroadcastToSubject sends data to every connection whose Principal.Subject matches the supplied subject. Subject "" returns ErrWebsocketBroadcastSubjectMissing (fail-closed: callers must declare an explicit identity). An unknown subject (no matching connections) is a no-op and returns nil.

O(1) lookup via the hub's subject index (centrifuge-style). Connection eviction on full send buffer is the same as BroadcastFilter.

ref: centrifugal/centrifuge hub.go connShard.users

func (*Hub) Close

func (h *Hub) Close(ctx context.Context) error

Close implements lifecycle.ManagedResource. It delegates to Stop(ctx) and is idempotent: if the Hub is already stopped, returns nil instead of the ErrWSAlreadyStopped error that Stop would return.

Bootstrap wiring: bootstrap.WithManagedResource(hub) calls Close in LIFO order during phase10 shutdown.

ref: adapters/rabbitmq/connection.go Close — idempotent + ctx-bounded.

func (*Hub) Config

func (h *Hub) Config() HubConfig

Config returns the Hub's configuration.

func (*Hub) ConnCount

func (h *Hub) ConnCount() int

ConnCount returns the number of active connections.

func (*Hub) IsRunning

func (h *Hub) IsRunning() bool

IsRunning reports whether the Hub is in the running state and can accept new connections. Use this to guard HTTP upgrade handlers.

func (*Hub) Probes

func (h *Hub) Probes() []healthz.Probe

Probes implements lifecycle.ManagedResource. It returns a single typed probe ProbeReady ("websocket_hub_ready") that reports nil (healthy) only when the Hub is in the running state. Any other state (idle/stopping/ stopped) returns errHubNotRunning.

probe name "websocket_hub_ready" follows the observability rule: snake_case + "_ready" suffix.

ref: adapters/rabbitmq/connection.go Probes — probe name + pre-allocated error pattern.

func (*Hub) Register

func (h *Hub) Register(ctx context.Context, conn Conn) error

Register adds a connection to the Hub and starts reading from it. ctx is used as the parent for per-connection values; cancellation is controlled by the Hub so request-scope cancellation does not immediately tear down accepted WebSocket connections after the HTTP handler returns. The read loop uses a per-connection context that is canceled when the connection is unregistered or the Hub shuts down.

The Hub must be in the running state (Start called). Register on an idle, stopping, or stopped Hub returns an error and closes the conn.

The per-connection context carries the connection's Principal (if present) so that MessageHandler can call auth.FromContext(ctx) for inbound ACL.

func (*Hub) Send

func (h *Hub) Send(ctx context.Context, connID string, data []byte) error

Send sends a text message to a specific connection. If the connection's send buffer is full, the connection is evicted and ErrWebsocketSlowClient is returned. If ctx is already canceled, returns ctx.Err() without enqueuing.

func (*Hub) Start

func (h *Hub) Start(ctx context.Context) error

Start begins the Hub's ping loop. It blocks until Stop is called or ctx is canceled.

If the caller's context is canceled, Start runs a full shutdown (drain + close) automatically. Stop may still be called afterwards but will return immediately since the Hub is already stopped.

func (*Hub) Stop

func (h *Hub) Stop(ctx context.Context) error

Stop gracefully shuts down the Hub. It rejects new connections, cancels the ping loop, closes all connections (breaking readLoops), and waits for goroutines to exit. The provided ctx bounds the entire operation.

After Stop the Hub is in a terminal state and cannot be restarted.

func (*Hub) Unregister

func (h *Hub) Unregister(connID string)

Unregister removes a connection from the Hub by ID and closes it. This is a force-remove: it always deletes the entry regardless of which conn object is registered.

func (*Hub) Worker

func (h *Hub) Worker() worker.Worker

Worker returns a worker.Worker that drives Start/Stop so bootstrap can auto-manage the hub lifecycle.

Composition root (PR #393, post-/fix): one of the two equivalent patterns

hub := websocket.NewHub(clk, cfg, handler)
bootstrap.New(..., bootstrap.WithManagedResource(hub))
// hub.Start is invoked by the bootstrap WorkerGroup; Close runs LIFO
// during phase10. Do NOT also run `go hub.Start(ctx)` manually — the
// duplicate Start would return ErrWSAlreadyStarted.

Manual mode (legacy / tests):

hub := websocket.NewHub(clk, cfg, handler)
go func() { _ = hub.Start(ctx) }()
// caller is responsible for hub.Stop(ctx) on shutdown.

ref: kernel/lifecycle/managed_resource.go::Worker — non-nil documents "auto-managed background worker" ref: uber-go/fx app.go Lifecycle — managed surface owns start+stop ref: go-kratos/kratos transport/transport.go Server — same contract

type HubConfig

type HubConfig struct {
	// PingInterval is the interval between ping sweeps. Default: 30s.
	PingInterval time.Duration
	// PingTimeout is the deadline for a single ping. Default: 5s.
	PingTimeout time.Duration
	// ReadLimit is the maximum message size in bytes. Default: 64KB.
	// The adapter applies this when creating a connection.
	ReadLimit int64
	// PingMissMax is the number of consecutive ping failures before a
	// connection is evicted. Default: 2.
	PingMissMax int
	// MaxConnections is the maximum number of concurrent connections.
	// 0 means unlimited. Default: 0.
	//
	// SECURITY: Production deployments MUST set an explicit cap to prevent
	// goroutine exhaustion (each connection runs 2 goroutines: readLoop +
	// writeLoop). Recommended starting value: expected concurrent sessions
	// + 20% headroom. Token-authenticated unlimited capacity is a DoS path.
	MaxConnections int
	// SendBufferSize is the per-connection send channel capacity used by the
	// writeLoop. When the channel is full the connection is evicted (slow
	// client; gorilla/websocket select-default-drop). Default 32; zero value
	// is replaced with the default at construction time. Must be > 0 if
	// explicitly set.
	SendBufferSize int

	// ShutdownTimeout limits the total time allowed for the external-cancel
	// shutdown path inside Start (caller's ctx is canceled with no deadline).
	// Stop(ctx) paths are not affected — caller's ctx governs directly.
	// 0 → defaultShutdownTimeout (10s).
	//
	// To bound Stop(ctx), supply a deadline on the ctx passed to Stop, or
	// configure bootstrap.WithShutdownTimeout for bootstrap-managed shutdown
	// — ShutdownTimeout only governs the external-cancel path inside Start.
	ShutdownTimeout time.Duration

	// ConcurrentCloseLimit bounds the semaphore pool used during shutdown drain
	// — at most this many conn.Close() calls run concurrently. Tuning guidance:
	// keep below system fd-limit / goroutine budget. Default 64 (matches
	// centrifuge's hubShutdownSemaphoreSize). Background goroutines spawned by
	// shutdown's outer wg.Wait/done-close pattern exit when readLoops/writeLoops
	// drain via Phase 1 context cancellation; max lifetime = connection teardown
	// latency.
	//
	// ref: centrifugal/centrifuge hub.go — bounded concurrent close
	ConcurrentCloseLimit int
}

HubConfig configures the Hub.

func DefaultHubConfig

func DefaultHubConfig() HubConfig

DefaultHubConfig returns a HubConfig with sensible defaults. Pass the clock as the first positional argument to NewHub.

type MessageHandler

type MessageHandler func(ctx context.Context, connID string, data []byte)

MessageHandler is called when a message is received from a client.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL