326 lines
11 KiB
Go
326 lines
11 KiB
Go
// 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"
|
|
"strings"
|
|
"time"
|
|
|
|
"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"
|
|
)
|
|
|
|
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
|
|
CtxHeader = "X-Crossbar-Ctx" // "moved:<old>new" when the context was relocated
|
|
)
|
|
|
|
// errBodyTooLarge is returned when a request body exceeds MaxBody during the model peek.
|
|
var errBodyTooLarge = errors.New("body too large")
|
|
|
|
// Health is what the proxy needs from the health table.
|
|
type Health interface {
|
|
Get(name string) (health.Status, bool)
|
|
MarkDown(name, reason string)
|
|
}
|
|
|
|
// Recorder is what the proxy needs to write an accounting row; *store.Store satisfies it.
|
|
type Recorder interface {
|
|
RecordRequest(store.Request) error
|
|
}
|
|
|
|
// Waker rouses a sleeping host. *wake.Waker satisfies it.
|
|
type Waker interface {
|
|
Wake(ctx context.Context, host string) bool
|
|
}
|
|
|
|
// 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
|
|
waker Waker
|
|
}
|
|
|
|
// SetWaker installs the waker the consults when a route has no healthy host left. A nil waker
|
|
// (the default) leaves the ErrNoHost answer as it was in v0: a plain 503.
|
|
func (p *Handler) SetWaker(w Waker) {
|
|
p.waker = w
|
|
}
|
|
|
|
// 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, 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
|
|
// "/a/" -> ("a", "/", true); "/", "//x", "", "noslash/v1" -> ("", "", false). The query string, if
|
|
// present, is kept in rest.
|
|
func SplitRoute(path string) (route, rest string, ok bool) {
|
|
if path == "" || path[0] != '/' {
|
|
return "", "", false
|
|
}
|
|
after := path[1:]
|
|
slash := strings.IndexByte(after, '/')
|
|
if slash == -1 {
|
|
if after == "" {
|
|
return "", "", false
|
|
}
|
|
return after, "/", true
|
|
}
|
|
if after[:slash] == "" {
|
|
return "", "", false
|
|
}
|
|
return after[:slash], after[slash:], true
|
|
}
|
|
|
|
// Choose returns the first host in order that is healthy and lists model in Loaded; failing that,
|
|
// the first healthy host. ok is false if none. model may be "".
|
|
func Choose(hosts []string, model string, h Health) (string, bool) {
|
|
for _, name := range hosts {
|
|
s, ok := h.Get(name)
|
|
if !ok || !s.Healthy || !hasModel(s.Loaded, model) {
|
|
continue
|
|
}
|
|
return name, true
|
|
}
|
|
for _, name := range hosts {
|
|
s, ok := h.Get(name)
|
|
if ok && s.Healthy {
|
|
return name, true
|
|
}
|
|
}
|
|
return "", false
|
|
}
|
|
|
|
func hasModel(loaded []string, model string) bool {
|
|
if model == "" {
|
|
return false
|
|
}
|
|
for _, m := range loaded {
|
|
if m == model {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// allowedPath reports whether rest may be proxied: under /v1/, or the two admin paths.
|
|
func allowedPath(rest string) bool {
|
|
return strings.HasPrefix(rest, "/v1/") || rest == "/health" || rest == "/props"
|
|
}
|
|
|
|
// 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. Compare concrete names.
|
|
if rname, _, ok := SplitRoute(rest); ok {
|
|
if _, _, rok := p.cfg.Route(rname); rok && rname != hdr {
|
|
return "", "", http.StatusBadRequest, "conflicting route"
|
|
}
|
|
}
|
|
if _, _, ok := p.cfg.Route(hdr); !ok {
|
|
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 _, _, ok := p.cfg.Route(route); !ok {
|
|
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, nil
|
|
}
|
|
if r.Body == nil || r.Body == http.NoBody {
|
|
return "", nil, nil
|
|
}
|
|
body, err := io.ReadAll(io.LimitReader(r.Body, MaxBody+1))
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
if len(body) > MaxBody {
|
|
return "", nil, errBodyTooLarge
|
|
}
|
|
r.Body = io.NopCloser(bytes.NewReader(body))
|
|
r.ContentLength = int64(len(body))
|
|
|
|
var req struct {
|
|
Model string `json:"model"`
|
|
}
|
|
_ = json.Unmarshal(body, &req)
|
|
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, code, msg := p.route(r)
|
|
if code != 0 {
|
|
p.writeError(w, code, msg)
|
|
return
|
|
}
|
|
routeCfg, _, _ := p.cfg.Route(route)
|
|
|
|
model, body, err := peekModel(r)
|
|
if err != nil {
|
|
p.writeError(w, http.StatusRequestEntityTooLarge, "body too large")
|
|
return
|
|
}
|
|
if model == "" {
|
|
model = routeCfg.DefaultModel
|
|
}
|
|
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, 0, "")
|
|
return
|
|
}
|
|
|
|
// 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 {
|
|
switch {
|
|
case errors.Is(err, lease.ErrNoHost):
|
|
// No host healthy. Ask a waker to rouse a sleeping one; it answers
|
|
// (served or 503) when it has had a turn, else falls through to the
|
|
// plain 503.
|
|
if p.waker != nil && p.wakeOnErrNoHost(w, r, route, routeCfg, rest, model, fp, started, lease.Key{Route: route, FP: fp, Model: model}) {
|
|
return
|
|
}
|
|
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
|
|
}
|
|
|
|
// Slot, context guard and forward, holding the slot for the leased host.
|
|
p.serveLeased(w, r, route, routeCfg, rest, model, fp, started, host, reused)
|
|
}
|
|
|
|
// serveLeased queues the request against the leased host's limiter, runs the
|
|
// context guard, and forwards. The slot is held for the originally leased host
|
|
// even if the guard relocates the lease: the guard already moved it.
|
|
func (p *Handler) serveLeased(w http.ResponseWriter, r *http.Request, route string, routeCfg config.Route, rest, model, fp string, started time.Time, host string, reused bool) {
|
|
// 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)
|
|
p.writeRecord(store.Request{
|
|
Route: route,
|
|
FP: fp,
|
|
Model: model,
|
|
Host: host,
|
|
Started: started,
|
|
TotalMs: time.Since(started).Milliseconds(),
|
|
Status: 499,
|
|
Err: "client cancelled while queued",
|
|
})
|
|
return
|
|
}
|
|
defer release()
|
|
|
|
// Context guard: if the prompt does not fit the leased host's per-slot
|
|
// context, move the conversation to a host where it fits, else answer 400.
|
|
now := time.Now()
|
|
host, header, _, done := p.guard(w, r, routeCfg.Hosts, host, route, model, fp, started)
|
|
if done {
|
|
return
|
|
}
|
|
p.forward(w, r, route, host, leaseState(reused), rest, fp, model, now, waited, 0, header)
|
|
}
|
|
|
|
// wakeOnErrNoHost answers the request when no host was healthy. It asks, in
|
|
// route order, each candidate with a wake target to rouse itself; a host that
|
|
// wakes is leased once more and then served. When none wakes, it answers 503
|
|
// with the hosts it tried. It returns true when the request has been answered.
|
|
func (p *Handler) wakeOnErrNoHost(w http.ResponseWriter, r *http.Request, route string, routeCfg config.Route, rest, model, fp string, started time.Time, key lease.Key) bool {
|
|
var tried []string
|
|
for _, name := range routeCfg.Hosts {
|
|
if p.cfg.Hosts[name].Wake == nil {
|
|
continue
|
|
}
|
|
tried = append(tried, name)
|
|
if !p.waker.Wake(r.Context(), name) {
|
|
continue
|
|
}
|
|
if newHost, _, err := p.leases.Acquire(key, routeCfg.Hosts, time.Now()); err == nil {
|
|
p.serveLeased(w, r, route, routeCfg, rest, model, fp, started, newHost, true)
|
|
return true
|
|
}
|
|
}
|
|
p.writeNoHealthyHost(w, route, model, fp, started, tried)
|
|
return true
|
|
}
|