Route by lease, queue per host and model, record every request

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)

internal/proxy/proxy.go

@@ -1,31 +1,34 @@
-// Package proxy is the routing reverse proxy. It takes /{route}/v1/…, picks a host from the
-// route's ordered list using the health table, forwards the request, streams the answer back as it
-// arrives, and tells the health table when a host fails.
+// Package proxy is the routing reverse proxy. It takes /{route}/v1/…, picks a host for the
+// conversation from its ordered list using a lease table (or the health table alone), queues per
+// (host, model), forwards the request streaming the answer back as it arrives, tees the response to
+// read usage/timings, marks a host down when a forward fails, and records one accounting row per
+// request.
 package proxy
 
 import (
 	"bytes"
-	"context"
 	"encoding/json"
 	"errors"
 	"io"
+	"log/slog"
 	"net/http"
-	"net/http/httputil"
-	"net/url"
 	"strings"
 	"time"
 
-	"log/slog"
-
 	"git.wntrmute.dev/kyle/crossbar/internal/config"
+	"git.wntrmute.dev/kyle/crossbar/internal/fingerprint"
 	"git.wntrmute.dev/kyle/crossbar/internal/health"
+	"git.wntrmute.dev/kyle/crossbar/internal/lease"
+	"git.wntrmute.dev/kyle/crossbar/internal/limiter"
+	"git.wntrmute.dev/kyle/crossbar/internal/store"
 )
 
-// MaxBody is the largest request body we look at for a top-level "model" field.
-const MaxBody = 16 << 20
-
-// HostHeader is set on every proxied response: the name of the host that answered.
-const HostHeader = "X-Crossbar-Host"
+const (
+	MaxBody     = 16 << 20
+	HostHeader  = "X-Crossbar-Host"
+	LeaseHeader = "X-Crossbar-Lease" // "new" or "reused"
+	RouteHeader = "X-Crossbar-Route" // client may name the route here instead of the path
+)
 
 // errBodyTooLarge is returned when a request body exceeds MaxBody during the model peek.
 var errBodyTooLarge = errors.New("body too large")
@@ -36,19 +39,30 @@ type Health interface {
 	MarkDown(name, reason string)
 }
 
-// Handler forwards requests for a route to one of the route's healthy hosts.
+// Recorder is what the proxy needs to write an accounting row; *store.Store satisfies it.
+type Recorder interface {
+	RecordRequest(store.Request) error
+}
+
+// Handler forwards requests for a route to one of the route's healthy hosts, choosing by lease when
+// one is configured and by health alone otherwise.
 type Handler struct {
 	cfg    *config.Config
 	health Health
+	leases *lease.Table
+	lim    *limiter.Limiter
+	rec    Recorder
 	log    *slog.Logger
 }
 
