kyle/crossbar
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge v1: leases, limiter, chooser, fingerprint, SQLite store, accounting, admin (8/8 tasks by Ornith)

Browse Source
This commit is contained in:
kyle
2026-09-25 07:00:34 -07:00
parent c6d115a8ff d7d8fcfa3b
commit bebee332a5
33 changed files with 4287 additions and 560 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+81 -22
View File
@@ -1,8 +1,10 @@
# crossbar
crossbar is an affinity router in front of several `llama-server` routers. A client's identity is
the first path segment of its base URL; v0 routes each request to the first healthy host on that
route's list and streams the answer back unbuffered.
the first path segment of its base URL — its route. Each conversation takes a sticky lease on one
host, chosen for the most free slots for its model times weight, and streams the answer back
incrementally with the usage chunk intact. Pins, drains, queueing, leases and accounting are all
new in v1.
## Build
@@ -16,22 +18,26 @@ streaming over real HTTP.
crossbar reads one TOML file. This is `example.toml`:
```toml
# crossbar example configuration. Replace <tailnet> and the addresses with your own.
# crossbar example configuration (v1). Replace <tailnet> and the addresses with your own.
listen = "127.0.0.1:17777" # never 0.0.0.0 — bind the tailnet address in production
db = "crossbar.db" # SQLite: leases + accounting (WAL). /var/lib/crossbar/crossbar.db under systemd
poll_interval = "1s" # 60s in production; 1s makes the smoke run quick
queue_max = 8
lease_idle = "30m" # a conversation idle this long loses its host
retention = "180d" # per-request rows older than this are rolled up daily
queue_max = 1 # waiting places per (host, model) beyond `parallel`; 503 past that
[hosts.alpha]
base_url = "http://127.0.0.1:18081" # e.g. http://straylight.<tailnet>:11434
weight = 1.0
models = { "ornith-1.5-35b-a3b" = { parallel = 4 }, "small-9b" = { parallel = 6 } }
models = { "ornith-1.5-35b-a3b" = { parallel = 1 }, "small-9b" = { parallel = 6 } }
[hosts.beta]
base_url = "http://127.0.0.1:18082" # e.g. http://titan.<tailnet>:8081
weight = 2.0
models = { "ornith-1.5-35b-a3b" = { parallel = 4 } }
models = { "ornith-1.5-35b-a3b" = { parallel = 2 } }
# v0: a route is a preference list; the first healthy host that has the model wins.
# v1: a route is a set of candidate hosts; each conversation gets a sticky lease on the host with
# the most free slots × weight at the time it starts. Pins and drains come from the admin API.
[routes.opencode-a]
hosts = ["alpha", "beta"]
default_model = "ornith-1.5-35b-a3b"
@@ -43,12 +49,15 @@ hosts = ["beta", "alpha"]
| Key | Meaning |
| --- | --- |
| `listen` | Where crossbar binds. A tailnet address, never `0.0.0.0`. |
| `db` | SQLite file holding leases and the accounting rows. |
| `lease_idle` | A conversation idle this long loses its host. |
| `retention` | Per-request rows older than this are rolled up daily. |
| `poll_interval` | How often each host is health-checked. 60s in production; 1s makes the smoke run quick. |
| `queue_max` | Reserved for v1 queueing; no effect in v0. |
| `queue_max` | Waiting places per (host, model) beyond `parallel`; a full queue returns 503. |
| `hosts.<name>.base_url` | The llama-server base URL this host serves. |
| `hosts.<name>.weight` | Relative share of new routes this host receives. |
| `hosts.<name>.weight` | Relative share of new requests this host receives. |
| `hosts.<name>.models` | The models this host serves, with per-model parallel tuning. |
| `routes.<name>.hosts` | Preference order: the first healthy host that serves the model wins. |
| `routes.<name>.hosts` | Candidate hosts, tried in order until one is healthy; a conversation leases one of them. |
| `routes.<name>.default_model` | Model used when a request omits one; must be served by a host in the route. |
## Run
@@ -84,23 +93,73 @@ custom_providers:
models: { ornith-1.5-35b-a3b: {} }
```
The route name in the URL must exist in `[routes]`; unknown routes are 404.
The route name in the URL must exist in `[routes]`; unknown routes are 404. A client may instead
name the route on an `X-Crossbar-Route` header and point at the bare `/v1` base:
## Inspect
`GET /_crossbar/hosts` reports every host's health and loaded models:
```json
{"alpha":{"healthy":true,"loaded":["ornith-1.5-35b-a3b","small-9b"],"last_ok":"2026-09-25T09:34:18Z","last_err":""},"beta":{"healthy":true,"loaded":["ornith-1.5-35b-a3b"],"last_ok":"2026-09-25T09:34:18Z","last_err":""}}
```sh
curl -H 'X-Crossbar-Route: opencode-a' \
https://crossbar.<tailnet>:7777/v1/chat/completions
```
`GET /_crossbar/routes` reports each route's preference order and default model:
## Operate
The operator's API lives under `/_crossbar/`. Every call returns 200 with a small JSON body unless
stated otherwise.
`GET /_crossbar/hosts` reports every host's health, loaded models, live concurrency from the
limiter and drain state:
```json
{"hermes-x":{"hosts":["beta","alpha"],"default_model":""},"opencode-a":{"hosts":["alpha","beta"],"default_model":"ornith-1.5-35b-a3b"}}
{"alpha":{"healthy":true,"loaded":["ornith-1.5-35b-a3b","small-9b"],"last_ok":"2026-09-25T13:53:25Z","last_err":"","free_slots":7,"in_flight":0,"queued":0,"draining":false},"beta":{"healthy":true,"loaded":["ornith-1.5-35b-a3b"],"last_ok":"2026-09-25T13:53:25Z","last_err":"","free_slots":2,"in_flight":0,"queued":0,"draining":false}}
```
## What v0 does not do
`GET /_crossbar/routes` reports each route's candidate hosts, default model, any pin and its live
leases:
Leases and stickiness, SQLite, `/slots`, queueing and wake-on-LAN are out of scope for v0; see
`PLAN.md`.
```json
{"hermes-x":{"hosts":["beta","alpha"],"default_model":"","pinned":"","leases":[]},"opencode-a":{"hosts":["alpha","beta"],"default_model":"ornith-1.5-35b-a3b","pinned":"","leases":[]}}
```
`POST /_crossbar/routes/{route}` pins a route to a host (`{"host":"alpha","pin":true}`) or releases
it and clears the pin (`{"release":true}`):
```json
{"ok":true}
```
`POST /_crossbar/hosts/{host}` sets or clears drain (`{"drain":true}`); a draining host takes no
new conversations but keeps its existing leases:
```json
{"ok":true}
```
`GET /_crossbar/usage` summarizes the accounting rows, grouped by `by=host`, `by=model` or
`by=route` (the default). Ask for JSON, or a fixed-width table with `Accept: text/plain`:
```json
[{"key":"beta","requests":2,"errors":0,"busy_ms":4,"queued_ms":0,"prompt_tokens":200,"cached_tokens":180,"completion_tokens":20}]
```
```
key requests errors busy_ms queued_ms prompt cached completion cache_hit
hermes-x 1 0 1 0 100 90 10 0.90
opencode-a 1 0 3 0 100 90 10 0.90
```
`GET /_crossbar/metrics` emits the Prometheus text exposition for request counts, token totals,
queue wait, host health and live slots:
```
# TYPE crossbar_requests_total counter
crossbar_requests_total{route="hermes-x",host="beta",status="200"} 1
crossbar_requests_total{route="opencode-a",host="beta",status="200"} 1
# TYPE crossbar_host_healthy gauge
crossbar_host_healthy{host="alpha"} 1
crossbar_host_healthy{host="beta"} 1
```
## What v1 does not do
The context-size guard, wake-on-LAN, Tailscale identity and `/slots` are out of scope for v1; see
`PLAN.md` v2.
cmd/crossbar/main.go
+81 -2
View File
@@ -17,7 +17,10 @@ import (
"git.wntrmute.dev/kyle/crossbar/internal/admin"
"git.wntrmute.dev/kyle/crossbar/internal/config"
"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/proxy"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
func main() {
@@ -38,6 +41,12 @@ func run() error {
log := slog.New(slog.NewTextHandler(os.Stderr, nil))
st, err := store.Open(cfg.DB)
if err != nil {
return err
}
defer st.Close()
baseURLs := make(map[string]string, len(cfg.Hosts))
for name, host := range cfg.Hosts {
baseURLs[name] = host.BaseURL
@@ -48,9 +57,79 @@ func run() error {
defer stop()
go table.Run(ctx)
hosts := proxy.HostView(table, cfg)
lim := limiter.New()
for name, h := range cfg.Hosts {
for model, m := range h.Models {
lim.Configure(name, model, m.Parallel, cfg.QueueMax)
}
}
leases, err := lease.New(st, hosts, proxy.Chooser(cfg, table, lim), cfg.LeaseIdle.Duration)
if err != nil {
return err
}
for name, rt := range cfg.Routes {
leases.Candidates(name, rt.Hosts)
}
mux := http.NewServeMux()
mux.Handle("/_crossbar/", admin.Handler(cfg, table))
mux.Handle("/", proxy.New(cfg, table, log))
mux.Handle("/_crossbar/", admin.Handler(cfg, table, leases, lim, st, hosts))
mux.Handle("/", proxy.New(cfg, table, leases, lim, st, log))
// Background maintenance until ctx is done. Errors are logged, never fatal.
go func() {
ticker := time.NewTicker(time.Minute)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
leases.ExpireIdle(time.Now())
}
}
}()
go func() {
ticker := time.NewTicker(time.Hour)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
n, err := st.Prune(time.Now(), cfg.Retention.Duration)
if err != nil {
log.Error("prune", "err", err)
continue
}
log.Info("pruned request rows", "rows", n)
}
}
}()
go func() {
ticker := time.NewTicker(cfg.PollInterval.Duration)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
for name, s := range table.All() {
if err := st.RecordHostHealth(store.HostHealth{
TS: time.Now(),
Host: name,
Healthy: s.Healthy,
Loaded: s.Loaded,
}); err != nil {
log.Error("record host health", "host", name, "err", err)
}
}
}
}
}()
srv := &http.Server{
Addr: cfg.Listen,
cmd/fakeupstream/main.go
+22 -8
View File
@@ -1,11 +1,13 @@
// fakeupstream stands in for a llama-server router in tests and the smoke run. Do not edit.
//
// fakeupstream -listen 127.0.0.1:18081 -name alpha -models a,b -down-file /tmp/alpha.down
// fakeupstream -listen 127.0.0.1:18081 -name alpha -models a,b -down-file /tmp/alpha.down -slow 0
//
// /health answers 503 while the down file exists, 200 otherwise. /v1/models lists -models.
// /props answers a small JSON object. /v1/chat/completions echoes: a streamed answer of five
// SSE chunks 200 ms apart when the body has "stream": true, one JSON answer otherwise. Every
// response carries X-Upstream: <name>.
// SSE chunks 200 ms apart when the body has "stream": true, then a final chunk carrying
// "usage" and llama-server style "timings", then [DONE]; one JSON answer with usage and
// timings otherwise. -slow adds that many milliseconds before answering (for queue tests).
// Every response carries X-Upstream: <name>.
package main
import (
@@ -25,11 +27,14 @@ func main() {
name := flag.String("name", "fake", "name reported in X-Upstream and answers")
models := flag.String("models", "m", "comma-separated model ids for /v1/models")
downFile := flag.String("down-file", "", "while this file exists, /health answers 503")
slow := flag.Int("slow", 0, "milliseconds to wait before answering a completion")
flag.Parse()
ids := strings.Split(*models, ",")
mux := http.NewServeMux()
stamp := func(w http.ResponseWriter) { w.Header().Set("X-Upstream", *name) }
usage := map[string]any{"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110}
timings := map[string]any{"prompt_n": 100, "cache_n": 90, "predicted_n": 10, "predicted_ms": 50.0}
mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
stamp(w)
@@ -61,11 +66,12 @@ func main() {
Stream bool `json:"stream"`
}
_ = json.Unmarshal(body, &req)
time.Sleep(time.Duration(*slow) * time.Millisecond)
if !req.Stream {
writeJSON(w, map[string]any{
"id": "chatcmpl-fake", "object": "chat.completion", "model": req.Model,
"choices": []map[string]any{{"index": 0, "message": map[string]string{"role": "assistant", "content": "hello from " + *name}, "finish_reason": "stop"}},
"usage": map[string]int{"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6},
"usage": usage, "timings": timings,
})
return
}
@@ -73,16 +79,24 @@ func main() {
w.Header().Set("Cache-Control", "no-cache")
w.WriteHeader(http.StatusOK)
fl, _ := w.(http.Flusher)
flush := func() {
if fl != nil {
fl.Flush()
}
}
for i := 1; i <= 5; i++ {
chunk := map[string]any{"id": "chatcmpl-fake", "object": "chat.completion.chunk", "model": req.Model,
"choices": []map[string]any{{"index": 0, "delta": map[string]string{"content": fmt.Sprintf("%s chunk %d ", *name, i)}}}}
b, _ := json.Marshal(chunk)
fmt.Fprintf(w, "data: %s\n\n", b)
if fl != nil {
fl.Flush()
}
flush()
time.Sleep(200 * time.Millisecond)
}
final := map[string]any{"id": "chatcmpl-fake", "object": "chat.completion.chunk", "model": req.Model,
"choices": []map[string]any{}, "usage": usage, "timings": timings}
b, _ := json.Marshal(final)
fmt.Fprintf(w, "data: %s\n\n", b)
flush()
fmt.Fprint(w, "data: [DONE]\n\n")
})
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
@@ -90,7 +104,7 @@ func main() {
http.Error(w, `{"error":"not found"}`, http.StatusNotFound)
})
log.Printf("fakeupstream %s listening on %s models=%v", *name, *listen, ids)
log.Printf("fakeupstream %s listening on %s models=%v slow=%dms", *name, *listen, ids, *slow)
srv := &http.Server{Addr: *listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
log.Fatal(srv.ListenAndServe())
}
docs/implementer-log.md
+39
View File
@@ -5,12 +5,21 @@ owner fills in the Model column. The reviewer adds findings under "Reviews" once
| Task | Date | Status | Gate runs | First gate | Deviations | Notes | Model |
|---|---|---|---|---|---|---|---|
| v1/08-smoke-readme | 2026-09-25 | done | 1 | pass | owner-directed fix to `Free` in `proxy.Chooser` | Changed `Free` from `c.lim.FreeSlots(host)` (sum over every model) to per-model free slots, `freeForModel(cfg.Hosts[host], model, c.lim.InFlight(host, model))`, floored at 0 and 0 when the host does not list the model (new helper in hosts.go); the one code change the task directs. `go test -race ./internal/proxy/` and `make gate` pass on the first run; `make smoke` → `smoke: ok (stream spread 1006 ms)`. README intro, `## Configure` (added db/lease_idle/retention, rewrote queue_max and hosts.<name>.hosts) and `## Inspect`→`## Operate` (all six endpoints, examples taken from the smoke run) updated. | llama.cpp/ornith-1.5-35b-a3b |
| v1/07-main | 2026-09-25 | done | 1 | pass | none | Wired store, limiter and lease table into `cmd/crossbar/main.go`: `store.Open` before the health table, `limiter.Configure` per (host, model) from `cfg.Hosts`, `lease.New` with `proxy.Chooser`, `Candidates` for every route, three background goroutines (idle expiry per minute, prune per hour logging the count, host-health recording per `poll_interval`), and `st.Close` via `defer`. The 3s SIGTERM run exits 0 with `listening`/`shutting down`; the missing-config run exits 1. | llama.cpp/ornith-1.5-35b-a3b |
| v1/06-admin | 2026-09-25 | done | 2 | fail | Split `internal/admin/admin.go` (196 lines) + `admin_ops.go` (366 lines) to stay under 400. Updated `cmd/crossbar/main.go`'s `admin.Handler` call from the committed 2-arg `(cfg, table)` to the task's 6-arg signature, passing the health table for `hosts` and `nil` for the not-yet-wired `leases`/`limiter`/`store`/`drainer` (task 07 wires them); this was a compile fix required for `go vet`/`go test ./...` on `cmd/crossbar` to pass — the full wiring is task 07. | First `make gate` failed on `go vet` (`admin.Handler` called with 2 args in `main.go` after the signature changed); fixed `main.go` and the gate passed on the second run. `admin_test.go` and `example.toml` verified byte-identical to `docs/plans/v1/_files/`; `internal/lease` and `internal/store` left untouched except the already-present `Candidates`/`StatusCounts`. | llama.cpp/ornith-1.5-35b-a3b |
| v1/05-proxy | 2026-09-25 | done | 2 | fail | Split `internal/proxy/proxy.go` (411 lines) into `proxy.go` + `forward.go` by moving `forward`, `newReverseProxy`, `forwardState`, `statusRecorder`, `leaseState`, `ttfbMs` and the `writeError`/`writeRecord` helpers to `forward.go`; the one `recorder_test.go` `proxy.New` call changed to `proxy.New(cfg, h, nil, nil, nil, nil)` per the task; `cmd/crossbar/main.go` passes `nil, nil, nil` for the new `leases`/`lim`/`rec` args (task 06 wires them). | The tee in `tee.go` already read the final SSE chunk's (streamed) and the JSON body's (non-streamed) usage/timings, so `TestAccountingRowsFromUsageAndTimings` passed on the first run — the only gate blocker was `proxy.go` at 411 lines. | llama.cpp/ornith-1.5-35b-a3b |
| v1/04-lease | 2026-09-25 | done | 1 | pass | The given `TestPinAndUnpin` was wrong and replaced by the owner mid-task; the corrected `internal/lease/lease_test.go` is byte-identical to `docs/plans/v1/_files/internal/lease/lease_test.go`. A `fmt.Printf("DEBUG …")` line the prior session left in `event` was removed before the gate. | `Acquire` order (pinned, existing, inherit, choose) with memory rolled back only after a successful save; `Pin` writes a pin event, then the pin row, then deletes other-host leases, so the pin event always precedes the unpin's release event in the log. | llama.cpp/ornith-1.5-35b-a3b |
| v1/02-fingerprint-config | 2026-09-25 | done | 1 | pass | Switched the existing `TestBadFiles` unknown-key example from `lease_idle` to `bogus_key`, and updated `testdata/bad-unknown-key.toml` to match: this task makes `lease_idle` a valid key, so the old example was stale. `config_test.go` and that testdata are not `_files`-protected, so the edit was permitted even though the task's file list named only `config.go` and `implementer-log.md`; the unknown-key rejection is still covered. | fingerprint.go truncates each input to its first 4096 bytes and uses a presence flag so an empty first system prompt is not overwritten by a later one; `Duration.UnmarshalText` matches `^[0-9]+d$` (regexp) before falling to `time.ParseDuration`. | llama.cpp/ornith-1.5-35b-a3b |
| v1/01-store | 2026-09-25 | done | 1 | pass | none | Gate passed on the first run once the owner gofmt'd the three previously-un-clean _files plan-tests under docs/plans/v1/_files/; the blocker in the stopped row no longer applies. | llama.cpp/ornith-1.5-35b-a3b |
| v1/01-store | 2026-09-25 | stopped | 2 | fail | none | Store implemented in `internal/store/store.go` + `schema.go`; `go test -race -count=1 ./internal/store/` is ok and `go vet`/`check-lines` pass. `make gate` cannot print `gate: ok` here: its `gofmt -l .` step flags three committed plan-tests under `docs/plans/v1/_files/` (admin, choose, proxy) that are not gofmt-clean under Go 1.26.7 (formatted by a gofmt that aligns one-line function bodies two columns wider; same diff on a pristine master). They live under `docs/plans/` (must not edit) and the gate covers them; the check cannot be scoped down without weakening it. Code left uncommitted for review. | llama.cpp/ornith-1.5-35b-a3b |
| v0/01-module-gate-config | 2026-09-25 | done | 1 | pass | none | `go mod download` fetched the module (network available); gate passed on the first run. | llama.cpp/ornith-1.5-35b-a3b |
| v0/02-health | 2026-09-25 | done | 1 | pass | none | First gate run passed. `MarkDown` initially forgot to write the entry back; caught by `TestMarkDown`. | llama.cpp/ornith-1.5-35b-a3b |
| v0/03-proxy | 2026-09-25 | done | 1 | pass | none | `SplitRoute` must reject an empty first segment (`/`, `//x`) as `ok=false`; the model peek restores the body and leaves non-JSON/empty as `""`. | llama.cpp/ornith-1.5-35b-a3b |
| v0/04-admin-main | 2026-09-25 | done | 1 | pass | none | `timeout --signal=TERM 3` exits 124 on a timed-out child on this GNU system, so the task's `exit=0` is not observable through it; sent SIGTERM directly and confirmed crossbar's own exit code is 0 with both log lines. | llama.cpp/ornith-1.5-35b-a3b |
| v0/05-smoke-readme-deploy | 2026-09-25 | done | 1 | pass | none | `README.md` `## Run` uses `install -m` instead of `cp` and adds `systemctl daemon-reload` before `enable --now`, which is required for systemd to see the new unit; the task said only "copy … then enable --now". | llama.cpp/ornith-1.5-35b-a3b |
| v0/01-review-fixes | 2026-09-25 | done | 1 | pass | none | `Flush` now two-value. Assertion inventory (`grep -n '\.(' internal/*/*.go`): proxy.go:150 fixed to two-value; proxy_test.go:274 net/http guarantees the server writer is a Flusher. No other unchecked outside assertion. Recorder test panicked before the fix, passed after; config tests passed as-is. | llama.cpp/ornith-1.5-35b-a3b |
| v1/03-limiter-choose | 2026-09-25 | done | 1 | pass | none | One mutex, a per-(host,model) pair with a FIFO waiter slice; release hands the slot to the head waiter by closing its channel without decrementing inflight, else frees it. A waiter whose ctx ends removes itself and, if the slot was handed in that same instant, gives it back so neither a slot nor a queue place leaks. FreeSlots counts only configured models so an unconfigured pair created by an Acquire does not add a phantom slot. | llama.cpp/ornith-1.5-35b-a3b |
## Reviews
@@ -41,3 +50,33 @@ Follow-ups for a `v0.1` task: fix 1 (`if f, ok := …; ok { f.Flush() }`), add t
test for 6, and make the log-row rule in `AGENTS.md` say that anything the Notes describe as a
change belongs in Deviations (finding 3).
### v1 review — 2026-09-25 (reviewer: claude, as owner for the night)
Checked: eight task commits `816614d`, `463cea1`, `7e0dbb4`, `9133240`, `97f7cdf`, `32ac7f5`,
`d82bfba`, `cf2aa24` with the trailer (plus one `stopped` commit and the owner's merges); every
given file byte-identical to its plan copy (v1 set, v0.1 set, and the v0 files not replaced;
`recorder_test.go` against the one owner-permitted edit); protected files untouched against the
merge base; `make gate` → `gate: ok`; `make smoke` → `smoke: ok (stream spread 1006 ms)`.
Probed outside the tests: a body whose `messages` is a string → 200 on the route lease; a leased
host drained *and* killed → 502 once with the host marked down, next turn moves with `lease=new`;
a second crossbar on the same `db` file → serves the same conversation on the leased host
(`lease=reused`) with no error; metrics carry the 502; SIGTERM mid-stream lets the stream finish
(7 SSE lines) and exits 0.
Tally: 8 tasks, 8 committed; first-run gate on 6 of the 8 sessions that reached the gate; 1
correct `stopped` (task 01, owner's gofmt fault); 3 owner-caused resumes (tasks 01, 04, 05) and 3
owner-caused restarts (tasks 05, 06 split, 08); 2 model-side process findings (below).
Wall time ~4 h including the owner's turnaround.
| # | Finding | Severity | Fault |
|---|---|---|---|
| 1 | A request whose client disconnects mid-stream writes **no accounting row** (`/usage` stays empty after a cut stream). Task rule 5 said the `ErrorHandler` does nothing on `context.Canceled`; rule 6 said "record what you have when `ServeHTTP` returns" — the second was not applied on that path. Cancelled requests are invisible to usage and error rate. | medium | task (ambiguous) + model (rule not applied everywhere) |
| 2 | `GET /_crossbar/usage` with no rows answers `null`, not `[]` (spec: a JSON array). | low | model |
| 3 | Task 02 edited two protected v0 files (fixture invalidated by the new key) with an honest deviation row instead of stopping. Content right, process wrong; the conflict itself was the owner's. | process | model + task |
| 4 | Task 05's first session ended its turn with a plan and no tool call after the sandbox refused a `/tmp` write (I9). | process | model |
| 5 | Owner faults, all recorded under "Changes during the run" in the plan README: given files not gofmt-clean; v0 fixture invalidated; pin-event positions; spread tie-break; queue-test read race; dropped test helper; unrecorded `/v1/models`; 433-line given test; task 06 oversized; task 06 text on the gate; chooser free slots summed across models. | — | task/test |
Follow-ups for `v1.1`: fix 1 (record the row on the cancel path with status 499 and `err`), fix 2
(`[]`), and an acceptance test for each; consider `lease_idle` expiry while a request is in flight
and `Prune` under concurrent writes, which this review did not probe.
example.toml
+9 -5
View File
@@ -1,19 +1,23 @@
# crossbar example configuration. Replace <tailnet> and the addresses with your own.
# crossbar example configuration (v1). Replace <tailnet> and the addresses with your own.
listen = "127.0.0.1:17777" # never 0.0.0.0 — bind the tailnet address in production
db = "crossbar.db" # SQLite: leases + accounting (WAL). /var/lib/crossbar/crossbar.db under systemd
poll_interval = "1s" # 60s in production; 1s makes the smoke run quick
queue_max = 8
lease_idle = "30m" # a conversation idle this long loses its host
retention = "180d" # per-request rows older than this are rolled up daily
queue_max = 1 # waiting places per (host, model) beyond `parallel`; 503 past that
[hosts.alpha]
base_url = "http://127.0.0.1:18081" # e.g. http://straylight.<tailnet>:11434
weight = 1.0
models = { "ornith-1.5-35b-a3b" = { parallel = 4 }, "small-9b" = { parallel = 6 } }
models = { "ornith-1.5-35b-a3b" = { parallel = 1 }, "small-9b" = { parallel = 6 } }
[hosts.beta]
base_url = "http://127.0.0.1:18082" # e.g. http://titan.<tailnet>:8081
weight = 2.0
models = { "ornith-1.5-35b-a3b" = { parallel = 4 } }
models = { "ornith-1.5-35b-a3b" = { parallel = 2 } }
# v0: a route is a preference list; the first healthy host that has the model wins.
# v1: a route is a set of candidate hosts; each conversation gets a sticky lease on the host with
# the most free slots × weight at the time it starts. Pins and drains come from the admin API.
[routes.opencode-a]
hosts = ["alpha", "beta"]
default_model = "ornith-1.5-35b-a3b"
go.mod
+16 -1
View File
@@ -2,4 +2,19 @@ module git.wntrmute.dev/kyle/crossbar
go 1.26
require github.com/BurntSushi/toml v1.6.0
require (
github.com/BurntSushi/toml v1.6.0
modernc.org/sqlite v1.59.0
)
require (
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/mattn/go-isatty v0.0.24 // indirect
github.com/ncruces/go-strftime v1.0.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
golang.org/x/sys v0.47.0 // indirect
modernc.org/libc v1.75.7 // indirect
modernc.org/mathutil v1.7.1 // indirect
modernc.org/memory v1.12.1 // indirect
)
go.sum
+51 -1
View File
@@ -1,2 +1,52 @@
github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk=
github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3/go.mod h1:jl5iWTm0/hd5PjEYEOuwAJ57L/CibdZfrqZ5XA5GrCk=
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3 h1:LMLX+LgTNWpfvCBdFebv6EsYotImrt/Ppc5cXIriCSo=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40=
golang.org/x/mod v0.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk=
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk=
golang.org/x/tools v0.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE=
modernc.org/ccgo/v4 v4.35.0/go.mod h1:qrVGs9S3Sr2Ztcg9ve+kTAYMp5a3YvWjo+SoN06kJ5I=
modernc.org/ccgo/v4 v4.35.0 h1:F+TUsmw09QxLzmi3aeYYGxjAXarmZaKgj3mKQHNaA8w=
modernc.org/cc/v4 v4.29.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
modernc.org/cc/v4 v4.29.2 h1:h6+9ciCnPKutf4I03CvheAvDLX7+IHlqR6Iy6J+cgd8=
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
modernc.org/gc/v3 v3.1.5/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
modernc.org/gc/v3 v3.1.5 h1:21ldfPfRYE31Tb7B3mwAK8gy1AxP4+dKjrOQPfqakoc=
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
modernc.org/libc v1.75.7/go.mod h1:bO5o2ztHxBb2rjz0PgdHN0sSMw57CgxGFLZ3Qd/QpVQ=
modernc.org/libc v1.75.7 h1:o3DTP9/0p9pKmY2WCKQaySW6wIiZhNM7wc2lUoyhfew=
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
modernc.org/memory v1.12.1/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
modernc.org/memory v1.12.1 h1:nFMiWrpStgZczNl6XI9GnIk/rWhYIyHGUaR04pGbp9g=
modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
modernc.org/sqlite v1.59.0/go.mod h1:+paeT2A3iPRHkQDwG7oA6Tk0zQd5woMEI8q7orfry8k=
modernc.org/sqlite v1.59.0 h1:X1es1GpqBlS/5T+vbM4HLUdaa8OtQx468DF2vrx+38A=
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
internal/admin/admin.go
+141 -47
View File
@@ -1,15 +1,19 @@
// Package admin serves the operator's view of crossbar: the health table and the routes as JSON,
// mounted at /_crossbar/ on the same listener as the proxy. The shape of /_crossbar/hosts is
// fixed so operators can read why a request went where it went.
// Package admin serves the operator's view of crossbar: the hosts view with
// slots and drain state, the routes view with leases and pins, the pin/release
// and drain controls, usage accounting, and Prometheus metrics, all under
// /_crossbar/.
package admin
import (
"encoding/json"
"net/http"
"sort"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/config"
"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"
)
// Hosts is what the admin handler needs from the health table.
@@ -17,59 +21,93 @@ type Hosts interface {
All() map[string]health.Status
}
type HostView struct {
Healthy bool `json:"healthy"`
Loaded []string `json:"loaded"` // never null: an empty slice when nothing is loaded
LastOK string `json:"last_ok"` // time.RFC3339 in UTC, or "" if never
LastErr string `json:"last_err"`
// Drainer is what the admin handler needs to steer draining; *proxy.Hosts
// satisfies it.
type Drainer interface {
Draining(name string) bool
SetDraining(name string, on bool)
}
// HostView is one host's row in the hosts view.
type HostView struct {
Healthy bool `json:"healthy"`
Loaded []string `json:"loaded"` // never null
LastOK string `json:"last_ok"` // RFC 3339 UTC or ""
LastErr string `json:"last_err"`
FreeSlots int `json:"free_slots"` // lim.FreeSlots(host)
InFlight int `json:"in_flight"` // sum over the host's configured models
Queued int `json:"queued"` // same
Draining bool `json:"draining"`
}
// LeaseView is one lease's row in a route's leases.
type LeaseView struct {
FP string `json:"fp"`
Model string `json:"model"`
Host string `json:"host"`
State string `json:"state"`
Created string `json:"created"` // RFC 3339 UTC
LastUsed string `json:"last_used"` // RFC 3339 UTC
}
// RouteView is one route's row in the routes view.
type RouteView struct {
Hosts []string `json:"hosts"`
DefaultModel string `json:"default_model"`
Pinned string `json:"pinned"` // "" when not pinned
Leases []LeaseView `json:"leases"` // never null
}
// Handler serves GET /_crossbar/hosts and GET /_crossbar/routes. Any other method on those paths is
// a 405 with an Allow: GET header; anything else under the handler is a 404.
func Handler(cfg *config.Config, h Hosts) http.Handler {
// handler implements the operator's endpoints under /_crossbar/.
type handler struct {
cfg *config.Config
h Hosts
lt *lease.Table
lim *limiter.Limiter
st *store.Store
d Drainer
}
// Handler builds the operator's HTTP handler.
func Handler(cfg *config.Config, h Hosts, lt *lease.Table, lim *limiter.Limiter, st *store.Store, d Drainer) http.Handler {
hx := &handler{cfg: cfg, h: h, lt: lt, lim: lim, st: st, d: d}
mux := http.NewServeMux()
mux.HandleFunc("/_crossbar/hosts", hostsHandler(h))
mux.HandleFunc("/_crossbar/routes", routesHandler(cfg))
mux.HandleFunc("/", notFound)
mux.HandleFunc("/_crossbar/hosts", hx.hostsGet)
mux.HandleFunc("/_crossbar/hosts/{host}", hx.hostDrain)
mux.HandleFunc("/_crossbar/routes", hx.routesGet)
mux.HandleFunc("/_crossbar/routes/{route}", hx.routePin)
mux.HandleFunc("/_crossbar/usage", hx.usageGet)
mux.HandleFunc("/_crossbar/metrics", hx.metricsGet)
mux.HandleFunc("/_crossbar/", hx.unknown)
return mux
}
func hostsHandler(h Hosts) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
func (hx *handler) hostsGet(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
wrongMethod(w)
wrongMethod(w, "GET")
return
}
views := make(map[string]HostView, len(h.All()))
for name, s := range h.All() {
views[name] = hostView(s)
}
writeJSON(w, http.StatusOK, views)
}
writeJSON(w, http.StatusOK, hx.hostViews())
}
func routesHandler(cfg *config.Config) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
wrongMethod(w)
return
}
views := make(map[string]RouteView, len(cfg.Routes))
for name, route := range cfg.Routes {
hosts := make([]string, len(route.Hosts))
copy(hosts, route.Hosts)
views[name] = RouteView{Hosts: hosts, DefaultModel: route.DefaultModel}
}
writeJSON(w, http.StatusOK, views)
// hostViews builds every host's row, keyed by host name.
func (hx *handler) hostViews() map[string]HostView {
all := hx.h.All()
out := make(map[string]HostView, len(all))
for name, s := range all {
out[name] = hx.hostView(name, s)
}
return out
}
func hostView(s health.Status) HostView {
// hostView builds one host's row: concurrency from the limiter summed over the
// models the host serves, draining from the drainer, and the health snapshot.
func (hx *handler) hostView(name string, s health.Status) HostView {
inflight, queued := 0, 0
for _, m := range configuredModels(hx.cfg, name) {
inflight += hx.lim.InFlight(name, m)
queued += hx.lim.Queued(name, m)
}
loaded := s.Loaded
if loaded == nil {
loaded = []string{}
@@ -83,20 +121,76 @@ func hostView(s health.Status) HostView {
Loaded: loaded,
LastOK: lastOK,
LastErr: s.LastErr,
FreeSlots: hx.lim.FreeSlots(name),
InFlight: inflight,
Queued: queued,
Draining: hx.d.Draining(name),
}
}
func wrongMethod(w http.ResponseWriter) {
w.Header().Set("Allow", "GET")
writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
// configuredModels returns the sorted model ids the host serves, or nil when
// the host is unknown.
func configuredModels(cfg *config.Config, name string) []string {
h, ok := cfg.Hosts[name]
if !ok {
return nil
}
models := make([]string, 0, len(h.Models))
for m := range h.Models {
models = append(models, m)
}
sort.Strings(models)
return models
}
func notFound(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusNotFound, map[string]string{"error": "not found"})
func (hx *handler) routesGet(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
wrongMethod(w, "GET")
return
}
snap := hx.lt.Snapshot()
views := make(map[string]RouteView, len(hx.cfg.Routes))
for name, route := range hx.cfg.Routes {
hosts := make([]string, len(route.Hosts))
copy(hosts, route.Hosts)
views[name] = hx.routeView(name, hosts, route.DefaultModel, snap)
}
writeJSON(w, http.StatusOK, views)
}
func writeJSON(w http.ResponseWriter, status int, v any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(v)
// routeView builds one route's row: the pinned host (empty if none) and the
// active leases on it, in snapshot order.
func (hx *handler) routeView(route string, hosts []string, defaultModel string, snap []lease.Lease) RouteView {
leases := make([]LeaseView, 0)
for _, l := range snap {
if l.Route == route {
leases = append(leases, leaseView(l))
}
}
return RouteView{
Hosts: hosts,
DefaultModel: defaultModel,
Pinned: hx.lt.Pinned(route),
Leases: leases,
}
}
// leaseView maps a lease to its operator view.
func leaseView(l lease.Lease) LeaseView {
return LeaseView{
FP: l.FP,
Model: l.Model,
Host: l.Host,
State: string(l.State),
Created: formatTime(l.Created),
LastUsed: formatTime(l.LastUsed),
}
}
// formatTime renders t as RFC 3339 in UTC, or "" for the zero time.
func formatTime(t time.Time) string {
if t.IsZero() {
return ""
}
return t.UTC().Format(time.RFC3339)
}
internal/admin/admin_ops.go
+366
View File
@@ -0,0 +1,366 @@
package admin
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"sort"
"strconv"
"strings"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/config"
"git.wntrmute.dev/kyle/crossbar/internal/lease"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
// routePin handles POST /_crossbar/routes/{route}: pin the route to a host or
// release and unpin it.
func (hx *handler) routePin(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
wrongMethod(w, "POST")
return
}
route := r.PathValue("route")
routeCfg, ok := hx.cfg.Routes[route]
if !ok {
writeError(w, http.StatusNotFound, "unknown route")
return
}
var raw struct {
Host string `json:"host"`
Pin *bool `json:"pin"`
Release *bool `json:"release"`
}
if err := decodeJSON(r, &raw); err != nil {
writeError(w, http.StatusBadRequest, "invalid body")
return
}
pinSet := raw.Pin != nil
releaseSet := raw.Release != nil
switch {
case pinSet && releaseSet:
writeError(w, http.StatusBadRequest, "pin and release at once")
case !pinSet && !releaseSet:
writeError(w, http.StatusBadRequest, "pin or release required")
case pinSet:
hx.pin(w, route, routeCfg, raw.Host)
default:
hx.release(w, route)
}
}
// pin validates the host, records candidates, and pins the route.
func (hx *handler) pin(w http.ResponseWriter, route string, routeCfg config.Route, host string) {
if host == "" {
writeError(w, http.StatusBadRequest, "pin requires host")
return
}
if !containsHost(routeCfg.Hosts, host) {
writeError(w, http.StatusNotFound, "host not in route")
return
}
// Record the route's hosts as candidates so Pin accepts a host no request
// has used yet.
hx.lt.Candidates(route, routeCfg.Hosts)
if err := hx.lt.Pin(route, host, time.Now()); err != nil {
if errors.Is(err, lease.ErrUnknownHost) {
writeError(w, http.StatusNotFound, "unknown host")
return
}
writeError(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
}
// release drops the route's leases and clears its pin.
func (hx *handler) release(w http.ResponseWriter, route string) {
n := hx.lt.Release(route)
hx.lt.Unpin(route)
writeJSON(w, http.StatusOK, map[string]any{"ok": true, "released": n})
}
// hostDrain handles POST /_crossbar/hosts/{host}: set or clear draining.
func (hx *handler) hostDrain(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
wrongMethod(w, "POST")
return
}
host := r.PathValue("host")
if _, ok := hx.cfg.Hosts[host]; !ok {
writeError(w, http.StatusNotFound, "unknown host")
return
}
var body struct {
Drain *bool `json:"drain"`
}
if err := decodeJSON(r, &body); err != nil {
writeError(w, http.StatusBadRequest, "invalid body")
return
}
if body.Drain == nil {
writeError(w, http.StatusBadRequest, "drain required")
return
}
hx.d.SetDraining(host, *body.Drain)
writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
}
// usageGet handles GET /_crossbar/usage: usage rows as JSON, or a fixed-width
// table when Accept is text/plain.
func (hx *handler) usageGet(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
wrongMethod(w, "GET")
return
}
q := r.URL.Query()
var byv store.By
switch q.Get("by") {
case "", "route":
byv = store.ByRoute
case "model":
byv = store.ByModel
case "host":
byv = store.ByHost
default:
writeError(w, http.StatusBadRequest, "invalid by")
return
}
since, err := parseSince(q.Get("since"))
if err != nil {
writeError(w, http.StatusBadRequest, "invalid since")
return
}
rows, err := hx.st.Usage(since, byv)
if err != nil {
writeError(w, http.StatusInternalServerError, "usage: "+err.Error())
return
}
if r.Header.Get("Accept") == "text/plain" {
writeUsageTable(w, rows)
return
}
writeJSON(w, http.StatusOK, rows)
}
// parseSince resolves the since query value: absent means all time, otherwise
// an RFC 3339 instant or a duration (which may end in "d" for days) meaning
// now - d.
func parseSince(s string) (time.Time, error) {
if s == "" {
return time.Time{}, nil
}
if t, err := time.Parse(time.RFC3339, s); err == nil {
return t.UTC(), nil
}
d, err := parseWindow(s)
if err != nil {
return time.Time{}, err
}
return time.Now().Add(-d), nil
}
// parseWindow parses a duration, accepting a trailing "d" for whole days.
func parseWindow(s string) (time.Duration, error) {
if n, ok := splitDays(s); ok {
return time.Duration(n) * 24 * time.Hour, nil
}
return time.ParseDuration(s)
}
// splitDays reports whether s is an integer number of days ("Nd").
func splitDays(s string) (int, bool) {
if len(s) < 2 || s[len(s)-1] != 'd' {
return 0, false
}
n, err := strconv.Atoi(s[:len(s)-1])
if err != nil || n < 0 {
return 0, false
}
return n, true
}
// writeUsageTable renders the rows as a fixed-width table with a header line,
// one row per entry, no trailing spaces.
func writeUsageTable(w http.ResponseWriter, rows []store.UsageRow) {
headers := []string{"key", "requests", "errors", "busy_ms", "queued_ms", "prompt", "cached", "completion", "cache_hit"}
lines := make([][]string, 0, len(rows)+1)
lines = append(lines, headers)
for _, u := range rows {
lines = append(lines, []string{
u.Key,
strconv.FormatInt(u.Requests, 10),
strconv.FormatInt(u.Errors, 10),
strconv.FormatInt(u.BusyMs, 10),
strconv.FormatInt(u.QueuedMs, 10),
strconv.FormatInt(u.PromptTokens, 10),
strconv.FormatInt(u.CachedTokens, 10),
strconv.FormatInt(u.CompletionTokens, 10),
strconv.FormatFloat(u.CacheHitRatio(), 'f', 2, 64),
})
}
widths := columnWidths(lines)
var b strings.Builder
for _, line := range lines {
for i, f := range line {
if i < len(line)-1 {
b.WriteString(fmt.Sprintf("%-*s ", widths[i], f))
} else {
b.WriteString(f)
}
}
b.WriteByte('\n')
}
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(b.String()))
}
// columnWidths returns the widest rendered field in each column.
func columnWidths(lines [][]string) []int {
widths := make([]int, len(lines[0]))
for _, line := range lines {
for i, f := range line {
if len(f) > widths[i] {
widths[i] = len(f)
}
}
}
return widths
}
// metricsGet handles GET /_crossbar/metrics, emitting the Prometheus text
// exposition format computed on request.
func (hx *handler) metricsGet(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
wrongMethod(w, "GET")
return
}
counts, err := hx.st.StatusCounts(time.Time{})
if err != nil {
writeError(w, http.StatusInternalServerError, "metrics: "+err.Error())
return
}
usage, err := hx.st.Usage(time.Time{}, store.ByRoute)
if err != nil {
writeError(w, http.StatusInternalServerError, "metrics: "+err.Error())
return
}
var reqSamples []string
for _, c := range counts {
reqSamples = append(reqSamples, fmt.Sprintf(
"crossbar_requests_total{route=\"%s\",host=\"%s\",status=\"%s\"} %d",
esc(c.Route), esc(c.Host), esc(strconv.Itoa(c.Status)), c.Count))
}
var prompt, cached, queue []string
for _, u := range usage {
prompt = append(prompt, fmt.Sprintf("crossbar_prompt_tokens_total{route=\"%s\"} %d", esc(u.Key), u.PromptTokens))
cached = append(cached, fmt.Sprintf("crossbar_cached_tokens_total{route=\"%s\"} %d", esc(u.Key), u.CachedTokens))
queue = append(queue, fmt.Sprintf("crossbar_queue_wait_ms_total{route=\"%s\"} %d", esc(u.Key), u.QueuedMs))
}
all := hx.h.All()
names := make([]string, 0, len(all))
for name := range all {
names = append(names, name)
}
sort.Strings(names)
var healthy, free, inflight, queued []string
for _, name := range names {
s := all[name]
healthy = append(healthy, fmt.Sprintf("crossbar_host_healthy{host=\"%s\"} %d", esc(name), btoi(s.Healthy)))
free = append(free, fmt.Sprintf("crossbar_host_free_slots{host=\"%s\"} %d", esc(name), hx.lim.FreeSlots(name)))
fi, q := 0, 0
for _, m := range configuredModels(hx.cfg, name) {
fi += hx.lim.InFlight(name, m)
q += hx.lim.Queued(name, m)
}
inflight = append(inflight, fmt.Sprintf("crossbar_host_in_flight{host=\"%s\"} %d", esc(name), fi))
queued = append(queued, fmt.Sprintf("crossbar_host_queued{host=\"%s\"} %d", esc(name), q))
}
var b strings.Builder
appendFamily(&b, "crossbar_requests_total", "counter", reqSamples)
appendFamily(&b, "crossbar_prompt_tokens_total", "counter", prompt)
appendFamily(&b, "crossbar_cached_tokens_total", "counter", cached)
appendFamily(&b, "crossbar_queue_wait_ms_total", "counter", queue)
appendFamily(&b, "crossbar_host_healthy", "gauge", healthy)
appendFamily(&b, "crossbar_host_free_slots", "gauge", free)
appendFamily(&b, "crossbar_host_in_flight", "gauge", inflight)
appendFamily(&b, "crossbar_host_queued", "gauge", queued)
w.Header().Set("Content-Type", "text/plain; version=0.0.4")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(b.String()))
}
// appendFamily writes a metric family: its TYPE line followed by the sorted
// sample lines.
func appendFamily(b *strings.Builder, name, typ string, samples []string) {
fmt.Fprintf(b, "# TYPE %s %s\n", name, typ)
sort.Strings(samples)
for _, s := range samples {
b.WriteString(s)
b.WriteByte('\n')
}
}
// esc escapes a label value for the Prometheus text format.
func esc(s string) string {
s = strings.ReplaceAll(s, `\`, `\\`)
s = strings.ReplaceAll(s, `"`, `\"`)
return s
}
// btoi converts a bool to 0/1 for a gauge.
func btoi(v bool) int {
if v {
return 1
}
return 0
}
// containsHost reports whether hosts contains h.
func containsHost(hosts []string, h string) bool {
for _, x := range hosts {
if x == h {
return true
}
}
return false
}
// decodeJSON decodes a bounded JSON body.
func decodeJSON(r *http.Request, v any) error {
dec := json.NewDecoder(io.LimitReader(r.Body, 4096))
return dec.Decode(v)
}
func writeJSON(w http.ResponseWriter, status int, v any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(v)
}
func writeError(w http.ResponseWriter, status int, msg string) {
writeJSON(w, status, map[string]string{"error": msg})
}
// wrongMethod answers 405 with the allowed method in the Allow header.
func wrongMethod(w http.ResponseWriter, allow string) {
w.Header().Set("Allow", allow)
writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
}
func (hx *handler) unknown(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusNotFound, map[string]string{"error": "not found"})
}
internal/admin/admin_test.go
+231 -37
View File
@@ -1,9 +1,12 @@
package admin_test
// v1 admin: read the tables, pin/release a route, drain a host, usage rollups, metrics.
import (
"encoding/json"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"testing"
"time"
@@ -11,21 +14,45 @@ import (
"git.wntrmute.dev/kyle/crossbar/internal/admin"
"git.wntrmute.dev/kyle/crossbar/internal/config"
"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"
)
type fakeHosts map[string]health.Status
type fakeHosts struct {
st map[string]health.Status
draining map[string]bool
}
func (f fakeHosts) All() map[string]health.Status { return f }
func (f *fakeHosts) All() map[string]health.Status { return f.st }
func (f *fakeHosts) Healthy(n string) bool { return f.st[n].Healthy }
func (f *fakeHosts) Draining(n string) bool { return f.draining[n] }
func (f *fakeHosts) SetDraining(n string, on bool) { f.draining[n] = on }
func (f *fakeHosts) Choose(c []string, model string) (string, bool) {
for _, h := range c {
if f.st[h].Healthy && !f.draining[h] {
return h, true
}
}
return "", false
}
func testConfig(t *testing.T) *config.Config {
c, err := config.Parse(strings.NewReader(`
type rig struct {
h http.Handler
store *store.Store
leases *lease.Table
hosts *fakeHosts
}
func newRig(t *testing.T) *rig {
cfg, err := config.Parse(strings.NewReader(`
listen = "127.0.0.1:1"
[hosts.alpha]
base_url = "http://alpha:1"
models = { "m" = { } }
models = { "m" = { parallel = 2 } }
[hosts.beta]
base_url = "http://beta:1"
models = { "m" = { } }
models = { "m" = { parallel = 4 } }
[routes.r]
hosts = ["alpha", "beta"]
default_model = "m"
@@ -33,67 +60,234 @@ default_model = "m"
if err != nil {
t.Fatal(err)
}
return c
st, err := store.Open(filepath.Join(t.TempDir(), "x.db"))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = st.Close() })
hosts := &fakeHosts{
st: map[string]health.Status{
"alpha": {Healthy: true, Loaded: []string{"m"}, LastOK: time.Date(2026, 9, 25, 8, 0, 0, 0, time.UTC)},
"beta": {Healthy: false, LastErr: "HTTP 503"},
},
draining: map[string]bool{},
}
lt, err := lease.New(st, hosts, hosts, 30*time.Minute)
if err != nil {
t.Fatal(err)
}
lim := limiter.New()
lim.Configure("alpha", "m", 2, 8)
lim.Configure("beta", "m", 4, 8)
return &rig{h: admin.Handler(cfg, hosts, lt, lim, st, hosts), store: st, leases: lt, hosts: hosts}
}
func TestHosts(t *testing.T) {
when := time.Date(2026, 9, 25, 8, 0, 0, 0, time.UTC)
h := admin.Handler(testConfig(t), fakeHosts{
"alpha": {Healthy: true, Loaded: []string{"m"}, LastOK: when},
"beta": {Healthy: false, LastErr: "HTTP 503"},
})
func (r *rig) do(t *testing.T, method, path, body string, hdr ...string) *httptest.ResponseRecorder {
req := httptest.NewRequest(method, path, strings.NewReader(body))
if body != "" {
req.Header.Set("Content-Type", "application/json")
}
for i := 0; i+1 < len(hdr); i += 2 {
req.Header.Set(hdr[i], hdr[i+1])
}
rec := httptest.NewRecorder()
h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/_crossbar/hosts", nil))
if rec.Code != 200 || !strings.HasPrefix(rec.Header().Get("Content-Type"), "application/json") {
t.Fatalf("status %d, content-type %q", rec.Code, rec.Header().Get("Content-Type"))
r.h.ServeHTTP(rec, req)
return rec
}
func TestHostsShowsSlotsAndDrain(t *testing.T) {
r := newRig(t)
rec := r.do(t, "GET", "/_crossbar/hosts", "")
if rec.Code != 200 {
t.Fatalf("%d %s", rec.Code, rec.Body.String())
}
var out map[string]admin.HostView
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
t.Fatal(err)
}
if a := out["alpha"]; !a.Healthy || len(a.Loaded) != 1 || a.LastOK != "2026-09-25T08:00:00Z" || a.LastErr != "" {
a := out["alpha"]
if !a.Healthy || a.FreeSlots != 2 || a.InFlight != 0 || a.Queued != 0 || a.Draining || a.LastOK != "2026-09-25T08:00:00Z" {
t.Errorf("alpha = %+v", a)
}
if b := out["beta"]; b.Healthy || b.LastOK != "" || b.LastErr != "HTTP 503" || b.Loaded == nil {
if b := out["beta"]; b.Healthy || b.LastErr != "HTTP 503" || b.FreeSlots != 4 || b.Loaded == nil {
t.Errorf("beta = %+v (loaded must be [] not null)", b)
}
if !strings.Contains(rec.Body.String(), `"loaded":[]`) {
t.Errorf("beta.loaded must encode as []: %s", rec.Body.String())
}
}
func TestRoutes(t *testing.T) {
h := admin.Handler(testConfig(t), fakeHosts{})
rec := httptest.NewRecorder()
h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/_crossbar/routes", nil))
func TestRoutesShowsLeases(t *testing.T) {
r := newRig(t)
now := time.Date(2026, 9, 25, 9, 0, 0, 0, time.UTC)
if _, _, err := r.leases.Acquire(lease.Key{Route: "r", FP: "abc", Model: "m"}, []string{"alpha", "beta"}, now); err != nil {
t.Fatal(err)
}
rec := r.do(t, "GET", "/_crossbar/routes", "")
var out map[string]admin.RouteView
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
t.Fatalf("%v: %s", err, rec.Body.String())
}
r := out["r"]
if len(r.Hosts) != 2 || r.Hosts[0] != "alpha" || r.DefaultModel != "m" {
t.Errorf("routes = %+v", out)
rv := out["r"]
if len(rv.Hosts) != 2 || rv.DefaultModel != "m" || rv.Pinned != "" {
t.Errorf("route view = %+v", rv)
}
if len(rv.Leases) != 1 || rv.Leases[0].FP != "abc" || rv.Leases[0].Host != "alpha" || rv.Leases[0].State != "active" || rv.Leases[0].LastUsed != "2026-09-25T09:00:00Z" {
t.Errorf("leases = %+v", rv.Leases)
}
}
func TestPinReleaseDrain(t *testing.T) {
r := newRig(t)
rec := r.do(t, "POST", "/_crossbar/routes/r", `{"host":"beta","pin":true}`)
if rec.Code != 200 {
t.Fatalf("pin: %d %s", rec.Code, rec.Body.String())
}
if h, _, err := r.leases.Acquire(lease.Key{Route: "r", FP: "x", Model: "m"}, []string{"alpha", "beta"}, time.Now()); err == nil || h != "" {
// beta is unhealthy in the rig: a pin to a down host is honoured, not silently moved
t.Errorf("acquire on a route pinned to a down host: %q %v, want ErrPinnedDown", h, err)
}
rec = r.do(t, "GET", "/_crossbar/routes", "")
var out map[string]admin.RouteView
_ = json.Unmarshal(rec.Body.Bytes(), &out)
if out["r"].Pinned != "beta" {
t.Errorf("Pinned = %q after pin", out["r"].Pinned)
}
rec = r.do(t, "POST", "/_crossbar/routes/r", `{"release":true}`)
if rec.Code != 200 {
t.Fatalf("release: %d %s", rec.Code, rec.Body.String())
}
if h, _, err := r.leases.Acquire(lease.Key{Route: "r", FP: "x", Model: "m"}, []string{"alpha", "beta"}, time.Now()); err != nil || h != "alpha" {
t.Errorf("after release: %q %v, want alpha (the only healthy host)", h, err)
}
for _, tc := range []struct {
body string
want int
}{
{`{"host":"nobody","pin":true}`, 404},
{`{"pin":true}`, 400},
{`not json`, 400},
{`{"release":true,"pin":true,"host":"alpha"}`, 400},
} {
if rec := r.do(t, "POST", "/_crossbar/routes/r", tc.body); rec.Code != tc.want {
t.Errorf("POST %s: %d, want %d (%s)", tc.body, rec.Code, tc.want, rec.Body.String())
}
}
if rec := r.do(t, "POST", "/_crossbar/routes/nope", `{"release":true}`); rec.Code != 404 {
t.Errorf("unknown route: %d", rec.Code)
}
rec = r.do(t, "POST", "/_crossbar/hosts/alpha", `{"drain":true}`)
if rec.Code != 200 || !r.hosts.Draining("alpha") {
t.Fatalf("drain: %d %s draining=%v", rec.Code, rec.Body.String(), r.hosts.Draining("alpha"))
}
rec = r.do(t, "GET", "/_crossbar/hosts", "")
var hv map[string]admin.HostView
_ = json.Unmarshal(rec.Body.Bytes(), &hv)
if !hv["alpha"].Draining {
t.Errorf("hosts view must show draining")
}
if rec := r.do(t, "POST", "/_crossbar/hosts/alpha", `{"drain":false}`); rec.Code != 200 || r.hosts.Draining("alpha") {
t.Errorf("undrain: %d draining=%v", rec.Code, r.hosts.Draining("alpha"))
}
if rec := r.do(t, "POST", "/_crossbar/hosts/nobody", `{"drain":true}`); rec.Code != 404 {
t.Errorf("unknown host: %d", rec.Code)
}
}
func seedUsage(t *testing.T, st *store.Store) {
t0 := time.Now().UTC().Add(-time.Hour)
for i, r := range []store.Request{
{Route: "r", FP: "a", Model: "m", Host: "alpha", Status: 200, TotalMs: 1000, PromptTokens: 100, CachedTokens: 80, CompletionTokens: 10},
{Route: "r", FP: "a", Model: "m", Host: "alpha", Status: 200, TotalMs: 500, QueuedMs: 30, PromptTokens: 100, CachedTokens: 100, CompletionTokens: 5},
{Route: "r2", FP: "b", Model: "m", Host: "beta", Status: 503, TotalMs: 1, Err: "queue full"},
} {
r.Started = t0.Add(time.Duration(i) * time.Minute)
if err := st.RecordRequest(r); err != nil {
t.Fatal(err)
}
}
}
func TestUsageJSONAndText(t *testing.T) {
r := newRig(t)
seedUsage(t, r.store)
rec := r.do(t, "GET", "/_crossbar/usage?by=route", "")
if rec.Code != 200 || !strings.HasPrefix(rec.Header().Get("Content-Type"), "application/json") {
t.Fatalf("%d %q", rec.Code, rec.Header().Get("Content-Type"))
}
var rows []store.UsageRow
if err := json.Unmarshal(rec.Body.Bytes(), &rows); err != nil {
t.Fatalf("%v: %s", err, rec.Body.String())
}
if len(rows) != 2 {
t.Fatalf("rows = %+v", rows)
}
for _, row := range rows {
if row.Key == "r" && (row.Requests != 2 || row.CachedTokens != 180 || row.QueuedMs != 30) {
t.Errorf("r = %+v", row)
}
if row.Key == "r2" && (row.Requests != 1 || row.Errors != 1) {
t.Errorf("r2 = %+v", row)
}
}
rec = r.do(t, "GET", "/_crossbar/usage?by=host&since=24h", "", "Accept", "text/plain")
if rec.Code != 200 || !strings.HasPrefix(rec.Header().Get("Content-Type"), "text/plain") {
t.Fatalf("text: %d %q", rec.Code, rec.Header().Get("Content-Type"))
}
body := rec.Body.String()
if !strings.Contains(body, "alpha") || !strings.Contains(body, "beta") || !strings.Contains(strings.ToLower(body), "cache") {
t.Errorf("text table = %q", body)
}
if rec := r.do(t, "GET", "/_crossbar/usage?by=colour", ""); rec.Code != 400 {
t.Errorf("bad by: %d", rec.Code)
}
if rec := r.do(t, "GET", "/_crossbar/usage?since=yesterday", ""); rec.Code != 400 {
t.Errorf("bad since: %d", rec.Code)
}
rec = r.do(t, "GET", "/_crossbar/usage?since=2026-09-25T00:00:00Z&by=model", "")
if rec.Code != 200 {
t.Errorf("RFC3339 since: %d %s", rec.Code, rec.Body.String())
}
}
func TestMetrics(t *testing.T) {
r := newRig(t)
seedUsage(t, r.store)
rec := r.do(t, "GET", "/_crossbar/metrics", "")
if rec.Code != 200 || !strings.HasPrefix(rec.Header().Get("Content-Type"), "text/plain") {
t.Fatalf("%d %q", rec.Code, rec.Header().Get("Content-Type"))
}
body := rec.Body.String()
for _, want := range []string{
`# TYPE crossbar_requests_total counter`,
`crossbar_requests_total{route="r",host="alpha",status="200"} 2`,
`crossbar_requests_total{route="r2",host="beta",status="503"} 1`,
`crossbar_host_healthy{host="alpha"} 1`,
`crossbar_host_healthy{host="beta"} 0`,
`crossbar_host_free_slots{host="alpha"} 2`,
`crossbar_prompt_tokens_total{route="r"} 200`,
`crossbar_cached_tokens_total{route="r"} 180`,
`crossbar_queue_wait_ms_total{route="r"} 30`,
} {
if !strings.Contains(body, want) {
t.Errorf("metrics missing %q\n%s", want, body)
}
}
}
func TestMethodsAndUnknown(t *testing.T) {
h := admin.Handler(testConfig(t), fakeHosts{})
r := newRig(t)
for _, tc := range []struct {
method, path string
want int
}{
{http.MethodPost, "/_crossbar/hosts", 405},
{http.MethodDelete, "/_crossbar/routes", 405},
{http.MethodGet, "/_crossbar/routes/r", 405},
{http.MethodGet, "/_crossbar/nope", 404},
{http.MethodGet, "/_crossbar/", 404},
{http.MethodPut, "/_crossbar/usage", 405},
} {
rec := httptest.NewRecorder()
h.ServeHTTP(rec, httptest.NewRequest(tc.method, tc.path, nil))
if rec.Code != tc.want {
t.Errorf("%s %s = %d, want %d", tc.method, tc.path, rec.Code, tc.want)
}
if !strings.HasPrefix(rec.Header().Get("Content-Type"), "application/json") {
t.Errorf("%s %s: errors are JSON too", tc.method, tc.path)
rec := r.do(t, tc.method, tc.path, "")
if rec.Code != tc.want || !strings.HasPrefix(rec.Header().Get("Content-Type"), "application/json") {
t.Errorf("%s %s = %d %q, want %d JSON", tc.method, tc.path, rec.Code, rec.Header().Get("Content-Type"), tc.want)
}
}
}
internal/choose/choose.go
+63
View File
@@ -0,0 +1,63 @@
// Package choose picks the host for a new lease: among the healthy, non-draining, known hosts it
// prefers those that have the model loaded over those that would only be able to serve it, then the
// one with the most free slots times weight, breaking ties by shortest queue and finally list
// order.
package choose
// Info is one candidate host's view of itself for a single model.
type Info struct {
Healthy bool // answered its last poll
Draining bool // operator is draining it; no new leases
Loaded bool // the model is resident here
CanServe bool // config lists the model, so we may load it
Free int
Queued int
Weight float64
}
// Best returns the best host for a new lease, or ok=false when nothing is eligible. It runs two
// passes over the candidates in order: the first over those that have the model loaded, the second,
// only if the first found nothing, over those that can serve it. A host must be healthy, not
// draining, and known (info reported ok) to be eligible in either pass; a host with zero free slots
// is still eligible, since it will queue. Among the eligible ones the highest Free*Weight wins, ties
// go to the lowest queue, and a remaining tie keeps the earlier candidate.
func Best(candidates []string, info func(host string) (Info, bool)) (string, bool) {
const (
passLoaded = 0
passCanServe = 1
)
best := ""
var (
bestScore float64
bestQueued int
found bool
)
for pass := passLoaded; pass <= passCanServe; pass++ {
for _, host := range candidates {
v, ok := info(host)
if !ok || !v.Healthy || v.Draining {
continue
}
switch pass {
case passLoaded:
if !v.Loaded {
continue
}
case passCanServe:
if !v.CanServe {
continue
}
}
score := float64(v.Free) * v.Weight
// Replace only when strictly better on score, or equal score with a shorter queue; a
// further tie keeps the earlier candidate because we scan in order and use "<".
if !found || score > bestScore || (score == bestScore && v.Queued < bestQueued) {
best, bestScore, bestQueued, found = host, score, v.Queued, true
}
}
if found {
break
}
}
return best, found
}
internal/choose/choose_test.go
+83
View File
@@ -0,0 +1,83 @@
package choose_test
import (
"testing"
"git.wntrmute.dev/kyle/crossbar/internal/choose"
)
func infoFor(m map[string]choose.Info) func(string) (choose.Info, bool) {
return func(name string) (choose.Info, bool) { i, ok := m[name]; return i, ok }
}
func TestMostFreeSlotsTimesWeightWins(t *testing.T) {
info := infoFor(map[string]choose.Info{
"alpha": {Healthy: true, Loaded: true, CanServe: true, Free: 3, Weight: 1.0},
"beta": {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 2.0}, // 4 > 3
"gamma": {Healthy: true, Loaded: true, CanServe: true, Free: 4, Weight: 0.5}, // 2
})
got, ok := choose.Best([]string{"alpha", "beta", "gamma"}, info)
if !ok || got != "beta" {
t.Errorf("got %q %v, want beta", got, ok)
}
}
func TestTieGoesToShortestQueueThenListOrder(t *testing.T) {
info := infoFor(map[string]choose.Info{
"alpha": {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 1, Queued: 3},
"beta": {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 1, Queued: 1},
"gamma": {Healthy: true, Loaded: true, CanServe: true, Free: 2, Weight: 1, Queued: 1},
})
if got, _ := choose.Best([]string{"alpha", "beta", "gamma"}, info); got != "beta" {
t.Errorf("tie on score: shortest queue wins, then list order; got %q", got)
}
if got, _ := choose.Best([]string{"gamma", "beta"}, info); got != "gamma" {
t.Errorf("full tie: first in list order wins; got %q", got)
}
}
func TestLoadedBeatsMerelyCapable(t *testing.T) {
info := infoFor(map[string]choose.Info{
"alpha": {Healthy: true, Loaded: false, CanServe: true, Free: 8, Weight: 4},
"beta": {Healthy: true, Loaded: true, CanServe: true, Free: 1, Weight: 1},
})
got, ok := choose.Best([]string{"alpha", "beta"}, info)
if !ok || got != "beta" {
t.Errorf("a host that has the model loaded wins over one that would have to load it; got %q", got)
}
}
func TestFallsBackToCapableHost(t *testing.T) {
info := infoFor(map[string]choose.Info{
"alpha": {Healthy: true, Loaded: false, CanServe: true, Free: 1, Weight: 1},
"beta": {Healthy: true, Loaded: false, CanServe: false, Free: 9, Weight: 9},
})
got, ok := choose.Best([]string{"beta", "alpha"}, info)
if !ok || got != "alpha" {
t.Errorf("only a host configured to serve the model may load it; got %q %v", got, ok)
}
}
func TestSkipsUnhealthyDrainingUnknownAndFull(t *testing.T) {
info := infoFor(map[string]choose.Info{
"down": {Healthy: false, Loaded: true, CanServe: true, Free: 9, Weight: 9},
"drain": {Healthy: true, Draining: true, Loaded: true, CanServe: true, Free: 9, Weight: 9},
"full": {Healthy: true, Loaded: true, CanServe: true, Free: 0, Weight: 9, Queued: 0},
"ok": {Healthy: true, Loaded: true, CanServe: true, Free: 1, Weight: 1},
})
got, ok := choose.Best([]string{"down", "drain", "missing", "full", "ok"}, info)
if !ok || got != "ok" {
t.Errorf("got %q %v, want ok", got, ok)
}
// A full host is still better than nothing: it gets the request (it will queue).
got, ok = choose.Best([]string{"down", "full"}, info)
if !ok || got != "full" {
t.Errorf("with only a full host left it must still be chosen; got %q %v", got, ok)
}
if _, ok := choose.Best([]string{"down", "drain", "missing"}, info); ok {
t.Errorf("nothing usable must give ok=false")
}
if _, ok := choose.Best(nil, info); ok {
t.Errorf("empty candidates must give ok=false")
}
}
internal/config/config.go
+75 -3
View File
@@ -15,18 +15,25 @@ import (
"os"
"regexp"
"sort"
"strconv"
"strings"
"time"
"github.com/BurntSushi/toml"
)
// Duration is a time.Duration that TOML reads from a string such as "60s" or
// "30m".
// Duration is a time.Duration that TOML reads from a string such as "60s",
// "30m", or "7d" (an integer number of days).
type Duration struct{ time.Duration }
// UnmarshalText implements encoding.TextUnmarshaler via time.ParseDuration.
// UnmarshalText implements encoding.TextUnmarshaler. It accepts the "Nd" form
// — an integer number of days, so "7d" is 7 × 24h — in addition to
// time.ParseDuration syntax.
func (d *Duration) UnmarshalText(text []byte) error {
if days, ok := parseDays(text); ok {
d.Duration = days
return nil
}
dt, err := time.ParseDuration(string(text))
if err != nil {
return err
@@ -35,6 +42,22 @@ func (d *Duration) UnmarshalText(text []byte) error {
return nil
}
// dayPattern matches a run of digits followed by "d", e.g. "7d".
var dayPattern = regexp.MustCompile(`^[0-9]+d$`)
// parseDays reports whether text is the "Nd" day form and returns that many
// hours. The regex guarantees the prefix is a base-10 integer.
func parseDays(text []byte) (time.Duration, bool) {
if !dayPattern.MatchString(string(text)) {
return 0, false
}
n, err := strconv.Atoi(string(text[:len(text)-1]))
if err != nil {
return 0, false
}
return time.Duration(n) * 24 * time.Hour, true
}
// Model is the per-model tuning carried by a host entry.
type Model struct {
Parallel int `toml:"parallel"`
@@ -58,6 +81,9 @@ type Config struct {
Listen string `toml:"listen"`
PollInterval Duration `toml:"poll_interval"`
QueueMax int `toml:"queue_max"`
DB string `toml:"db"`
LeaseIdle Duration `toml:"lease_idle"`
Retention Duration `toml:"retention"`
Hosts map[string]Host `toml:"hosts"`
Routes map[string]Route `toml:"routes"`
}
@@ -76,6 +102,13 @@ const (
DefaultPollInterval = 60 * time.Second
DefaultQueueMax = 8
MinPollInterval = time.Second
DefaultDB = "crossbar.db"
DefaultLeaseIdle = 30 * time.Minute
DefaultRetention = 180 * 24 * time.Hour
MinLeaseIdle = time.Minute
MinRetention = 24 * time.Hour
)
var routeName = regexp.MustCompile(`^[a-z0-9][a-z0-9-]*$`)
@@ -109,6 +142,15 @@ func Parse(r io.Reader) (*Config, error) {
return nil, &Error{Field: keys[0], Msg: "unknown key"}
}
if !md.IsDefined("db") {
c.DB = DefaultDB
}
if !md.IsDefined("lease_idle") {
c.LeaseIdle.Duration = DefaultLeaseIdle
}
if !md.IsDefined("retention") {
c.Retention.Duration = DefaultRetention
}
if c.PollInterval.Duration == 0 {
c.PollInterval.Duration = DefaultPollInterval
}
@@ -152,6 +194,15 @@ func (c *Config) validate() *Error {
if e := c.checkQueue(); e != nil {
return e
}
if e := c.checkDB(); e != nil {
return e
}
if e := c.checkLeaseIdle(); e != nil {
return e
}
if e := c.checkRetention(); e != nil {
return e
}
if e := c.checkHosts(); e != nil {
return e
}
@@ -193,6 +244,27 @@ func (c *Config) checkQueue() *Error {
return nil
}
func (c *Config) checkDB() *Error {
if c.DB == "" {
return &Error{Field: "db", Msg: "required"}
}
return nil
}
func (c *Config) checkLeaseIdle() *Error {
if c.LeaseIdle.Duration < MinLeaseIdle {
return &Error{Field: "lease_idle", Msg: "must be at least 1m"}
}
return nil
}
func (c *Config) checkRetention() *Error {
if c.Retention.Duration < MinRetention {
return &Error{Field: "retention", Msg: "must be at least 1d"}
}
return nil
}
func (c *Config) checkHosts() *Error {
if len(c.Hosts) == 0 {
return &Error{Field: "hosts", Msg: "at least one required"}
internal/config/config_test.go
+1 -1
View File
@@ -81,7 +81,7 @@ func TestBadFiles(t *testing.T) {
{"bad-listen.toml", "listen"},
{"bad-unknown-host.toml", "routes.r.hosts"},
{"bad-default-model.toml", "routes.r.default_model"},
{"bad-unknown-key.toml", "lease_idle"},
{"bad-unknown-key.toml", "bogus_key"},
}
for _, tc := range cases {
t.Run(tc.file, func(t *testing.T) {
internal/config/config_v1_test.go
+81
View File
@@ -0,0 +1,81 @@
package config_test
import (
"strings"
"testing"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/config"
)
const v1Base = `
listen = "127.0.0.1:1"
[hosts.a]
base_url = "http://a:1"
models = { "m" = { } }
[routes.r]
hosts = ["a"]
`
func TestV1Defaults(t *testing.T) {
c, err := config.Parse(strings.NewReader(v1Base))
if err != nil {
t.Fatal(err)
}
if c.DB != "crossbar.db" {
t.Errorf("DB default = %q", c.DB)
}
if c.LeaseIdle.Duration != 30*time.Minute {
t.Errorf("LeaseIdle default = %v", c.LeaseIdle.Duration)
}
if c.Retention.Duration != 180*24*time.Hour {
t.Errorf("Retention default = %v", c.Retention.Duration)
}
}
func TestV1Values(t *testing.T) {
c, err := config.Parse(strings.NewReader(`
db = "/var/lib/crossbar/crossbar.db"
lease_idle = "45m"
retention = "30d"
` + v1Base))
if err != nil {
t.Fatal(err)
}
if c.DB != "/var/lib/crossbar/crossbar.db" || c.LeaseIdle.Duration != 45*time.Minute || c.Retention.Duration != 30*24*time.Hour {
t.Errorf("got db %q idle %v retention %v", c.DB, c.LeaseIdle.Duration, c.Retention.Duration)
}
}
func TestDurationAcceptsDays(t *testing.T) {
var d config.Duration
for _, tc := range []struct {
in string
want time.Duration
}{
{"1d", 24 * time.Hour}, {"7d", 7 * 24 * time.Hour}, {"90m", 90 * time.Minute}, {"2h30m", 150 * time.Minute},
} {
if err := d.UnmarshalText([]byte(tc.in)); err != nil || d.Duration != tc.want {
t.Errorf("UnmarshalText(%q) = %v %v, want %v", tc.in, d.Duration, err, tc.want)
}
}
for _, bad := range []string{"1.5d", "d", "3 days", "1d2h"} {
if err := d.UnmarshalText([]byte(bad)); err == nil {
t.Errorf("UnmarshalText(%q) must fail", bad)
}
}
}
func TestV1Validation(t *testing.T) {
for _, tc := range []struct{ name, text, field string }{
{"empty db", "db = \"\"\n" + v1Base, "db"},
{"lease_idle too short", "lease_idle = \"10s\"\n" + v1Base, "lease_idle"},
{"retention too short", "retention = \"12h\"\n" + v1Base, "retention"},
} {
_, err := config.Parse(strings.NewReader(tc.text))
e, ok := config.IsError(err)
if !ok || e.Field != tc.field {
t.Errorf("%s: %v, want *Error on %s", tc.name, err, tc.field)
}
}
}
internal/config/testdata/bad-unknown-key.toml
Vendored
+1 -1
View File
@@ -1,5 +1,5 @@
listen = "127.0.0.1:7777"
lease_idle = "30m"
bogus_key = 1
[hosts.alpha]
base_url = "http://alpha.example:11434"
internal/fingerprint/fingerprint.go
+88
View File
@@ -0,0 +1,88 @@
// Package fingerprint identifies a chat-completions conversation without a
// session id: the system prompt and the first user message never change from
// turn to turn, so hashing them pins the conversation.
package fingerprint
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"strings"
)
// maxPart is the number of bytes of each input string that contributes to the
// key: 4 KiB keeps a huge first message from slowing every turn.
const maxPart = 4096
// Of returns the lowercase hex SHA-256 of the system prompt and the first user
// message of a chat-completions body (first 4 KiB of each, joined with "\n"),
// or "" when the body is not a JSON object with a "messages" array containing a
// user message.
func Of(body []byte) string {
var doc struct {
Messages []json.RawMessage `json:"messages"`
}
if err := json.Unmarshal(body, &doc); err != nil || doc.Messages == nil {
return ""
}
var system, user string
seenSystem, seenUser := false, false
for _, raw := range doc.Messages {
var msg struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
}
if err := json.Unmarshal(raw, &msg); err != nil {
return ""
}
switch msg.Role {
case "system":
if !seenSystem {
seenSystem = true
system = contentText(msg.Content)
}
case "user":
if !seenUser {
seenUser = true
user = contentText(msg.Content)
}
}
}
if !seenUser {
return ""
}
if len(system) > maxPart {
system = system[:maxPart]
}
if len(user) > maxPart {
user = user[:maxPart]
}
sum := sha256.Sum256([]byte(system + "\n" + user))
return hex.EncodeToString(sum[:])
}
// contentText renders a message content value: a JSON string is returned as-is,
// an array of parts is the concatenation of its text parts (other types
// ignored), and anything else is "".
func contentText(raw json.RawMessage) string {
var text string
if err := json.Unmarshal(raw, &text); err == nil {
return text
}
var parts []struct {
Type string `json:"type"`
Text string `json:"text"`
}
if err := json.Unmarshal(raw, &parts); err != nil {
return ""
}
var b strings.Builder
for _, p := range parts {
if p.Type == "text" {
b.WriteString(p.Text)
}
}
return b.String()
}
internal/fingerprint/fingerprint_test.go
+70
View File
@@ -0,0 +1,70 @@
package fingerprint_test
import (
"strings"
"testing"
"git.wntrmute.dev/kyle/crossbar/internal/fingerprint"
)
const conv1 = `{"model":"m","messages":[{"role":"system","content":"You are the project A assistant."},{"role":"user","content":"Add a config loader."},{"role":"assistant","content":"Sure."},{"role":"user","content":"Now tests."}]}`
const conv1later = `{"model":"m","messages":[{"role":"system","content":"You are the project A assistant."},{"role":"user","content":"Add a config loader."},{"role":"assistant","content":"Sure."},{"role":"user","content":"Now tests."},{"role":"assistant","content":"Done."},{"role":"user","content":"And docs."}]}`
const conv2 = `{"model":"m","messages":[{"role":"system","content":"You are the project A assistant."},{"role":"user","content":"Fix the flaky test."}]}`
const conv3 = `{"model":"m","messages":[{"role":"system","content":"You are the project B assistant."},{"role":"user","content":"Add a config loader."}]}`
func TestSameConversationSameKey(t *testing.T) {
a := fingerprint.Of([]byte(conv1))
b := fingerprint.Of([]byte(conv1later))
if a == "" || a != b {
t.Errorf("later turns of one conversation must keep the key: %q vs %q", a, b)
}
if len(a) != 64 || strings.Trim(a, "0123456789abcdef") != "" {
t.Errorf("key must be lowercase hex sha256 (64 chars), got %q", a)
}
}
func TestDifferentConversationsDifferentKeys(t *testing.T) {
a, b, c := fingerprint.Of([]byte(conv1)), fingerprint.Of([]byte(conv2)), fingerprint.Of([]byte(conv3))
if a == b {
t.Errorf("different first user message must change the key")
}
if a == c {
t.Errorf("different system prompt must change the key")
}
}
func TestNoUserMessageIsEmpty(t *testing.T) {
for _, body := range []string{
`{"model":"m","messages":[{"role":"system","content":"only a system prompt"}]}`,
`{"model":"m","messages":[]}`,
`{"model":"m"}`,
`{"input":"an embeddings request"}`,
`not json at all`,
``,
} {
if got := fingerprint.Of([]byte(body)); got != "" {
t.Errorf("Of(%q) = %q, want empty", body, got)
}
}
}
func TestOnlyTheFirstFourKiBCount(t *testing.T) {
long := strings.Repeat("x", 5000)
a := `{"messages":[{"role":"user","content":"` + long + `A"}]}`
b := `{"messages":[{"role":"user","content":"` + long + `B"}]}`
if fingerprint.Of([]byte(a)) != fingerprint.Of([]byte(b)) {
t.Errorf("bytes after the first 4 KiB of a message must not change the key")
}
c := `{"messages":[{"role":"user","content":"A` + long + `"}]}`
if fingerprint.Of([]byte(a)) == fingerprint.Of([]byte(c)) {
t.Errorf("bytes inside the first 4 KiB must change the key")
}
}
func TestContentPartsAreFlattened(t *testing.T) {
plain := `{"messages":[{"role":"user","content":"hello world"}]}`
parts := `{"messages":[{"role":"user","content":[{"type":"text","text":"hello world"}]}]}`
if fingerprint.Of([]byte(plain)) != fingerprint.Of([]byte(parts)) {
t.Errorf("a content array of text parts must fingerprint like the joined text")
}
}
internal/lease/lease.go
+330
View File
@@ -0,0 +1,330 @@
// Package lease is crossbar's sticky placement table. It remembers which host each
// conversation (and each route) is on and keeps it there unless the host is
// unhealthy, the lease has been idle past lease_idle, or an operator releases or
// pins the route. Every change is written through to a Persister and replayed back
// at start so a restart does not reshuffle sessions.
package lease
import (
"errors"
"fmt"
"sort"
"sync"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
var (
ErrNoHost = errors.New("lease: no usable host")
ErrPinnedDown = errors.New("lease: pinned host is not healthy")
ErrUnknownHost = errors.New("lease: unknown host")
)
// Key identifies one conversation: a route, the fingerprint of its first request
// (empty for a request with no user message, which leases the route itself), and
// the model it wants.
type Key struct {
Route, FP, Model string
}
// Lease is one routed model on one host, in memory.
type Lease struct {
Key
Host string
State store.State
Created, LastUsed time.Time
}
// Persister is the durable half of the table: leases, the pin row, and the event
// log. *store.Store satisfies it.
type Persister interface {
SaveLease(store.Lease) error
DeleteLease(route, fp, model string) error
ListLeases() ([]store.Lease, error)
RecordEvent(store.LeaseEvent) error
}
// Hosts is the live health view the table consults before placing a lease.
type Hosts interface {
Healthy(name string) bool
Draining(name string) bool
}
// Chooser picks a host for a new lease among the eligible candidates.
type Chooser interface {
Choose(candidates []string, model string) (string, bool)
}
type Table struct {
mu sync.Mutex
leases map[Key]*Lease // active leases, keyed by (route, fp, model)
pins map[string]string // route -> pinned host
seen map[string]map[string]bool // route -> candidate hosts ever asked for or stored
p Persister
hosts Hosts
choose Chooser
idle time.Duration
}
// New builds a table and loads its state. Rows with FP=="" && Model=="" &&
// State==Pinned are pins; every other row becomes an active lease and its host
// counts as a candidate already seen for the route.
func New(p Persister, h Hosts, c Chooser, idle time.Duration) (*Table, error) {
loads, err := p.ListLeases()
if err != nil {
return nil, fmt.Errorf("lease: list leases: %w", err)
}
t := &Table{
leases: make(map[Key]*Lease),
pins: make(map[string]string),
seen: make(map[string]map[string]bool),
p: p,
hosts: h,
choose: c,
idle: idle,
}
for _, l := range loads {
if l.FP == "" && l.Model == "" && l.State == store.Pinned {
t.pins[l.Route] = l.Host
} else {
t.leases[Key{l.Route, l.FP, l.Model}] = &Lease{Key{l.Route, l.FP, l.Model}, l.Host, l.State, l.Created, l.LastUsed}
}
if t.seen[l.Route] == nil {
t.seen[l.Route] = make(map[string]bool)
}
t.seen[l.Route][l.Host] = true
}
return t, nil
}
func (l *Lease) store() store.Lease {
return store.Lease{Route: l.Route, FP: l.FP, Model: l.Model, Host: l.Host, State: l.State, Created: l.Created, LastUsed: l.LastUsed}
}
// Acquire places k, keeping it sticky. See the task's Acquire ordering: pinned
// route, existing lease, inherit the route's host, else choose. Only one lease is
// created per call, and a failed save is rolled back in memory.
func (t *Table) Acquire(k Key, candidates []string, now time.Time) (host string, reused bool, err error) {
t.mu.Lock()
defer t.mu.Unlock()
if t.seen[k.Route] == nil {
t.seen[k.Route] = make(map[string]bool)
}
for _, c := range candidates {
t.seen[k.Route][c] = true
}
// unhealthyFrom is set when an existing lease sat on a dead host; the next
// create records an unhealthy move instead of a fresh new.
var unhealthyFrom string
reason := store.ReasonNew
// 1. Pinned route.
if pin, ok := t.pins[k.Route]; ok {
if !t.hosts.Healthy(pin) {
return "", false, ErrPinnedDown
}
if _, exists := t.leases[k]; exists {
return pin, true, nil
}
l := &Lease{k, pin, store.Active, now, now}
t.leases[k] = l
if err := t.save(l); err != nil {
delete(t.leases, k)
return "", false, err
}
t.event(now, k, store.ReasonNew, "", pin)
return pin, false, nil
}
// 2. Existing lease for k.
if l, exists := t.leases[k]; exists {
if t.hosts.Healthy(l.Host) {
l.LastUsed = now
if err := t.save(l); err != nil {
return "", false, err
}
return l.Host, true, nil
}
unhealthyFrom = l.Host
}
// 3. Inherit the route's own host when a fingerprinted request can start
// where the route already lives.
if k.FP != "" {
rk := Key{k.Route, "", k.Model}
if rl, exists := t.leases[rk]; exists && t.hosts.Healthy(rl.Host) {
l := &Lease{k, rl.Host, store.Active, now, now}
t.leases[k] = l
if err := t.save(l); err != nil {
delete(t.leases, k)
return "", false, err
}
if unhealthyFrom != "" {
reason = store.ReasonUnhealthy
}
t.event(now, k, reason, unhealthyFrom, rl.Host)
return rl.Host, true, nil
}
}
// 4. Choose among healthy, non-draining candidates.
filtered := make([]string, 0, len(candidates))
for _, c := range candidates {
if t.hosts.Healthy(c) && !t.hosts.Draining(c) {
filtered = append(filtered, c)
}
}
host, ok := t.choose.Choose(filtered, k.Model)
if !ok {
return "", false, ErrNoHost
}
l := &Lease{k, host, store.Active, now, now}
t.leases[k] = l
if err := t.save(l); err != nil {
delete(t.leases, k)
return "", false, err
}
if unhealthyFrom != "" {
reason = store.ReasonUnhealthy
}
t.event(now, k, reason, unhealthyFrom, host)
return host, false, nil
}
// Candidates records hosts as seen for route (idempotent), so Pin can accept a host the route
// is configured for before any request has used it. cmd/crossbar calls it for every route at
// start; the admin handler calls it before Pin.
func (t *Table) Candidates(route string, hosts []string) {
t.mu.Lock()
defer t.mu.Unlock()
if t.seen[route] == nil {
t.seen[route] = make(map[string]bool)
}
for _, h := range hosts {
t.seen[route][h] = true
}
}
// save writes a lease through, failing the call on a persister error.
func (t *Table) save(l *Lease) error {
if err := t.p.SaveLease(l.store()); err != nil {
return fmt.Errorf("lease: %w", err)
}
return nil
}
// event appends a lease event. from is empty for a fresh placement.
func (t *Table) event(now time.Time, k Key, reason, from, to string) {
_ = t.p.RecordEvent(store.LeaseEvent{TS: now, Route: k.Route, Model: k.Model, FromHost: from, ToHost: to, Reason: reason})
}
// ExpireIdle removes leases idle longer than lease_idle (never pins, which are
// not in the lease map) and records an idle event for each. It returns how many
// it removed.
func (t *Table) ExpireIdle(now time.Time) int {
t.mu.Lock()
defer t.mu.Unlock()
var idle []Key
for k, l := range t.leases {
if now.Sub(l.LastUsed) > t.idle {
idle = append(idle, k)
}
}
for _, k := range idle {
l := t.leases[k]
delete(t.leases, k)
_ = t.p.DeleteLease(k.Route, k.FP, k.Model)
t.event(now, k, store.ReasonIdle, l.Host, "")
}
return len(idle)
}
// Pin routes route to host. host must be a candidate the table has seen for the
// route, else ErrUnknownHost. It records a pin event, stores the pin row, and
// deletes the route's existing leases on other hosts so the next turn lands on
// the pin.
func (t *Table) Pin(route, host string, now time.Time) error {
t.mu.Lock()
defer t.mu.Unlock()
if t.seen[route] == nil || !t.seen[route][host] {
return ErrUnknownHost
}
t.event(now, Key{Route: route}, store.ReasonPin, "", host)
if err := t.p.SaveLease(store.Lease{Route: route, FP: "", Model: "", Host: host, State: store.Pinned, Created: now, LastUsed: now}); err != nil {
return fmt.Errorf("lease: %w", err)
}
t.pins[route] = host
for k, l := range t.leases {
if k.Route == route && l.Host != host {
delete(t.leases, k)
_ = t.p.DeleteLease(k.Route, k.FP, k.Model)
t.event(now, k, store.ReasonPin, l.Host, host)
}
}
return nil
}
// Unpin clears the pin for route and records a release. Existing leases stay put.
func (t *Table) Unpin(route string) {
t.mu.Lock()
defer t.mu.Unlock()
delete(t.pins, route)
_ = t.p.DeleteLease(route, "", "")
t.event(time.Now(), Key{Route: route}, store.ReasonRelease, "", "")
}
// Pinned returns the pinned host for route, or "" if it is not pinned.
func (t *Table) Pinned(route string) string {
t.mu.Lock()
defer t.mu.Unlock()
return t.pins[route]
}
// Release drops every lease (not the pin) of route and records a release event
// per dropped lease. It returns how many were removed.
func (t *Table) Release(route string) int {
t.mu.Lock()
defer t.mu.Unlock()
var keys []Key
for k := range t.leases {
if k.Route == route {
keys = append(keys, k)
}
}
for _, k := range keys {
delete(t.leases, k)
_ = t.p.DeleteLease(k.Route, k.FP, k.Model)
t.event(time.Now(), k, store.ReasonRelease, "", "")
}
return len(keys)
}
// Snapshot returns copies of the active leases, sorted by route, fp, model. Pins
// are excluded.
func (t *Table) Snapshot() []Lease {
t.mu.Lock()
defer t.mu.Unlock()
out := make([]Lease, 0, len(t.leases))
for _, l := range t.leases {
out = append(out, *l)
}
sort.Slice(out, func(i, j int) bool {
if out[i].Route != out[j].Route {
return out[i].Route < out[j].Route
}
if out[i].FP != out[j].FP {
return out[i].FP < out[j].FP
}
return out[i].Model < out[j].Model
})
return out
}
internal/lease/lease_test.go
+308
View File
@@ -0,0 +1,308 @@
package lease_test
import (
"errors"
"sync"
"testing"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/lease"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
// memPersister is an in-memory Persister that also counts writes.
type memPersister struct {
mu sync.Mutex
leases map[[3]string]store.Lease
events []store.LeaseEvent
saves int
}
func newPersister() *memPersister { return &memPersister{leases: map[[3]string]store.Lease{}} }
func (m *memPersister) SaveLease(l store.Lease) error {
m.mu.Lock()
defer m.mu.Unlock()
m.saves++
m.leases[[3]string{l.Route, l.FP, l.Model}] = l
return nil
}
func (m *memPersister) DeleteLease(route, fp, model string) error {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.leases, [3]string{route, fp, model})
return nil
}
func (m *memPersister) ListLeases() ([]store.Lease, error) {
m.mu.Lock()
defer m.mu.Unlock()
out := []store.Lease{}
for _, l := range m.leases {
out = append(out, l)
}
return out, nil
}
func (m *memPersister) RecordEvent(e store.LeaseEvent) error {
m.mu.Lock()
defer m.mu.Unlock()
m.events = append(m.events, e)
return nil
}
func (m *memPersister) reasons() []string {
m.mu.Lock()
defer m.mu.Unlock()
var r []string
for _, e := range m.events {
r = append(r, e.Reason)
}
return r
}
// world is a hand-set view of hosts plus a chooser that returns a fixed answer.
type world struct {
mu sync.Mutex
healthy map[string]bool
draining map[string]bool
pick string
picks []string // candidates seen by Choose, for assertions
}
func (w *world) Healthy(name string) bool { w.mu.Lock(); defer w.mu.Unlock(); return w.healthy[name] }
func (w *world) Draining(name string) bool { w.mu.Lock(); defer w.mu.Unlock(); return w.draining[name] }
func (w *world) Choose(candidates []string, model string) (string, bool) {
w.mu.Lock()
defer w.mu.Unlock()
w.picks = append([]string{}, candidates...)
for _, c := range candidates {
if c == w.pick {
return c, true
}
}
if len(candidates) > 0 {
return candidates[0], true
}
return "", false
}
var t0 = time.Date(2026, 9, 25, 10, 0, 0, 0, time.UTC)
func newTable(t *testing.T, p *memPersister, w *world) *lease.Table {
tbl, err := lease.New(p, w, w, 30*time.Minute)
if err != nil {
t.Fatal(err)
}
return tbl
}
func TestNewLeaseThenSticky(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "beta"}
tbl := newTable(t, p, w)
k := lease.Key{Route: "r", FP: "conv1", Model: "m"}
host, reused, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0)
if err != nil || host != "beta" || reused {
t.Fatalf("first: %q %v %v", host, reused, err)
}
w.pick = "alpha" // the chooser would now prefer alpha; the lease must hold
for i := 1; i <= 5; i++ {
host, reused, err = tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(time.Duration(i)*time.Minute))
if err != nil || host != "beta" || !reused {
t.Fatalf("turn %d: %q reused=%v %v, want beta reused", i, host, reused, err)
}
}
if got := p.reasons(); len(got) != 1 || got[0] != store.ReasonNew {
t.Errorf("events = %v, want one 'new'", got)
}
snap := tbl.Snapshot()
if len(snap) != 1 || snap[0].Host != "beta" || !snap[0].LastUsed.Equal(t0.Add(5*time.Minute)) {
t.Errorf("snapshot = %+v", snap)
}
if p.saves < 2 {
t.Errorf("LastUsed must be written through (saves=%d)", p.saves)
}
}
func TestUnhealthyHostMovesTheLease(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
tbl := newTable(t, p, w)
k := lease.Key{Route: "r", FP: "c", Model: "m"}
if host, _, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0); host != "alpha" {
t.Fatalf("first: %q", host)
}
w.mu.Lock()
w.healthy["alpha"] = false
w.pick = "beta"
w.mu.Unlock()
host, reused, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(time.Minute))
if err != nil || host != "beta" || reused {
t.Fatalf("after alpha down: %q reused=%v %v", host, reused, err)
}
if got := p.reasons(); len(got) != 2 || got[1] != store.ReasonUnhealthy {
t.Errorf("events = %v, want [new unhealthy]", got)
}
if len(w.picks) != 1 || w.picks[0] != "beta" {
t.Errorf("Choose must not see the unhealthy host: %v", w.picks)
}
}
func TestIdleExpiry(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
tbl := newTable(t, p, w)
k := lease.Key{Route: "r", FP: "c", Model: "m"}
tbl.Acquire(k, []string{"alpha", "beta"}, t0)
if n := tbl.ExpireIdle(t0.Add(29 * time.Minute)); n != 0 {
t.Errorf("expired %d before lease_idle", n)
}
if n := tbl.ExpireIdle(t0.Add(31 * time.Minute)); n != 1 {
t.Errorf("expired %d after lease_idle, want 1", n)
}
if got := p.reasons(); got[len(got)-1] != store.ReasonIdle {
t.Errorf("events = %v, want idle last", got)
}
w.pick = "beta"
if host, reused, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(32*time.Minute)); host != "beta" || reused {
t.Errorf("after expiry a new lease is chosen: %q reused=%v", host, reused)
}
if l, _ := p.ListLeases(); len(l) != 1 {
t.Errorf("persister holds %d leases, want 1", len(l))
}
}
func TestFingerprintInheritsRouteLease(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "beta"}
tbl := newTable(t, p, w)
// A request without a fingerprint (no user message) leases the route itself…
if host, _, _ := tbl.Acquire(lease.Key{Route: "r", FP: "", Model: "m"}, []string{"alpha", "beta"}, t0); host != "beta" {
t.Fatalf("route lease: %q", host)
}
w.pick = "alpha"
// …and a new conversation on that route starts where the route already is.
host, reused, err := tbl.Acquire(lease.Key{Route: "r", FP: "conv", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Second))
if err != nil || host != "beta" || !reused {
t.Errorf("fingerprint lease must inherit the route's host: %q reused=%v %v", host, reused, err)
}
if len(tbl.Snapshot()) != 2 {
t.Errorf("both the route lease and the conversation lease exist: %+v", tbl.Snapshot())
}
}
func TestPinAndUnpin(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
tbl := newTable(t, p, w)
k := lease.Key{Route: "r", FP: "c", Model: "m"}
tbl.Acquire(k, []string{"alpha", "beta"}, t0)
if err := tbl.Pin("r", "beta", t0.Add(time.Minute)); err != nil {
t.Fatal(err)
}
host, _, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(2*time.Minute))
if err != nil || host != "beta" {
t.Fatalf("pinned route must go to beta: %q %v", host, err)
}
host, _, err = tbl.Acquire(lease.Key{Route: "r", FP: "other", Model: "m"}, []string{"alpha", "beta"}, t0.Add(2*time.Minute))
if err != nil || host != "beta" {
t.Fatalf("new conversations on a pinned route go to the pin too: %q %v", host, err)
}
w.mu.Lock()
w.healthy["beta"] = false
w.mu.Unlock()
if _, _, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(3*time.Minute)); !errors.Is(err, lease.ErrPinnedDown) {
t.Errorf("a pinned host that is down is ErrPinnedDown, never a silent move: %v", err)
}
if err := tbl.Pin("r", "nobody", t0); !errors.Is(err, lease.ErrUnknownHost) {
t.Errorf("pinning to a host not in the candidates of any lease: %v, want ErrUnknownHost", err)
}
tbl.Unpin("r")
w.mu.Lock()
w.healthy["beta"] = true
w.mu.Unlock()
if host, _, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(4*time.Minute)); host != "beta" {
t.Errorf("after unpin the existing lease (on beta) simply continues: %q", host)
}
// Events: a pin event naming beta must exist, and the unpin's release event must come after it.
// Acquires under the pin may record their own events in between; their number is not fixed here.
got := p.reasons()
pinAt, releaseAt := -1, -1
for i, r := range got {
if r == store.ReasonPin && pinAt < 0 {
pinAt = i
}
if r == store.ReasonRelease {
releaseAt = i
}
}
if pinAt < 0 || releaseAt < pinAt {
t.Errorf("events = %v, want a pin event followed later by a release event", got)
}
p.mu.Lock()
if pinAt >= 0 && p.events[pinAt].ToHost != "beta" {
t.Errorf("pin event = %+v, want ToHost beta", p.events[pinAt])
}
p.mu.Unlock()
}
func TestDrainKeepsExistingRefusesNew(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, draining: map[string]bool{}, pick: "alpha"}
tbl := newTable(t, p, w)
k := lease.Key{Route: "r", FP: "c", Model: "m"}
tbl.Acquire(k, []string{"alpha", "beta"}, t0)
w.mu.Lock()
w.draining["alpha"] = true
w.mu.Unlock()
if host, reused, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(time.Minute)); host != "alpha" || !reused {
t.Errorf("an existing lease on a draining host continues: %q reused=%v", host, reused)
}
host, _, err := tbl.Acquire(lease.Key{Route: "r2", FP: "x", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Minute))
if err != nil || host != "beta" {
t.Errorf("a new lease avoids the draining host: %q %v", host, err)
}
if len(w.picks) != 1 || w.picks[0] != "beta" {
t.Errorf("Choose must not see the draining host: %v", w.picks)
}
if _, _, err := tbl.Acquire(lease.Key{Route: "r3", FP: "y", Model: "m"}, []string{"alpha"}, t0); !errors.Is(err, lease.ErrNoHost) {
t.Errorf("only draining candidates: %v, want ErrNoHost", err)
}
}
func TestReleaseRoute(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
tbl := newTable(t, p, w)
tbl.Acquire(lease.Key{Route: "r", FP: "a", Model: "m"}, []string{"alpha", "beta"}, t0)
tbl.Acquire(lease.Key{Route: "r", FP: "b", Model: "m"}, []string{"alpha", "beta"}, t0)
tbl.Acquire(lease.Key{Route: "other", FP: "c", Model: "m"}, []string{"alpha", "beta"}, t0)
if n := tbl.Release("r"); n != 2 {
t.Errorf("Release removed %d, want 2", n)
}
if n := tbl.Release("r"); n != 0 {
t.Errorf("second Release removed %d", n)
}
if l, _ := p.ListLeases(); len(l) != 1 || l[0].Route != "other" {
t.Errorf("persister after release: %+v", l)
}
w.pick = "beta"
if host, reused, _ := tbl.Acquire(lease.Key{Route: "r", FP: "a", Model: "m"}, []string{"alpha", "beta"}, t0); host != "beta" || reused {
t.Errorf("after release the route is re-chosen: %q reused=%v", host, reused)
}
}
func TestLoadsFromPersister(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
_ = p.SaveLease(store.Lease{Route: "r", FP: "c", Model: "m", Host: "beta", State: store.Active, Created: t0, LastUsed: t0})
_ = p.SaveLease(store.Lease{Route: "pinned", FP: "", Model: "", Host: "beta", State: store.Pinned, Created: t0, LastUsed: t0})
tbl := newTable(t, p, w)
if host, reused, _ := tbl.Acquire(lease.Key{Route: "r", FP: "c", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Second)); host != "beta" || !reused {
t.Errorf("a restart must not reshuffle: %q reused=%v", host, reused)
}
if host, _, _ := tbl.Acquire(lease.Key{Route: "pinned", FP: "new", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Second)); host != "beta" {
t.Errorf("a pin survives a restart: %q", host)
}
}
func TestNoCandidates(t *testing.T) {
p, w := newPersister(), &world{healthy: map[string]bool{}, pick: ""}
tbl := newTable(t, p, w)
if _, _, err := tbl.Acquire(lease.Key{Route: "r", FP: "c", Model: "m"}, []string{"alpha"}, t0); !errors.Is(err, lease.ErrNoHost) {
t.Errorf("no healthy host: %v, want ErrNoHost", err)
}
if len(tbl.Snapshot()) != 0 {
t.Errorf("a failed acquire must not create a lease")
}
}
internal/limiter/limiter.go
+180
View File
@@ -0,0 +1,180 @@
// Package limiter hands out at most `parallel` concurrent slots per (host, model) and lets at
// most `queue_max` requests wait in a FIFO. A request that finds the queue full is refused at
// once so the caller can retry elsewhere; a waiting request can cancel and leave without leaking
// a slot or a queue place.
package limiter
import (
"context"
"errors"
"sync"
"time"
)
// ErrQueueFull is returned by Acquire when the queue is already at queue_max; the caller may try a
// different host.
var ErrQueueFull = errors.New("queue full")
// defaults is what an unconfigured (host, model) behaves as: one slot, no waiting room.
const (
defaultParallel = 1
defaultQueueMax = 0
)
// pair holds the live state for one (host, model): how many slots exist, how many are taken, and
// the FIFO of waiters. All fields are guarded by Limiter.mu.
type pair struct {
parallel int
queueMax int
inflight int
waiters []chan struct{}
configured bool
}
// Limiter tracks one pair per (host, model). Safe for concurrent use.
type Limiter struct {
mu sync.Mutex
pairs map[pairKey]*pair
}
type pairKey struct {
host string
model string
}
// New returns an empty Limiter.
func New() *Limiter {
return &Limiter{pairs: make(map[pairKey]*pair)}
}
// Configure sets the slot and queue limits for one (host, model). It may be called before any
// request or after one has created the pair with the defaults; either way the limits apply.
func (l *Limiter) Configure(host, model string, parallel, queueMax int) {
l.mu.Lock()
defer l.mu.Unlock()
p := l.pairLocked(host, model)
p.parallel = parallel
p.queueMax = queueMax
p.configured = true
}
// pairLocked returns the pair for (host, model), creating it with the unconfigured defaults if it
// does not exist yet. The caller holds l.mu.
func (l *Limiter) pairLocked(host, model string) *pair {
k := pairKey{host, model}
p := l.pairs[k]
if p == nil {
p = &pair{parallel: defaultParallel, queueMax: defaultQueueMax}
l.pairs[k] = p
}
return p
}
// Acquire blocks until a slot is held. It returns a release that gives the slot back exactly once
// (a second call is a no-op), how long the caller spent in the queue, and an error: ErrQueueFull
// when the queue is already full (returned immediately, without waiting), or ctx.Err() when the
// context ends while waiting.
func (l *Limiter) Acquire(ctx context.Context, host, model string) (release func(), waited time.Duration, err error) {
l.mu.Lock()
start := time.Now()
p := l.pairLocked(host, model)
if p.inflight < p.parallel {
p.inflight++
l.mu.Unlock()
return l.release(p), time.Since(start), nil
}
if len(p.waiters) >= p.queueMax {
l.mu.Unlock()
return nil, time.Since(start), ErrQueueFull
}
waiter := make(chan struct{})
p.waiters = append(p.waiters, waiter)
l.mu.Unlock()
select {
case <-ctx.Done():
// The slot may have been handed to us the instant the context ended; release it either
// way so neither a slot nor a queue place leaks.
l.mu.Lock()
if !p.dropWaiter(waiter) {
l.mu.Unlock()
l.release(p)
return nil, time.Since(start), ctx.Err()
}
l.mu.Unlock()
return nil, time.Since(start), ctx.Err()
case <-waiter:
return l.release(p), time.Since(start), nil
}
}
// release returns the function the caller holds for a slot: it hands the slot to the next waiter
// if one is waiting, otherwise it frees the slot. It is safe to call through the sync.Once that
// Acquire wrapped it in.
func (l *Limiter) release(p *pair) func() {
var once sync.Once
return func() {
once.Do(func() {
l.mu.Lock()
defer l.mu.Unlock()
if len(p.waiters) > 0 {
next := p.waiters[0]
p.waiters = p.waiters[1:]
close(next)
return
}
p.inflight--
})
}
}
// dropWaiter removes w from the middle of the queue. It reports whether w was there; false means
// the slot was already handed to w (its channel closed) and the caller must give it back.
func (p *pair) dropWaiter(w chan struct{}) bool {
for i, cw := range p.waiters {
if cw == w {
p.waiters = append(p.waiters[:i], p.waiters[i+1:]...)
return true
}
}
return false
}
// InFlight returns the number of held slots for (host, model).
func (l *Limiter) InFlight(host, model string) int {
l.mu.Lock()
defer l.mu.Unlock()
p := l.pairs[pairKey{host, model}]
if p == nil {
return 0
}
return p.inflight
}
// Queued returns the number of requests waiting for (host, model).
func (l *Limiter) Queued(host, model string) int {
l.mu.Lock()
defer l.mu.Unlock()
p := l.pairs[pairKey{host, model}]
if p == nil {
return 0
}
return len(p.waiters)
}
// FreeSlots sums the unused slots over the host's configured models, never counting below zero for
// one. An unknown host has none.
func (l *Limiter) FreeSlots(host string) int {
l.mu.Lock()
defer l.mu.Unlock()
total := 0
for k, p := range l.pairs {
if k.host != host || !p.configured {
continue
}
if free := p.parallel - p.inflight; free > 0 {
total += free
}
}
return total
}
internal/limiter/limiter_test.go
+178
View File
@@ -0,0 +1,178 @@
package limiter_test
import (
"context"
"errors"
"sync"
"testing"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/limiter"
)
func TestParallelAndQueue(t *testing.T) {
l := limiter.New()
l.Configure("alpha", "m", 2, 1) // two slots, one waiting place
ctx := context.Background()
rel1, w1, err := l.Acquire(ctx, "alpha", "m")
if err != nil || w1 > 50*time.Millisecond {
t.Fatalf("first acquire: err %v waited %v", err, w1)
}
rel2, _, err := l.Acquire(ctx, "alpha", "m")
if err != nil {
t.Fatalf("second acquire: %v", err)
}
if l.InFlight("alpha", "m") != 2 || l.FreeSlots("alpha") != 0 {
t.Errorf("in flight %d free %d, want 2 and 0", l.InFlight("alpha", "m"), l.FreeSlots("alpha"))
}
// Third waits in the queue.
got3 := make(chan error, 1)
go func() {
rel, waited, err := l.Acquire(ctx, "alpha", "m")
if err == nil {
defer rel()
if waited < 40*time.Millisecond {
err = errors.New("third acquire did not wait")
}
}
got3 <- err
}()
time.Sleep(20 * time.Millisecond)
if l.Queued("alpha", "m") != 1 {
t.Errorf("queued = %d, want 1", l.Queued("alpha", "m"))
}
// Fourth finds the queue full and is refused at once.
start := time.Now()
_, _, err = l.Acquire(ctx, "alpha", "m")
if !errors.Is(err, limiter.ErrQueueFull) {
t.Fatalf("fourth acquire: %v, want ErrQueueFull", err)
}
if time.Since(start) > 50*time.Millisecond {
t.Errorf("a full queue must refuse immediately, took %v", time.Since(start))
}
time.Sleep(30 * time.Millisecond)
rel1() // frees a slot: the queued third proceeds
select {
case err := <-got3:
if err != nil {
t.Fatalf("third: %v", err)
}
case <-time.After(time.Second):
t.Fatal("queued acquire did not proceed after a release")
}
rel2()
if l.InFlight("alpha", "m") != 0 || l.Queued("alpha", "m") != 0 {
t.Errorf("after releases: inflight %d queued %d", l.InFlight("alpha", "m"), l.Queued("alpha", "m"))
}
}
func TestReleaseIsIdempotent(t *testing.T) {
l := limiter.New()
l.Configure("h", "m", 1, 0)
rel, _, err := l.Acquire(context.Background(), "h", "m")
if err != nil {
t.Fatal(err)
}
rel()
rel() // a second call must not free a slot that was never taken
if l.InFlight("h", "m") != 0 {
t.Errorf("in flight %d after double release", l.InFlight("h", "m"))
}
if _, _, err := l.Acquire(context.Background(), "h", "m"); err != nil {
t.Errorf("slot must be free again: %v", err)
}
}
func TestCancelWhileQueuedLeaksNothing(t *testing.T) {
l := limiter.New()
l.Configure("h", "m", 1, 2)
rel, _, err := l.Acquire(context.Background(), "h", "m")
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { _, _, err := l.Acquire(ctx, "h", "m"); done <- err }()
time.Sleep(20 * time.Millisecond)
cancel()
select {
case err := <-done:
if !errors.Is(err, context.Canceled) {
t.Fatalf("cancelled acquire returned %v", err)
}
case <-time.After(time.Second):
t.Fatal("cancelled acquire did not return")
}
if l.Queued("h", "m") != 0 {
t.Errorf("queued = %d after cancel", l.Queued("h", "m"))
}
rel()
if l.InFlight("h", "m") != 0 {
t.Errorf("in flight %d, the cancelled waiter must not have taken the slot", l.InFlight("h", "m"))
}
}
func TestQueueIsFIFO(t *testing.T) {
l := limiter.New()
l.Configure("h", "m", 1, 8)
rel, _, err := l.Acquire(context.Background(), "h", "m")
if err != nil {
t.Fatal(err)
}
var mu sync.Mutex
var order []int
var wg sync.WaitGroup
for i := 1; i <= 4; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
r, _, err := l.Acquire(context.Background(), "h", "m")
if err != nil {
t.Errorf("waiter %d: %v", i, err)
return
}
mu.Lock()
order = append(order, i)
mu.Unlock()
time.Sleep(5 * time.Millisecond)
r()
}(i)
time.Sleep(15 * time.Millisecond) // stagger arrivals so the order is defined
}
rel()
wg.Wait()
if len(order) != 4 || order[0] != 1 || order[1] != 2 || order[2] != 3 || order[3] != 4 {
t.Errorf("waiters proceeded in order %v, want [1 2 3 4]", order)
}
}
func TestUnconfiguredPairIsOneSlotNoQueue(t *testing.T) {
l := limiter.New()
rel, _, err := l.Acquire(context.Background(), "x", "y")
if err != nil {
t.Fatal(err)
}
defer rel()
if _, _, err := l.Acquire(context.Background(), "x", "y"); !errors.Is(err, limiter.ErrQueueFull) {
t.Errorf("second acquire on an unconfigured pair: %v, want ErrQueueFull", err)
}
}
func TestFreeSlotsSumsModels(t *testing.T) {
l := limiter.New()
l.Configure("h", "a", 4, 0)
l.Configure("h", "b", 2, 0)
if got := l.FreeSlots("h"); got != 6 {
t.Fatalf("free = %d, want 6", got)
}
rel, _, _ := l.Acquire(context.Background(), "h", "a")
defer rel()
if got := l.FreeSlots("h"); got != 5 {
t.Errorf("free = %d, want 5", got)
}
if l.FreeSlots("nobody") != 0 {
t.Errorf("unknown host has no slots")
}
}
internal/proxy/forward.go
+166
View File
@@ -0,0 +1,166 @@
package proxy
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httputil"
"net/url"
"strings"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
// forward builds the reverse proxy for one host, tees the response, records the accounting row, and
// logs. leaseState is "new" or "reused"; waited is the time spent in the queue.
func (p *Handler) forward(w http.ResponseWriter, r *http.Request, route, host, leaseState, rest, fp, model string, started time.Time, waited time.Duration) {
hostCfg, ok := p.cfg.Hosts[host]
if !ok {
p.writeError(w, http.StatusBadGateway, "upstream failed")
return
}
target, err := url.Parse(hostCfg.BaseURL)
if err != nil {
p.writeError(w, http.StatusBadGateway, "upstream failed")
return
}
rev := &forwardState{started: started}
rp := newReverseProxy(p.health, host, leaseState, target, rest, rev)
rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
rp.ServeHTTP(rec, r)
total := time.Since(started)
req := store.Request{
Route: route,
FP: fp,
Model: model,
Host: host,
Started: started,
QueuedMs: waited.Milliseconds(),
TTFBMs: ttfbMs(rev),
TotalMs: total.Milliseconds(),
Status: rec.status,
Streamed: rev.streamed,
}
if rev.tee != nil {
prompt, cached, completion := rev.tee.tokens()
req.PromptTokens = int64(prompt)
req.CachedTokens = int64(cached)
req.CompletionTokens = int64(completion)
}
p.writeRecord(req)
fp8 := fp
if len(fp8) > 8 {
fp8 = fp8[:8]
}
p.log.Info("request",
"route", route,
"host", host,
"method", r.Method,
"path", rest,
"status", rec.status,
"lease", leaseState,
"queued_ms", waited.Milliseconds(),
"fp", fp8,
"ms", total.Milliseconds(),
)
}
// leaseState is "reused" when the lease already held the conversation, else "new".
func leaseState(reused bool) string {
if reused {
return "reused"
}
return "new"
}
// ttfbMs is the time from request start to the response head; zero when the head never arrived.
func ttfbMs(rev *forwardState) int64 {
if rev.ttfb.IsZero() || rev.ttfb.Before(rev.started) {
return 0
}
return rev.ttfb.Sub(rev.started).Milliseconds()
}
// forwardState carries, across one forward, when the request started, when the head arrived, whether
// the response streamed, and the tee that scanned it.
type forwardState struct {
started time.Time
ttfb time.Time
streamed bool
tee *tee
}
// 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()
}
}
// writeError answers with a JSON {"error":"…"} body.
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})
}
// writeRecord writes one accounting row, logging (never returning) a recorder error.
func (p *Handler) writeRecord(req store.Request) {
if p.rec == nil {
return
}
if err := p.rec.RecordRequest(req); err != nil {
p.log.Error("record request", "err", err)
}
}
// 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, tees the response for usage/timings, and marks the host down on any transport error
// other than a client disconnect.
func newReverseProxy(h Health, host, leaseState string, target *url.URL, rest string, rev *forwardState) *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, host)
resp.Header.Set(LeaseHeader, leaseState)
rev.ttfb = time.Now()
rev.streamed = strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
t := newTee(resp.Body, rev.streamed)
resp.Body = t
rev.tee = t
return nil
},
ErrorHandler: func(w http.ResponseWriter, req *http.Request, err error) {
if errors.Is(err, context.Canceled) {
return
}
h.MarkDown(host, err.Error())
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadGateway)
_ = json.NewEncoder(w).Encode(map[string]string{"error": "upstream failed", "host": host})
},
}
}
internal/proxy/helpers_test.go
+211
View File
@@ -0,0 +1,211 @@
package proxy_test
// Test scaffolding shared by proxy_test.go and recorder_test.go: the fake health table, the fake
// llama-server upstream, and the rig that builds a whole crossbar over real HTTP.
import (
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/config"
"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/proxy"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
// fakeHealth is a hand-set health table that also records MarkDown calls. It lived in the v0
// proxy_test.go; the v1 given test replaces that file, so recorder_test.go (which still exercises
// the nil-lease path through proxy.New) needs it here.
type fakeHealth struct {
mu sync.Mutex
st map[string]health.Status
marked []string
}
func (f *fakeHealth) Get(name string) (health.Status, bool) {
f.mu.Lock()
defer f.mu.Unlock()
s, ok := f.st[name]
return s, ok
}
func (f *fakeHealth) MarkDown(name, reason string) {
f.mu.Lock()
defer f.mu.Unlock()
f.marked = append(f.marked, name)
s := f.st[name]
s.Healthy = false
s.LastErr = reason
f.st[name] = s
}
func (f *fakeHealth) markedHosts() []string {
f.mu.Lock()
defer f.mu.Unlock()
return append([]string{}, f.marked...)
}
// upstream is a llama-server stand-in: streams N chunks with a delay, reports usage/timings in
// the final chunk, counts requests, and can be slowed down or killed.
type upstream struct {
name string
srv *httptest.Server
hits atomic.Int32
delay time.Duration
mu sync.Mutex
last recorded
}
type recorded struct{ method, path, host, xff, body string }
func newUpstream(t *testing.T, name string) *upstream {
u := &upstream{name: name}
mux := http.NewServeMux()
mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, `{"status":"ok"}`) })
mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
u.mu.Lock()
u.last = recorded{r.Method, r.URL.RequestURI(), r.Host, r.Header.Get("X-Forwarded-For"), ""}
u.mu.Unlock()
fmt.Fprint(w, `{"object":"list","data":[{"id":"shared"},{"id":"`+name+`-only"}]}`)
})
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
u.hits.Add(1)
b, _ := io.ReadAll(r.Body)
u.mu.Lock()
u.last = recorded{r.Method, r.URL.RequestURI(), r.Host, r.Header.Get("X-Forwarded-For"), string(b)}
u.mu.Unlock()
var req struct {
Stream bool `json:"stream"`
}
_ = json.Unmarshal(b, &req)
w.Header().Set("X-Upstream", name)
time.Sleep(u.delay)
if !req.Stream {
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{"choices":[{"message":{"role":"assistant","content":"hi from %s"}}],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110},"timings":{"prompt_n":100,"cache_n":90,"predicted_n":10,"predicted_ms":50.0}}`, name)
return
}
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(200)
fl := w.(http.Flusher)
for i := 0; i < 3; i++ {
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"%s %d \"}}]}\n\n", name, i)
fl.Flush()
time.Sleep(10 * time.Millisecond)
}
fmt.Fprint(w, `data: {"choices":[],"usage":{"prompt_tokens":200,"completion_tokens":20,"total_tokens":220},"timings":{"prompt_n":200,"cache_n":150,"predicted_n":20,"predicted_ms":80.0}}`+"\n\n")
fl.Flush()
fmt.Fprint(w, "data: [DONE]\n\n")
})
u.srv = httptest.NewServer(mux)
t.Cleanup(u.srv.Close)
return u
}
func (u *upstream) lastReq() recorded { u.mu.Lock(); defer u.mu.Unlock(); return u.last }
// rig is one crossbar: config, real health table (polled once), real lease table over a real
// SQLite store, real limiter, the proxy handler served by httptest.
type rig struct {
t *testing.T
cfg *config.Config
health *health.Table
store *store.Store
leases *lease.Table
lim *limiter.Limiter
front *httptest.Server
}
// newRig builds crossbar from a config text where %s placeholders are the upstream base URLs.
func newRig(t *testing.T, cfgText string, ups ...*upstream) *rig {
urls := make([]any, len(ups))
for i, u := range ups {
urls[i] = u.srv.URL
}
cfg, err := config.Parse(strings.NewReader(fmt.Sprintf(cfgText, urls...)))
if err != nil {
t.Fatal(err)
}
bases := map[string]string{}
for name, h := range cfg.Hosts {
bases[name] = h.BaseURL
}
ht := health.New(bases, time.Hour, nil)
ht.PollOnce(t.Context())
st, err := store.Open(filepath.Join(t.TempDir(), "crossbar.db"))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = st.Close() })
lim := limiter.New()
for name, h := range cfg.Hosts {
for model, m := range h.Models {
lim.Configure(name, model, m.Parallel, cfg.QueueMax)
}
}
lt, err := lease.New(st, proxy.HostView(ht, cfg), proxy.Chooser(cfg, ht, lim), cfg.LeaseIdle.Duration)
if err != nil {
t.Fatal(err)
}
p := proxy.New(cfg, ht, lt, lim, st, nil)
front := httptest.NewServer(p)
t.Cleanup(front.Close)
return &rig{t: t, cfg: cfg, health: ht, store: st, leases: lt, lim: lim, front: front}
}
const twoHosts = `
listen = "127.0.0.1:1"
queue_max = 1
lease_idle = "30m"
[hosts.alpha]
base_url = %q
weight = 1.0
models = { "shared" = { parallel = 2 }, "alpha-only" = { } }
[hosts.beta]
base_url = %q
weight = 2.0
models = { "shared" = { parallel = 2 }, "beta-only" = { } }
[routes.r]
hosts = ["alpha", "beta"]
default_model = "shared"
[routes.other]
hosts = ["alpha"]
`
func conversation(id, turn int) string {
msgs := fmt.Sprintf(`{"role":"system","content":"project"},{"role":"user","content":"conversation %d opening"}`, id)
for i := 1; i < turn; i++ {
msgs += fmt.Sprintf(`,{"role":"assistant","content":"ok"},{"role":"user","content":"turn %d"}`, i)
}
return `{"model":"shared","stream":false,"messages":[` + msgs + `]}`
}
func (r *rig) post(path, body string, hdr ...string) *http.Response {
req, _ := http.NewRequest(http.MethodPost, r.front.URL+path, strings.NewReader(body))
req.Header.Set("Content-Type", "application/json")
for i := 0; i+1 < len(hdr); i += 2 {
req.Header.Set(hdr[i], hdr[i+1])
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
r.t.Fatal(err)
}
return resp
}
func drain(resp *http.Response) string {
b, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return string(b)
}
internal/proxy/hosts.go
+107
View File
@@ -0,0 +1,107 @@
package proxy
import (
"sync"
"git.wntrmute.dev/kyle/crossbar/internal/choose"
"git.wntrmute.dev/kyle/crossbar/internal/config"
"git.wntrmute.dev/kyle/crossbar/internal/health"
"git.wntrmute.dev/kyle/crossbar/internal/lease"
"git.wntrmute.dev/kyle/crossbar/internal/limiter"
)
// Hosts adapts the health table and config for the lease table, and holds the operator's drain set.
// The lease table filters candidates by Healthy && !Draining: Healthy is the health table's view of
// a host, Draining is an operator flag that keeps new leases off a host being taken out of service.
type Hosts struct {
health *health.Table
drain map[string]bool
mu sync.Mutex
}
// HostView builds the lease-table view of the health table and config.
func HostView(h *health.Table, cfg *config.Config) *Hosts {
return &Hosts{health: h, drain: make(map[string]bool)}
}
// Healthy reports whether the health table says the host answered its last poll; an unknown host is
// not healthy.
func (h *Hosts) Healthy(name string) bool {
s, ok := h.health.Get(name)
if !ok {
return false
}
return s.Healthy
}
// Draining reports whether an operator is draining the host.
func (h *Hosts) Draining(name string) bool {
h.mu.Lock()
defer h.mu.Unlock()
return h.drain[name]
}
// SetDraining turns draining on or off for a host.
func (h *Hosts) SetDraining(name string, on bool) {
h.mu.Lock()
defer h.mu.Unlock()
if on {
h.drain[name] = true
return
}
delete(h.drain, name)
}
// hostChooser adapts config, health and limiter to lease.Chooser via choose.Best.
type hostChooser struct {
cfg *config.Config
health *health.Table
lim *limiter.Limiter
}
// Chooser adapts config, health and limiter to lease.Chooser using choose.Best: among the candidates
// it prefers the hosts that have the model loaded over those that can only serve it, then the one
// with the most free slots times weight, breaking ties by shortest queue. Draining is reported false
// because the lease table already filtered draining hosts out before calling Choose.
func Chooser(cfg *config.Config, h *health.Table, l *limiter.Limiter) lease.Chooser {
return &hostChooser{cfg: cfg, health: h, lim: l}
}
func (c *hostChooser) Choose(candidates []string, model string) (string, bool) {
return choose.Best(candidates, func(host string) (choose.Info, bool) {
s, ok := c.health.Get(host)
info := choose.Info{
Healthy: ok && s.Healthy,
Draining: false,
Loaded: contains(s.Loaded, model),
CanServe: c.cfg.Serves(host, model),
Free: freeForModel(c.cfg.Hosts[host], model, c.lim.InFlight(host, model)),
Queued: c.lim.Queued(host, model),
Weight: c.cfg.Hosts[host].Weight,
}
return info, ok
})
}
func contains(list []string, v string) bool {
for _, s := range list {
if s == v {
return true
}
}
return false
}
// freeForModel returns the free slots for one model on one host: its parallel minus the in-flight
// count, floored at zero, and zero when the host does not list that model.
func freeForModel(h config.Host, model string, inflight int) int {
m, ok := h.Models[model]
if !ok {
return 0
}
free := m.Parallel - inflight
if free < 0 {
return 0
}
return free
}
internal/proxy/proxy.go
+120 -105
View File
@@ -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
}
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
}
host, ok := p.cfg.Hosts[name]
if !ok {
p.writeError(w, http.StatusBadGateway, "upstream failed")
p.forward(w, r, route, name, "", rest, fp, model, started, 0)
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 {
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) {
// 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
}
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})
},
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)
}
internal/proxy/proxy_test.go
+231 -278
View File
@@ -1,318 +1,271 @@
package proxy_test
// v1 acceptance tests for the proxy: leases, queueing, accounting, header route override.
// They drive the whole handler over real HTTP against fake upstreams; only what a client or an
// operator can observe is asserted (status codes, headers, the accounting rows, the health table).
// The rig, the fake upstream and the request helpers live in helpers_test.go.
import (
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/config"
"git.wntrmute.dev/kyle/crossbar/internal/health"
"git.wntrmute.dev/kyle/crossbar/internal/proxy"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
// fakeHealth is a hand-set health table that also records MarkDown calls.
type fakeHealth struct {
mu sync.Mutex
st map[string]health.Status
marked []string
func TestConversationIsStickyAndLeaseHeaderTellsWhy(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
r := newRig(t, twoHosts, alpha, beta)
first := r.post("/r/v1/chat/completions", conversation(1, 1))
drain(first)
host := first.Header.Get(proxy.HostHeader)
if first.StatusCode != 200 || host != "beta" { // beta: same free slots, double weight
t.Fatalf("first turn: %d from %q, want 200 from beta", first.StatusCode, host)
}
if got := first.Header.Get(proxy.LeaseHeader); got != "new" {
t.Errorf("%s = %q on the first turn, want new", proxy.LeaseHeader, got)
}
// Take alpha's slots away as a "better host" signal: it must not matter, the lease holds.
for turn := 2; turn <= 6; turn++ {
resp := r.post("/r/v1/chat/completions", conversation(1, turn))
drain(resp)
if resp.Header.Get(proxy.HostHeader) != host || resp.Header.Get(proxy.LeaseHeader) != "reused" {
t.Fatalf("turn %d: host %q lease %q, want %q reused", turn, resp.Header.Get(proxy.HostHeader), resp.Header.Get(proxy.LeaseHeader), host)
}
}
if alpha.hits.Load() != 0 || beta.hits.Load() != 6 {
t.Errorf("hits alpha=%d beta=%d, want 0 and 6", alpha.hits.Load(), beta.hits.Load())
}
}
func (f *fakeHealth) Get(name string) (health.Status, bool) {
f.mu.Lock()
defer f.mu.Unlock()
s, ok := f.st[name]
return s, ok
}
func (f *fakeHealth) MarkDown(name, reason string) {
f.mu.Lock()
defer f.mu.Unlock()
f.marked = append(f.marked, name)
s := f.st[name]
s.Healthy = false
s.LastErr = reason
f.st[name] = s
}
func (f *fakeHealth) markedHosts() []string {
f.mu.Lock()
defer f.mu.Unlock()
return append([]string{}, f.marked...)
}
// upstream records what it received and answers with its name.
type upstream struct {
name string
srv *httptest.Server
mu sync.Mutex
reqs []recorded
}
type recorded struct {
method, path, host, xff string
body string
}
func newUpstream(t *testing.T, name string) *upstream {
u := &upstream{name: name}
u.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
b, _ := io.ReadAll(r.Body)
u.mu.Lock()
u.reqs = append(u.reqs, recorded{r.Method, r.URL.RequestURI(), r.Host, r.Header.Get("X-Forwarded-For"), string(b)})
u.mu.Unlock()
w.Header().Set("Content-Type", "application/json")
fmt.Fprintf(w, `{"from":%q}`, name)
}))
t.Cleanup(u.srv.Close)
return u
}
func (u *upstream) last(t *testing.T) recorded {
u.mu.Lock()
defer u.mu.Unlock()
if len(u.reqs) == 0 {
t.Fatalf("%s: no request received", u.name)
}
return u.reqs[len(u.reqs)-1]
}
func cfgFor(t *testing.T, alpha, beta string) *config.Config {
c, err := config.Parse(strings.NewReader(fmt.Sprintf(`
// spreadHosts: beta is preferred (weight 10) until both of its "shared" slots are busy; then
// alpha (2 free × 1) beats beta (0 free × 10), and a new conversation must start on alpha.
const spreadHosts = `
listen = "127.0.0.1:1"
queue_max = 4
lease_idle = "30m"
[hosts.alpha]
base_url = %q
models = { "shared" = { }, "alpha-only" = { } }
weight = 1.0
models = { "shared" = { parallel = 2 } }
[hosts.beta]
base_url = %q
models = { "shared" = { }, "beta-only" = { } }
weight = 10.0
models = { "shared" = { parallel = 2 } }
[routes.r]
hosts = ["alpha", "beta"]
default_model = "shared"
[routes.beta-first]
hosts = ["beta", "alpha"]
`, alpha, beta)))
if err != nil {
t.Fatal(err)
`
func TestDifferentConversationsSpreadByFreeSlots(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
beta.delay = 400 * time.Millisecond
r := newRig(t, spreadHosts, alpha, beta)
// Two slow conversations occupy beta's two "shared" slots…
var wg sync.WaitGroup
for i := 1; i <= 2; i++ {
wg.Add(1)
go func(i int) { defer wg.Done(); drain(r.post("/r/v1/chat/completions", conversation(i, 1))) }(i)
time.Sleep(50 * time.Millisecond) // arrive one after the other so both pick beta (10 > 2)
}
time.Sleep(50 * time.Millisecond)
// …so a third conversation starting now is sent to alpha (beta has 0 free slots, alpha 2).
resp := r.post("/r/v1/chat/completions", conversation(3, 1))
drain(resp)
if resp.Header.Get(proxy.HostHeader) != "alpha" {
t.Errorf("third conversation went to %q, want alpha (free slots beat weight)", resp.Header.Get(proxy.HostHeader))
}
wg.Wait()
if beta.hits.Load() != 2 || alpha.hits.Load() != 1 {
t.Errorf("hits beta=%d alpha=%d, want 2 and 1", beta.hits.Load(), alpha.hits.Load())
}
return c
}
func healthy(loaded ...string) health.Status {
return health.Status{Healthy: true, Loaded: loaded, Consecutive: 1}
func TestQueueFullIs503(t *testing.T) {
alpha := newUpstream(t, "alpha")
alpha.delay = 400 * time.Millisecond
r := newRig(t, `
listen = "127.0.0.1:1"
queue_max = 1
[hosts.alpha]
base_url = %q
models = { "shared" = { parallel = 1 } }
[routes.r]
hosts = ["alpha"]
default_model = "shared"
`, alpha)
codes := make(chan int, 3)
for i := 1; i <= 3; i++ {
go func(i int) {
resp := r.post("/r/v1/chat/completions", conversation(i, 1))
drain(resp)
codes <- resp.StatusCode
}(i)
time.Sleep(30 * time.Millisecond) // arrival order: 1 runs, 2 queues, 3 finds the queue full
}
got := map[int]int{}
for i := 0; i < 3; i++ {
got[<-codes]++
}
if got[200] != 2 || got[503] != 1 {
t.Fatalf("status counts = %v, want two 200 and one 503", got)
}
// Rows are written after each response completes; allow the store a moment to catch up.
var rows []store.UsageRow
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
rows, _ = r.store.Usage(time.Time{}, store.ByRoute)
if len(rows) == 1 && rows[0].Requests == 3 {
break
}
time.Sleep(20 * time.Millisecond)
}
if len(rows) != 1 || rows[0].Requests != 3 || rows[0].Errors != 1 {
t.Fatalf("usage = %+v, want 3 requests, 1 error (the 503 is recorded too)", rows)
}
if rows[0].QueuedMs <= 0 {
t.Errorf("the queued request must record its wait: %+v", rows[0])
}
}
func TestSplitRoute(t *testing.T) {
func TestUnhealthyHostReleasesAndMoves(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
r := newRig(t, twoHosts, alpha, beta)
drain(r.post("/r/v1/chat/completions", conversation(1, 1))) // lands on beta
beta.srv.Close()
resp := r.post("/r/v1/chat/completions", conversation(1, 2))
drain(resp)
if resp.StatusCode != http.StatusBadGateway {
t.Fatalf("first request after beta died: %d, want 502", resp.StatusCode)
}
if s, _ := r.health.Get("beta"); s.Healthy {
t.Fatalf("beta must be marked down after the 502")
}
resp = r.post("/r/v1/chat/completions", conversation(1, 3))
drain(resp)
if resp.StatusCode != 200 || resp.Header.Get(proxy.HostHeader) != "alpha" || resp.Header.Get(proxy.LeaseHeader) != "new" {
t.Errorf("after the move: %d from %q lease %q, want 200 alpha new", resp.StatusCode, resp.Header.Get(proxy.HostHeader), resp.Header.Get(proxy.LeaseHeader))
}
ev, _ := r.store.Events(time.Time{}, 10)
var reasons []string
for _, e := range ev {
reasons = append(reasons, e.Reason)
}
if len(reasons) != 2 || reasons[0] != store.ReasonNew || reasons[1] != store.ReasonUnhealthy {
t.Errorf("lease events = %v, want [new unhealthy]", reasons)
}
}
func TestAccountingRowsFromUsageAndTimings(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
r := newRig(t, twoHosts, alpha, beta)
drain(r.post("/r/v1/chat/completions", conversation(1, 1))) // non-streamed
drain(r.post("/r/v1/chat/completions", strings.Replace(conversation(1, 2), `"stream":false`, `"stream":true`, 1))) // streamed
deadline := time.Now().Add(2 * time.Second)
var rows []store.UsageRow
for time.Now().Before(deadline) {
rows, _ = r.store.Usage(time.Time{}, store.ByHost)
if len(rows) == 1 && rows[0].Requests == 2 {
break
}
time.Sleep(20 * time.Millisecond)
}
if len(rows) != 1 || rows[0].Requests != 2 {
t.Fatalf("usage by host = %+v, want one host with 2 requests (rows may be written after the response completes, within 2 s)", rows)
}
u := rows[0]
if u.PromptTokens != 300 || u.CachedTokens != 240 || u.CompletionTokens != 30 {
t.Errorf("tokens = prompt %d cached %d completion %d, want 300/240/30 (100+200, 90+150, 10+20)", u.PromptTokens, u.CachedTokens, u.CompletionTokens)
}
if u.BusyMs <= 0 || u.Errors != 0 {
t.Errorf("busy %d errors %d", u.BusyMs, u.Errors)
}
if got := u.CacheHitRatio(); got < 0.79 || got > 0.81 {
t.Errorf("cache hit ratio = %v, want 0.8", got)
}
}
func TestStreamIsUnalteredWhileTeed(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
r := newRig(t, twoHosts, alpha, beta)
resp := r.post("/r/v1/chat/completions", strings.Replace(conversation(9, 1), `"stream":false`, `"stream":true`, 1))
body := drain(resp)
want := 0
for _, line := range strings.Split(body, "\n") {
if strings.HasPrefix(line, "data: ") {
want++
}
}
if want != 5 || !strings.HasSuffix(strings.TrimSpace(body), "data: [DONE]") {
t.Errorf("client must receive every SSE line untouched (3 deltas, usage, DONE); got %d data lines:\n%s", want, body)
}
}
func TestHeaderRouteOverride(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
r := newRig(t, twoHosts, alpha, beta)
// The header names the route; the path has none.
resp := r.post("/v1/chat/completions", conversation(1, 1), proxy.RouteHeader, "other")
drain(resp)
if resp.StatusCode != 200 || resp.Header.Get(proxy.HostHeader) != "alpha" {
t.Errorf("header route 'other' (alpha only): %d from %q", resp.StatusCode, resp.Header.Get(proxy.HostHeader))
}
if alpha.lastReq().path != "/v1/chat/completions" {
t.Errorf("upstream path = %q", alpha.lastReq().path)
}
// A path route and a header route that disagree: the header is the operator's intent → 400.
resp = r.post("/r/v1/chat/completions", conversation(1, 1), proxy.RouteHeader, "other")
if drain(resp); resp.StatusCode != 400 {
t.Errorf("conflicting route in path and header: %d, want 400", resp.StatusCode)
}
resp = r.post("/v1/chat/completions", conversation(1, 1), proxy.RouteHeader, "nope")
if drain(resp); resp.StatusCode != 404 {
t.Errorf("unknown header route: %d, want 404", resp.StatusCode)
}
}
func TestV0BehaviourStillHolds(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
r := newRig(t, twoHosts, alpha, beta)
for _, tc := range []struct {
path, route, rest string
ok bool
}{
{"/a/v1/x", "a", "/v1/x", true},
{"/a/v1/x?q=1", "a", "/v1/x?q=1", true},
{"/a", "a", "/", true},
{"/a/", "a", "/", true},
{"/opencode-a/v1/chat/completions", "opencode-a", "/v1/chat/completions", true},
{"/", "", "", false},
{"//x", "", "", false},
{"", "", "", false},
{"noslash/v1", "", "", false},
} {
route, rest, ok := proxy.SplitRoute(tc.path)
if route != tc.route || rest != tc.rest || ok != tc.ok {
t.Errorf("SplitRoute(%q) = %q %q %v, want %q %q %v", tc.path, route, rest, ok, tc.route, tc.rest, tc.ok)
}
}
}
func TestChoose(t *testing.T) {
h := &fakeHealth{st: map[string]health.Status{
"down": {Healthy: false, Loaded: []string{"m"}},
"alpha": healthy("shared", "alpha-only"),
"beta": healthy("shared", "beta-only"),
}}
hosts := []string{"down", "alpha", "beta"}
if got, ok := proxy.Choose(hosts, "", h); !ok || got != "alpha" {
t.Errorf("no model: %q %v, want alpha (first healthy)", got, ok)
}
if got, ok := proxy.Choose(hosts, "beta-only", h); !ok || got != "beta" {
t.Errorf("beta-only: %q %v, want beta (has the model loaded)", got, ok)
}
if got, ok := proxy.Choose(hosts, "nobody-has-it", h); !ok || got != "alpha" {
t.Errorf("unknown model falls back to the first healthy host: %q %v", got, ok)
}
if got, ok := proxy.Choose([]string{"down", "missing"}, "m", h); ok {
t.Errorf("no healthy host must give ok=false, got %q", got)
}
if got, ok := proxy.Choose(nil, "m", h); ok {
t.Errorf("empty hosts: %q %v", got, ok)
}
}
func TestRoutesToFirstHealthyAndRewrites(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
h := &fakeHealth{st: map[string]health.Status{"alpha": healthy("shared"), "beta": healthy("shared")}}
p := proxy.New(cfgFor(t, alpha.srv.URL, beta.srv.URL), h, nil)
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "http://crossbar.local:7777/r/v1/models?x=1", nil)
req.RemoteAddr = "10.9.8.7:5555"
p.ServeHTTP(rec, req)
if rec.Code != 200 || rec.Header().Get(proxy.HostHeader) != "alpha" {
t.Fatalf("status %d host %q body %s", rec.Code, rec.Header().Get(proxy.HostHeader), rec.Body.String())
}
got := alpha.last(t)
if got.path != "/v1/models?x=1" {
t.Errorf("upstream path = %q, want route stripped and query kept", got.path)
}
if got.host != strings.TrimPrefix(alpha.srv.URL, "http://") {
t.Errorf("Host header = %q, want the upstream's %q", got.host, strings.TrimPrefix(alpha.srv.URL, "http://"))
}
if got.xff != "10.9.8.7" {
t.Errorf("X-Forwarded-For = %q, want the client address", got.xff)
}
if !strings.Contains(rec.Body.String(), `"from":"alpha"`) {
t.Errorf("body = %s", rec.Body.String())
}
}
func TestModelPreferenceAndBodyPassThrough(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
h := &fakeHealth{st: map[string]health.Status{"alpha": healthy("shared", "alpha-only"), "beta": healthy("shared", "beta-only")}}
p := proxy.New(cfgFor(t, alpha.srv.URL, beta.srv.URL), h, nil)
body := `{"model":"beta-only","messages":[{"role":"user","content":"hi"}],"stream":false}`
rec := httptest.NewRecorder()
p.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/r/v1/chat/completions", strings.NewReader(body)))
if rec.Code != 200 || rec.Header().Get(proxy.HostHeader) != "beta" {
t.Fatalf("status %d host %q", rec.Code, rec.Header().Get(proxy.HostHeader))
}
if got := beta.last(t); got.body != body || got.method != http.MethodPost {
t.Errorf("upstream got %+v; the body must arrive unchanged after the model peek", got)
}
// Not JSON: no model, the route default ("shared") applies, first healthy wins.
rec = httptest.NewRecorder()
p.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/r/v1/embeddings", strings.NewReader("plain text")))
if rec.Header().Get(proxy.HostHeader) != "alpha" {
t.Errorf("non-JSON body: host %q, want alpha", rec.Header().Get(proxy.HostHeader))
}
if got := alpha.last(t); got.body != "plain text" {
t.Errorf("non-JSON body must pass through unchanged, got %q", got.body)
}
}
func TestFailoverOnUpstreamError(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
h := &fakeHealth{st: map[string]health.Status{"alpha": healthy("shared"), "beta": healthy("shared")}}
p := proxy.New(cfgFor(t, alpha.srv.URL, beta.srv.URL), h, nil)
alpha.srv.Close() // health still believes alpha is up
rec := httptest.NewRecorder()
p.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/r/v1/models", nil))
if rec.Code != http.StatusBadGateway {
t.Fatalf("first request after alpha died: %d, want 502", rec.Code)
}
var e map[string]string
if err := json.Unmarshal(rec.Body.Bytes(), &e); err != nil || e["error"] != "upstream failed" || e["host"] != "alpha" {
t.Errorf("502 body = %s", rec.Body.String())
}
if m := h.markedHosts(); len(m) != 1 || m[0] != "alpha" {
t.Errorf("MarkDown calls = %v, want [alpha]", m)
}
rec = httptest.NewRecorder()
p.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/r/v1/models", nil))
if rec.Code != 200 || rec.Header().Get(proxy.HostHeader) != "beta" {
t.Errorf("second request: %d %q, want 200 from beta", rec.Code, rec.Header().Get(proxy.HostHeader))
}
}
func TestErrors(t *testing.T) {
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
h := &fakeHealth{st: map[string]health.Status{"alpha": {Healthy: false}, "beta": {Healthy: false}}}
p := proxy.New(cfgFor(t, alpha.srv.URL, beta.srv.URL), h, nil)
for _, tc := range []struct {
name, method, path string
body io.Reader
method, path string
want int
msg string
}{
{"bare slash", http.MethodGet, "/", nil, 400, "missing route"},
{"double slash", http.MethodGet, "//v1/models", nil, 400, "missing route"},
{"unknown route", http.MethodGet, "/nope/v1/models", nil, 404, "unknown route"},
{"disallowed path", http.MethodGet, "/r/slots", nil, 404, "not found"},
{"admin through proxy", http.MethodGet, "/r/_crossbar/hosts", nil, 404, "not found"},
{"no healthy host", http.MethodGet, "/r/v1/models", nil, 503, "no healthy host"},
{"body too large", http.MethodPost, "/r/v1/chat/completions", strings.NewReader(strings.Repeat("x", proxy.MaxBody+1)), 413, "body too large"},
{http.MethodGet, "/", 400, "missing route"},
{http.MethodGet, "/nope/v1/models", 404, "unknown route"},
{http.MethodGet, "/r/slots", 404, "not found"},
{http.MethodGet, "/r/_crossbar/hosts", 404, "not found"},
} {
rec := httptest.NewRecorder()
p.ServeHTTP(rec, httptest.NewRequest(tc.method, tc.path, tc.body))
if rec.Code != tc.want {
t.Errorf("%s: status %d, want %d", tc.name, rec.Code, tc.want)
}
var e map[string]string
if err := json.Unmarshal(rec.Body.Bytes(), &e); err != nil || e["error"] != tc.msg {
t.Errorf("%s: body %s, want error %q", tc.name, rec.Body.String(), tc.msg)
}
if !strings.HasPrefix(rec.Header().Get("Content-Type"), "application/json") {
t.Errorf("%s: errors are JSON", tc.name)
}
}
if len(h.markedHosts()) != 0 {
t.Errorf("errors before choosing a host must not mark anything down: %v", h.markedHosts())
}
}
// TestStreamingIsNotBuffered: the upstream writes one chunk, flushes, and then waits until the
// test has *read* that chunk. If the proxy buffered, the read would never complete.
func TestStreamingIsNotBuffered(t *testing.T) {
release := make(chan struct{})
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(200)
fmt.Fprint(w, "data: first\n\n")
w.(http.Flusher).Flush()
select {
case <-release:
case <-time.After(5 * time.Second):
}
fmt.Fprint(w, "data: second\n\n")
}))
t.Cleanup(up.Close)
beta := newUpstream(t, "beta")
h := &fakeHealth{st: map[string]health.Status{"alpha": healthy("shared"), "beta": healthy("shared")}}
front := httptest.NewServer(proxy.New(cfgFor(t, up.URL, beta.srv.URL), h, nil))
t.Cleanup(front.Close)
resp, err := http.Post(front.URL+"/r/v1/chat/completions", "application/json", strings.NewReader(`{"model":"shared","stream":true}`))
req, _ := http.NewRequest(tc.method, r.front.URL+tc.path, nil)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
buf := make([]byte, 64)
done := make(chan string, 1)
go func() {
n, err := resp.Body.Read(buf)
body := drain(resp)
var e map[string]string
if resp.StatusCode != tc.want || json.Unmarshal([]byte(body), &e) != nil || e["error"] != tc.msg {
t.Errorf("%s: %d %s, want %d %q", tc.path, resp.StatusCode, body, tc.want, tc.msg)
}
}
big := strings.Repeat("x", proxy.MaxBody+1)
resp := r.post("/r/v1/chat/completions", big)
if drain(resp); resp.StatusCode != 413 {
t.Errorf("oversize body: %d, want 413", resp.StatusCode)
}
// GET pass-through with query string, Host and X-Forwarded-For as in v0.
resp, err := http.Get(r.front.URL + "/r/v1/models?x=1")
if err != nil {
done <- "read error: " + err.Error()
return
t.Fatal(err)
}
done <- string(buf[:n])
}()
select {
case got := <-done:
if !strings.HasPrefix(got, "data: first") {
t.Fatalf("first read = %q", got)
}
case <-time.After(2 * time.Second):
t.Fatal("the first chunk did not arrive before the upstream finished: the proxy buffers")
}
close(release)
rest, _ := io.ReadAll(resp.Body)
if !strings.Contains(string(rest), "data: second") {
t.Errorf("rest = %q", rest)
}
if resp.Header.Get(proxy.HostHeader) != "alpha" {
t.Errorf("host header %q", resp.Header.Get(proxy.HostHeader))
drain(resp)
host := resp.Header.Get(proxy.HostHeader)
u := map[string]*upstream{"alpha": alpha, "beta": beta}[host]
if u == nil || u.lastReq().path != "/v1/models?x=1" || u.lastReq().host != strings.TrimPrefix(u.srv.URL, "http://") || u.lastReq().xff == "" {
t.Errorf("GET pass-through: host %q last %+v", host, u.lastReq())
}
}
internal/proxy/recorder_test.go
+1 -1
View File
@@ -42,7 +42,7 @@ hosts = ["alpha"]
t.Fatal(err)
}
h := &fakeHealth{st: map[string]health.Status{"alpha": {Healthy: true, Loaded: []string{"m"}}}}
p := proxy.New(cfg, h, nil)
p := proxy.New(cfg, h, nil, nil, nil, nil)
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/r/v1/chat/completions", strings.NewReader(`{"model":"m","stream":true}`))
internal/proxy/tee.go
+147
View File
@@ -0,0 +1,147 @@
package proxy
import (
"bytes"
"encoding/json"
"io"
"sync"
)
// maxParseBody bounds the non-streamed JSON body accumulated to extract usage/timings: beyond it we
// record no tokens rather than hold an unbounded body in memory.
const maxParseBody = 1 << 20 // 1 MiB
// usageTimings is the last usage/timings object seen in a stream, or the single object parsed from a
// non-streamed JSON body.
type usageTimings struct {
hasUsage bool
usage struct{ prompt, completion int }
hasTimings bool
timings struct{ promptN, cacheN, predictedN int }
}
// tokens resolves the recorded usage and timings to the three counts the accounting row needs.
// prompt and completion come from usage when present, else from timings; cached comes only from
// timings.
func (u *usageTimings) tokens() (prompt, cached, completion int) {
if u.hasUsage {
prompt = u.usage.prompt
completion = u.usage.completion
} else if u.hasTimings {
prompt = u.timings.promptN
completion = u.timings.predictedN
}
if u.hasTimings {
cached = u.timings.cacheN
}
return
}
// tee wraps a response body, passing every byte through unchanged while scanning for usage/timings.
// For a text/event-stream body it scans complete data: lines and remembers the last object seen; for
// anything else it accumulates the body (bounded) and parses it once at Close.
type tee struct {
body io.Reader
closed bool
streamed bool
pending []byte
buf bytes.Buffer
mu sync.Mutex
last usageTimings
}
func newTee(r io.Reader, streamed bool) *tee {
return &tee{body: r, streamed: streamed}
}
// Read reads from the upstream body and feeds the bytes to the scanner without holding any back.
func (t *tee) Read(b []byte) (int, error) {
n, err := t.body.Read(b)
if n > 0 {
t.ingest(b[:n])
}
return n, err
}
// ingest routes a fresh chunk to the streamed scanner or the non-streamed accumulator.
func (t *tee) ingest(p []byte) {
t.mu.Lock()
defer t.mu.Unlock()
if t.streamed {
t.scanSSE(p)
return
}
if t.buf.Len() < maxParseBody {
t.buf.Write(p)
}
}
// scanSSE splits complete lines off the pending buffer and parses each "data: " line.
func (t *tee) scanSSE(p []byte) {
t.pending = append(t.pending, p...)
for {
i := bytes.IndexByte(t.pending, '\n')
if i < 0 {
return
}
line := t.pending[:i]
t.pending = t.pending[i+1:]
if bytes.HasPrefix(line, []byte("data: ")) {
t.parseData(line[len("data: "):])
}
}
}
// parseData decodes one data line's JSON object and remembers its usage and/or timings.
func (t *tee) parseData(data []byte) {
var doc struct {
Usage *struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
} `json:"usage"`
Timings *struct {
PromptN int `json:"prompt_n"`
CacheN int `json:"cache_n"`
PredictedN int `json:"predicted_n"`
} `json:"timings"`
}
if err := json.Unmarshal(data, &doc); err != nil {
return
}
if doc.Usage != nil {
t.last.hasUsage = true
t.last.usage.prompt = doc.Usage.PromptTokens
t.last.usage.completion = doc.Usage.CompletionTokens
}
if doc.Timings != nil {
t.last.hasTimings = true
t.last.timings.promptN = doc.Timings.PromptN
t.last.timings.cacheN = doc.Timings.CacheN
t.last.timings.predictedN = doc.Timings.PredictedN
}
}
// Close closes the underlying body, parsing a non-streamed JSON body once at the end.
func (t *tee) Close() error {
t.mu.Lock()
if t.closed {
t.mu.Unlock()
return nil
}
t.closed = true
if !t.streamed && t.buf.Len() > 0 {
t.parseData(t.buf.Bytes())
}
t.mu.Unlock()
if c, ok := t.body.(io.Closer); ok {
return c.Close()
}
return nil
}
// tokens returns the last usage/timings seen, if any.
func (t *tee) tokens() (prompt, cached, completion int) {
t.mu.Lock()
defer t.mu.Unlock()
return t.last.tokens()
}
internal/store/schema.go
+168
View File
@@ -0,0 +1,168 @@
package store
import (
"database/sql"
"encoding/json"
"fmt"
"time"
)
// schema is the durable layout: the lease table, the lease-event log, the
// per-request accounting rows, the host-health poll log, and the daily
// rollup that Prune writes into. Times are Unix milliseconds.
const schema = `
CREATE TABLE IF NOT EXISTS leases (
route TEXT NOT NULL,
fp TEXT NOT NULL,
model TEXT NOT NULL,
host TEXT NOT NULL,
state TEXT NOT NULL,
created INTEGER NOT NULL,
last_used INTEGER NOT NULL,
PRIMARY KEY (route, fp, model)
);
CREATE TABLE IF NOT EXISTS lease_events (
ts INTEGER NOT NULL,
route TEXT NOT NULL,
model TEXT NOT NULL,
from_host TEXT NOT NULL,
to_host TEXT NOT NULL,
reason TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS requests (
id INTEGER PRIMARY KEY,
route TEXT NOT NULL,
fp TEXT NOT NULL,
model TEXT NOT NULL,
host TEXT NOT NULL,
started INTEGER NOT NULL,
queued_ms INTEGER NOT NULL,
ttfb_ms INTEGER NOT NULL,
total_ms INTEGER NOT NULL,
status INTEGER NOT NULL,
streamed INTEGER NOT NULL,
prompt_tokens INTEGER NOT NULL,
cached_tokens INTEGER NOT NULL,
completion_tokens INTEGER NOT NULL,
err TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS host_health (
ts INTEGER NOT NULL,
host TEXT NOT NULL,
healthy INTEGER NOT NULL,
loaded_models TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS requests_daily (
day INTEGER NOT NULL,
route TEXT NOT NULL,
model TEXT NOT NULL,
host TEXT NOT NULL,
requests INTEGER NOT NULL,
errors INTEGER NOT NULL,
busy_ms INTEGER NOT NULL,
queued_ms INTEGER NOT NULL,
prompt_tokens INTEGER NOT NULL,
cached_tokens INTEGER NOT NULL,
completion_tokens INTEGER NOT NULL,
PRIMARY KEY (day, route, model, host)
);
`
// byColumn maps a grouping to its table column, rejecting anything else.
func byColumn(b By) (string, error) {
switch b {
case ByRoute:
return "route", nil
case ByModel:
return "model", nil
case ByHost:
return "host", nil
default:
return "", fmt.Errorf("store: unknown grouping %q", b)
}
}
// midnight returns UTC midnight of the day holding ms.
func midnight(ms int64) int64 {
return time.UnixMilli(ms).Truncate(24 * time.Hour).UnixMilli()
}
// btoi converts a bool to 0/1 for storage.
func btoi(b bool) int64 {
if b {
return 1
}
return 0
}
// encodeLoaded serialises a Loaded slice as JSON, writing nil as [].
func encodeLoaded(loaded []string) ([]byte, error) {
if loaded == nil {
loaded = []string{}
}
return json.Marshal(loaded)
}
// wrap prefixes a driver error with the package name.
func wrap(err error) error {
if err == nil {
return nil
}
return fmt.Errorf("store: %w", err)
}
func scanLeases(rows *sql.Rows) ([]Lease, error) {
var out []Lease
for rows.Next() {
var (
l Lease
created, last int64
)
if err := rows.Scan(&l.Route, &l.FP, &l.Model, &l.Host, (*string)(&l.State), &created, &last); err != nil {
return nil, wrap(err)
}
l.Created = time.UnixMilli(created).UTC()
l.LastUsed = time.UnixMilli(last).UTC()
out = append(out, l)
}
return out, rows.Err()
}
func scanUsage(rows *sql.Rows) ([]UsageRow, error) {
var out []UsageRow
for rows.Next() {
var u UsageRow
if err := rows.Scan(&u.Key, &u.Requests, &u.Errors, &u.BusyMs, &u.QueuedMs,
&u.PromptTokens, &u.CachedTokens, &u.CompletionTokens); err != nil {
return nil, wrap(err)
}
out = append(out, u)
}
return out, rows.Err()
}
func scanStatusCounts(rows *sql.Rows) ([]StatusCount, error) {
var out []StatusCount
for rows.Next() {
var c StatusCount
if err := rows.Scan(&c.Route, &c.Host, &c.Status, &c.Count); err != nil {
return nil, wrap(err)
}
out = append(out, c)
}
return out, rows.Err()
}
func scanEvents(rows *sql.Rows) ([]LeaseEvent, error) {
var out []LeaseEvent
for rows.Next() {
var e LeaseEvent
var ts int64
if err := rows.Scan(&ts, &e.Route, &e.Model, &e.FromHost, &e.ToHost, (*string)(&e.Reason)); err != nil {
return nil, wrap(err)
}
e.TS = time.UnixMilli(ts).UTC()
out = append(out, e)
}
return out, rows.Err()
}
internal/store/store.go
+373
View File
@@ -0,0 +1,373 @@
// Package store is crossbar's durable state: the lease table and the
// accounting log (requests, lease events, host-health polls) with rollup
// queries that answer per-route / per-model / per-host usage. All times are
// stored as Unix milliseconds (INTEGER) and returned as time.Time in UTC.
package store
import (
"database/sql"
"fmt"
"os"
"path/filepath"
"time"
_ "modernc.org/sqlite"
)
// State is the lease's placement state.
type State string
const (
Active State = "active"
Pinned State = "pinned"
)
// Reasons a lease event carries.
const (
ReasonNew = "new"
ReasonUnhealthy = "unhealthy"
ReasonIdle = "idle"
ReasonPin = "pin"
ReasonRelease = "release"
ReasonDrain = "drain"
)
// By selects the grouping column of a Usage query.
type By string
const (
ByRoute By = "route"
ByModel By = "model"
ByHost By = "host"
)
// Lease is one routed model on one host.
type Lease struct {
Route, FP, Model, Host string
State State
Created, LastUsed time.Time
}
// LeaseEvent records a change to a lease.
type LeaseEvent struct {
TS time.Time
Route, Model, FromHost, ToHost string
Reason string
}
// Request is one proxied request, for the accounting log.
type Request struct {
Route, FP, Model, Host string
Started time.Time
QueuedMs, TTFBMs, TotalMs int64
Status int
Streamed bool
PromptTokens, CachedTokens, CompletionTokens int64
Err string
}
// HostHealth is one poller observation of a host.
type HostHealth struct {
TS time.Time
Host string
Healthy bool
Loaded []string
}
// UsageRow is one group of a Usage query.
type UsageRow struct {
Key string `json:"key"`
Requests int64 `json:"requests"`
Errors int64 `json:"errors"`
BusyMs int64 `json:"busy_ms"`
QueuedMs int64 `json:"queued_ms"`
PromptTokens int64 `json:"prompt_tokens"`
CachedTokens int64 `json:"cached_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
}
// CacheHitRatio is CachedTokens / PromptTokens, or 0 when there were no prompt
// tokens to spend.
func (u UsageRow) CacheHitRatio() float64 {
if u.PromptTokens == 0 {
return 0
}
return float64(u.CachedTokens) / float64(u.PromptTokens)
}
// StatusCount is one (route, host, status) group of live requests, for the metrics endpoint, which
// needs the per-status breakdown Usage cannot give.
type StatusCount struct {
Route, Host string
Status int
Count int64
}
// Store holds the SQLite connection to crossbar's durable state.
type Store struct {
db *sql.DB
}
// Open connects to path in WAL mode and creates the tables if they are missing.
// It fails if the directory that holds the file does not exist.
func Open(path string) (*Store, error) {
if dir := filepath.Dir(path); dir != "" {
if _, err := os.Stat(dir); err != nil {
return nil, fmt.Errorf("store: %s: %w", dir, err)
}
}
db, err := sql.Open("sqlite", "file:"+path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)")
if err != nil {
return nil, wrap(err)
}
if _, err := db.Exec(schema); err != nil {
_ = db.Close()
return nil, wrap(err)
}
return &Store{db: db}, nil
}
// Close releases the connection.
func (s *Store) Close() error {
return wrap(s.db.Close())
}
// JournalMode reports the active journal mode ("wal").
func (s *Store) JournalMode() string {
var mode string
_ = s.db.QueryRow(`PRAGMA journal_mode`).Scan(&mode)
return mode
}
// SaveLease inserts or replaces a lease on (route, fp, model).
func (s *Store) SaveLease(l Lease) error {
_, err := s.db.Exec(`
INSERT OR REPLACE INTO leases (route, fp, model, host, state, created, last_used)
VALUES (?, ?, ?, ?, ?, ?, ?)`,
l.Route, l.FP, l.Model, l.Host, string(l.State),
l.Created.UnixMilli(), l.LastUsed.UnixMilli())
return wrap(err)
}
// DeleteLease removes the lease identified by (route, fp, model).
func (s *Store) DeleteLease(route, fp, model string) error {
_, err := s.db.Exec(`DELETE FROM leases WHERE route = ? AND fp = ? AND model = ?`, route, fp, model)
return wrap(err)
}
// ListLeases returns every lease, ordered by (route, fp, model).
func (s *Store) ListLeases() ([]Lease, error) {
rows, err := s.db.Query(`
SELECT route, fp, model, host, state, created, last_used
FROM leases ORDER BY route, fp, model`)
if err != nil {
return nil, wrap(err)
}
defer rows.Close()
return scanLeases(rows)
}
// RecordEvent appends a lease event.
func (s *Store) RecordEvent(e LeaseEvent) error {
_, err := s.db.Exec(`
INSERT INTO lease_events (ts, route, model, from_host, to_host, reason)
VALUES (?, ?, ?, ?, ?, ?)`,
e.TS.UnixMilli(), e.Route, e.Model, e.FromHost, e.ToHost, e.Reason)
return wrap(err)
}
// RecordRequest appends one request to the accounting log.
func (s *Store) RecordRequest(r Request) error {
_, err := s.db.Exec(`
INSERT INTO requests (route, fp, model, host, started, queued_ms, ttfb_ms, total_ms, status, streamed, prompt_tokens, cached_tokens, completion_tokens, err)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
r.Route, r.FP, r.Model, r.Host,
r.Started.UnixMilli(), r.QueuedMs, r.TTFBMs, r.TotalMs,
r.Status, btoi(r.Streamed),
r.PromptTokens, r.CachedTokens, r.CompletionTokens, r.Err)
return wrap(err)
}
// RecordHostHealth appends one host-health observation.
func (s *Store) RecordHostHealth(h HostHealth) error {
b, err := encodeLoaded(h.Loaded)
if err != nil {
return wrap(err)
}
_, err = s.db.Exec(`
INSERT INTO host_health (ts, host, healthy, loaded_models)
VALUES (?, ?, ?, ?)`,
h.TS.UnixMilli(), h.Host, btoi(h.Healthy), string(b))
return wrap(err)
}
// Usage sums requests at or after since, grouped by by. A zero since means all
// time. It counts the live requests table plus the requests_daily rollup whose
// day is at or after since. Results are ordered by Key.
func (s *Store) Usage(since time.Time, by By) ([]UsageRow, error) {
col, err := byColumn(by)
if err != nil {
return nil, err
}
sinceMs := since.UnixMilli()
query := `
SELECT key,
SUM(requests) AS requests,
SUM(errors) AS errors,
SUM(busy_ms) AS busy_ms,
SUM(queued_ms) AS queued_ms,
SUM(prompt_tokens) AS prompt_tokens,
SUM(cached_tokens) AS cached_tokens,
SUM(completion_tokens) AS completion_tokens
FROM (
SELECT ` + col + ` AS key, 1 AS requests,
CASE WHEN status >= 400 THEN 1 ELSE 0 END AS errors,
total_ms AS busy_ms, queued_ms,
prompt_tokens, cached_tokens, completion_tokens
FROM requests WHERE started >= ?
UNION ALL
SELECT ` + col + ` AS key, requests, errors, busy_ms, queued_ms,
prompt_tokens, cached_tokens, completion_tokens
FROM requests_daily WHERE day >= ?
)
GROUP BY key
ORDER BY key`
rows, err := s.db.Query(query, sinceMs, sinceMs)
if err != nil {
return nil, wrap(err)
}
defer rows.Close()
return scanUsage(rows)
}
// StatusCounts groups the live requests at or after since by (route, host, status). It reads the
// requests table only; the rolled-up requests_daily rows are not in it (Prune has moved them out of
// requests), so counts cover only traffic still in the live table.
func (s *Store) StatusCounts(since time.Time) ([]StatusCount, error) {
sinceMs := since.UnixMilli()
rows, err := s.db.Query(`
SELECT route, host, status, COUNT(*)
FROM requests WHERE started >= ?
GROUP BY route, host, status
ORDER BY route, host, status`, sinceMs)
if err != nil {
return nil, wrap(err)
}
defer rows.Close()
return scanStatusCounts(rows)
}
// Events returns lease events at or after since, oldest first, at most limit.
func (s *Store) Events(since time.Time, limit int) ([]LeaseEvent, error) {
rows, err := s.db.Query(`
SELECT ts, route, model, from_host, to_host, reason
FROM lease_events WHERE ts >= ?
ORDER BY ts ASC LIMIT ?`, since.UnixMilli(), limit)
if err != nil {
return nil, wrap(err)
}
defer rows.Close()
return scanEvents(rows)
}
// Prune moves every request older than now-retention into requests_daily (adding
// into the existing daily row for that day/route/model/host), deletes the live
// rows, and returns how many were removed. All in one transaction.
func (s *Store) Prune(now time.Time, retention time.Duration) (int64, error) {
threshold := now.Add(-retention).UnixMilli()
tx, err := s.db.Begin()
if err != nil {
return 0, wrap(err)
}
defer func() { _ = tx.Rollback() }()
rows, err := tx.Query(`
SELECT route, model, host, started, total_ms, queued_ms,
status, prompt_tokens, cached_tokens, completion_tokens
FROM requests WHERE started < ?`, threshold)
if err != nil {
_ = tx.Rollback()
return 0, wrap(err)
}
type bucket struct {
day int64
route string
model string
host string
}
type agg struct {
requests int64
errors int64
busyMs int64
queuedMs int64
prompt int64
cached int64
done int64
}
aggs := map[bucket]*agg{}
count := int64(0)
for rows.Next() {
var (
route, model, host string
started int64
totalMs, queuedMs int64
status int
prompt, cached, done int64
)
if err := rows.Scan(&route, &model, &host, &started, &totalMs, &queuedMs, &status, &prompt, &cached, &done); err != nil {
_ = rows.Close()
_ = tx.Rollback()
return 0, wrap(err)
}
count++
k := bucket{day: midnight(started), route: route, model: model, host: host}
a := aggs[k]
if a == nil {
a = &agg{}
aggs[k] = a
}
a.requests++
if status >= 400 {
a.errors++
}
a.busyMs += totalMs
a.queuedMs += queuedMs
a.prompt += prompt
a.cached += cached
a.done += done
}
if err := rows.Err(); err != nil {
_ = rows.Close()
_ = tx.Rollback()
return 0, wrap(err)
}
_ = rows.Close()
upsert := `
INSERT INTO requests_daily (day, route, model, host, requests, errors, busy_ms, queued_ms, prompt_tokens, cached_tokens, completion_tokens)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT (day, route, model, host) DO UPDATE SET
requests = requests + excluded.requests,
errors = errors + excluded.errors,
busy_ms = busy_ms + excluded.busy_ms,
queued_ms = queued_ms + excluded.queued_ms,
prompt_tokens = prompt_tokens + excluded.prompt_tokens,
cached_tokens = cached_tokens + excluded.cached_tokens,
completion_tokens = completion_tokens + excluded.completion_tokens`
for k, a := range aggs {
if _, err := tx.Exec(upsert, k.day, k.route, k.model, k.host,
a.requests, a.errors, a.busyMs, a.queuedMs, a.prompt, a.cached, a.done); err != nil {
_ = tx.Rollback()
return 0, wrap(err)
}
}
if _, err := tx.Exec(`DELETE FROM requests WHERE started < ?`, threshold); err != nil {
_ = tx.Rollback()
return 0, wrap(err)
}
if err := tx.Commit(); err != nil {
return 0, wrap(err)
}
return count, nil
}
internal/store/store_test.go
+185
View File
@@ -0,0 +1,185 @@
package store_test
import (
"path/filepath"
"testing"
"time"
"git.wntrmute.dev/kyle/crossbar/internal/store"
)
func open(t *testing.T, dir string) *store.Store {
s, err := store.Open(filepath.Join(dir, "crossbar.db"))
if err != nil {
t.Fatalf("Open: %v", err)
}
t.Cleanup(func() { _ = s.Close() })
return s
}
func TestOpenIsIdempotentAndWAL(t *testing.T) {
dir := t.TempDir()
s := open(t, dir)
if got := s.JournalMode(); got != "wal" {
t.Errorf("journal_mode = %q, want wal", got)
}
if err := s.Close(); err != nil {
t.Fatal(err)
}
open(t, dir) // second open on the same file must not fail on existing tables
}
func TestLeasesSurviveReopen(t *testing.T) {
dir := t.TempDir()
s := open(t, dir)
now := time.Date(2026, 9, 25, 10, 0, 0, 0, time.UTC)
l := store.Lease{Route: "opencode-a", FP: "abc", Model: "m", Host: "alpha", State: store.Active, Created: now, LastUsed: now}
if err := s.SaveLease(l); err != nil {
t.Fatal(err)
}
l2 := l
l2.FP = "def"
l2.Host = "beta"
l2.State = store.Pinned
if err := s.SaveLease(l2); err != nil {
t.Fatal(err)
}
// Saving the same key again replaces, not duplicates.
l.Host = "beta"
l.LastUsed = now.Add(time.Minute)
if err := s.SaveLease(l); err != nil {
t.Fatal(err)
}
if err := s.Close(); err != nil {
t.Fatal(err)
}
s = open(t, dir)
got, err := s.ListLeases()
if err != nil {
t.Fatal(err)
}
if len(got) != 2 {
t.Fatalf("ListLeases = %d rows, want 2: %+v", len(got), got)
}
byFP := map[string]store.Lease{}
for _, x := range got {
byFP[x.FP] = x
}
if a := byFP["abc"]; a.Host != "beta" || !a.LastUsed.Equal(now.Add(time.Minute)) || a.State != store.Active {
t.Errorf("abc = %+v", a)
}
if d := byFP["def"]; d.State != store.Pinned || d.Host != "beta" {
t.Errorf("def = %+v", d)
}
if err := s.DeleteLease("opencode-a", "abc", "m"); err != nil {
t.Fatal(err)
}
got, _ = s.ListLeases()
if len(got) != 1 || got[0].FP != "def" {
t.Errorf("after delete: %+v", got)
}
}
func TestEventsAndRequestsAndUsage(t *testing.T) {
s := open(t, t.TempDir())
t0 := time.Date(2026, 9, 25, 10, 0, 0, 0, time.UTC)
must := func(err error) {
if err != nil {
t.Fatal(err)
}
}
must(s.RecordEvent(store.LeaseEvent{TS: t0, Route: "r1", Model: "m", FromHost: "", ToHost: "alpha", Reason: store.ReasonNew}))
must(s.RecordEvent(store.LeaseEvent{TS: t0.Add(time.Hour), Route: "r1", Model: "m", FromHost: "alpha", ToHost: "beta", Reason: store.ReasonUnhealthy}))
reqs := []store.Request{
{Route: "r1", FP: "a", Model: "m", Host: "alpha", Started: t0, QueuedMs: 0, TTFBMs: 100, TotalMs: 1000, Status: 200, Streamed: true, PromptTokens: 1000, CachedTokens: 900, CompletionTokens: 50},
{Route: "r1", FP: "a", Model: "m", Host: "alpha", Started: t0.Add(time.Minute), QueuedMs: 40, TTFBMs: 120, TotalMs: 2000, Status: 200, Streamed: true, PromptTokens: 1100, CachedTokens: 1000, CompletionTokens: 60},
{Route: "r2", FP: "b", Model: "m", Host: "beta", Started: t0.Add(2 * time.Minute), TotalMs: 500, Status: 502, Err: "upstream failed"},
{Route: "r2", FP: "b", Model: "m", Host: "beta", Started: t0.Add(-48 * time.Hour), TotalMs: 300, Status: 200, PromptTokens: 10, CompletionTokens: 5},
}
for _, r := range reqs {
must(s.RecordRequest(r))
}
must(s.RecordHostHealth(store.HostHealth{TS: t0, Host: "alpha", Healthy: true, Loaded: []string{"m"}}))
rows, err := s.Usage(t0.Add(-time.Hour), store.ByRoute)
must(err)
if len(rows) != 2 {
t.Fatalf("Usage by route since t0-1h: %d rows, want 2 (r1, r2): %+v", len(rows), rows)
}
byKey := map[string]store.UsageRow{}
for _, r := range rows {
byKey[r.Key] = r
}
r1 := byKey["r1"]
if r1.Requests != 2 || r1.Errors != 0 || r1.BusyMs != 3000 || r1.QueuedMs != 40 {
t.Errorf("r1 = %+v", r1)
}
if r1.PromptTokens != 2100 || r1.CachedTokens != 1900 || r1.CompletionTokens != 110 {
t.Errorf("r1 tokens = %+v", r1)
}
if got := r1.CacheHitRatio(); got < 0.904 || got > 0.905 {
t.Errorf("r1 cache hit ratio = %v, want 1900/2100", got)
}
r2 := byKey["r2"]
if r2.Requests != 1 || r2.Errors != 1 || r2.BusyMs != 500 {
t.Errorf("r2 = %+v (the 48h-old request is outside since)", r2)
}
if r2.CacheHitRatio() != 0 {
t.Errorf("no prompt tokens: ratio must be 0, got %v", r2.CacheHitRatio())
}
byHost, err := s.Usage(time.Time{}, store.ByHost)
must(err)
if len(byHost) != 2 {
t.Errorf("by host, all time: %+v", byHost)
}
for _, r := range byHost {
if r.Key == "beta" && r.Requests != 2 {
t.Errorf("beta all-time requests = %d, want 2", r.Requests)
}
}
byModel, err := s.Usage(time.Time{}, store.ByModel)
must(err)
if len(byModel) != 1 || byModel[0].Key != "m" || byModel[0].Requests != 4 {
t.Errorf("by model: %+v", byModel)
}
ev, err := s.Events(t0.Add(-time.Minute), 10)
must(err)
if len(ev) != 2 || ev[0].Reason != store.ReasonNew || ev[1].ToHost != "beta" {
t.Errorf("events = %+v", ev)
}
}
func TestPruneRollsUpOldRequests(t *testing.T) {
s := open(t, t.TempDir())
t0 := time.Date(2026, 9, 25, 10, 0, 0, 0, time.UTC)
old := t0.Add(-200 * 24 * time.Hour)
for i := 0; i < 3; i++ {
if err := s.RecordRequest(store.Request{Route: "r", Model: "m", Host: "h", Started: old.Add(time.Duration(i) * time.Minute), TotalMs: 100, Status: 200, PromptTokens: 10, CachedTokens: 5, CompletionTokens: 1}); err != nil {
t.Fatal(err)
}
}
if err := s.RecordRequest(store.Request{Route: "r", Model: "m", Host: "h", Started: t0, TotalMs: 100, Status: 200}); err != nil {
t.Fatal(err)
}
n, err := s.Prune(t0, 180*24*time.Hour)
if err != nil {
t.Fatal(err)
}
if n != 3 {
t.Errorf("Prune removed %d rows, want 3", n)
}
rows, _ := s.Usage(time.Time{}, store.ByRoute)
if len(rows) != 1 || rows[0].Requests != 4 || rows[0].PromptTokens != 30 {
t.Errorf("usage must still include pruned traffic through the daily rollup: %+v", rows)
}
live, _ := s.Usage(old.Add(24*time.Hour), store.ByRoute)
if len(live) != 1 || live[0].Requests != 1 {
t.Errorf("recent-only usage = %+v", live)
}
}
func TestBadPath(t *testing.T) {
if _, err := store.Open(filepath.Join(t.TempDir(), "no", "such", "dir", "x.db")); err == nil {
t.Fatal("Open must fail when the directory does not exist")
}
}
tools/smoke.sh
+59 -24
View File
@@ -1,45 +1,80 @@
#!/bin/sh
# Smoke run: two fake upstreams, one crossbar, real HTTP. Prints "smoke: ok" or fails.
# Needs: bin/crossbar and bin/fakeupstream (make build), curl.
# Smoke run (v1): two fake upstreams, one crossbar with a fresh SQLite file, real HTTP.
# Checks routing, leases (sticky + header), failover, recovery, streaming, queueing, pin, drain,
# usage and metrics. Prints "smoke: ok" or fails with the crossbar log.
set -eu
cd "$(dirname "$0")/.."
tmp=$(mktemp -d); trap 'kill $pids 2>/dev/null; rm -rf "$tmp"' EXIT INT TERM
pids=""
bin/fakeupstream -listen 127.0.0.1:18081 -name alpha -models ornith-1.5-35b-a3b,small-9b -down-file "$tmp/alpha.down" >"$tmp/alpha.log" 2>&1 & pids="$pids $!"
sed "s#^db .*#db = \"$tmp/crossbar.db\"#" example.toml > "$tmp/crossbar.toml"
bin/fakeupstream -listen 127.0.0.1:18081 -name alpha -models ornith-1.5-35b-a3b,small-9b -down-file "$tmp/alpha.down" -slow 600 >"$tmp/alpha.log" 2>&1 & pids="$pids $!"
bin/fakeupstream -listen 127.0.0.1:18082 -name beta -models ornith-1.5-35b-a3b -down-file "$tmp/beta.down" >"$tmp/beta.log" 2>&1 & pids="$pids $!"
bin/crossbar -config example.toml >"$tmp/crossbar.log" 2>&1 & pids="$pids $!"
bin/crossbar -config "$tmp/crossbar.toml" >"$tmp/crossbar.log" 2>&1 & pids="$pids $!"
sleep 1.5
fail() { echo "smoke: FAIL: $*" >&2; echo "--- crossbar.log"; cat "$tmp/crossbar.log"; exit 1; }
base=http://127.0.0.1:17777
conv() { printf '{"model":"ornith-1.5-35b-a3b","stream":false,"messages":[{"role":"system","content":"smoke"},{"role":"user","content":"conversation %s"}]}' "$1"; }
hdrs() { curl -s -o /dev/null -w '%{http_code} %header{X-Crossbar-Host} %header{X-Crossbar-Lease}' "$@"; }
h=$(curl -s -o /dev/null -w '%{http_code} %header{X-Crossbar-Host}' "$base/opencode-a/v1/models")
[ "$h" = "200 alpha" ] || fail "opencode-a should go to alpha, got '$h'"
h=$(curl -s -o /dev/null -w '%{http_code} %header{X-Crossbar-Host}' "$base/hermes-x/v1/models")
[ "$h" = "200 beta" ] || fail "hermes-x should go to beta, got '$h'"
h=$(curl -s -o /dev/null -w '%{http_code}' "$base/nope/v1/models")
[ "$h" = "404" ] || fail "unknown route should be 404, got '$h'"
# 1. a conversation gets a lease and keeps it; beta wins (2 slots × weight 2 vs 1 × 1)
h=$(hdrs -X POST -H 'Content-Type: application/json' -d "$(conv A)" "$base/opencode-a/v1/chat/completions")
[ "$h" = "200 beta new" ] || fail "first turn should be '200 beta new', got '$h'"
h=$(hdrs -X POST -H 'Content-Type: application/json' -d "$(conv A)" "$base/opencode-a/v1/chat/completions")
[ "$h" = "200 beta reused" ] || fail "second turn should reuse beta, got '$h'"
touch "$tmp/alpha.down"; sleep 2.5 # poll_interval is 1s in example.toml
h=$(curl -s -o /dev/null -w '%{http_code} %header{X-Crossbar-Host}' "$base/opencode-a/v1/models")
[ "$h" = "200 beta" ] || fail "with alpha down, opencode-a should fail over to beta, got '$h'"
curl -s "$base/_crossbar/hosts" | grep -q '"alpha":{"healthy":false' || fail "/_crossbar/hosts does not show alpha unhealthy: $(curl -s $base/_crossbar/hosts)"
# 2. header route
h=$(hdrs -X POST -H 'Content-Type: application/json' -H 'X-Crossbar-Route: hermes-x' -d "$(conv B)" "$base/v1/chat/completions")
case "$h" in "200 beta new") ;; *) fail "header route hermes-x should be '200 beta new', got '$h'";; esac
h=$(curl -s -o /dev/null -w '%{http_code}' "$base/nope/v1/models"); [ "$h" = "404" ] || fail "unknown route 404, got $h"
rm "$tmp/alpha.down"; sleep 3.5 # recovery needs two good polls
h=$(curl -s -o /dev/null -w '%header{X-Crossbar-Host}' "$base/opencode-a/v1/models")
[ "$h" = "alpha" ] || fail "alpha should be back after two good polls, got '$h'"
# 3. pin opencode-a to alpha: conversation A's next turn moves (an operator pin outranks the lease)
h=$(curl -s -o /dev/null -w '%{http_code}' -X POST -H 'Content-Type: application/json' -d '{"host":"alpha","pin":true}' "$base/_crossbar/routes/opencode-a")
[ "$h" = "200" ] || fail "pin returned $h"
h=$(hdrs -X POST -H 'Content-Type: application/json' -d "$(conv A)" "$base/opencode-a/v1/chat/completions")
[ "$h" = "200 alpha new" ] || fail "after pin, conversation A should be '200 alpha new', got '$h'"
curl -s "$base/_crossbar/routes" | grep -q '"pinned":"alpha"' || fail "routes view does not show the pin: $(curl -s $base/_crossbar/routes)"
# Streaming: five chunks 200 ms apart must arrive over >= 0.6 s, not all at once at the end.
# 4. queue: alpha has parallel 1, queue_max 1, and answers in 600 ms → of three concurrent, one is 503
for i in 1 2 3; do (curl -s -o /dev/null -w '%{http_code}\n' -X POST -H 'Content-Type: application/json' -d "$(conv Q$i)" "$base/opencode-a/v1/chat/completions" >> "$tmp/codes") & sleep 0.1; done; wait $! 2>/dev/null || true
sleep 2.5
sort "$tmp/codes" | uniq -c | tr -s ' ' > "$tmp/counts"
grep -q '2 200' "$tmp/counts" && grep -q '1 503' "$tmp/counts" || fail "queue test wanted two 200 and one 503, got: $(cat "$tmp/counts")"
# 5. release the pin, drain alpha: new conversations go to beta, A stays on alpha
curl -s -o /dev/null -X POST -H 'Content-Type: application/json' -d '{"release":true}' "$base/_crossbar/routes/opencode-a"
h=$(curl -s -o /dev/null -w '%{http_code}' -X POST -H 'Content-Type: application/json' -d '{"drain":true}' "$base/_crossbar/hosts/alpha"); [ "$h" = "200" ] || fail "drain returned $h"
h=$(hdrs -X POST -H 'Content-Type: application/json' -d "$(conv C)" "$base/opencode-a/v1/chat/completions")
[ "$h" = "200 beta new" ] || fail "with alpha draining a new conversation should go to beta, got '$h'"
curl -s "$base/_crossbar/hosts" | grep -q '"alpha":{[^}]*"draining":true' || fail "hosts view does not show alpha draining"
curl -s -o /dev/null -X POST -H 'Content-Type: application/json' -d '{"drain":false}' "$base/_crossbar/hosts/alpha"
# 6. failover + recovery
touch "$tmp/beta.down"; sleep 2.5
h=$(hdrs -X POST -H 'Content-Type: application/json' -d "$(conv C)" "$base/opencode-a/v1/chat/completions")
[ "$h" = "200 alpha new" ] || fail "with beta down conversation C should move to alpha, got '$h'"
curl -s "$base/_crossbar/hosts" | grep -q '"beta":{"healthy":false' || fail "hosts view does not show beta unhealthy"
rm "$tmp/beta.down"; sleep 3.5
curl -s "$base/_crossbar/hosts" | grep -q '"beta":{"healthy":true' || fail "beta did not recover after two good polls"
# 7. streaming still arrives incrementally, and the final usage chunk is untouched
start=$(date +%s%N)
first=""
curl -sN -X POST -H 'Content-Type: application/json' -d '{"model":"ornith-1.5-35b-a3b","stream":true,"messages":[]}' \
curl -sN -X POST -H 'Content-Type: application/json' -d '{"model":"ornith-1.5-35b-a3b","stream":true,"messages":[{"role":"user","content":"stream me"}]}' \
"$base/opencode-a/v1/chat/completions" | while IFS= read -r line; do
[ -n "$line" ] || continue
now=$(date +%s%N); echo "$(( (now - start) / 1000000 )) $line"
done > "$tmp/stream.txt"
firstms=$(head -1 "$tmp/stream.txt" | cut -d' ' -f1); lastms=$(tail -1 "$tmp/stream.txt" | cut -d' ' -f1)
[ -n "$firstms" ] && [ "$((lastms - firstms))" -ge 600 ] || fail "stream arrived in one burst (first ${firstms:-?} ms, last ${lastms:-?} ms):
$(cat "$tmp/stream.txt")"
grep -q 'DONE' "$tmp/stream.txt" || fail "stream did not end with [DONE]"
[ -n "$firstms" ] && [ "$((lastms - firstms))" -ge 600 ] || fail "stream arrived in one burst: $(cat "$tmp/stream.txt")"
grep -q '"usage"' "$tmp/stream.txt" && grep -q 'DONE' "$tmp/stream.txt" || fail "stream lost the usage chunk or DONE"
grep -q 'route=opencode-a host=alpha' "$tmp/crossbar.log" || fail "no request log line"
# 8. accounting and metrics
sleep 1
u=$(curl -s "$base/_crossbar/usage?by=host")
echo "$u" | grep -q '"key":"alpha"' && echo "$u" | grep -q '"key":"beta"' || fail "usage by host: $u"
echo "$u" | grep -q '"cached_tokens":[1-9]' || fail "usage has no cached tokens (SSE/JSON usage not captured): $u"
curl -s -H 'Accept: text/plain' "$base/_crossbar/usage?by=route" | grep -qi 'cache' || fail "text usage table missing"
m=$(curl -s "$base/_crossbar/metrics")
echo "$m" | grep -q 'crossbar_requests_total{route="opencode-a",host="alpha",status="503"} 1' || fail "metrics missing the 503: $m"
echo "$m" | grep -q 'crossbar_host_healthy{host="beta"} 1' || fail "metrics missing host health"
grep -q 'route=opencode-a host=' "$tmp/crossbar.log" || fail "no request log line"
echo "smoke: ok (stream spread $((lastms - firstms)) ms)"