puma-plus-core 0.1.0

data/go/internal/realtime/hub.go

@@ -0,0 +1,456 @@
+// Package realtime holds the transport-agnostic half of puma-plus's persistent
+// connections: the connection registry, topic membership, fan-out, and the
+// sticky/flexible routing model.
+//
+// Both WebSocket and WebTransport plug into it by implementing Transport. That
+// is not incidental tidiness — it is what lets one topic contain clients of both
+// kinds, so an application broadcasts once and reaches everybody regardless of
+// which protocol their browser supports. During the years that support is
+// uneven, dual-publishing logic in the app is exactly the thing worth avoiding.
+package realtime
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// Errors returned by transports that cannot honour a capability.
+var (
+	// ErrNoDatagrams means the transport has no unreliable channel. WebSocket
+	// runs over TCP, so every message is reliable and ordered whether you wanted
+	// that or not.
+	ErrNoDatagrams = errors.New("realtime: transport does not support datagrams")
+	// ErrNoStreams means the transport has a single message channel rather than
+	// independent streams.
+	ErrNoStreams = errors.New("realtime: transport does not support streams")
+	// ErrClosed means the connection is gone.
+	ErrClosed = errors.New("realtime: connection closed")
+)
+
+// Transport is how a connection actually delivers bytes.
+//
+// The asymmetry between the methods is deliberate and load-bearing: WebSocket
+// cannot do datagrams or independent streams, and rather than pretend otherwise
+// those methods return an error naming the reason. An application that reaches
+// for a datagram on a WebSocket client should find out, not silently get
+// reliable delivery it did not ask for and will not get on the next client.
+type Transport interface {
+	// Kind is "websocket" or "webtransport".
+	Kind() string
+
+	// SendMessage delivers reliably and in order on the connection's default
+	// channel.
+	SendMessage(payload []byte, binary bool) error
+
+	// SendOnStream delivers on a specific stream. ErrNoStreams if unsupported.
+	SendOnStream(streamID string, payload []byte, binary bool) error
+
+	// OpenStream initiates a server-side stream and returns its id.
+	// ErrNoStreams if unsupported.
+	OpenStream() (string, error)
+
+	// SendDatagram delivers best-effort: it may be dropped, duplicated or
+	// reordered. ErrNoDatagrams if unsupported.
+	SendDatagram(payload []byte) error
+
+	// Close terminates the connection.
+	Close(code int, reason string)
+
+	// RemoteAddr for logging and the Rack env.
+	RemoteAddr() string
+}
+
+// preframer is an optional optimization for fan-out.
+//
+// A WebSocket message has to be wrapped in a frame; doing that once per publish
+// instead of once per subscriber is most of why fan-out belongs in Go. A
+// transport that can reuse pre-encoded bytes implements this; one that cannot
+// just gets SendMessage.
+type preframer interface {
+	// PreframeMessage encodes a message once so it can be written to many
+	// connections of the same kind.
+	PreframeMessage(payload []byte, binary bool) []byte
+	// SendPreframed writes bytes produced by PreframeMessage.
+	SendPreframed(frame []byte) error
+}
+
+// Mode is how a connection's messages are routed to Ruby workers.
+type Mode string
+
+const (
+	// ModeFlexible delivers each message to whichever worker is free.
+	//
+	// Maximum throughput, and one chatty connection can use the whole pool. The
+	// cost is ordering: two messages from the same client may be handled
+	// concurrently by different workers.
+	ModeFlexible Mode = "flexible"
+
+	// ModeSticky pins a connection to one worker process for its lifetime, so
+	// that process can keep per-connection state in ordinary memory and sees the
+	// connection's messages in order.
+	//
+	// Worth noting the tension with WebTransport: independent streams exist
+	// precisely so they do not block each other, and funnelling them all through
+	// one worker gives some of that back. Sticky is still the right choice for a
+	// stateful session; it is the wrong one for parallel bulk transfer.
+	ModeSticky Mode = "sticky"
+)
+
+// Conn is one persistent client connection, of either transport.
+type Conn struct {
+	ID   string
+	Mode Mode
+	Path string
+	Meta map[string]string
+
+	// WorkerID is the worker a sticky connection is pinned to, valid only once
+	// Pinned is true. Pinned is a separate flag because 0 is a legitimate worker
+	// id: a single worker running without a shepherd has exactly that id.
+	WorkerID int
+	Pinned   bool
+
+	tr     Transport
+	hub    *Hub
+	closed atomic.Bool
+	opened time.Time
+}
+
+// NewConn wraps a transport as a hub connection.
+func NewConn(h *Hub, tr Transport, mode Mode, path string, meta map[string]string) *Conn {
+	if mode != ModeSticky {
+		mode = ModeFlexible
+	}
+	if meta == nil {
+		meta = map[string]string{}
+	}
+	return &Conn{
+		ID: NewID(), Mode: mode, Path: path, Meta: meta,
+		tr: tr, hub: h, opened: time.Now(),
+	}
+}
+
+// Transport exposes the underlying transport.
+func (c *Conn) Transport() Transport { return c.tr }
+
+// Kind is "websocket" or "webtransport".
+func (c *Conn) Kind() string { return c.tr.Kind() }
+
+// RemoteAddr reports the peer address.
+func (c *Conn) RemoteAddr() string { return c.tr.RemoteAddr() }
+
+// IsClosed reports whether the connection has been torn down.
+func (c *Conn) IsClosed() bool { return c.closed.Load() }
+
+// Send delivers a reliable message on the default channel.
+func (c *Conn) Send(payload []byte, binary bool) error {
+	if c.closed.Load() {
+		return ErrClosed
+	}
+	return c.tr.SendMessage(payload, binary)
+}
+
+// SendOnStream delivers on a specific stream.
+func (c *Conn) SendOnStream(streamID string, payload []byte, binary bool) error {
+	if c.closed.Load() {
+		return ErrClosed
+	}
+	return c.tr.SendOnStream(streamID, payload, binary)
+}
+
+// OpenStream initiates a server-side stream.
+func (c *Conn) OpenStream() (string, error) {
+	if c.closed.Load() {
+		return "", ErrClosed
+	}
+	return c.tr.OpenStream()
+}
+
+// SendDatagram delivers best-effort.
+func (c *Conn) SendDatagram(payload []byte) error {
+	if c.closed.Load() {
+		return ErrClosed
+	}
+	return c.tr.SendDatagram(payload)
+}
+
+// CloseWith terminates the connection.
+func (c *Conn) CloseWith(code int, reason string) {
+	if c.closed.Swap(true) {
+		return
+	}
+	c.tr.Close(code, reason)
+}
+
+// Hub owns every live connection and the topic membership between them.
+type Hub struct {
+	mu     sync.RWMutex
+	conns  map[string]*Conn
+	topics map[string]map[string]*Conn
+
+	// connTopics is the reverse index, so closing a connection does not have to
+	// scan every topic to unsubscribe it.
+	connTopics map[string]map[string]struct{}
+}
+
+// NewHub creates an empty hub.
+func NewHub() *Hub {
+	return &Hub{
+		conns:      map[string]*Conn{},
+		topics:     map[string]map[string]*Conn{},
+		connTopics: map[string]map[string]struct{}{},
+	}
+}
+
+// Add registers a connection.
+func (h *Hub) Add(c *Conn) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.conns[c.ID] = c
+	h.connTopics[c.ID] = map[string]struct{}{}
+}
+
+// Remove deregisters a connection and drops all of its subscriptions.
+func (h *Hub) Remove(id string) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	delete(h.conns, id)
+	for topic := range h.connTopics[id] {
+		if subs := h.topics[topic]; subs != nil {
+			delete(subs, id)
+			if len(subs) == 0 {
+				delete(h.topics, topic)
+			}
+		}
+	}
+	delete(h.connTopics, id)
+}
+
+// Get looks up a connection.
+func (h *Hub) Get(id string) *Conn {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return h.conns[id]
+}
+
+// Count is the number of live connections.
+func (h *Hub) Count() int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return len(h.conns)
+}
+
+// Subscribe adds a connection to a topic.
+func (h *Hub) Subscribe(id, topic string) bool {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	c := h.conns[id]
+	if c == nil {
+		return false
+	}
+	if h.topics[topic] == nil {
+		h.topics[topic] = map[string]*Conn{}
+	}
+	h.topics[topic][id] = c
+	h.connTopics[id][topic] = struct{}{}
+	return true
+}
+
+// Unsubscribe removes a connection from a topic.
+func (h *Hub) Unsubscribe(id, topic string) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	if subs := h.topics[topic]; subs != nil {
+		delete(subs, id)
+		if len(subs) == 0 {
+			delete(h.topics, topic)
+		}
+	}
+	if ct := h.connTopics[id]; ct != nil {
+		delete(ct, topic)
+	}
+}
+
+// Subscribers counts a topic's members.
+func (h *Hub) Subscribers(topic string) int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return len(h.topics[topic])
+}
+
+// targets snapshots a topic's members. Taken under lock and written outside it,
+// so one stalled peer cannot freeze every other connection's routing.
+func (h *Hub) targets(topic string) []*Conn {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	subs := h.topics[topic]
+	out := make([]*Conn, 0, len(subs))
+	for _, c := range subs {
+		out = append(out, c)
+	}
+	return out
+}
+
+// Publish fans a reliable message out to every subscriber of a topic, across
+// both transports, and reports how many received it.
+//
+// A topic may hold WebSocket and WebTransport clients at once; each gets the
+// message in its own native form. That is the point of a shared hub: the
+// application publishes once and does not track who is on which protocol.
+func (h *Hub) Publish(topic string, payload []byte, binary bool) int {
+	targets := h.targets(topic)
+	if len(targets) == 0 {
+		return 0
+	}
+
+	// Pre-encode once per transport kind rather than once per recipient. For a
+	// WebSocket broadcast to ten thousand subscribers that is one framing
+	// operation instead of ten thousand.
+	preframed := map[string][]byte{}
+	for _, c := range targets {
+		k := c.tr.Kind()
+		if _, done := preframed[k]; done {
+			continue
+		}
+		if pf, ok := c.tr.(preframer); ok {
+			preframed[k] = pf.PreframeMessage(payload, binary)
+		} else {
+			preframed[k] = nil
+		}
+	}
+
+	var wg sync.WaitGroup
+	var delivered atomic.Int64
+	for _, c := range targets {
+		wg.Add(1)
+		go func(c *Conn) {
+			defer wg.Done()
+			if c.closed.Load() {
+				return
+			}
+			var err error
+			if frame := preframed[c.tr.Kind()]; frame != nil {
+				err = c.tr.(preframer).SendPreframed(frame)
+			} else {
+				err = c.tr.SendMessage(payload, binary)
+			}
+			if err == nil {
+				delivered.Add(1)
+			}
+		}(c)
+	}
+	wg.Wait()
+	return int(delivered.Load())
+}
+
+// PublishDatagram fans a message out unreliably.
+//
+// Deliberately a separate method from Publish rather than a flag on it.
+// Datagrams may be dropped, duplicated or reordered, which is a fundamentally
+// different contract, and a one-word keyword argument is far too easy to copy
+// between call sites without noticing the guarantee changed. Subscribers whose
+// transport has no datagram channel are skipped and counted separately, so a
+// caller can tell the difference between "delivered" and "could not possibly be
+// delivered".
+func (h *Hub) PublishDatagram(topic string, payload []byte) (delivered, unsupported int) {
+	targets := h.targets(topic)
+	if len(targets) == 0 {
+		return 0, 0
+	}
+
+	var wg sync.WaitGroup
+	var ok, no atomic.Int64
+	for _, c := range targets {
+		wg.Add(1)
+		go func(c *Conn) {
+			defer wg.Done()
+			if c.closed.Load() {
+				return
+			}
+			switch err := c.tr.SendDatagram(payload); {
+			case err == nil:
+				ok.Add(1)
+			case errors.Is(err, ErrNoDatagrams):
+				no.Add(1)
+			}
+		}(c)
+	}
+	wg.Wait()
+	return int(ok.Load()), int(no.Load())
+}
+
+// DropWorker closes every sticky connection pinned to a departing worker.
+//
+// Sticky means the application is keeping per-connection state inside that
+// process. When the process goes, so does the state, and re-pinning would hand a
+// live conversation to a worker that has never heard of it. Flexible connections
+// are untouched: by definition they hold no per-worker state.
+func (h *Hub) DropWorker(workerID int, code int, reason string) []string {
+	h.mu.RLock()
+	var doomed []*Conn
+	for _, c := range h.conns {
+		if c.Mode == ModeSticky && c.Pinned && c.WorkerID == workerID {
+			doomed = append(doomed, c)
+		}
+	}
+	h.mu.RUnlock()
+
+	ids := make([]string, 0, len(doomed))
+	for _, c := range doomed {
+		ids = append(ids, c.ID)
+		c.CloseWith(code, reason)
+	}
+	return ids
+}
+
+// Stats summarises hub state for /stats.
+type Stats struct {
+	Connections  int `json:"connections"`
+	Topics       int `json:"topics"`
+	Sticky       int `json:"sticky"`
+	Flexible     int `json:"flexible"`
+	WebSocket    int `json:"websocket"`
+	WebTransport int `json:"webtransport"`
+}
+
+// Stats snapshots the hub.
+func (h *Hub) Stats() Stats {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	s := Stats{Connections: len(h.conns), Topics: len(h.topics)}
+	for _, c := range h.conns {
+		if c.Mode == ModeSticky {
+			s.Sticky++
+		} else {
+			s.Flexible++
+		}
+		switch c.tr.Kind() {
+		case "webtransport":
+			s.WebTransport++
+		default:
+			s.WebSocket++
+		}
+	}
+	return s
+}
+
+// CloseAll shuts every connection down, for server shutdown.
+func (h *Hub) CloseAll(code int, reason string) {
+	h.mu.RLock()
+	all := make([]*Conn, 0, len(h.conns))
+	for _, c := range h.conns {
+		all = append(all, c)
+	}
+	h.mu.RUnlock()
+	for _, c := range all {
+		c.CloseWith(code, reason)
+	}
+}
+
+// Hub returns the hub this connection belongs to.
+func (c *Conn) Hub() *Hub { return c.hub }
+
+// MarkClosed records that the transport has gone away, without trying to send a
+// close frame over it. Used by read loops that have already seen EOF.
+func (c *Conn) MarkClosed() { c.closed.Store(true) }