job-search/internal/websocket/hub.go

package websocket

import (
	"context"
	"encoding/json"
	"log/slog"
	"sync"

	ws "github.com/coder/websocket"

	"git.avsdev.uk/Go/job-search/internal/config"
)

// Hub tracks connected clients, WebSocket (ServeHTTP) and plain TCP
// (ServeTCP) alike, and sends pushed messages to all of them.
type Hub struct {
	conf config.JSCServerWebSocket
	log  *slog.Logger
	// trustedOrigins are the origins, other than the server's own, that may
	// open a WebSocket (see WithTrustedOrigins).
	trustedOrigins []string

	mu      sync.RWMutex
	clients map[*client]struct{}
	closed  bool
	wg      sync.WaitGroup
}

// defaultMaxMessageSize is the largest message read from a client when
// max_message_size is not set.
const defaultMaxMessageSize = 32 << 10

// Option configures a Hub.
type Option func(*Hub)

// WithLogger makes the hub log to logger instead of slog.Default().
func WithLogger(logger *slog.Logger) Option {
	return func(h *Hub) {
		h.log = logger
	}
}

// WithTrustedOrigins lets pages from origins (such as
// "https://admin.example.com") open a WebSocket (ServeHTTP) as well as pages
// from the server's own origin, as the API's cross-origin protection does
// for csrf.cross_origin_protection.trusted_origins.
func WithTrustedOrigins(origins []string) Option {
	return func(h *Hub) {
		h.trustedOrigins = origins
	}
}

// NewHub creates a hub. It logs to slog.Default() unless given WithLogger,
// and accepts WebSocket upgrades only from the server's own origin unless
// given WithTrustedOrigins.
func NewHub(conf config.JSCServerWebSocket, opts ...Option) *Hub {
	hub := &Hub{
		conf:    conf,
		clients: make(map[*client]struct{}),
	}
	for _, opt := range opts {
		opt(hub)
	}
	return hub
}

// logger returns the logger the hub logs to, tagged with its component.
func (h *Hub) logger() *slog.Logger {
	log := h.log
	if log == nil {
		log = slog.Default()
	}
	return log.With("component", "ws")
}

// Push sends a message of type msgType to every client. The payload is
// marshalled once and the same bytes are queued for every client; a client
// whose queue is full is closed as a slow consumer. A payload that can't be
// marshalled is logged and nothing is sent.
func (h *Hub) Push(msgType string, payload any) {
	var rawPayload json.RawMessage
	if payload != nil {
		var err error
		if rawPayload, err = json.Marshal(payload); err != nil {
			h.logger().Error("marshal push payload", "type", msgType, "error", err)
			return
		}
	}

	msg, err := json.Marshal(Envelope{
		Type:    msgType,
		Payload: rawPayload,
	})
	if err != nil {
		h.logger().Error("marshal push", "type", msgType, "error", err)
		return
	}

	var slow []*client
	h.mu.RLock()
	for client := range h.clients {
		if !client.trySend(msg) {
			slow = append(slow, client)
		}
	}
	h.mu.RUnlock()

	for _, client := range slow {
		h.logger().Warn(
			"dropping slow client",
			"client", client.id,
			"ip", client.ip,
		)
		go client.close(ws.StatusPolicyViolation, "slow consumer")
	}
}

// Count returns the number of connected clients.
func (h *Hub) Count() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients)
}

// Shutdown closes every connection with 1001 (going away) and waits for the
// handlers to return. If ctx expires first, the remaining connections are
// closed at once and ctx's error is returned. New connections are refused
// afterwards.
func (h *Hub) Shutdown(ctx context.Context) error {
	h.mu.Lock()
	h.closed = true
	clientShutdowns := make([]*client, 0, len(h.clients))
	for client := range h.clients {
		clientShutdowns = append(clientShutdowns, client)
	}
	h.mu.Unlock()

	if len(clientShutdowns) == 0 {
		// Nothing to close. Clients that just left are only finishing their
		// unregister, so this wait is short and ctx doesn't apply.
		h.wg.Wait()
		return nil
	}

	for _, client := range clientShutdowns {
		go client.close(ws.StatusGoingAway, "server shutting down")
	}

	done := make(chan struct{})
	go func() {
		h.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		return nil
	case <-ctx.Done():
		// Both may be ready; a finished shutdown is not a timeout.
		select {
		case <-done:
			return nil
		default:
		}
		for _, client := range clientShutdowns {
			client.conn.CloseNow() //nolint:errcheck,gosec // out of time; dropping it
		}
		return ctx.Err()
	}
}

// maxMessageSize is the largest message read from a client: max_message_size,
// or defaultMaxMessageSize if that is 0 or less. The configuration rejects -1
// (no limit) for this setting, as a client could then exhaust the server's
// memory; a -1 set in code also gets the default, never no limit.
func (h *Hub) maxMessageSize() int {
	if h.conf.MaxMessageSize <= 0 {
		return defaultMaxMessageSize
	}
	return int(h.conf.MaxMessageSize)
}

// register adds a client to the hub. Hub clients receive all push messages and
// have their lifecycles tracked by the hub.
func (h *Hub) register(client *client) bool {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.closed {
		return false
	}

	h.clients[client] = struct{}{}
	h.wg.Add(1)

	return true
}

// unregister removes a client from the hub.
func (h *Hub) unregister(client *client) {
	h.mu.Lock()
	delete(h.clients, client)
	h.mu.Unlock()
	h.wg.Done()
}