-// New builds a Handler. A nil logger becomes slog.Default().
-func New(cfg *config.Config, h Health, log *slog.Logger) *Handler {
+// New builds a Handler. A nil logger becomes slog.Default(). With a nil lease table it behaves like
+// the v0 proxy: first healthy host, no queueing, no recording; nil limiter and recorder are likewise
+// no-ops.
+func New(cfg *config.Config, h Health, leases *lease.Table, lim *limiter.Limiter, rec Recorder, log *slog.Logger) *Handler {
 	if log == nil {
 		log = slog.Default()
 	}
-	return &Handler{cfg: cfg, health: h, log: log}
+	return &Handler{cfg: cfg, health: h, leases: leases, lim: lim, rec: rec, log: log}
 }
 
 // SplitRoute takes the first path segment as the route. "/a/v1/x" -> ("a", "/v1/x", true); "/a" and
@@ -108,22 +122,57 @@ func allowedPath(rest string) bool {
 	return strings.HasPrefix(rest, "/v1/") || rest == "/health" || rest == "/props"
 }
 
-// peekModel reads a non-GET/HEAD body up to MaxBody+1 bytes, restores it on the request, and
-// returns the top-level "model". A non-JSON body or one without a model gives "". A body larger
-// than MaxBody returns errBodyTooLarge.
-func peekModel(r *http.Request) (string, error) {
+// route resolves the route name and the upstream path (rest) from the request, honouring the
+// optional route header. code is non-zero when the request must be answered; msg is the JSON error
+// text for that code.
+func (p *Handler) route(r *http.Request) (route, rest string, code int, msg string) {
+	hdr := r.Header.Get(RouteHeader)
+	if hdr != "" {
+		rest := r.URL.Path
+		// A path that also carries a (different) route name is a client mistake: the header is the
+		// operator's intent, but the path disagrees.
+		if rname, _, ok := SplitRoute(rest); ok {
+			if _, known := p.cfg.Routes[rname]; known && rname != hdr {
+				return "", "", http.StatusBadRequest, "conflicting route"
+			}
+		}
+		if _, known := p.cfg.Routes[hdr]; !known {
+			return "", "", http.StatusNotFound, "unknown route"
+		}
+		if !allowedPath(rest) {
+			return "", "", http.StatusNotFound, "not found"
+		}
+		return hdr, rest, 0, ""
+	}
+	route, rest, ok := SplitRoute(r.URL.Path)
+	if !ok {
+		return "", "", http.StatusBadRequest, "missing route"
+	}
+	if _, known := p.cfg.Routes[route]; !known {
+		return "", "", http.StatusNotFound, "unknown route"
+	}
+	if !allowedPath(rest) {
+		return "", "", http.StatusNotFound, "not found"
+	}
+	return route, rest, 0, ""
+}
+
+// peekModel reads a non-GET/HEAD body up to MaxBody+1 bytes, restores it on the request, and returns
+// the top-level "model" and the body itself (for fingerprinting). A non-JSON body or one without a
+// model gives "". A body larger than MaxBody returns errBodyTooLarge.
+func peekModel(r *http.Request) (string, []byte, error) {
 	if r.Method == http.MethodGet || r.Method == http.MethodHead {
-		return "", nil
+		return "", nil, nil
 	}
 	if r.Body == nil || r.Body == http.NoBody {
-		return "", nil
+		return "", nil, nil
 	}
 	body, err := io.ReadAll(io.LimitReader(r.Body, MaxBody+1))
 	if err != nil {
-		return "", err
+		return "", nil, err
 	}
 	if len(body) > MaxBody {
-		return "", errBodyTooLarge
+		return "", nil, errBodyTooLarge
 	}
 	r.Body = io.NopCloser(bytes.NewReader(body))
 	r.ContentLength = int64(len(body))
@@ -132,42 +181,21 @@ func peekModel(r *http.Request) (string, error) {
 		Model string `json:"model"`
 	}
 	_ = json.Unmarshal(body, &req)
-	return req.Model, nil
-}
-
-// statusRecorder records the status written and forwards Flush so the reverse proxy can stream.
-type statusRecorder struct {
-	http.ResponseWriter
-	status int
-}
-
-func (r *statusRecorder) WriteHeader(code int) {
-	r.status = code
-	r.ResponseWriter.WriteHeader(code)
-}
-
-func (r *statusRecorder) Flush() {
-	if f, ok := r.ResponseWriter.(http.Flusher); ok {
-		f.Flush()
-	}
+	return req.Model, body, nil
 }
 
+// ServeHTTP routes, fingerprints, leases a host, queues per (host, model), forwards with streaming,
+// tees the response for usage/timings, and records one accounting row. Every error answer is JSON
+// {"error":"…"}.
 func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	route, rest, ok := SplitRoute(r.URL.Path)
-	if !ok {
-		p.writeError(w, http.StatusBadRequest, "missing route")
+	route, rest, code, msg := p.route(r)
+	if code != 0 {
+		p.writeError(w, code, msg)
 		return
 	}
-	routeCfg, ok := p.cfg.Routes[route]
-	if !ok {
-		p.writeError(w, http.StatusNotFound, "unknown route")
-		return
-	}
-	if !allowedPath(rest) {
-		p.writeError(w, http.StatusNotFound, "not found")
-		return
-	}
-	model, err := peekModel(r)
+	routeCfg := p.cfg.Routes[route]
+
+	model, body, err := peekModel(r)
 	if err != nil {
 		p.writeError(w, http.StatusRequestEntityTooLarge, "body too large")
 		return
@@ -175,68 +203,55 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if model == "" {
 		model = routeCfg.DefaultModel
 	}
-	name, ok := Choose(routeCfg.Hosts, model, p.health)
-	if !ok {
-		p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
+	fp := fingerprint.Of(body)
+	started := time.Now()
+
+	// v0 compatibility path: no lease table, no limiter, no recording.
+	if p.leases == nil {
+		name, ok := Choose(routeCfg.Hosts, model, p.health)
+		if !ok {
+			p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
+			return
+		}
+		p.forward(w, r, route, name, "", rest, fp, model, started, 0)
 		return
 	}
 
-	host, ok := p.cfg.Hosts[name]
-	if !ok {
-		p.writeError(w, http.StatusBadGateway, "upstream failed")
-		return
-	}
-	target, err := url.Parse(host.BaseURL)
+	// Lease. The route's ordered host list is the candidate set.
+	host, reused, err := p.leases.Acquire(lease.Key{Route: route, FP: fp, Model: model}, routeCfg.Hosts, time.Now())
 	if err != nil {
-		p.writeError(w, http.StatusBadGateway, "upstream failed")
+		switch {
+		case errors.Is(err, lease.ErrNoHost):
+			p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
+		case errors.Is(err, lease.ErrPinnedDown):
+			p.writeError(w, http.StatusServiceUnavailable, "pinned host down")
+		default:
+			p.writeError(w, http.StatusBadGateway, "upstream failed")
+		}
 		return
 	}
 
-	pr := newReverseProxy(p.health, name, target, rest)
-	rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
-	start := time.Now()
-	pr.ServeHTTP(rec, r)
-	p.log.Info("request",
-		"route", route,
-		"host", name,
-		"method", r.Method,
-		"path", rest,
-		"status", rec.status,
-		"ms", time.Since(start).Milliseconds(),
-	)
-}
-
-func (p *Handler) writeError(w http.ResponseWriter, status int, msg string) {
-	w.Header().Set("Content-Type", "application/json")
-	w.WriteHeader(status)
-	_ = json.NewEncoder(w).Encode(map[string]string{"error": msg})
-}
-
-// newReverseProxy forwards to a single host, rewriting the path to target.Path+rest and keeping the
-// original query string. It flushes after every write so long server-sent-event streams are not
-// buffered, and marks the host down on any transport error other than a client disconnect.
-func newReverseProxy(h Health, name string, target *url.URL, rest string) *httputil.ReverseProxy {
-	return &httputil.ReverseProxy{
-		Rewrite: func(pr *httputil.ProxyRequest) {
-			pr.SetURL(target)
-			pr.Out.URL.Path = target.Path + rest
-			pr.Out.URL.RawPath = ""
-			pr.Out.Host = target.Host
-			pr.SetXForwarded()
-		},
-		FlushInterval: -1,
-		ModifyResponse: func(resp *http.Response) error {
-			resp.Header.Set(HostHeader, name)
-			return nil
-		},
-		ErrorHandler: func(w http.ResponseWriter, req *http.Request, err error) {
-			if errors.Is(err, context.Canceled) {
-				return
-			}
-			h.MarkDown(name, err.Error())
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusBadGateway)
-			_ = json.NewEncoder(w).Encode(map[string]string{"error": "upstream failed", "host": name})
-		},
+	// Slot. A full queue is a 503; a context done while waiting means the client left.
+	release, waited, err := p.lim.Acquire(r.Context(), host, model)
+	if err != nil {
+		if errors.Is(err, limiter.ErrQueueFull) {
+			p.writeRecord(store.Request{
+				Route:   route,
+				FP:      fp,
+				Model:   model,
+				Host:    host,
+				Started: started,
+				TotalMs: time.Since(started).Milliseconds(),
+				Status:  http.StatusServiceUnavailable,
+				Err:     "queue full",
+			})
+			p.writeError(w, http.StatusServiceUnavailable, "queue full")
+			return
+		}
+		p.log.Warn("request", "route", route, "host", host, "method", r.Method, "path", rest, "status", 499)
+		return
 	}
+	defer release()
+
+	p.forward(w, r, route, host, leaseState(reused), rest, fp, model, started, waited)
 }