puma-plus-core 0.1.0

data/go/internal/frontend/rtbridge.go

@@ -0,0 +1,410 @@
+package frontend
+
+import (
+	"bufio"
+	"context"
+	"errors"
+	"fmt"
+	"log"
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/evanphx/puma-plus/internal/metrics"
+	"github.com/evanphx/puma-plus/internal/queue"
+	"github.com/evanphx/puma-plus/internal/realtime"
+	"github.com/evanphx/puma-plus/internal/websocket"
+	"github.com/evanphx/puma-plus/internal/webtransport"
+	"github.com/evanphx/puma-plus/internal/wire"
+
+	wt "github.com/quic-go/webtransport-go"
+)
+
+// RTBridge connects the realtime hub to the request queue, for both WebSocket
+// and WebTransport.
+//
+// The key decision, and the reason both transports share this code: an event is
+// dispatched as an ordinary REQUEST frame with method WEBSOCKET. It travels the
+// same queue, gets the same queue-time measurement, appears in the same
+// histograms, and therefore drives the same autoscaling as an HTTP request --
+// with no parallel code path and no second set of metrics.
+//
+// That is what puma structurally cannot do: there a persistent connection's load
+// is invisible to every scaling signal, because its thread is parked rather than
+// queued.
+type RTBridge struct {
+	hub    *realtime.Hub
+	reg    *queue.Registry
+	col    *metrics.Collector
+	logger *log.Logger
+	opts   Options
+
+	wtServer *wt.Server
+}
+
+// NewRTBridge wires a hub to a registry.
+func NewRTBridge(hub *realtime.Hub, reg *queue.Registry, col *metrics.Collector,
+	opts Options, logger *log.Logger) *RTBridge {
+	if logger == nil {
+		logger = log.Default()
+	}
+	return &RTBridge{hub: hub, reg: reg, col: col, opts: opts, logger: logger}
+}
+
+// AttachWebTransport enables WebTransport upgrades over the HTTP/3 listener.
+func (b *RTBridge) AttachWebTransport(s *wt.Server) { b.wtServer = s }
+
+// Hub exposes the connection registry.
+func (b *RTBridge) Hub() *realtime.Hub { return b.hub }
+
+// Dispatch delivers one event to Ruby and applies any reply.
+//
+// It blocks while no worker is free, which is the whole point: the wait is
+// measured as queue time and shows up in the scaling signal.
+func (b *RTBridge) Dispatch(ev realtime.Event) error {
+	runnable := time.Now()
+
+	ctx := context.Background()
+	var cancel context.CancelFunc
+	if b.opts.AcquireTimeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, b.opts.AcquireTimeout)
+		defer cancel()
+	}
+
+	var (
+		conn      *queue.Conn
+		queueWait time.Duration
+		err       error
+	)
+
+	// A sticky connection is pinned by its FIRST event: until then there is no
+	// worker to prefer, so take whoever is free and remember them.
+	if ev.Conn.Mode == realtime.ModeSticky && ev.Conn.Pinned {
+		conn, queueWait, err = b.reg.AcquirePinned(ctx, ev.Conn.WorkerID)
+		if errors.Is(err, queue.ErrWorkerGone) {
+			// The process holding this connection's state is gone. Closing is
+			// the honest answer; re-pinning would hand a live conversation to a
+			// worker that has never heard of it.
+			b.logger.Printf("realtime %s: pinned worker %d is gone, closing",
+				ev.Conn.ID, ev.Conn.WorkerID)
+			ev.Conn.CloseWith(websocket.CloseServiceRestarting, "worker retired")
+			return err
+		}
+	} else {
+		conn, queueWait, err = b.reg.Acquire(ctx)
+	}
+	if err != nil {
+		if b.col != nil {
+			b.col.ObserveShed()
+		}
+		return err
+	}
+
+	if ev.Conn.Mode == realtime.ModeSticky && !ev.Conn.Pinned {
+		ev.Conn.WorkerID = conn.WID
+		ev.Conn.Pinned = true
+	}
+
+	healthy, service := b.relay(conn, ev, queueWait, runnable)
+	b.reg.Release(conn, healthy)
+
+	if b.col != nil {
+		b.col.ObserveRequest(queueWait, service, queueWait+service, 0)
+		if !healthy {
+			b.col.ObserveError()
+		}
+	}
+	if !healthy {
+		return fmt.Errorf("realtime %s: worker failed handling %s", ev.Conn.ID, ev.Kind)
+	}
+	return nil
+}
+
+// relay writes the event as a REQUEST frame and consumes the reply.
+func (b *RTBridge) relay(conn *queue.Conn, ev realtime.Event,
+	queueWait time.Duration, runnable time.Time) (bool, time.Duration) {
+
+	env := b.buildEnv(ev)
+
+	mode := wire.BodyNone
+	if len(ev.Payload) > 0 {
+		mode = wire.BodyInline
+	}
+	meta := wire.RequestMeta{
+		QueueNs:        uint64(queueWait.Nanoseconds()),
+		RunnableWallUs: uint64(runnable.UnixMicro()),
+		HTTPVersion:    wire.HTTP11,
+		BodyMode:       mode,
+	}
+
+	payload := wire.AppendRequestMeta(nil, meta)
+	payload = wire.AppendKV(payload, env)
+	if mode == wire.BodyInline {
+		payload = append(payload, ev.Payload...)
+	}
+
+	if err := conn.W.WriteFrame(wire.TypeRequest, payload); err != nil {
+		return false, 0
+	}
+
+	var reply []byte
+	var status int
+	var service time.Duration
+	for {
+		typ, frame, err := conn.R.ReadFrame()
+		if err != nil {
+			return false, service
+		}
+		switch typ {
+		case wire.TypeResponse:
+			if len(frame) >= 2 {
+				status = int(frame[0])<<8 | int(frame[1])
+			}
+		case wire.TypeRespChunk:
+			reply = append(reply, frame...)
+		case wire.TypeRespEnd:
+			if em, _, err := wire.DecodeRespEndMeta(frame); err == nil {
+				service = time.Duration(em.ServiceNs)
+			}
+			// A response body is shorthand for "send this back where it came
+			// from", which covers echo handlers without the outbound channel.
+			// It replies on the originating stream, so a WebTransport answer
+			// lands on the stream that asked.
+			if len(reply) > 0 && status != http.StatusNoContent && ev.Kind != realtime.EventClose {
+				if ev.StreamID != "" {
+					_ = ev.Conn.SendOnStream(ev.StreamID, reply, ev.Binary)
+				} else {
+					_ = ev.Conn.Send(reply, ev.Binary)
+				}
+			}
+			return true, service
+		case wire.TypeHijack:
+			b.logger.Printf("realtime %s: worker attempted HIJACK on an event", ev.Conn.ID)
+			return false, service
+		default:
+			return false, service
+		}
+	}
+}
+
+// buildEnv constructs the Rack env for an event.
+//
+// A normal-looking env with REQUEST_METHOD set to WEBSOCKET, so an existing Rack
+// app routes it by PATH_INFO exactly as it routes anything else and middleware
+// that inspects the env keeps working.
+func (b *RTBridge) buildEnv(ev realtime.Event) []wire.Pair {
+	c := ev.Conn
+	env := []wire.Pair{
+		{Key: []byte("REQUEST_METHOD"), Value: []byte("WEBSOCKET")},
+		{Key: []byte("PATH_INFO"), Value: []byte(c.Path)},
+		{Key: []byte("SCRIPT_NAME"), Value: nil},
+		{Key: []byte("QUERY_STRING"), Value: nil},
+		{Key: []byte("SERVER_PROTOCOL"), Value: []byte("HTTP/1.1")},
+		{Key: []byte("SERVER_SOFTWARE"), Value: serverSoftware},
+		{Key: []byte("REMOTE_ADDR"), Value: []byte(c.RemoteAddr())},
+		{Key: []byte("rack.url_scheme"), Value: []byte(b.opts.Scheme)},
+
+		{Key: []byte("puma_plus.ws.event"), Value: []byte(ev.Kind)},
+		{Key: []byte("puma_plus.ws.conn_id"), Value: []byte(c.ID)},
+		{Key: []byte("puma_plus.ws.mode"), Value: []byte(c.Mode)},
+		// Which protocol the client used. The application usually should not
+		// care -- that is the point of a shared hub -- but it must be able to
+		// find out before reaching for a datagram.
+		{Key: []byte("puma_plus.ws.transport"), Value: []byte(c.Kind())},
+	}
+
+	if ev.StreamID != "" {
+		env = append(env, wire.Pair{
+			Key: []byte("puma_plus.ws.stream_id"), Value: []byte(ev.StreamID),
+		})
+	}
+	if ev.Kind == realtime.EventMessage || ev.Kind == realtime.EventDatagram {
+		op := "text"
+		if ev.Binary {
+			op = "binary"
+		}
+		env = append(env, wire.Pair{Key: []byte("puma_plus.ws.opcode"), Value: []byte(op)})
+	}
+	if ev.Kind == realtime.EventClose {
+		env = append(env,
+			wire.Pair{Key: []byte("puma_plus.ws.close_code"),
+				Value: []byte(strconv.Itoa(ev.Code))},
+			wire.Pair{Key: []byte("puma_plus.ws.close_reason"),
+				Value: []byte(ev.Reason)},
+		)
+	}
+
+	for k, v := range c.Meta {
+		env = append(env, wire.Pair{
+			Key: []byte("puma_plus.ws.meta." + k), Value: []byte(v),
+		})
+	}
+	return env
+}
+
+// UpgradeSpec is what the Rack app asked for when it returned 101.
+type UpgradeSpec struct {
+	Mode      realtime.Mode
+	Topics    []string
+	Meta      map[string]string
+	Transport string // "websocket" or "webtransport"
+}
+
+// parseUpgradeSpec reads the puma-plus directives the app set on its response.
+//
+// Encoded as ordinary response headers because that is what a Rack app can
+// already return, and because it makes the upgrade a value the app produces
+// rather than a side effect it performs -- middleware can see it, log it, or
+// override it.
+func parseUpgradeSpec(h http.Header) (UpgradeSpec, bool) {
+	mode := h.Get("Puma-Plus-Websocket")
+	if mode == "" {
+		return UpgradeSpec{}, false
+	}
+	spec := UpgradeSpec{
+		Mode: realtime.Mode(mode), Meta: map[string]string{},
+		Transport: h.Get("Puma-Plus-Transport"),
+	}
+	if spec.Mode != realtime.ModeSticky {
+		spec.Mode = realtime.ModeFlexible
+	}
+	for _, t := range h.Values("Puma-Plus-Websocket-Topic") {
+		if t != "" {
+			spec.Topics = append(spec.Topics, t)
+		}
+	}
+	for _, m := range h.Values("Puma-Plus-Websocket-Meta") {
+		if k, v, ok := strings.Cut(m, "="); ok {
+			spec.Meta[k] = v
+		}
+	}
+	h.Del("Puma-Plus-Websocket")
+	h.Del("Puma-Plus-Websocket-Topic")
+	h.Del("Puma-Plus-Websocket-Meta")
+	h.Del("Puma-Plus-Transport")
+	return spec, true
+}
+
+// Upgrade completes a handshake and hands the connection to the hub.
+//
+// Which protocol depends on the request: an HTTP/3 Extended CONNECT with
+// :protocol=webtransport becomes a WebTransport session, anything else is a
+// WebSocket upgrade over HTTP/1.1.
+func (b *RTBridge) Upgrade(w http.ResponseWriter, r *http.Request, spec UpgradeSpec) error {
+	if isWebTransportRequest(r) || spec.Transport == "webtransport" {
+		return b.upgradeWebTransport(w, r, spec)
+	}
+	return b.upgradeWebSocket(w, r, spec)
+}
+
+// isWebTransportRequest detects an Extended CONNECT carrying :protocol.
+//
+// The exact string has moved with the drafts -- quic-go currently reports
+// "webtransport-h3" where earlier revisions used plain "webtransport" -- so this
+// matches on the prefix rather than an exact value. Getting it wrong is silent
+// and confusing: the request falls through to the ordinary HTTP path, whose body
+// read then blocks forever on a CONNECT stream that never ends.
+func isWebTransportRequest(r *http.Request) bool {
+	return r.Method == http.MethodConnect &&
+		strings.HasPrefix(r.Proto, "webtransport")
+}
+
+func (b *RTBridge) upgradeWebTransport(w http.ResponseWriter, r *http.Request,
+	spec UpgradeSpec) error {
+	if b.wtServer == nil {
+		return errors.New("webtransport: not enabled (requires an HTTP/3 listener)")
+	}
+	sess, err := b.wtServer.Upgrade(w, r)
+	if err != nil {
+		return fmt.Errorf("webtransport upgrade: %w", err)
+	}
+
+	tr := webtransport.NewTransport(sess)
+	c := realtime.NewConn(b.hub, tr, spec.Mode, r.URL.Path, spec.Meta)
+	b.hub.Add(c)
+	for _, t := range spec.Topics {
+		b.hub.Subscribe(c.ID, t)
+	}
+
+	b.logger.Printf("webtransport %s open mode=%s path=%s (%d live)",
+		c.ID, c.Mode, c.Path, b.hub.Count())
+
+	go webtransport.Serve(c, b)
+	return nil
+}
+
+func (b *RTBridge) upgradeWebSocket(w http.ResponseWriter, r *http.Request,
+	spec UpgradeSpec) error {
+	key := r.Header.Get("Sec-WebSocket-Key")
+	if key == "" {
+		return errors.New("websocket: missing Sec-WebSocket-Key")
+	}
+	if v := r.Header.Get("Sec-WebSocket-Version"); v != "13" {
+		return fmt.Errorf("websocket: unsupported version %q", v)
+	}
+
+	hj, ok := w.(http.Hijacker)
+	if !ok {
+		return fmt.Errorf("%w: WebSocket upgrade requires HTTP/1.1 "+
+			"(over h2/h3 use WebTransport instead)", ErrHijackUnsupported)
+	}
+	netConn, buf, err := hj.Hijack()
+	if err != nil {
+		return err
+	}
+
+	resp := "HTTP/1.1 101 Switching Protocols\r\n" +
+		"Upgrade: websocket\r\n" +
+		"Connection: Upgrade\r\n" +
+		"Sec-WebSocket-Accept: " + websocket.AcceptKey(key) + "\r\n\r\n"
+	if _, err := netConn.Write([]byte(resp)); err != nil {
+		netConn.Close()
+		return err
+	}
+
+	tr := websocket.NewTransport(netConn)
+	c := realtime.NewConn(b.hub, tr, spec.Mode, r.URL.Path, spec.Meta)
+	b.hub.Add(c)
+	for _, t := range spec.Topics {
+		b.hub.Subscribe(c.ID, t)
+	}
+
+	b.logger.Printf("websocket %s open mode=%s path=%s (%d live)",
+		c.ID, c.Mode, c.Path, b.hub.Count())
+
+	go websocket.Serve(c, b, websocket.DefaultServeOptions(),
+		bufio.NewReader(&joinedReader{buffered: buf.Reader, conn: netConn}))
+	return nil
+}
+
+// joinedReader replays bytes net/http had already buffered before the raw
+// connection. Dropping them would corrupt the first frame.
+type joinedReader struct {
+	buffered *bufio.Reader
+	conn     interface{ Read([]byte) (int, error) }
+}
+
+func (j *joinedReader) Read(p []byte) (int, error) {
+	if j.buffered != nil {
+		if n := j.buffered.Buffered(); n > 0 {
+			return j.buffered.Read(p)
+		}
+		j.buffered = nil
+	}
+	return j.conn.Read(p)
+}
+
+// drainToRespEnd consumes the rest of a response so the worker connection is
+// left at a frame boundary and can be reused.
+func drainToRespEnd(conn *queue.Conn) bool {
+	for {
+		typ, _, err := conn.R.ReadFrame()
+		if err != nil {
+			return false
+		}
+		if typ == wire.TypeRespEnd {
+			return true
+		}
+	}
+}

