puma-plus-core 0.1.0

data/go/internal/frontend/handler.go

@@ -0,0 +1,543 @@
+package frontend
+
+import (
+	"context"
+	"io"
+	"log"
+	"net/http"
+	"strconv"
+	"time"
+
+	"github.com/evanphx/puma-plus/internal/metrics"
+	"github.com/evanphx/puma-plus/internal/queue"
+	"github.com/evanphx/puma-plus/internal/wire"
+)
+
+// Options configures the frontend handler.
+type Options struct {
+	// BodyBufferMax is the largest request body Go buffers before marking the
+	// request runnable. Buffering first is what keeps a slow uploader's time in
+	// *body wait* rather than leaking into queue time or service time.
+	//
+	// Past this limit we dispatch anyway and stream the remainder, because
+	// waiting for a multi-gigabyte upload before starting work would be absurd.
+	// In that case the worker blocks on rack.input and reports body_read_ns,
+	// which the controller subtracts from service time -- so client slowness
+	// still never gets mistaken for the application being slow.
+	BodyBufferMax int64
+
+	// AcquireTimeout bounds how long a request may wait for capacity before it
+	// is shed. Shedding converts an unbounded wait into a fast 503.
+	AcquireTimeout time.Duration
+
+	// DebugHeaders adds X-PumaPlus-* timing headers to every response. Useful
+	// during development and for the benchmark harness's cross-check; off in any
+	// serious measurement run because it perturbs response size.
+	DebugHeaders bool
+
+	Scheme      string
+	DefaultPort string
+	Logger      *log.Logger
+
+	// Collector receives per-request timings. Optional; nil disables metrics.
+	Collector *metrics.Collector
+}
+
+// Handler serves HTTP requests by relaying them to Ruby worker connections.
+type Handler struct {
+	reg  *queue.Registry
+	opts Options
+	col  *metrics.Collector
+	ws   *RTBridge
+}
+
+// AttachRealtime enables Go-managed WebSocket upgrades.
+func (h *Handler) AttachRealtime(b *RTBridge) { h.ws = b }
+
+// NewHandler wires a registry to an http.Handler.
+func NewHandler(reg *queue.Registry, opts Options) *Handler {
+	if opts.BodyBufferMax == 0 {
+		opts.BodyBufferMax = 1 << 20
+	}
+	if opts.AcquireTimeout == 0 {
+		opts.AcquireTimeout = 30 * time.Second
+	}
+	if opts.Scheme == "" {
+		opts.Scheme = "http"
+	}
+	if opts.DefaultPort == "" {
+		opts.DefaultPort = "80"
+	}
+	if opts.Logger == nil {
+		opts.Logger = log.Default()
+	}
+	return &Handler{reg: reg, opts: opts, col: opts.Collector}
+}
+
+// observe records a completed request, if metrics are enabled.
+func (h *Handler) observe(queueWait, service, ttfb time.Duration, bytesOut int) {
+	if h.col != nil {
+		h.col.ObserveRequest(queueWait, service, ttfb, bytesOut)
+	}
+}
+
+func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// A WebTransport CONNECT must never have its body read.
+	//
+	// The CONNECT stream IS the session: it stays open for the connection's
+	// whole life and never sends EOF, so io.ReadAll below would block forever
+	// and the request would never reach Ruby. It arrives with an unknown
+	// content length, which is exactly what the buffering path treats as "read
+	// until EOF".
+	if isWebTransportRequest(r) {
+		h.serveNoBody(w, r)
+		return
+	}
+
+	// --- T2: headers parsed. Now buffer the body so that client slowness lands
+	// in body wait rather than polluting queue time or service time.
+	bodyStart := time.Now()
+	body, moreBody, err := h.readBody(r)
+	if err != nil {
+		// The client vanished or stalled mid-upload. Nothing to serve.
+		http.Error(w, "error reading request body", http.StatusBadRequest)
+		return
+	}
+	bodyWait := time.Since(bodyStart)
+
+	// --- T2b: runnable. The request can no longer be delayed by anything the
+	// client does. Everything from here to dispatch is our fault, and only that
+	// interval is queue time.
+	runnable := time.Now()
+	runnableWallUs := uint64(runnable.UnixMicro())
+
+	ctx := r.Context()
+	if h.opts.AcquireTimeout > 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, h.opts.AcquireTimeout)
+		defer cancel()
+	}
+
+	// Acquire a worker and dispatch. A worker process that exits between our
+	// acquiring its connection and our writing to it produces a broken pipe --
+	// unavoidable, because there is no way to reserve a process against dying.
+	//
+	// But a REQUEST that failed to *write* was never delivered: the app did not
+	// see it, nothing was executed, no side effect occurred. So it is always
+	// safe to retry it on a different connection, and that is strictly better
+	// than returning 502 to a client whose request we simply failed to hand
+	// over. Only write failures are retried; once a byte of the request is in
+	// the worker's hands we must not send it twice.
+	var (
+		conn      *queue.Conn
+		queueWait time.Duration
+		res       relayResult
+	)
+	const maxDispatchAttempts = 3
+
+	for attempt := 0; attempt < maxDispatchAttempts; attempt++ {
+		var acquireWait time.Duration
+		conn, acquireWait, err = h.reg.Acquire(ctx)
+		if err != nil {
+			// Either the client hung up while queued, or we hit the acquire
+			// timeout. A queued-then-abandoned request costs us nothing further;
+			// a timeout is a shed, and the client gets told to come back.
+			if r.Context().Err() != nil {
+				return
+			}
+			if h.col != nil {
+				h.col.ObserveShed()
+			}
+			w.Header().Set("Retry-After", "1")
+			http.Error(w, "server busy", http.StatusServiceUnavailable)
+			return
+		}
+		// Queue time accumulates across attempts: from the client's point of
+		// view it waited for all of them.
+		queueWait += acquireWait
+		// --- T3: dispatched. queueWait is Q, the product metric.
+
+		res = h.relay(w, r, conn, body, moreBody, wire.RequestMeta{
+			QueueNs:        uint64(queueWait.Nanoseconds()),
+			RunnableWallUs: runnableWallUs,
+			BodyWaitNs:     uint64(bodyWait.Nanoseconds()),
+			HTTPVersion:    protoVersion(r),
+			BodyMode:       bodyMode(body, moreBody),
+		})
+		if res.hijacked {
+			// Already retired by spliceHijack; releasing would hand a tunnelled
+			// connection back to the pool as if it were idle capacity.
+			return
+		}
+		h.reg.Release(conn, res.healthy)
+		if res.upgraded {
+			// Worker returned to the pool above; the client socket is the hub's.
+			return
+		}
+
+		if !res.undelivered {
+			break
+		}
+		// The worker died before receiving the request. Its connection has been
+		// removed by Release; try another one.
+		h.opts.Logger.Printf("worker pid=%d died before receiving the request; retrying (attempt %d)",
+			conn.Pid, attempt+1)
+		if h.col != nil {
+			h.col.ObserveError()
+		}
+	}
+
+	if res.undelivered {
+		// Every attempt hit a dying worker. Something is badly wrong with the
+		// worker pool, and a 502 is now the honest answer.
+		http.Error(w, "no healthy worker", http.StatusBadGateway)
+		return
+	}
+	if !res.healthy && h.col != nil {
+		h.col.ObserveError()
+	}
+	// TTFB is measured from runnable, not from accept: the client's own slowness
+	// getting the request to us is not something this server can affect, and
+	// folding it in would make the number unusable as a server-side signal.
+	h.observe(queueWait, res.service, res.ttfb, res.bytesOut)
+}
+
+// serveNoBody dispatches a request whose body must not be consumed, then hands
+// the result to the upgrade path. Used for WebTransport's CONNECT.
+func (h *Handler) serveNoBody(w http.ResponseWriter, r *http.Request) {
+	runnable := time.Now()
+
+	ctx := r.Context()
+	if h.opts.AcquireTimeout > 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, h.opts.AcquireTimeout)
+		defer cancel()
+	}
+
+	conn, queueWait, err := h.reg.Acquire(ctx)
+	if err != nil {
+		if r.Context().Err() != nil {
+			return
+		}
+		if h.col != nil {
+			h.col.ObserveShed()
+		}
+		w.Header().Set("Retry-After", "1")
+		http.Error(w, "server busy", http.StatusServiceUnavailable)
+		return
+	}
+
+	res := h.relay(w, r, conn, nil, false, wire.RequestMeta{
+		QueueNs:        uint64(queueWait.Nanoseconds()),
+		RunnableWallUs: uint64(runnable.UnixMicro()),
+		HTTPVersion:    protoVersion(r),
+		BodyMode:       wire.BodyNone,
+	})
+	if res.hijacked || res.upgraded {
+		if !res.hijacked {
+			h.reg.Release(conn, res.healthy)
+		}
+		return
+	}
+	h.reg.Release(conn, res.healthy)
+	h.observe(queueWait, res.service, res.ttfb, res.bytesOut)
+}
+
+// relayResult carries what the relay learned back to the metrics path.
+type relayResult struct {
+	// healthy reports whether the worker connection is reusable.
+	healthy bool
+	// service is the Rack app's own timing, reported by Ruby in RESP_END. Taken
+	// from the worker rather than measured in Go so it excludes transport and
+	// scheduling, which is what makes it a usable service-time input for
+	// Little's Law.
+	service  time.Duration
+	ttfb     time.Duration
+	bytesOut int
+	// undelivered means the REQUEST frame never reached the worker, so the app
+	// never ran and the request may safely be retried on another connection.
+	// Distinct from a failure *after* delivery, which must never be retried.
+	undelivered bool
+	// hijacked means the connection became a raw tunnel and has already been
+	// retired from the pool. It must not be released, and no response may be
+	// written -- the application owns the socket now.
+	hijacked bool
+	// upgraded means the CLIENT socket was handed to the WebSocket hub, but the
+	// worker connection is untouched and must be released normally.
+	//
+	// Distinct from hijacked, and the distinction is load-bearing: conflating
+	// them leaked a worker thread on every WebSocket upgrade, because the conn
+	// was neither released nor retired. Under Go-managed WebSockets the worker
+	// finished an ordinary request and is immediately free -- that is the entire
+	// point of the feature.
+	upgraded bool
+}
+
+// readBody buffers the request body up to BodyBufferMax and reports whether
+// more remains.
+//
+// Buffering before dispatch is what keeps the queue clock honest: a slow
+// uploader's time lands in body wait, not in queue time and not in service time.
+// But we cannot wait for an arbitrarily large upload before starting work, so
+// past the buffer limit we dispatch and stream the remainder. The worker then
+// blocks on rack.input and reports body_read_ns, which the controller subtracts
+// from service time -- so even in streaming mode client slowness never gets
+// mistaken for the application being slow.
+func (h *Handler) readBody(r *http.Request) (buffered []byte, more bool, err error) {
+	if r.Body == nil || r.ContentLength == 0 {
+		return nil, false, nil
+	}
+	// Read one byte past the limit so "exactly at the limit" and "there is more"
+	// are distinguishable without a second read.
+	buf, err := io.ReadAll(io.LimitReader(r.Body, h.opts.BodyBufferMax+1))
+	if err != nil {
+		return nil, false, err
+	}
+	// Return everything read, including the probe byte. Returning only the first
+	// BodyBufferMax bytes would silently drop that byte: it has already been
+	// consumed from r.Body, so the pump would never see it and every streamed
+	// body would arrive exactly one byte short.
+	return buf, int64(len(buf)) > h.opts.BodyBufferMax, nil
+}
+
+// relay writes the REQUEST frame and pumps response frames back to the client.
+// It returns whether the connection is still healthy and reusable; a protocol or
+// IO error means the worker is not safe to hand another request.
+func (h *Handler) relay(w http.ResponseWriter, r *http.Request, conn *queue.Conn,
+	body []byte, moreBody bool, meta wire.RequestMeta) relayResult {
+
+	var res relayResult
+	dispatched := time.Now()
+	flusher, _ := w.(http.Flusher)
+
+	payload := wire.AppendRequestMeta(nil, meta)
+	payload = wire.AppendKV(payload, BuildEnv(r, h.opts.Scheme, h.opts.DefaultPort, nil))
+	if meta.BodyMode == wire.BodyInline {
+		payload = append(payload, body...)
+	}
+
+	if err := conn.W.WriteFrame(wire.TypeRequest, payload); err != nil {
+		// Nothing was delivered. Do not write a response here -- the caller will
+		// retry on another connection, and writing now would commit us to an
+		// error the client need not see.
+		res.undelivered = true
+		return res
+	}
+
+	// A streamed body is pumped concurrently with the worker running the app,
+	// because the worker may block reading rack.input while we are still
+	// sending. A unix socket is full duplex, so this and the response loop below
+	// coexist. conn.W is not safe for concurrent use, but it does not need to
+	// be: the pump is the only writer for the rest of this request, since the
+	// loop below only reads.
+	var pump *bodyPump
+	if meta.BodyMode == wire.BodyStream {
+		pump = h.startBodyPump(conn, r.Body, body)
+	}
+
+	wroteHeader := false
+	for {
+		typ, frame, err := conn.R.ReadFrame()
+		if err != nil {
+			if !wroteHeader {
+				// The worker died holding our request. We cannot know whether it
+				// ran the app before dying, so this is NOT safe to retry -- a
+				// non-idempotent request could execute twice. Report it.
+				h.opts.Logger.Printf("worker pid=%d died while holding a request: %v", conn.Pid, err)
+				http.Error(w, "backend error", http.StatusBadGateway)
+			} else {
+				h.opts.Logger.Printf("worker pid=%d read: %v", conn.Pid, err)
+			}
+			return res
+		}
+
+		switch typ {
+		case wire.TypeResponse:
+			if wroteHeader {
+				h.opts.Logger.Printf("worker pid=%d sent a second RESPONSE", conn.Pid)
+				return res
+			}
+			status, err := h.writeHeaders(w, frame, meta)
+			if err != nil {
+				h.opts.Logger.Printf("worker pid=%d bad RESPONSE: %v", conn.Pid, err)
+				http.Error(w, "backend error", http.StatusBadGateway)
+				return res
+			}
+
+			// A 101 carrying puma-plus.websocket means the app authorised a
+			// WebSocket upgrade. Go takes the connection from here: the app has
+			// already done auth, routing and middleware as an ordinary request,
+			// and never touches a socket.
+			if status == http.StatusSwitchingProtocols && h.ws != nil {
+				if spec, ok := parseUpgradeSpec(w.Header()); ok {
+					// Drain RESP_END before hijacking, or the worker connection is
+					// left mid-frame and cannot be reused.
+					if !drainToRespEnd(conn) {
+						return res
+					}
+					res.healthy = true
+					if err := h.ws.Upgrade(w, r, spec); err != nil {
+						h.opts.Logger.Printf("websocket upgrade failed: %v", err)
+						return res
+					}
+					// The CLIENT socket belongs to the hub now, but this worker
+					// just finished an ordinary request and goes straight back to
+					// the pool. It does NOT stay bound to the connection.
+					res.upgraded = true
+					return res
+				}
+			}
+			w.WriteHeader(status)
+			wroteHeader = true
+			// Headers are on the wire: this is T5 from the client's point of
+			// view. Adding the queue wait gives TTFB measured from runnable.
+			res.ttfb = time.Duration(meta.QueueNs) + time.Since(dispatched)
+
+		case wire.TypeRespChunk:
+			if !wroteHeader {
+				h.opts.Logger.Printf("worker pid=%d sent RESP_CHUNK before RESPONSE", conn.Pid)
+				return res
+			}
+			if len(frame) == 0 {
+				continue
+			}
+			if n, err := w.Write(frame); err != nil {
+				// The client went away mid-body. The worker is fine; we just
+				// stop copying. Keep draining so the conn stays in protocol
+				// sync for the next request.
+				continue
+			} else {
+				res.bytesOut += n
+			}
+			// Flush every chunk. This is the puma-dev bug we are deliberately
+			// not repeating: its reverse proxy uses FlushInterval: 1s
+			// (puma-dev/dev/http.go:63), which batches server-sent events into
+			// one-second hiccups. A chunk is a unit the app chose to emit, so
+			// emitting it is the correct behavior; for a large buffered body the
+			// extra flushes are amortised across 64 KiB frames and cost nothing
+			// measurable.
+			if flusher != nil {
+				flusher.Flush()
+			}
+
+		case wire.TypeRespEnd:
+			if !wroteHeader {
+				h.opts.Logger.Printf("worker pid=%d sent RESP_END before RESPONSE", conn.Pid)
+				return res
+			}
+			// The worker is required to drain the request body before RESP_END,
+			// so by now the pump has finished. Confirm it, because a pump that
+			// died leaves the connection in an unknown protocol state.
+			if pump != nil {
+				if err := pump.wait(); err != nil {
+					h.opts.Logger.Printf("worker pid=%d body pump: %v", conn.Pid, err)
+					return res
+				}
+			}
+			// Ruby's own timing rides along on the response, so per-request
+			// service and CPU cost arrive with no extra frame and no extra
+			// syscall. body_read_ns is subtracted so a slow uploader on a
+			// streamed body cannot inflate service time and fool the
+			// controller into thinking the app got slower.
+			if em, _, err := wire.DecodeRespEndMeta(frame); err == nil {
+				if em.ServiceNs > em.BodyReadNs {
+					res.service = time.Duration(em.ServiceNs - em.BodyReadNs)
+				} else {
+					res.service = time.Duration(em.ServiceNs)
+				}
+				// Where the service time actually went, which is what decides
+				// threads versus processes.
+				if h.col != nil {
+					now := time.Now()
+					h.col.ObserveWork(res.service, time.Duration(em.CPUNs), now)
+					// The GVL counter is process-cumulative, and its denominator
+					// is this worker's thread count.
+					h.col.ObserveGVL(conn.WID, em.GVLWaitNs,
+						h.reg.ThreadsFor(conn.WID), now)
+				}
+			}
+			res.healthy = true
+			return res
+
+		case wire.TypeHijack:
+			if wroteHeader {
+				// The app already committed to a normal response; it cannot then
+				// take the socket.
+				h.opts.Logger.Printf("worker pid=%d sent HIJACK after RESPONSE", conn.Pid)
+				return res
+			}
+			if err := h.spliceHijack(w, conn); err != nil {
+				h.opts.Logger.Printf("worker pid=%d hijack failed: %v", conn.Pid, err)
+				http.Error(w, "hijack unsupported on this transport", http.StatusBadRequest)
+				return res
+			}
+			// The connection now belongs to the tunnel. Report healthy so the
+			// caller does not treat this as an error, but hijacked so it does not
+			// release a conn that no longer exists.
+			res.healthy = true
+			res.hijacked = true
+			return res
+
+		default:
+			h.opts.Logger.Printf("worker pid=%d unexpected frame %s", conn.Pid, wire.TypeName(typ))
+			return res
+		}
+	}
+}
+
+// writeHeaders decodes a RESPONSE frame onto the ResponseWriter and returns the
+// status code.
+func (h *Handler) writeHeaders(w http.ResponseWriter, frame []byte, meta wire.RequestMeta) (int, error) {
+	if len(frame) < 2 {
+		return 0, wire.ErrTruncated
+	}
+	status := int(frame[0])<<8 | int(frame[1])
+	pairs, _, err := wire.DecodeKV(frame[2:])
+	if err != nil {
+		return 0, err
+	}
+
+	hdr := w.Header()
+	for _, p := range pairs {
+		// Repeated keys are Rack 3 array-valued headers; Add preserves them as
+		// separate header lines rather than collapsing them, which matters for
+		// Set-Cookie above all.
+		hdr.Add(string(p.Key), string(p.Value))
+	}
+
+	if h.opts.DebugHeaders {
+		hdr.Set("X-PumaPlus-Queue-Ns", strconv.FormatUint(meta.QueueNs, 10))
+		hdr.Set("X-PumaPlus-Body-Wait-Ns", strconv.FormatUint(meta.BodyWaitNs, 10))
+	}
+	return status, nil
+}
+
+// bodyMode picks the carriage for a request body.
+//
+// INLINE is capped at wire.InlineBodyMax because it rides inside the REQUEST
+// frame; anything larger travels as BODY_CHUNK frames even when it is fully
+// buffered, so the frame stays a bounded size and Ruby's StringIO path stays
+// on the small-body fast track.
+func bodyMode(body []byte, more bool) uint8 {
+	switch {
+	case len(body) == 0 && !more:
+		return wire.BodyNone
+	case !more && len(body) <= wire.InlineBodyMax:
+		return wire.BodyInline
+	default:
+		return wire.BodyStream
+	}
+}
+
+func protoVersion(r *http.Request) uint8 {
+	switch {
+	case r.ProtoMajor == 3:
+		return wire.HTTP3
+	case r.ProtoMajor == 2:
+		return wire.HTTP2
+	case r.ProtoMajor == 1 && r.ProtoMinor == 0:
+		return wire.HTTP10
+	default:
+		return wire.HTTP11
+	}
+}