gw-05 step 01 — A WebSocket server and a push registry

Goal

Stand up a WebSocket server that accepts device connections, registers each deviceId → node in a push registry, and keeps connections alive with ping/pong. This is the connection-holding half of Pushy.

Code — src/go/node.go

package pushy

import (
	"context"
	"sync"
	"time"

	"nhooyr.io/websocket" // go get nhooyr.io/websocket
)

// Node is one Pushy instance: it holds live device connections and can
// write a message down any of them.
type Node struct {
	ID       string
	Registry Registry

	mu    sync.RWMutex
	conns map[string]*deviceConn // deviceId -> connection
}

type deviceConn struct {
	deviceID string
	ws       *websocket.Conn
	out      chan []byte // bounded per-connection write queue (backpressure)
}

func NewNode(id string, r Registry) *Node {
	return &Node{ID: id, Registry: r, conns: map[string]*deviceConn{}}
}

// Handle runs for the life of one device connection.
func (n *Node) Handle(ctx context.Context, ws *websocket.Conn, deviceID string) {
	dc := &deviceConn{deviceID: deviceID, ws: ws, out: make(chan []byte, 16)}

	n.mu.Lock()
	n.conns[deviceID] = dc
	n.mu.Unlock()
	n.Registry.Register(deviceID, n.ID, 60*time.Second) // device D lives on this node

	defer func() {
		n.mu.Lock()
		delete(n.conns, deviceID)
		n.mu.Unlock()
		n.Registry.Unregister(deviceID, n.ID)
		ws.Close(websocket.StatusNormalClosure, "bye")
	}()

	go n.writePump(ctx, dc) // drains dc.out to the socket
	n.readPump(ctx, dc)     // reads client frames + keeps registry TTL fresh
}

// writePump is the ONLY goroutine that writes to the socket — serializes
// writes and applies a per-connection deadline.
func (n *Node) writePump(ctx context.Context, dc *deviceConn) {
	ping := time.NewTicker(20 * time.Second) // app-layer keepalive
	defer ping.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case msg := <-dc.out:
			wctx, cancel := context.WithTimeout(ctx, 5*time.Second)
			err := dc.ws.Write(wctx, websocket.MessageText, msg)
			cancel()
			if err != nil {
				return
			}
		case <-ping.C:
			pctx, cancel := context.WithTimeout(ctx, 5*time.Second)
			err := dc.ws.Ping(pctx)
			cancel()
			if err != nil {
				return // dead peer detected
			}
		}
	}
}

func (n *Node) readPump(ctx context.Context, dc *deviceConn) {
	for {
		_, _, err := dc.ws.Read(ctx) // client->server msgs / pongs
		if err != nil {
			return // connection closed
		}
		n.Registry.Touch(dc.deviceID, n.ID, 60*time.Second) // refresh TTL
	}
}

// Deliver enqueues a message for a locally-connected device. Returns
// false if the device isn't on this node (caller should consult registry).
func (n *Node) Deliver(deviceID string, msg []byte) bool {
	n.mu.RLock()
	dc, ok := n.conns[deviceID]
	n.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case dc.out <- msg:
		return true
	default:
		// Write queue full (slow client): drop or disconnect per policy.
		// Never block the delivery path on one slow device.
		return false
	}
}

Code — the registry interface

package pushy

import (
	"sync"
	"time"
)

// Registry maps deviceId -> owning node, with TTLs so a crashed node's
// entries expire. Backed by a map here; KeyValue/Redis in production.
type Registry interface {
	Register(deviceID, node string, ttl time.Duration)
	Touch(deviceID, node string, ttl time.Duration)
	Unregister(deviceID, node string)
	Lookup(deviceID string) (node string, ok bool)
}

type MemRegistry struct {
	mu sync.RWMutex
	m  map[string]entry
}
type entry struct {
	node    string
	expires time.Time
}

func NewMemRegistry() *MemRegistry { return &MemRegistry{m: map[string]entry{}} }

func (r *MemRegistry) Register(d, n string, ttl time.Duration) {
	r.mu.Lock(); defer r.mu.Unlock()
	r.m[d] = entry{node: n, expires: time.Now().Add(ttl)}
}
func (r *MemRegistry) Touch(d, n string, ttl time.Duration) { r.Register(d, n, ttl) }
func (r *MemRegistry) Unregister(d, n string) {
	r.mu.Lock(); defer r.mu.Unlock()
	if e, ok := r.m[d]; ok && e.node == n {
		delete(r.m, d)
	}
}
func (r *MemRegistry) Lookup(d string) (string, bool) {
	r.mu.RLock(); defer r.mu.RUnlock()
	e, ok := r.m[d]
	if !ok || time.Now().After(e.expires) {
		return "", false // expired entries are misses (crashed-node cleanup)
	}
	return e.node, true
}

Tasks

  1. Wire Node.Handle behind an http.HandlerFunc that accepts the upgrade (websocket.Accept) and reads a deviceId query param.
  2. Connect a few clients (a k6 script or a tiny Go client). Confirm the registry shows each deviceId → nodeID and that entries expire when a client disconnects.
  3. Kill a client uncleanly (SIGKILL the client process); confirm ping/pong detects the dead peer within the keepalive interval and the registry entry is removed/expires.

Acceptance

  • Multiple devices stay connected; registry reflects ownership with TTLs.
  • A dead peer is detected via failed ping and cleaned up.
  • A slow consumer (full out queue) does not block delivery to other devices (you can prove this in step 02).

Discussion prompts

  • Why is there exactly one writer goroutine per connection (the writePump)? (Concurrent writes to a WebSocket corrupt framing.)
  • Why must registry entries have a TTL even though you Unregister on clean close? (Crashes/network partitions never run your defer.)
  • At 200k connections/node, how big can out be? Do the memory math for the buffer size you chose.