data/go/internal/frontend/shepherd.go

@@ -0,0 +1,188 @@
+package frontend
+
+import (
+	"errors"
+	"log"
+	"net"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/evanphx/puma-plus/internal/wire"
+)
+
+// ShepherdConn is the Go side of the control connection to the Ruby shepherd.
+//
+// It carries commands down (SET_SLOTS, QUIESCE, SHUTDOWN) and heartbeats up.
+// Liveness is heartbeat recency at 1 Hz, which is an order of magnitude finer
+// than puma's 5s worker pipe ping measured against a 60s worker_timeout
+// (puma/lib/puma/cluster/worker.rb:125, worker_handle.rb:95).
+type ShepherdConn struct {
+	mu     sync.RWMutex
+	conn   net.Conn
+	w      *wire.Writer
+	logger *log.Logger
+
+	lastSeen time.Time
+	workers  int
+	pids     []int
+	seq      uint64
+}
+
+// ErrNoShepherd means no shepherd is attached.
+var ErrNoShepherd = errors.New("frontend: no shepherd connected")
+
+// NewShepherdConn wraps an accepted control connection and starts reading
+// heartbeats.
+func NewShepherdConn(c net.Conn, logger *log.Logger) *ShepherdConn {
+	if logger == nil {
+		logger = log.Default()
+	}
+	s := &ShepherdConn{
+		conn: c, w: wire.NewWriter(c), logger: logger, lastSeen: time.Now(),
+	}
+	go s.readLoop()
+	return s
+}
+
+func (s *ShepherdConn) readLoop() {
+	r := wire.NewReader(s.conn)
+	for {
+		typ, payload, err := r.ReadFrame()
+		if err != nil {
+			s.mu.Lock()
+			s.conn = nil
+			s.mu.Unlock()
+			s.logger.Printf("shepherd disconnected: %v", err)
+			return
+		}
+
+		switch typ {
+		case wire.TypeWorkerStatus:
+			pairs, _, err := wire.DecodeKV(payload)
+			if err != nil {
+				continue
+			}
+			s.mu.Lock()
+			s.lastSeen = time.Now()
+			s.seq++
+			for _, p := range pairs {
+				switch string(p.Key) {
+				case "workers":
+					if n, err := strconv.Atoi(string(p.Value)); err == nil {
+						s.workers = n
+					}
+				case "worker_pids":
+					s.pids = parsePIDs(string(p.Value))
+				}
+			}
+			s.mu.Unlock()
+		case wire.TypePong:
+			s.mu.Lock()
+			s.lastSeen = time.Now()
+			s.mu.Unlock()
+		}
+	}
+}
+
+func parsePIDs(s string) []int {
+	if s == "" {
+		return nil
+	}
+	parts := strings.Split(s, ",")
+	out := make([]int, 0, len(parts))
+	for _, p := range parts {
+		if n, err := strconv.Atoi(strings.TrimSpace(p)); err == nil {
+			out = append(out, n)
+		}
+	}
+	return out
+}
+
+// Connected reports whether a shepherd is attached and recently heard from.
+//
+// Two missed heartbeats means suspect. At 1 Hz that is a 2s detection window,
+// against puma's 60s worker_timeout.
+func (s *ShepherdConn) Connected() bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.conn != nil && time.Since(s.lastSeen) < 3*time.Second
+}
+
+// WorkerPIDs reports the shepherd's current worker processes, so the controller
+// can read their RSS from /proc rather than trusting a self-report.
+func (s *ShepherdConn) WorkerPIDs() []int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	out := make([]int, len(s.pids))
+	copy(out, s.pids)
+	return out
+}
+
+// WorkerCount is how many worker processes exist, which for a forking shepherd
+// is exactly how many pids it reports.
+func (s *ShepherdConn) WorkerCount() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return len(s.pids)
+}
+
+// Workers is the shepherd's reported worker count.
+func (s *ShepherdConn) Workers() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.workers
+}
+
+// SetWorkers asks the shepherd for n worker processes.
+func (s *ShepherdConn) SetWorkers(n int) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.conn == nil {
+		return ErrNoShepherd
+	}
+	payload := []byte{byte(n >> 24), byte(n >> 16), byte(n >> 8), byte(n)}
+	return s.w.WriteFrame(wire.TypeSetSlots, payload)
+}
+
+// Quiesce tells the shepherd to stop respawning dead workers.
+func (s *ShepherdConn) Quiesce() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.conn == nil {
+		return ErrNoShepherd
+	}
+	return s.w.WriteFrame(wire.TypeQuiesce, nil)
+}
+
+// Shutdown asks the shepherd to stop, allowing grace for in-flight work.
+func (s *ShepherdConn) Shutdown(grace time.Duration) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.conn == nil {
+		return ErrNoShepherd
+	}
+	ms := uint32(grace.Milliseconds())
+	return s.w.WriteFrame(wire.TypeShutdown,
+		[]byte{byte(ms >> 24), byte(ms >> 16), byte(ms >> 8), byte(ms)})
+}
+
+// Close tears down the control connection.
+func (s *ShepherdConn) Close() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.conn == nil {
+		return nil
+	}
+	err := s.conn.Close()
+	s.conn = nil
+	return err
+}
+
+// NoShepherd is a Shepherd that is never connected, for running without one.
+type NoShepherd struct{}
+
+func (NoShepherd) SetWorkers(int) error { return ErrNoShepherd }
+func (NoShepherd) WorkerPIDs() []int    { return nil }
+func (NoShepherd) Connected() bool      { return false }

data/go/internal/frontend/socket.go

@@ -0,0 +1,37 @@
+package frontend
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"os"
+)
+
+// removeStaleSocket clears a leftover unix socket at path so bind can succeed.
+//
+// It refuses to unlink a socket that something is still listening on: dialing it
+// successfully means a live server owns that path, and removing it would silently
+// steal traffic from a running process. A socket that exists but refuses
+// connections is genuinely stale (its owner crashed) and is safe to remove.
+func removeStaleSocket(path string) error {
+	fi, err := os.Stat(path)
+	if errors.Is(err, os.ErrNotExist) {
+		return nil
+	}
+	if err != nil {
+		return fmt.Errorf("stat worker socket %s: %w", path, err)
+	}
+	if fi.Mode()&os.ModeSocket == 0 {
+		return fmt.Errorf("worker socket path %s exists and is not a socket", path)
+	}
+
+	if c, err := net.Dial("unix", path); err == nil {
+		c.Close()
+		return fmt.Errorf("worker socket %s is already in use by a live process", path)
+	}
+
+	if err := os.Remove(path); err != nil {
+		return fmt.Errorf("remove stale worker socket %s: %w", path, err)
+	}
+	return nil
+}

data/go/internal/frontend/socketpair_test.go

@@ -0,0 +1,34 @@
+package frontend
+
+import (
+	"net"
+	"os"
+
+	"golang.org/x/sys/unix"
+)
+
+// makeSocketPair returns two connected unix sockets.
+//
+// net.Pipe is not a substitute here: it is fully synchronous and unbuffered, so
+// the body pump would deadlock against a worker that has not started reading
+// yet. Real request streaming depends on kernel socket buffering to make
+// progress, and the tests need to exercise that.
+func makeSocketPair() ([2]net.Conn, error) {
+	var out [2]net.Conn
+
+	fds, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM, 0)
+	if err != nil {
+		return out, err
+	}
+
+	for i, fd := range fds {
+		f := os.NewFile(uintptr(fd), "socketpair")
+		c, err := net.FileConn(f)
+		f.Close()
+		if err != nil {
+			return out, err
+		}
+		out[i] = c
+	}
+	return out, nil
+}