Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Ilker ManapandClaude Opus 5
|
3a4260c8c5 | ||
|
Ilker ManapandClaude Opus 5
|
b9c2be6cc2 | ||
|
Ilker ManapandClaude Opus 5
|
4ba84fa90e | ||
|
Ilker ManapandClaude Opus 5
|
ade689fc8e | ||
|
Ilker ManapandClaude Opus 5
|
5b5e61cb97 | ||
|
Ilker ManapandClaude Opus 5
|
fa60b81f37 | ||
|
Ilker ManapandClaude Opus 5
|
b6de469351 | ||
|
Ilker ManapandClaude Opus 5
|
a76ebdb79a | ||
|
Ilker ManapandClaude Opus 5
|
ccc9aa8b76 | ||
|
Ilker ManapandClaude Opus 5
|
539bdac152 | ||
|
Ilker ManapandClaude Opus 5
|
7dd1de6537 | ||
 ilker
|
c815b629e2 |
No files matched your search
+84
-1
@@ -11,12 +11,95 @@ is only released once its **exit criterion** is met, not when its code lands.
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | released as part of v0.5.0 |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler | released |
|
||||
| v0.6.0 | 4 | Multi-node replicas and fabric awareness | released |
|
||||
| v0.7.0 | — | Conversations, agent memory, disk management | released |
|
||||
| v0.8.0 | — | PostgreSQL, leader election, agent failover | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
Until v1.0.0 the minor version is the breaking one: configuration formats and Go
|
||||
APIs change between minors without a deprecation period.
|
||||
|
||||
## Unreleased
|
||||
## v0.8.0 — 2026-09-27
|
||||
|
||||
The store stops being a file on one machine's disk. PostgreSQL, a way to move
|
||||
an existing installation onto it without reissuing every key, and the two
|
||||
things a second control plane needs so that running one is safe rather than
|
||||
clever.
|
||||
|
||||
### A second control plane, and a way to fail over to it
|
||||
|
||||
Leadership is a lease in the store, so two control planes against one database
|
||||
cannot both place work — which they otherwise would, both seeing a model
|
||||
wanting replicas and none running, leaving twice the replicas each holding GPUs
|
||||
the other believes are free.
|
||||
|
||||
A standby also refuses agent streams. Commands travel down the stream the agent
|
||||
opened, so a standby holding one is a node the leader cannot command, and the
|
||||
leader's placements would fail with "no agent connected" for a node that is
|
||||
plainly connected. Agents therefore take a list of addresses rather than one,
|
||||
trying each in turn: a standby refusing, a stopped control plane refusing the
|
||||
connection and a crashed one never answering are the same thing from a node.
|
||||
|
||||
Measured on the cluster. Leadership moved three seconds after the leader was
|
||||
stopped — it released the lease on the way out rather than letting it lapse —
|
||||
and all five agents had registered with the standby six seconds later. The
|
||||
standby adopted every replica rather than restarting any: six replicas, all
|
||||
ready, engines that never stopped serving.
|
||||
|
||||
Everything is unset-means-leader, so a single control plane and every
|
||||
deployment that existed before this behave exactly as they did.
|
||||
|
||||
What this does not do is move client traffic. The standby served inference
|
||||
correctly on its own port while the public address returned 502, because
|
||||
nothing in front of them moved — that needs a load balancer or a shared
|
||||
address, and it is the remaining piece.
|
||||
|
||||
### PostgreSQL
|
||||
|
||||
The store runs on PostgreSQL, chosen by what the address looks like: a path is
|
||||
still SQLite, a `postgres://` URL is not, and `@` followed by a path reads the
|
||||
address out of a file so a DSN carrying a password stays out of a configuration
|
||||
that gets copied and pasted into issues. Every existing installation keeps
|
||||
working untouched.
|
||||
|
||||
One set of queries rather than one per database, because the SQL here is almost
|
||||
entirely portable and two copies would drift one query at a time. The rewrite
|
||||
of placeholders is a wrapper around the handle rather than a change at sixty
|
||||
call sites — which found six places reaching past it, including all three
|
||||
manual transactions and two prepared statements, every one of which worked on
|
||||
SQLite and would have failed on PostgreSQL.
|
||||
|
||||
Exactly one query is not the same sentence in both: the note search, MATCH
|
||||
against `@@`. SQLite gets full-text search from an FTS5 table kept in step by
|
||||
three triggers; PostgreSQL derives a tsvector from the row, so there is nothing
|
||||
to keep in step and no trigger to forget. It uses the 'simple' configuration
|
||||
rather than English — these are an agent's own notes, in whatever language it
|
||||
chose, and stemming English rules over Turkish would quietly stop finding
|
||||
things.
|
||||
|
||||
`llmgw keys move-store` copies an existing store rather than starting a fresh
|
||||
one, because keys are stored hashed: a migration that drops them cannot
|
||||
recreate them, and every person and every script would need a new one.
|
||||
|
||||
Testing it against a real database found three things nothing else would have.
|
||||
`schema_migrations` is created in Go rather than in a migration file, so it
|
||||
escaped the port and still said INTEGER — 64 bits in SQLite, 32 in PostgreSQL,
|
||||
and a millisecond timestamp does not fit. The copy order followed the order the
|
||||
migrations create tables in, and `users` is created after `api_keys` even
|
||||
though `api_keys` references it. And `key_claims.sealed` is emptied the moment
|
||||
a claim is used, so a claimed row holds an empty blob in a NOT NULL column — an
|
||||
empty blob scans as a typed nil slice and a real NULL as an untyped one, a
|
||||
difference pgx cannot see, so it arrived as NULL.
|
||||
|
||||
The driver is pgx, which is pure Go: the whole binary is built with
|
||||
CGO_ENABLED=0 so that it runs on a node whose libc is older than the builder's,
|
||||
and a cgo driver would take that away from every deployment, including the ones
|
||||
staying on SQLite.
|
||||
|
||||
## v0.7.0 — 2026-09-25
|
||||
|
||||
What a cluster needs to be lived with rather than demonstrated: somewhere to
|
||||
keep a conversation, somewhere to keep a note, and an honest answer about the
|
||||
disk underneath both.
|
||||
|
||||
### Memory notes for agents, and threads without API keys
|
||||
|
||||
|
||||
@@ -30,6 +30,8 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, engine supervision | released |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler, ctl | released |
|
||||
| v0.6.0 | 4 | Multi-node placement and fabric awareness | released |
|
||||
| v0.7.0 | — | Conversations, agent memory, disk management | released |
|
||||
| v0.8.0 | — | PostgreSQL, leader election, agent failover | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
|
||||
+53
-4
@@ -59,7 +59,7 @@ a certificate and changing it later means reissuing one.
|
||||
One static binary, no runtime dependencies, the same on every machine.
|
||||
|
||||
```
|
||||
VERSION=v0.6.0 # check the releases page for the current tag
|
||||
VERSION=v0.8.0 # check the releases page for the current tag
|
||||
curl -fsSLO https://gitea.mynodes.xyz/ilker/llmgw/releases/download/$VERSION/llmgw-$VERSION-linux-amd64
|
||||
curl -fsSL https://gitea.mynodes.xyz/ilker/llmgw/releases/download/$VERSION/checksums.txt | sha256sum -c --ignore-missing
|
||||
sudo install -m0755 llmgw-$VERSION-linux-amd64 /usr/local/bin/llmgw
|
||||
@@ -595,6 +595,52 @@ fleet means the measurement is missing, not that the fleet is idle. And
|
||||
`llmgw_gpu_nvlink_peers` at zero everywhere means the topology is unknown,
|
||||
which is what a binary built without the nvml tag reports.
|
||||
|
||||
## PostgreSQL, when SQLite is no longer the right shape
|
||||
|
||||
SQLite is the default and stays the right answer for most installations: one
|
||||
file, no server, and a request path that barely touches it — usage is written
|
||||
in batches, budgets are counted in memory, authentication is cached.
|
||||
|
||||
Two things change that. Conversations and agent notes read and write on every
|
||||
turn, which is the first per-request load this store has had; and SQLite takes
|
||||
one writer at a time, so those turns serialise. And a second control plane is
|
||||
impossible while the store is a file on one machine's disk.
|
||||
|
||||
The setting is the same one, and the address decides:
|
||||
|
||||
```yaml
|
||||
auth:
|
||||
store: "@/etc/llmgw/secrets/store-dsn"
|
||||
```
|
||||
|
||||
A path is SQLite, a `postgres://` URL is PostgreSQL, and `@` followed by a path
|
||||
reads the address out of that file — which is how a DSN carrying a password
|
||||
stays out of a configuration that gets copied, diffed and pasted into issues.
|
||||
Keep that file at `0600` and owned by the user llmgw runs as.
|
||||
|
||||
Moving an existing installation is a copy, not a fresh start. Keys are stored
|
||||
hashed, so starting fresh means reissuing every one of them and being unable
|
||||
to say which was which:
|
||||
|
||||
```
|
||||
sudo systemctl stop llmgw-server
|
||||
llmgw keys move-store --from /var/lib/llmgw/llmgw.db --to "@/etc/llmgw/secrets/store-dsn"
|
||||
# point auth.store at the new address, then
|
||||
sudo systemctl start llmgw-server
|
||||
```
|
||||
|
||||
Stop the gateway first. The copy makes no attempt to follow a store being
|
||||
written to: rows landing after the table they belong to has been passed would
|
||||
simply be absent, and nothing would say so. The destination must be empty, and
|
||||
the source is left untouched — so the way back is to point the configuration at
|
||||
the file again.
|
||||
|
||||
On a machine that also runs the control plane, PostgreSQL's defaults assume it
|
||||
owns the box. A gateway swapped out to give the database cache is a worse
|
||||
outcome than a smaller cache, so cap it: `shared_buffers = 96MB` and
|
||||
`max_connections = 40` are enough for a store whose working set is measured in
|
||||
megabytes.
|
||||
|
||||
## Conversations across models
|
||||
|
||||
Off everywhere until a project asks. llmgw otherwise stores no message content
|
||||
@@ -735,9 +781,12 @@ requests. On a slow model, set `timeouts.generation` higher.
|
||||
|
||||
## Known gaps, so you find them here rather than at 2am
|
||||
|
||||
- **A fetch is silent while it runs.** llmgw will fetch weights for you, but
|
||||
it reports no progress, so a first start on a large model looks like a hang.
|
||||
`ls -l` in the model directory is the only progress bar there is.
|
||||
- **A fetch reports its progress where nobody is looking.** The agent logs a
|
||||
percentage as it downloads, so `journalctl -u llmgw-agent` on that node knows
|
||||
exactly how far along it is. Nothing carries it to the control plane, so from
|
||||
`ctl` or the dashboard a replica pulling a hundred gigabytes and one that is
|
||||
wedged still look identical — which is the comparison the log line itself
|
||||
says matters.
|
||||
- **Every node fetches its own copy.** There is no seeding between nodes, so
|
||||
bringing up eight nodes on a new model is eight downloads from the same
|
||||
origin. On a fabric where the nodes can see each other this is wasteful, and
|
||||
|
||||
@@ -6,6 +6,7 @@ require (
|
||||
connectrpc.com/connect v1.21.0
|
||||
github.com/NVIDIA/go-nvml v0.13.4-0
|
||||
github.com/cespare/xxhash/v2 v2.3.0
|
||||
github.com/jackc/pgx/v5 v5.11.0
|
||||
github.com/knadh/koanf/parsers/yaml v1.1.1
|
||||
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
||||
github.com/knadh/koanf/providers/file v1.2.1
|
||||
@@ -14,6 +15,7 @@ require (
|
||||
github.com/prometheus/client_model v0.6.3
|
||||
github.com/prometheus/common v0.71.0
|
||||
github.com/spf13/cobra v1.10.2
|
||||
golang.org/x/crypto v0.57.0
|
||||
golang.org/x/net v0.59.0
|
||||
golang.org/x/sync v0.23.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
@@ -27,6 +29,9 @@ require (
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/knadh/koanf/maps v0.1.2 // indirect
|
||||
github.com/mattn/go-isatty v0.0.24 // indirect
|
||||
github.com/mitchellh/copystructure v1.2.0 // indirect
|
||||
@@ -37,7 +42,6 @@ require (
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/spf13/pflag v1.0.9 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/crypto v0.57.0 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
modernc.org/libc v1.75.7 // indirect
|
||||
|
||||
@@ -7,6 +7,7 @@ github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6r
|
||||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=
|
||||
@@ -23,6 +24,14 @@ github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.11.0 h1:IzBBtyK9AHqf98cctWFifYSci2hgQR/cd56wB4p+ogg=
|
||||
github.com/jackc/pgx/v5 v5.11.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk=
|
||||
github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
|
||||
github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo=
|
||||
@@ -47,6 +56,7 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU=
|
||||
github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE=
|
||||
github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo=
|
||||
@@ -62,6 +72,9 @@ github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
|
||||
github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
|
||||
github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY=
|
||||
github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE=
|
||||
github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg=
|
||||
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
|
||||
@@ -88,6 +101,7 @@ golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo
|
||||
google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc=
|
||||
google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
modernc.org/cc/v4 v4.29.2 h1:h6+9ciCnPKutf4I03CvheAvDLX7+IHlqR6Iy6J+cgd8=
|
||||
modernc.org/cc/v4 v4.29.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
|
||||
modernc.org/ccgo/v4 v4.35.0 h1:F+TUsmw09QxLzmi3aeYYGxjAXarmZaKgj3mKQHNaA8w=
|
||||
|
||||
+79
-24
@@ -34,8 +34,18 @@ import (
|
||||
|
||||
// Config configures the agent.
|
||||
type Config struct {
|
||||
NodeID string `koanf:"node_id"`
|
||||
NodeID string `koanf:"node_id"`
|
||||
// Server is the control plane, and Servers is the same thing when there is
|
||||
// more than one. Both are accepted: every configuration written before
|
||||
// there could be a second control plane names one address, and none of
|
||||
// them should have to change.
|
||||
//
|
||||
// The order is tried in order, and a refusal moves to the next. Only the
|
||||
// leader accepts an agent, so a standby is skipped in the time it takes to
|
||||
// be told so — which is what makes failover a reconnect rather than an
|
||||
// edit to every node's configuration.
|
||||
Server string `koanf:"server"`
|
||||
Servers []string `koanf:"servers"`
|
||||
Hostname string `koanf:"hostname"`
|
||||
MgmtAddr string `koanf:"mgmt_addr"`
|
||||
Labels map[string]string `koanf:"labels"`
|
||||
@@ -108,9 +118,6 @@ func (c *Config) normalize() error {
|
||||
if c.NodeID == "" {
|
||||
return errors.New("node_id is required")
|
||||
}
|
||||
if c.Server == "" {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
if err := c.TLS.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -140,26 +147,59 @@ func (c *Config) normalize() error {
|
||||
// that says plaintext while TLS is configured. Silently connecting anyway would
|
||||
// mean a cluster believing it has mutual authentication when it has none.
|
||||
func (c *Config) normalizeServerURL() error {
|
||||
secure := c.TLS.Enabled()
|
||||
switch {
|
||||
case strings.HasPrefix(c.Server, "https://"):
|
||||
if !secure {
|
||||
return errors.New("server is an https:// address but no tls certificates are configured")
|
||||
}
|
||||
case strings.HasPrefix(c.Server, "http://"):
|
||||
if secure {
|
||||
return errors.New("tls is configured but server is an http:// address; use https://")
|
||||
}
|
||||
case strings.Contains(c.Server, "://"):
|
||||
return fmt.Errorf("server %q must be an http:// or https:// address", c.Server)
|
||||
case secure:
|
||||
c.Server = "https://" + c.Server
|
||||
default:
|
||||
c.Server = "http://" + c.Server
|
||||
// One list, however it was written. Server first when both are given,
|
||||
// because an operator adding a standby to an existing file expects the
|
||||
// one already there to stay preferred.
|
||||
all := make([]string, 0, 1+len(c.Servers))
|
||||
if c.Server != "" {
|
||||
all = append(all, c.Server)
|
||||
}
|
||||
all = append(all, c.Servers...)
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
out := make([]string, 0, len(all))
|
||||
for _, s := range all {
|
||||
fixed, err := normalizeOne(s, c.TLS.Enabled())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// A duplicate would be tried twice in a row on failover, which reads
|
||||
// as the rotation not working.
|
||||
if _, dup := seen[fixed]; dup {
|
||||
continue
|
||||
}
|
||||
seen[fixed] = struct{}{}
|
||||
out = append(out, fixed)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
c.Servers = out
|
||||
c.Server = out[0]
|
||||
return nil
|
||||
}
|
||||
|
||||
// normalizeOne applies the scheme rules to a single address.
|
||||
func normalizeOne(addr string, secure bool) (string, error) {
|
||||
switch {
|
||||
case strings.HasPrefix(addr, "https://"):
|
||||
if !secure {
|
||||
return "", fmt.Errorf("%s is an https:// address but no tls certificates are configured", addr)
|
||||
}
|
||||
case strings.HasPrefix(addr, "http://"):
|
||||
if secure {
|
||||
return "", fmt.Errorf("tls is configured but %s is an http:// address; use https://", addr)
|
||||
}
|
||||
case strings.Contains(addr, "://"):
|
||||
return "", fmt.Errorf("server %q must be an http:// or https:// address", addr)
|
||||
case secure:
|
||||
return "https://" + addr, nil
|
||||
default:
|
||||
return "http://" + addr, nil
|
||||
}
|
||||
return addr, nil
|
||||
}
|
||||
|
||||
// Agent runs on a GPU node.
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
@@ -230,15 +270,30 @@ func (a *Agent) Run(ctx context.Context) error {
|
||||
a.recoverOrphans()
|
||||
|
||||
attempt := 0
|
||||
next := 0
|
||||
for ctx.Err() == nil {
|
||||
err := a.session(ctx)
|
||||
server := a.cfg.Servers[next%len(a.cfg.Servers)]
|
||||
err := a.session(ctx, server)
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
// On to the next address whatever went wrong. A standby refuses, a
|
||||
// stopped control plane refuses the connection, and a crashed one
|
||||
// never answers — from here they are the same thing, and trying
|
||||
// somewhere else is the answer to all three.
|
||||
next++
|
||||
attempt++
|
||||
delay := a.reconnectDelay(attempt)
|
||||
// No wait when there is somewhere else to try: the delay exists to
|
||||
// stop hammering one control plane, and the next address has not been
|
||||
// hammered. With a single address this is the old behaviour exactly.
|
||||
if len(a.cfg.Servers) > 1 && next%len(a.cfg.Servers) != 0 {
|
||||
delay = 0
|
||||
}
|
||||
a.log.Warn("disconnected from the control plane, retrying",
|
||||
"err", err, "in", delay.Round(time.Millisecond), "attempt", attempt)
|
||||
"err", err, "server", server,
|
||||
"next", a.cfg.Servers[next%len(a.cfg.Servers)],
|
||||
"in", delay.Round(time.Millisecond), "attempt", attempt)
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -469,8 +524,8 @@ func (a *Agent) scan(ctx context.Context) error {
|
||||
}
|
||||
|
||||
// session runs one connection until it breaks.
|
||||
func (a *Agent) session(ctx context.Context) error {
|
||||
client := llmgwv1connect.NewAgentServiceClient(a.http, a.cfg.Server)
|
||||
func (a *Agent) session(ctx context.Context, server string) error {
|
||||
client := llmgwv1connect.NewAgentServiceClient(a.http, server)
|
||||
|
||||
// parent is kept because the derived context is cancelled from inside this
|
||||
// function; asking the derived one whether we are shutting down would
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/h2c"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
type standby bool
|
||||
|
||||
func (s standby) IsLeader() bool { return !bool(s) }
|
||||
|
||||
// startPlain runs an agent service over h2c, which is enough to test who
|
||||
// accepts a connection and who refuses it.
|
||||
//
|
||||
// A nil leader means none was configured, which is every deployment that
|
||||
// existed before leadership did and has to keep accepting agents.
|
||||
func startPlain(t *testing.T, leader agentsvc.Leader) (*inventory.Store, string) {
|
||||
t.Helper()
|
||||
inv := inventory.NewStore()
|
||||
svc := agentsvc.New(inv, newRecorder(), false, slog.New(slog.DiscardHandler))
|
||||
if leader != nil {
|
||||
svc.SetLeadership(leader)
|
||||
}
|
||||
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
// h2c, because the agent speaks HTTP/2 with prior knowledge and a plain
|
||||
// server negotiates HTTP/1.1 — which turns a clear refusal into a frame
|
||||
// parsing error that says nothing about why.
|
||||
srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))
|
||||
t.Cleanup(srv.Close)
|
||||
return inv, srv.URL
|
||||
}
|
||||
|
||||
func TestAStandbyRefusesAnAgent(t *testing.T) {
|
||||
// Commands travel down the stream the agent opened, so whichever control
|
||||
// plane holds it is the only one that can start anything on that node. A
|
||||
// standby holding it means the leader's placements fail with "no agent
|
||||
// connected" for a node that is plainly connected.
|
||||
_, addr := startPlain(t, standby(true))
|
||||
|
||||
ag := runAgent(t, Config{NodeID: "n1", Server: addr, MgmtAddr: "127.0.0.1"})
|
||||
err := ag.session(context.Background(), ag.cfg.Server)
|
||||
if err == nil {
|
||||
t.Fatal("a standby accepted an agent")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "standby") {
|
||||
t.Errorf("the refusal does not say why: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAgentReachesTheLeaderPastAStandby(t *testing.T) {
|
||||
// The point of the list. The standby is first, so the agent only registers
|
||||
// if it moved on rather than sitting on the refusal.
|
||||
_, standbyAddr := startPlain(t, standby(true))
|
||||
leaderInv, leaderAddr := startPlain(t, standby(false))
|
||||
|
||||
ag := runAgent(t, Config{
|
||||
NodeID: "n1", MgmtAddr: "127.0.0.1",
|
||||
Servers: []string{standbyAddr, leaderAddr},
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
go ag.Run(ctx)
|
||||
|
||||
deadline := time.Now().Add(8 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, ok := leaderInv.Get("n1"); ok {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the agent never reached the leader past the standby")
|
||||
}
|
||||
|
||||
func TestALeaderStillAcceptsWithNoLeadershipConfigured(t *testing.T) {
|
||||
// Every deployment before this had no leadership at all, and must keep
|
||||
// accepting agents with nothing configured.
|
||||
inv, addr := startPlain(t, nil)
|
||||
|
||||
ag := runAgent(t, Config{NodeID: "n1", Server: addr, MgmtAddr: "127.0.0.1"})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
go ag.Run(ctx)
|
||||
|
||||
deadline := time.Now().Add(4 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, ok := inv.Get("n1"); ok {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("a service with no leadership configured refused an agent")
|
||||
}
|
||||
@@ -255,7 +255,7 @@ func TestRefusedAgentLearnsWhy(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
err = ag.session(context.Background())
|
||||
err = ag.session(context.Background(), ag.cfg.Server)
|
||||
if err == nil {
|
||||
t.Fatal("a refused session reported success")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
func normalized(t *testing.T, c Config) Config {
|
||||
t.Helper()
|
||||
if err := c.normalizeServerURL(); err != nil {
|
||||
t.Fatalf("normalising %+v: %v", c, err)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func TestOneServerStillWorksExactlyAsBefore(t *testing.T) {
|
||||
// Every agent configuration written before there could be a second control
|
||||
// plane names one address, and none of them should have to change.
|
||||
c := normalized(t, Config{Server: "10.0.0.140:7443"})
|
||||
if !slices.Equal(c.Servers, []string{"http://10.0.0.140:7443"}) {
|
||||
t.Errorf("servers = %v", c.Servers)
|
||||
}
|
||||
if c.Server != "http://10.0.0.140:7443" {
|
||||
t.Errorf("server = %q", c.Server)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheAddressAlreadyThereStaysPreferred(t *testing.T) {
|
||||
// An operator adding a standby to an existing file expects the control
|
||||
// plane that was already working to keep being tried first.
|
||||
c := normalized(t, Config{
|
||||
Server: "first:7443",
|
||||
Servers: []string{"second:7443", "third:7443"},
|
||||
})
|
||||
want := []string{"http://first:7443", "http://second:7443", "http://third:7443"}
|
||||
if !slices.Equal(c.Servers, want) {
|
||||
t.Errorf("servers = %v, want %v", c.Servers, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestADuplicateAddressIsNotTriedTwiceInARow(t *testing.T) {
|
||||
// Which would look exactly like the rotation not working.
|
||||
c := normalized(t, Config{
|
||||
Server: "a:7443",
|
||||
Servers: []string{"http://a:7443", "b:7443"},
|
||||
})
|
||||
want := []string{"http://a:7443", "http://b:7443"}
|
||||
if !slices.Equal(c.Servers, want) {
|
||||
t.Errorf("servers = %v, want %v", c.Servers, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEveryAddressIsCheckedAgainstTheTLSSetting(t *testing.T) {
|
||||
// The check existed for one address and has to hold for all of them: a
|
||||
// standby reached over plaintext while the leader is authenticated would
|
||||
// be a way around mutual TLS, not a fallback.
|
||||
c := Config{
|
||||
Servers: []string{"https://a:7443", "http://b:7443"},
|
||||
TLS: tlsconf.Config{CA: "ca.pem", Cert: "c.pem", Key: "k.pem"},
|
||||
}
|
||||
err := c.normalizeServerURL()
|
||||
if err == nil {
|
||||
t.Fatal("a plaintext standby was accepted alongside TLS")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "b:7443") {
|
||||
t.Errorf("the error does not say which address: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoServerAtAllIsStillAnError(t *testing.T) {
|
||||
c := Config{}
|
||||
if err := c.normalizeServerURL(); err == nil {
|
||||
t.Error("an agent with no control plane configured was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchemeIsAddedPerAddress(t *testing.T) {
|
||||
c := normalized(t, Config{
|
||||
Servers: []string{"a:7443", "b:7443"},
|
||||
TLS: tlsconf.Config{CA: "ca.pem", Cert: "c.pem", Key: "k.pem"},
|
||||
})
|
||||
for _, s := range c.Servers {
|
||||
if !strings.HasPrefix(s, "https://") {
|
||||
t.Errorf("%q did not get the secure scheme", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -27,7 +27,7 @@ func newTenancyCmd() *cobra.Command {
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store")
|
||||
|
||||
open := func(ctx context.Context) (*store.SQLite, error) { return store.OpenSQLite(ctx, dbPath) }
|
||||
open := func(ctx context.Context) (*store.SQLite, error) { return store.Open(ctx, dbPath) }
|
||||
|
||||
cmd.AddCommand(
|
||||
newKeyInitCmd(open),
|
||||
@@ -37,6 +37,7 @@ func newTenancyCmd() *cobra.Command {
|
||||
newLimitsCmd(open),
|
||||
newThreadsCmd(open),
|
||||
newMemoryCmd(open),
|
||||
newMoveStoreCmd(),
|
||||
)
|
||||
return cmd
|
||||
}
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// newMoveStoreCmd copies a store into another one.
|
||||
//
|
||||
// Written because keys are stored hashed. Moving to PostgreSQL without this
|
||||
// means every person and every script needs a new key issued, and the old ones
|
||||
// cannot even be looked up to say which they were.
|
||||
//
|
||||
// It reads and writes, and nothing else runs while it does: run it with the
|
||||
// gateway stopped. A copy taken while requests are landing would miss whatever
|
||||
// arrived after the table it had already passed.
|
||||
func newMoveStoreCmd() *cobra.Command {
|
||||
var from, to string
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "move-store",
|
||||
Short: "Copy a store into another one, for moving SQLite to PostgreSQL",
|
||||
Long: "Copies every table, parents first, into an empty destination.\n\n" +
|
||||
"Stop the gateway first. This makes no attempt to copy a store that is\n" +
|
||||
"being written to: rows landing after the table they belong to has been\n" +
|
||||
"passed would simply be absent, and nothing would say so.\n\n" +
|
||||
"The destination must be empty. Merging two stores means deciding which\n" +
|
||||
"of two projects with the same id was meant, and there is no answer to\n" +
|
||||
"that which is not a guess.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if from == "" || to == "" {
|
||||
return fmt.Errorf("--from and --to are both required")
|
||||
}
|
||||
src, err := store.Open(cmd.Context(), from)
|
||||
if err != nil {
|
||||
return fmt.Errorf("opening the source: %w", err)
|
||||
}
|
||||
defer src.Close()
|
||||
|
||||
dst, err := store.Open(cmd.Context(), to)
|
||||
if err != nil {
|
||||
return fmt.Errorf("opening the destination: %w", err)
|
||||
}
|
||||
defer dst.Close()
|
||||
|
||||
res, err := src.CopyTo(cmd.Context(), dst)
|
||||
// Printed even on failure: a partial copy is the one case where
|
||||
// knowing how far it got decides what to do next.
|
||||
out := cmd.OutOrStdout()
|
||||
var total int64
|
||||
for _, r := range res {
|
||||
if r.Rows > 0 {
|
||||
fmt.Fprintf(out, " %-16s %d\n", r.Table, r.Rows)
|
||||
}
|
||||
total += r.Rows
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(out, "\n%d rows copied. Point the gateway at the new store and start it.\n", total)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&from, "from", "", "source: a SQLite path, a postgres:// URL, or @file holding one")
|
||||
cmd.Flags().StringVar(&to, "to", "", "destination, same forms")
|
||||
return cmd
|
||||
}
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
|
||||
storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/web"
|
||||
)
|
||||
@@ -335,7 +336,29 @@ func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Live
|
||||
if u, ok := sink.(control.UsageSource); ok {
|
||||
ctl.SetUsageSource(u)
|
||||
}
|
||||
// With a shared store, two control planes can see each other, so this one
|
||||
// competes rather than assuming it is alone. Without one there is nobody
|
||||
// to compete with and it places unconditionally, exactly as before.
|
||||
// The same shape as the usage assertion above, and for the same reason:
|
||||
// the sink is the router when the router is serving, and nothing when a
|
||||
// control plane runs on its own.
|
||||
var lead *control.Leadership
|
||||
if h, ok := sink.(interface{ Store() *storepkg.SQLite }); ok {
|
||||
if st := h.Store(); st != nil {
|
||||
lead = control.NewLeadership(st, cfg.Control.LeaseTTL, log)
|
||||
ctl.SetLeadership(lead)
|
||||
log.Info("competing to be the control plane that places work",
|
||||
"owner", lead.Owner())
|
||||
}
|
||||
}
|
||||
svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), requireIdentity, log)
|
||||
// A standby must refuse agents as well as refusing to place. Commands
|
||||
// travel down the stream the agent opened, so a standby holding one is a
|
||||
// node the leader cannot command — and the leader's placements would fail
|
||||
// with "no agent connected" for a node that is plainly connected.
|
||||
if lead != nil {
|
||||
svc.SetLeadership(lead)
|
||||
}
|
||||
fwd.set(svc)
|
||||
|
||||
ctl.Start(ctx)
|
||||
|
||||
@@ -24,7 +24,7 @@ func newUsersCmd() *cobra.Command {
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store")
|
||||
open := func(cmd *cobra.Command) (*store.SQLite, error) {
|
||||
return store.OpenSQLite(cmd.Context(), dbPath)
|
||||
return store.Open(cmd.Context(), dbPath)
|
||||
}
|
||||
|
||||
cmd.AddCommand(newUserListCmd(open), newUserApproveCmd(open),
|
||||
|
||||
@@ -177,6 +177,15 @@ type Control struct {
|
||||
PortRange []int `koanf:"port_range"`
|
||||
// MgmtInterface is the ethernet interface NCCL bootstraps over.
|
||||
MgmtInterface string `koanf:"mgmt_interface"`
|
||||
|
||||
// LeaseTTL is how long this control plane's claim to be the one placing
|
||||
// work stands without being renewed.
|
||||
//
|
||||
// Only meaningful with a shared store, where a second control plane can
|
||||
// see the claim. Long enough to survive a slow query or a brief network
|
||||
// stall, short enough that a standby takes over before anyone has finished
|
||||
// reading the alert. Zero takes the default.
|
||||
LeaseTTL time.Duration `koanf:"lease_ttl"`
|
||||
}
|
||||
|
||||
// Router holds routing behaviour.
|
||||
|
||||
@@ -31,6 +31,9 @@ type StatusSink interface {
|
||||
// time the reconciler's Dispatcher: it holds the open streams, so it is the
|
||||
// only thing that can deliver a command to a node.
|
||||
type Service struct {
|
||||
// leader gates whether this process accepts agent streams at all.
|
||||
leader Leader
|
||||
|
||||
inv *inventory.Store
|
||||
status StatusSink
|
||||
log *slog.Logger
|
||||
@@ -86,6 +89,19 @@ func (s *Service) Connect(
|
||||
ctx context.Context,
|
||||
stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage],
|
||||
) error {
|
||||
// A standby must not hold an agent's stream. Commands travel down the
|
||||
// stream the agent opened, so whichever control plane holds it is the only
|
||||
// one that can start or stop anything on that node — and if that is not
|
||||
// the leader, the leader's placements fail with "no agent connected" for a
|
||||
// node that is plainly connected.
|
||||
//
|
||||
// Refused before the opening message is read, so the agent spends nothing
|
||||
// on a connection that cannot serve it.
|
||||
if s.leader != nil && !s.leader.IsLeader() {
|
||||
return connect.NewError(connect.CodeUnavailable,
|
||||
errors.New("this control plane is a standby; connect to the leader"))
|
||||
}
|
||||
|
||||
first, err := stream.Receive()
|
||||
if err != nil {
|
||||
return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("no opening message: %w", err))
|
||||
@@ -345,6 +361,16 @@ func (s *Service) ConnectedNodes() []string {
|
||||
|
||||
// --- Dispatcher ---
|
||||
|
||||
// Leader reports whether this process is the one that should be holding agent
|
||||
// streams.
|
||||
type Leader interface{ IsLeader() bool }
|
||||
|
||||
// SetLeadership makes this service refuse agents while it is a standby.
|
||||
//
|
||||
// Nil accepts them unconditionally, which is what a single control plane wants
|
||||
// and what every deployment before leadership had.
|
||||
func (s *Service) SetLeadership(l Leader) { s.leader = l }
|
||||
|
||||
// ErrNotConnected means no agent stream is open for that node.
|
||||
var ErrNotConnected = errors.New("no agent connected")
|
||||
|
||||
|
||||
@@ -75,6 +75,7 @@ type Controller struct {
|
||||
inv *inventory.Store
|
||||
cat *catalog.Live
|
||||
rec *reconciler.Reconciler
|
||||
lead *Leadership
|
||||
disp reconciler.Dispatcher
|
||||
sink EndpointSink
|
||||
usage UsageSource
|
||||
@@ -172,6 +173,22 @@ func (c *Controller) Reconciler() *reconciler.Reconciler { return c.rec }
|
||||
// Catalog exposes the model catalog.
|
||||
func (c *Controller) Catalog() *catalog.Catalog { return c.cat.Load() }
|
||||
|
||||
// SetLeadership makes this control plane compete for the right to place work,
|
||||
// instead of assuming it is the only one.
|
||||
//
|
||||
// Wired after construction because deciding leadership needs a store, and the
|
||||
// store is opened by the router. Left unset, this process places
|
||||
// unconditionally — which is what a single control plane wants, and what every
|
||||
// deployment that existed before this had.
|
||||
func (c *Controller) SetLeadership(l *Leadership) {
|
||||
c.lead = l
|
||||
c.rec.SetLeadership(l)
|
||||
}
|
||||
|
||||
// IsLeader reports whether this process is placing work. True when leadership
|
||||
// was never configured, because then there is nobody to compete with.
|
||||
func (c *Controller) IsLeader() bool { return c.lead == nil || c.lead.IsLeader() }
|
||||
|
||||
// WeightRemover can ask a node to delete cached model files.
|
||||
//
|
||||
// Separate from Dispatcher because it is not part of scheduling and not every
|
||||
@@ -189,6 +206,9 @@ func (c *Controller) WeightRemover() WeightRemover {
|
||||
|
||||
// Start runs the reconciler and begins publishing endpoints.
|
||||
func (c *Controller) Start(ctx context.Context) {
|
||||
if c.lead != nil {
|
||||
go c.lead.Run(ctx)
|
||||
}
|
||||
c.rec.Start(ctx)
|
||||
|
||||
go func() {
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"os"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// LeaseStore is the part of the store leadership needs.
|
||||
//
|
||||
// An interface so the controller does not take a dependency on a concrete
|
||||
// database for one feature, and so a deployment with no store at all — a
|
||||
// router-only run, or a single control plane on a static configuration — keeps
|
||||
// working with nothing to configure.
|
||||
type LeaseStore interface {
|
||||
AcquireLeadership(ctx context.Context, role, owner string, ttl time.Duration) (store.Lease, error)
|
||||
ReleaseLeadership(ctx context.Context, role, owner string) error
|
||||
}
|
||||
|
||||
// controlRole is the one thing there is to be leader of: placing work.
|
||||
const controlRole = "control-plane"
|
||||
|
||||
// Leadership decides whether this process is the one that places work.
|
||||
//
|
||||
// Two control planes against one store both reconcile: both see a model
|
||||
// wanting two replicas and none running, and both place. The result is twice
|
||||
// the replicas, each holding GPUs the other believes are free, and nothing
|
||||
// crashes to say so.
|
||||
//
|
||||
// A standby is not idle — it serves its own routing and keeps its connection
|
||||
// to the store — it simply does not place, evict or drain. That makes running
|
||||
// a second one safe rather than clever, which is the property worth having
|
||||
// first: the common failure is not a planned failover, it is somebody starting
|
||||
// a second control plane by accident.
|
||||
type Leadership struct {
|
||||
store LeaseStore
|
||||
log *slog.Logger
|
||||
// owner identifies this process. Random rather than the hostname, so a
|
||||
// restart on the same machine cannot be fooled by its own stale lease.
|
||||
owner string
|
||||
ttl time.Duration
|
||||
// renew is how often the holder pushes the expiry out. Well under ttl, so
|
||||
// one slow query does not cost the leadership.
|
||||
renew time.Duration
|
||||
|
||||
leader atomic.Bool
|
||||
}
|
||||
|
||||
// NewLeadership prepares leadership over a store.
|
||||
//
|
||||
// A nil store means there is nobody to coordinate with, and this process is
|
||||
// the leader by definition — which is every single-control-plane deployment,
|
||||
// and must stay the case without anyone opting in.
|
||||
func NewLeadership(st LeaseStore, ttl time.Duration, log *slog.Logger) *Leadership {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if ttl <= 0 {
|
||||
// Long enough to survive a slow database or a brief network stall,
|
||||
// short enough that a standby takes over before anyone has finished
|
||||
// reading the alert.
|
||||
ttl = 15 * time.Second
|
||||
}
|
||||
l := &Leadership{
|
||||
store: st,
|
||||
log: log,
|
||||
owner: newOwnerID(),
|
||||
ttl: ttl,
|
||||
renew: ttl / 3,
|
||||
}
|
||||
if st == nil {
|
||||
l.leader.Store(true)
|
||||
}
|
||||
return l
|
||||
}
|
||||
|
||||
// IsLeader reports whether this process should be placing work.
|
||||
func (l *Leadership) IsLeader() bool { return l.leader.Load() }
|
||||
|
||||
// Owner is this process's id, for logs and for ctl.
|
||||
func (l *Leadership) Owner() string { return l.owner }
|
||||
|
||||
// Run keeps the lease for as long as it can hold it, and keeps trying when it
|
||||
// cannot.
|
||||
//
|
||||
// A standby polls at the same interval the leader renews at, so a lapsed lease
|
||||
// is picked up within about a third of its life rather than at the end of it.
|
||||
func (l *Leadership) Run(ctx context.Context) {
|
||||
if l.store == nil {
|
||||
return
|
||||
}
|
||||
t := time.NewTicker(l.renew)
|
||||
defer t.Stop()
|
||||
|
||||
l.attempt(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
l.step(ctx)
|
||||
return
|
||||
case <-t.C:
|
||||
l.attempt(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// attempt takes or renews the lease, logging only transitions.
|
||||
//
|
||||
// Only transitions, because this runs every few seconds for the life of the
|
||||
// process: a line per attempt would bury everything else, and the thing an
|
||||
// operator needs to see is the moment leadership moved.
|
||||
func (l *Leadership) attempt(ctx context.Context) {
|
||||
lease, err := l.store.AcquireLeadership(ctx, controlRole, l.owner, l.ttl)
|
||||
switch {
|
||||
case err == nil:
|
||||
if l.leader.CompareAndSwap(false, true) {
|
||||
l.log.Info("this control plane is now placing work", "owner", l.owner)
|
||||
}
|
||||
case errors.Is(err, store.ErrNotLeader):
|
||||
if l.leader.CompareAndSwap(true, false) {
|
||||
// Losing a lease that was held is worth a warning: it means this
|
||||
// process was too slow to renew, and whatever made it slow is
|
||||
// still true.
|
||||
l.log.Warn("lost leadership; standing by",
|
||||
"owner", l.owner, "now", lease.Owner)
|
||||
} else if l.leader.Load() == false && lease.Owner != "" {
|
||||
l.log.Debug("standing by", "leader", lease.Owner,
|
||||
"until", lease.ExpiresAt.Format(time.RFC3339))
|
||||
}
|
||||
default:
|
||||
// A store that cannot be reached is not evidence that somebody else
|
||||
// is leading. Keeping the claim is the safer of two bad options: two
|
||||
// leaders is worse than one, but no leader at all means nothing is
|
||||
// placed and nothing is evicted, and a brief database blip should not
|
||||
// stop a cluster from healing itself.
|
||||
l.log.Warn("could not reach the store to renew leadership; keeping it for now",
|
||||
"err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// step gives the lease up on the way out, so a standby takes over in seconds
|
||||
// rather than waiting out the clock.
|
||||
func (l *Leadership) step(ctx context.Context) {
|
||||
if !l.leader.Load() {
|
||||
return
|
||||
}
|
||||
// A fresh context: the one that just ended is the reason this is running.
|
||||
ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
|
||||
defer cancel()
|
||||
if err := l.store.ReleaseLeadership(ctx, controlRole, l.owner); err != nil {
|
||||
l.log.Warn("could not release leadership; a standby will wait it out", "err", err)
|
||||
return
|
||||
}
|
||||
l.log.Info("released leadership", "owner", l.owner)
|
||||
}
|
||||
|
||||
// newOwnerID is random, not the hostname.
|
||||
//
|
||||
// A hostname would let a restarted process match its own stale lease and
|
||||
// believe it had been leading all along — including through whatever gap made
|
||||
// it restart, during which another process may have taken over and started
|
||||
// placing.
|
||||
func newOwnerID() string {
|
||||
var b [8]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
// Only reached if the system has no entropy, where a pid and the time
|
||||
// is still better than a fixed string.
|
||||
return "ctl-" + hex.EncodeToString([]byte(time.Now().Format(time.RFC3339Nano)))[:16]
|
||||
}
|
||||
host, _ := os.Hostname()
|
||||
if host == "" {
|
||||
host = "ctl"
|
||||
}
|
||||
return host + "-" + hex.EncodeToString(b[:])
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
type fixedLeader bool
|
||||
|
||||
func (f fixedLeader) IsLeader() bool { return bool(f) }
|
||||
|
||||
func TestAStandbyPlacesNothing(t *testing.T) {
|
||||
// The failure this prevents: two control planes against one store both see
|
||||
// a model wanting replicas and none running, and both place. The model
|
||||
// ends up with twice what it asked for, each replica holding GPUs the
|
||||
// other's inventory believes are free, and nothing crashes to say so.
|
||||
h := newHarness(t, testModel(1), gpuNode("n1", 8))
|
||||
h.r.SetLeadership(fixedLeader(false))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.r.Replicas()); got != 0 {
|
||||
t.Errorf("a standby placed %d replicas", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestALeaderPlacesAsBefore(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("n1", 8))
|
||||
h.r.SetLeadership(fixedLeader(true))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.r.Replicas()); got == 0 {
|
||||
t.Error("a leader placed nothing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithoutLeadershipConfiguredNothingChanges(t *testing.T) {
|
||||
// Every single-control-plane deployment, and every test written before
|
||||
// leadership existed. Unset must mean "act", not "wait for permission
|
||||
// from a thing that was never wired up".
|
||||
h := newHarness(t, testModel(1), gpuNode("n1", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.r.Replicas()); got == 0 {
|
||||
t.Error("a reconciler with no leadership configured placed nothing")
|
||||
}
|
||||
}
|
||||
@@ -166,6 +166,9 @@ func (c *Config) normalize() {
|
||||
|
||||
// Reconciler compares the catalog against what is running and closes the gap.
|
||||
type Reconciler struct {
|
||||
// leader gates every action; nil means act unconditionally.
|
||||
leader Leader
|
||||
|
||||
cfg Config
|
||||
cat *catalog.Live
|
||||
inv *inventory.Store
|
||||
@@ -235,6 +238,17 @@ func (r *Reconciler) SetClock(now func() time.Time) {
|
||||
r.startedAt = now()
|
||||
}
|
||||
|
||||
// Leader reports whether this process should be acting on the fleet.
|
||||
type Leader interface{ IsLeader() bool }
|
||||
|
||||
// SetLeadership gates every action on holding leadership.
|
||||
//
|
||||
// A setter rather than a constructor argument because the thing that decides
|
||||
// leadership needs a store, and the store is wired after the controller. Nil
|
||||
// leaves the reconciler acting unconditionally, which is what a single control
|
||||
// plane wants and what every test expects.
|
||||
func (r *Reconciler) SetLeadership(l Leader) { r.leader = l }
|
||||
|
||||
// Start runs the loop until Close.
|
||||
func (r *Reconciler) Start(ctx context.Context) {
|
||||
go func() {
|
||||
@@ -278,6 +292,16 @@ func (r *Reconciler) Notify() {
|
||||
// Reconcile performs one pass. It is exported so tests can drive it directly
|
||||
// instead of racing a ticker.
|
||||
func (r *Reconciler) Reconcile(ctx context.Context) {
|
||||
// A standby does nothing at all here. Everything below issues commands to
|
||||
// agents — placing, evicting, draining, reaping — and two control planes
|
||||
// doing that against one fleet is how a model ends up with twice the
|
||||
// replicas it asked for, each holding GPUs the other believes are free.
|
||||
//
|
||||
// Unset means leader, so a single control plane and every test behave
|
||||
// exactly as they did before leadership existed.
|
||||
if r.leader != nil && !r.leader.IsLeader() {
|
||||
return
|
||||
}
|
||||
r.expireNodes()
|
||||
r.expireStuckReplicas()
|
||||
r.drainCordonedNodes(ctx)
|
||||
|
||||
@@ -189,7 +189,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
}, log)
|
||||
|
||||
if cfg.Auth.Enabled {
|
||||
st, err := storepkg.OpenSQLite(context.Background(), cfg.Auth.Store)
|
||||
st, err := storepkg.Open(context.Background(), cfg.Auth.Store)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("auth store: %w", err)
|
||||
}
|
||||
@@ -234,7 +234,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
r.recorder, log)
|
||||
r.threads = r.store
|
||||
} else if cfg.Threads.Store != "" {
|
||||
st, err := storepkg.OpenSQLite(context.Background(), cfg.Threads.Store)
|
||||
st, err := storepkg.Open(context.Background(), cfg.Threads.Store)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("threads store: %w", err)
|
||||
}
|
||||
|
||||
@@ -92,7 +92,7 @@ func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error)
|
||||
hash := HashKey(token)
|
||||
now := time.Now()
|
||||
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
q, tx, err := s.begin(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -100,7 +100,7 @@ func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error)
|
||||
|
||||
var sealed []byte
|
||||
var expires int64
|
||||
err = tx.QueryRowContext(ctx,
|
||||
err = q.QueryRowContext(ctx,
|
||||
`SELECT sealed, expires_at FROM key_claims WHERE token_hash = ? AND claimed_at IS NULL`,
|
||||
hash).Scan(&sealed, &expires)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
@@ -120,7 +120,7 @@ func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error)
|
||||
|
||||
// Blank the sealed value as well as marking it claimed, so a used link
|
||||
// leaves nothing behind even before the row is purged.
|
||||
res, err := tx.ExecContext(ctx,
|
||||
res, err := q.ExecContext(ctx,
|
||||
`UPDATE key_claims SET claimed_at = ?, sealed = x'' WHERE token_hash = ? AND claimed_at IS NULL`,
|
||||
millis(now), hash)
|
||||
if err != nil {
|
||||
|
||||
@@ -67,7 +67,7 @@ func TestTheDatabaseNeverHoldsAUsableKey(t *testing.T) {
|
||||
s, token, secret := claimFixture(t)
|
||||
|
||||
var sealed []byte
|
||||
if err := s.db.QueryRow(`SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
if err := s.db.QueryRowContext(context.Background(), `SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sealed) == 0 {
|
||||
@@ -79,7 +79,7 @@ func TestTheDatabaseNeverHoldsAUsableKey(t *testing.T) {
|
||||
// And the token itself is not there either, so nothing in the database can
|
||||
// derive the key that opens it.
|
||||
var count int
|
||||
if err := s.db.QueryRow(`SELECT COUNT(*) FROM key_claims WHERE token_hash = ?`, []byte(token)).
|
||||
if err := s.db.QueryRowContext(context.Background(), `SELECT COUNT(*) FROM key_claims WHERE token_hash = ?`, []byte(token)).
|
||||
Scan(&count); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -98,7 +98,7 @@ func TestAClaimedLinkLeavesNothingBehind(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var sealed []byte
|
||||
if err := s.db.QueryRow(`SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
if err := s.db.QueryRowContext(context.Background(), `SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sealed) != 0 {
|
||||
@@ -149,7 +149,7 @@ func TestAnExpiredLinkIsRefused(t *testing.T) {
|
||||
// Age it directly. CreateClaim treats a non-positive TTL as "use the
|
||||
// default", which is the right thing for a caller and no way to build the
|
||||
// row this test needs.
|
||||
if _, err := s.db.Exec(`UPDATE key_claims SET expires_at = ?`,
|
||||
if _, err := s.db.ExecContext(context.Background(), `UPDATE key_claims SET expires_at = ?`,
|
||||
time.Now().Add(-time.Hour).UnixMilli()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// copyOrder is every table, parents before children.
|
||||
//
|
||||
// Written out rather than discovered, because the order is a fact about the
|
||||
// foreign keys and getting it wrong produces a constraint violation halfway
|
||||
// through — with some tables copied and some not, which is the worst place to
|
||||
// stop. A table added without being added here is copied by nothing, and the
|
||||
// destination comes up quietly missing it; the test below is what catches
|
||||
// that.
|
||||
var copyOrder = []string{
|
||||
"organizations",
|
||||
"teams",
|
||||
"projects",
|
||||
// Before api_keys, which references it. The first version of this list
|
||||
// followed the order the tables are created across the migrations, and
|
||||
// users is created in a later one than api_keys even though api_keys
|
||||
// points at it — so the copy failed on a foreign key three tables in.
|
||||
"users",
|
||||
"api_keys",
|
||||
"key_claims",
|
||||
"sessions",
|
||||
"limits",
|
||||
"usage",
|
||||
"models",
|
||||
"threads",
|
||||
"thread_messages",
|
||||
"memory_notes",
|
||||
}
|
||||
|
||||
// notCopied are tables deliberately left behind.
|
||||
//
|
||||
// Leadership is a claim about a process that is running right now. Copying it
|
||||
// would hand the new store a leader that does not exist, and the standby
|
||||
// waiting on that lease would wait out its full expiry before noticing — at
|
||||
// the exact moment somebody is watching a migration and wondering why nothing
|
||||
// is being placed.
|
||||
var notCopied = map[string]string{
|
||||
"leadership": "a lease on a running process, meaningless in a new store",
|
||||
}
|
||||
|
||||
// CopyResult is how many rows each table received.
|
||||
type CopyResult struct {
|
||||
Table string
|
||||
Rows int64
|
||||
}
|
||||
|
||||
// CopyTo copies every row into another store.
|
||||
//
|
||||
// For moving from SQLite to PostgreSQL, which otherwise means reissuing every
|
||||
// API key: the keys are stored hashed, so there is no way to recreate them and
|
||||
// no way for anyone to keep using the one they have.
|
||||
//
|
||||
// The destination must be empty. Merging two stores would mean deciding what
|
||||
// happens when both have a project with the same id and different contents,
|
||||
// and every answer to that is a guess about which one someone meant.
|
||||
//
|
||||
// Not a transaction across the whole copy. It is one transaction per table,
|
||||
// because a single transaction holding every row of a large usage table is how
|
||||
// a migration runs out of memory on the machine least able to spare it. A
|
||||
// failure therefore leaves a partial copy, which is why the destination has to
|
||||
// have been empty: starting again is dropping and recreating it.
|
||||
func (s *SQLite) CopyTo(ctx context.Context, dst *SQLite) ([]CopyResult, error) {
|
||||
empty, where, err := dst.isEmpty(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !empty {
|
||||
return nil, fmt.Errorf("the destination already has rows in %s; copy into an empty store", where)
|
||||
}
|
||||
|
||||
var out []CopyResult
|
||||
for _, table := range copyOrder {
|
||||
n, err := s.copyTable(ctx, dst, table)
|
||||
if err != nil {
|
||||
return out, fmt.Errorf("copying %s: %w", table, err)
|
||||
}
|
||||
out = append(out, CopyResult{Table: table, Rows: n})
|
||||
}
|
||||
// Part of the copy, not a step someone has to remember: a destination
|
||||
// whose sequences still start at 1 is not a destination anyone can use.
|
||||
if err := dst.resetIdentities(ctx); err != nil {
|
||||
return out, fmt.Errorf("pointing the identity columns past the copied rows: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// isEmpty reports whether any table has rows, and which one first does.
|
||||
func (s *SQLite) isEmpty(ctx context.Context) (bool, string, error) {
|
||||
for _, table := range copyOrder {
|
||||
var n int64
|
||||
//nolint:gosec // table comes from copyOrder, not from input
|
||||
if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM `+table).Scan(&n); err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
if n > 0 {
|
||||
return false, table, nil
|
||||
}
|
||||
}
|
||||
return true, "", nil
|
||||
}
|
||||
|
||||
// copyTable moves one table, column names taken from the source.
|
||||
//
|
||||
// Reading the columns rather than listing them keeps this working when a
|
||||
// migration adds one: a hand-written column list would copy the old shape and
|
||||
// silently drop the new field, and nothing would notice until someone looked
|
||||
// for a setting that was there before the move.
|
||||
func (s *SQLite) copyTable(ctx context.Context, dst *SQLite, table string) (int64, error) {
|
||||
//nolint:gosec // table comes from copyOrder, not from input
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT * FROM `+table)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
q, tx, err := dst.begin(ctx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
insert := `INSERT INTO ` + table + ` (` + strings.Join(quoteAll(cols), ", ") +
|
||||
`) VALUES (` + strings.Repeat("?, ", len(cols)-1) + `?)`
|
||||
|
||||
var n int64
|
||||
for rows.Next() {
|
||||
vals := make([]any, len(cols))
|
||||
ptrs := make([]any, len(cols))
|
||||
for i := range vals {
|
||||
ptrs[i] = &vals[i]
|
||||
}
|
||||
if err := rows.Scan(ptrs...); err != nil {
|
||||
return n, err
|
||||
}
|
||||
for i, v := range vals {
|
||||
vals[i] = keepEmptyBytes(v)
|
||||
}
|
||||
if _, err := q.ExecContext(ctx, insert, vals...); err != nil {
|
||||
return n, err
|
||||
}
|
||||
n++
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return n, err
|
||||
}
|
||||
return n, tx.Commit()
|
||||
}
|
||||
|
||||
// keepEmptyBytes stops an empty byte column becoming NULL on the way across.
|
||||
//
|
||||
// An empty BLOB scans as a typed nil slice — []byte(nil) inside a non-nil any
|
||||
// — while a real NULL scans as an untyped nil. The two are distinguishable
|
||||
// here and not in pgx, which encodes any nil slice as NULL, so the difference
|
||||
// has to be restored before the insert.
|
||||
//
|
||||
// It is not a detail. key_claims.sealed is emptied the moment a claim is used,
|
||||
// so every claimed row holds an empty blob in a NOT NULL column, and the copy
|
||||
// failed on the first one — after the gateway had been stopped and four tables
|
||||
// had already been written.
|
||||
func keepEmptyBytes(v any) any {
|
||||
if b, ok := v.([]byte); ok && b == nil {
|
||||
return []byte{}
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// quoteAll quotes column names, because `key` and `text` are reserved words in
|
||||
// one dialect or the other and both are column names here.
|
||||
func quoteAll(cols []string) []string {
|
||||
out := make([]string, len(cols))
|
||||
for i, c := range cols {
|
||||
out[i] = `"` + c + `"`
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// resetIdentities points PostgreSQL's identity columns past the copied rows.
|
||||
//
|
||||
// Without it the first insert after a migration collides with a row that was
|
||||
// copied in: the sequence still starts at 1 while the table already holds a
|
||||
// thousand notes, and the failure arrives later, to somebody else, looking
|
||||
// like a bug in the feature rather than in the move.
|
||||
func (s *SQLite) resetIdentities(ctx context.Context) error {
|
||||
if s.dialect != PostgresDialect {
|
||||
return nil
|
||||
}
|
||||
// Only the tables with a generated id. SQLite's rowid needs nothing.
|
||||
for _, t := range []string{"memory_notes", "usage"} {
|
||||
var has bool
|
||||
if err := s.db.QueryRowContext(ctx,
|
||||
`SELECT EXISTS (SELECT 1 FROM information_schema.columns
|
||||
WHERE table_name = ? AND column_name = 'id'
|
||||
AND is_identity = 'YES')`, t).Scan(&has); err != nil {
|
||||
return err
|
||||
}
|
||||
if !has {
|
||||
continue
|
||||
}
|
||||
//nolint:gosec // t is a literal from the list above
|
||||
if _, err := s.db.ExecContext(ctx,
|
||||
`SELECT setval(pg_get_serial_sequence('`+t+`', 'id'),
|
||||
GREATEST((SELECT COALESCE(MAX(id), 0) FROM `+t+`), 1))`); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestEveryTableIsInTheCopyOrder(t *testing.T) {
|
||||
// A table added by a migration and not added here is copied by nothing,
|
||||
// and the destination comes up quietly missing it — discovered when
|
||||
// somebody looks for a setting that was there before the move. Read out of
|
||||
// the migrations rather than listed twice, so the two cannot disagree.
|
||||
var declared []string
|
||||
entries, err := migrationFS.ReadDir("migrations/sqlite")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, e := range entries {
|
||||
body, err := migrationFS.ReadFile("migrations/sqlite/" + e.Name())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, line := range strings.Split(string(body), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
rest, ok := strings.CutPrefix(line, "CREATE TABLE ")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
name, _, _ := strings.Cut(rest, " ")
|
||||
name = strings.TrimSuffix(strings.TrimSpace(name), "(")
|
||||
declared = append(declared, name)
|
||||
}
|
||||
}
|
||||
if len(declared) == 0 {
|
||||
t.Fatal("found no tables in the migrations, so this test proves nothing")
|
||||
}
|
||||
for _, name := range declared {
|
||||
if slices.Contains(copyOrder, name) {
|
||||
continue
|
||||
}
|
||||
// Left behind on purpose, with the reason written down beside it.
|
||||
if why, ok := notCopied[name]; ok {
|
||||
if why == "" {
|
||||
t.Errorf("%q is excluded from the copy with no reason given", name)
|
||||
}
|
||||
continue
|
||||
}
|
||||
t.Errorf("%q is created by a migration but nothing copies it", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyRefusesADestinationThatAlreadyHasRows(t *testing.T) {
|
||||
// Merging two stores means deciding what happens when both hold a project
|
||||
// with the same id and different contents, and every answer is a guess
|
||||
// about which one somebody meant.
|
||||
ctx := context.Background()
|
||||
src, dst := openDB(t), openDB(t)
|
||||
tree(t, src)
|
||||
tree(t, dst)
|
||||
|
||||
_, err := src.CopyTo(ctx, dst)
|
||||
if err == nil || !strings.Contains(err.Error(), "empty store") {
|
||||
t.Errorf("copying into a populated store gave %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestACopiedStoreStillAuthenticatesTheSameKey(t *testing.T) {
|
||||
// The reason this exists. Keys are stored hashed, so a migration that
|
||||
// loses them cannot recreate them: every person and every script would
|
||||
// need a new one issued.
|
||||
ctx := context.Background()
|
||||
src := openDB(t)
|
||||
proj := tree(t, src)
|
||||
secret, k := mustKey(t, src, "k1", proj)
|
||||
|
||||
if err := src.WriteUsage(ctx, []Usage{{
|
||||
At: time.Now(), KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
|
||||
PromptTokens: 3, CompletionTokens: 4, CostMicros: 99, Status: 200,
|
||||
}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer dst.Close()
|
||||
|
||||
res, err := src.CopyTo(ctx, dst)
|
||||
if err != nil {
|
||||
t.Fatalf("copying: %v", err)
|
||||
}
|
||||
if len(res) != len(copyOrder) {
|
||||
t.Errorf("copied %d tables of %d", len(res), len(copyOrder))
|
||||
}
|
||||
|
||||
// The same secret the caller already holds still resolves, in the copy.
|
||||
got, err := dst.KeyByHash(ctx, HashKey(secret))
|
||||
if err != nil {
|
||||
t.Fatalf("the copied store does not know the key: %v", err)
|
||||
}
|
||||
if got.ID != k.ID || got.ProjectID != proj {
|
||||
t.Errorf("key came back as %+v", got)
|
||||
}
|
||||
|
||||
// And the spend it had, so a budget is not quietly reset by moving.
|
||||
spend, err := dst.SpendSince(ctx, time.Now().Add(-time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if spend[k.ID] != 99 {
|
||||
t.Errorf("spend after the copy = %d, want 99", spend[k.ID])
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyCarriesConversationsAndNotes(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
src := openDB(t)
|
||||
proj := tree(t, src)
|
||||
if err := src.SetThreadRetention(ctx, proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := src.AppendThread(ctx, proj, "th1", []Message{
|
||||
{Role: "user", Content: "remembered"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer dst.Close()
|
||||
if _, err := src.CopyTo(ctx, dst); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
msgs, err := dst.Thread(ctx, proj, "th1")
|
||||
if err != nil {
|
||||
t.Fatalf("the conversation did not survive: %v", err)
|
||||
}
|
||||
if len(msgs) != 1 || msgs[0].Content != "remembered" {
|
||||
t.Errorf("thread came back as %+v", msgs)
|
||||
}
|
||||
// The retention setting is a column on projects, so it rides along — a
|
||||
// destination that kept the conversations and forgot the promise about
|
||||
// how long would be worse than one that kept neither.
|
||||
if d, err := dst.ThreadRetention(ctx, proj); err != nil || d == 0 {
|
||||
t.Errorf("retention after the copy = %v (err %v)", d, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyingAnEmptyStoreIsFine(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
src := openDB(t)
|
||||
dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer dst.Close()
|
||||
if _, err := src.CopyTo(ctx, dst); err != nil {
|
||||
t.Errorf("copying an empty store failed: %v", err)
|
||||
}
|
||||
_ = os.Remove("")
|
||||
}
|
||||
|
||||
func TestTheCopyOrderSatisfiesTheForeignKeys(t *testing.T) {
|
||||
// Completeness is not enough: a table copied before the one it references
|
||||
// fails on a constraint partway through, leaving some tables copied and
|
||||
// some not. The first version of copyOrder followed the order the
|
||||
// migrations create tables in, and users is created after api_keys even
|
||||
// though api_keys points at it — so the copy died three tables in, on a
|
||||
// real database, after the gateway had already been stopped.
|
||||
//
|
||||
// Derived from the migrations rather than restated, so a new reference
|
||||
// fails here rather than during somebody's migration.
|
||||
deps := map[string]map[string]bool{}
|
||||
note := func(table, ref string) {
|
||||
if table == ref {
|
||||
return
|
||||
}
|
||||
if deps[table] == nil {
|
||||
deps[table] = map[string]bool{}
|
||||
}
|
||||
deps[table][ref] = true
|
||||
}
|
||||
|
||||
entries, err := migrationFS.ReadDir("migrations/sqlite")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body := ""
|
||||
for _, e := range entries {
|
||||
b, err := migrationFS.ReadFile("migrations/sqlite/" + e.Name())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body += string(b)
|
||||
}
|
||||
|
||||
// Table bodies, so a REFERENCES is attributed to the table holding it.
|
||||
for _, m := range regexpCreateTable.FindAllStringSubmatch(body, -1) {
|
||||
for _, r := range regexpReferences.FindAllStringSubmatch(m[2], -1) {
|
||||
note(m[1], r[1])
|
||||
}
|
||||
}
|
||||
// And columns added later, which is exactly how api_keys got its user_id.
|
||||
for _, m := range regexpAlterRef.FindAllStringSubmatch(body, -1) {
|
||||
note(m[1], m[2])
|
||||
}
|
||||
if len(deps) == 0 {
|
||||
t.Fatal("found no foreign keys, so this test proves nothing")
|
||||
}
|
||||
|
||||
position := map[string]int{}
|
||||
for i, t := range copyOrder {
|
||||
position[t] = i
|
||||
}
|
||||
for table, refs := range deps {
|
||||
for ref := range refs {
|
||||
if position[ref] > position[table] {
|
||||
t.Errorf("%s is copied before %s, which it references", table, ref)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
regexpCreateTable = regexp.MustCompile(`(?s)CREATE TABLE (\w+)\s*\((.*?)\n\);`)
|
||||
regexpReferences = regexp.MustCompile(`REFERENCES\s+(\w+)`)
|
||||
regexpAlterRef = regexp.MustCompile(`ALTER TABLE (\w+) ADD COLUMN[^;]*REFERENCES\s+(\w+)`)
|
||||
)
|
||||
|
||||
func TestAnEmptyBlobDoesNotBecomeNullInTheCopy(t *testing.T) {
|
||||
// key_claims.sealed is emptied the moment a claim is used, so a claimed
|
||||
// row holds an empty blob in a NOT NULL column. An empty blob scans as a
|
||||
// typed nil slice and a real NULL as an untyped one; pgx cannot tell them
|
||||
// apart, so the distinction is restored before the insert.
|
||||
if got := keepEmptyBytes([]byte(nil)); got == nil {
|
||||
t.Error("an empty byte column would be inserted as NULL")
|
||||
} else if b, ok := got.([]byte); !ok || len(b) != 0 {
|
||||
t.Errorf("empty bytes came out as %#v", got)
|
||||
}
|
||||
// A genuine NULL arrives as an untyped nil and must stay one.
|
||||
if got := keepEmptyBytes(nil); got != nil {
|
||||
t.Errorf("a NULL became %#v", got)
|
||||
}
|
||||
// Everything else passes through.
|
||||
if got := keepEmptyBytes(int64(7)); got != int64(7) {
|
||||
t.Errorf("a number was rewritten to %#v", got)
|
||||
}
|
||||
if got := keepEmptyBytes([]byte("x")); string(got.([]byte)) != "x" {
|
||||
t.Errorf("a non-empty blob was rewritten to %#v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Dialect is which SQL database is underneath.
|
||||
//
|
||||
// There is one set of queries, not one per database, because the SQL this
|
||||
// store writes is almost entirely portable — the only construct that genuinely
|
||||
// differs is the placeholder, and `ON CONFLICT` upserts work in both. Two
|
||||
// copies of every query would drift, and the drift would be discovered one
|
||||
// query at a time in production.
|
||||
type Dialect string
|
||||
|
||||
// The supported databases.
|
||||
const (
|
||||
SQLiteDialect Dialect = "sqlite"
|
||||
PostgresDialect Dialect = "postgres"
|
||||
)
|
||||
|
||||
// rebind rewrites `?` placeholders for the dialect in use.
|
||||
//
|
||||
// SQLite takes `?`, PostgreSQL takes `$1`, `$2` and so on. Queries are written
|
||||
// with `?` because that is what the majority of them were written with
|
||||
// already, and rewriting is cheaper and safer than maintaining two spellings
|
||||
// of each.
|
||||
//
|
||||
// A `?` inside a string literal would be rewritten too, which would be a bug —
|
||||
// so this skips over quoted sections rather than scanning blindly. No query
|
||||
// here contains one today; this exists so that adding one is not a silent
|
||||
// corruption.
|
||||
func (d Dialect) rebind(query string) string {
|
||||
if d != PostgresDialect || !strings.Contains(query, "?") {
|
||||
return query
|
||||
}
|
||||
var b strings.Builder
|
||||
b.Grow(len(query) + 8)
|
||||
|
||||
n := 0
|
||||
inSingle, inDouble := false, false
|
||||
for i := 0; i < len(query); i++ {
|
||||
c := query[i]
|
||||
switch {
|
||||
case c == '\'' && !inDouble:
|
||||
inSingle = !inSingle
|
||||
case c == '"' && !inSingle:
|
||||
inDouble = !inDouble
|
||||
case c == '?' && !inSingle && !inDouble:
|
||||
n++
|
||||
b.WriteByte('$')
|
||||
b.WriteString(strconv.Itoa(n))
|
||||
continue
|
||||
}
|
||||
b.WriteByte(c)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// conn is the part of *sql.DB and *sql.Tx this store uses.
|
||||
//
|
||||
// Named so the same query helpers work inside and outside a transaction, and
|
||||
// so every path goes through the rebinding wrapper rather than some of them
|
||||
// reaching the driver directly with the wrong placeholders.
|
||||
type conn interface {
|
||||
ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
|
||||
QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
|
||||
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
|
||||
// PrepareContext is here because a prepared statement carries the query
|
||||
// with it: rebinding at Exec time would be too late, the placeholders are
|
||||
// fixed when it is prepared.
|
||||
PrepareContext(ctx context.Context, query string) (*sql.Stmt, error)
|
||||
}
|
||||
|
||||
// bound is a conn that rewrites placeholders on the way through.
|
||||
//
|
||||
// A wrapper rather than a change at each call site: there are sixty of those,
|
||||
// and sixty chances to forget one. Here the only way to reach the driver is
|
||||
// through the rewrite.
|
||||
type bound struct {
|
||||
c Dialect
|
||||
// to is the underlying *sql.DB or *sql.Tx.
|
||||
to conn
|
||||
}
|
||||
|
||||
func (b bound) ExecContext(ctx context.Context, q string, args ...any) (sql.Result, error) {
|
||||
return b.to.ExecContext(ctx, b.c.rebind(q), args...)
|
||||
}
|
||||
|
||||
func (b bound) QueryContext(ctx context.Context, q string, args ...any) (*sql.Rows, error) {
|
||||
return b.to.QueryContext(ctx, b.c.rebind(q), args...)
|
||||
}
|
||||
|
||||
func (b bound) QueryRowContext(ctx context.Context, q string, args ...any) *sql.Row {
|
||||
return b.to.QueryRowContext(ctx, b.c.rebind(q), args...)
|
||||
}
|
||||
|
||||
func (b bound) PrepareContext(ctx context.Context, q string) (*sql.Stmt, error) {
|
||||
return b.to.PrepareContext(ctx, b.c.rebind(q))
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package store
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestSQLiteKeepsItsPlaceholders(t *testing.T) {
|
||||
q := "SELECT a FROM t WHERE b = ? AND c = ?"
|
||||
if got := SQLiteDialect.rebind(q); got != q {
|
||||
t.Errorf("sqlite query was rewritten: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresGetsNumberedPlaceholders(t *testing.T) {
|
||||
got := PostgresDialect.rebind("SELECT a FROM t WHERE b = ? AND c = ?")
|
||||
want := "SELECT a FROM t WHERE b = $1 AND c = $2"
|
||||
if got != want {
|
||||
t.Errorf("rebind = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAQuestionMarkInsideALiteralIsLeftAlone(t *testing.T) {
|
||||
// No query here contains one today. This exists so that adding one is a
|
||||
// test failure rather than a query that silently reads the wrong argument
|
||||
// — or, worse, one that still parses and returns the wrong rows.
|
||||
got := PostgresDialect.rebind(`SELECT a FROM t WHERE name = 'why?' AND b = ?`)
|
||||
want := `SELECT a FROM t WHERE name = 'why?' AND b = $1`
|
||||
if got != want {
|
||||
t.Errorf("rebind = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoubleQuotedIdentifiersAreAlsoSkipped(t *testing.T) {
|
||||
got := PostgresDialect.rebind(`SELECT "odd?name" FROM t WHERE b = ?`)
|
||||
want := `SELECT "odd?name" FROM t WHERE b = $1`
|
||||
if got != want {
|
||||
t.Errorf("rebind = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAQueryWithNoPlaceholdersIsUntouched(t *testing.T) {
|
||||
q := "DELETE FROM t"
|
||||
if got := PostgresDialect.rebind(q); got != q {
|
||||
t.Errorf("rebind = %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrNotLeader means somebody else holds the lease and it has not lapsed.
|
||||
var ErrNotLeader = errors.New("store: another control plane holds the lease")
|
||||
|
||||
// Lease is who is in charge and until when.
|
||||
type Lease struct {
|
||||
Owner string
|
||||
ExpiresAt time.Time
|
||||
AcquiredAt time.Time
|
||||
}
|
||||
|
||||
// AcquireLeadership takes or renews the lease for a role.
|
||||
//
|
||||
// Renewing and taking are the same statement on purpose: a leader that renews
|
||||
// takes the same path as one starting fresh, so there is no way for the two to
|
||||
// disagree about what a held lease looks like. It succeeds when the row is
|
||||
// absent, when this process already owns it, or when whoever owned it let it
|
||||
// lapse; otherwise it returns ErrNotLeader.
|
||||
//
|
||||
// The clock is the caller's rather than the database's, which is a compromise
|
||||
// worth naming. A database clock would be one clock for everyone; taking it
|
||||
// means a round trip that says now() inside every comparison, in SQL that is
|
||||
// spelled differently in each dialect. Two control planes on machines whose
|
||||
// clocks disagree by more than the lease will fight, so keep them on NTP —
|
||||
// which any cluster running InfiniBand already has to be.
|
||||
func (s *SQLite) AcquireLeadership(ctx context.Context, role, owner string, ttl time.Duration) (Lease, error) {
|
||||
now := time.Now()
|
||||
expires := now.Add(ttl)
|
||||
|
||||
// One statement, so two processes racing cannot both believe they won:
|
||||
// the WHERE decides, and the loser changes no rows.
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`INSERT INTO leadership (role, owner, expires_at, acquired_at) VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT (role) DO UPDATE SET
|
||||
owner = excluded.owner,
|
||||
expires_at = excluded.expires_at,
|
||||
acquired_at = CASE WHEN leadership.owner = excluded.owner
|
||||
THEN leadership.acquired_at
|
||||
ELSE excluded.acquired_at END
|
||||
WHERE leadership.owner = excluded.owner OR leadership.expires_at < ?`,
|
||||
role, owner, expires.UnixMilli(), now.UnixMilli(), now.UnixMilli())
|
||||
if err != nil {
|
||||
return Lease{}, err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return Lease{}, err
|
||||
}
|
||||
if n == 0 {
|
||||
// Lost, and the caller is told who to, because "not leader" without a
|
||||
// name is the start of a long afternoon. The lease is returned beside
|
||||
// the error rather than instead of it.
|
||||
held, lerr := s.Leadership(ctx, role)
|
||||
if lerr != nil {
|
||||
return Lease{}, lerr
|
||||
}
|
||||
return held, ErrNotLeader
|
||||
}
|
||||
return s.Leadership(ctx, role)
|
||||
}
|
||||
|
||||
// Leadership reports who holds a role's lease.
|
||||
func (s *SQLite) Leadership(ctx context.Context, role string) (Lease, error) {
|
||||
var owner string
|
||||
var expires, acquired int64
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT owner, expires_at, acquired_at FROM leadership WHERE role = ?`, role).
|
||||
Scan(&owner, &expires, &acquired)
|
||||
if err != nil {
|
||||
return Lease{}, err
|
||||
}
|
||||
return Lease{
|
||||
Owner: owner,
|
||||
ExpiresAt: fromMillis(expires),
|
||||
AcquiredAt: fromMillis(acquired),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ReleaseLeadership gives the lease up, so a standby takes over in seconds
|
||||
// rather than waiting out the clock.
|
||||
//
|
||||
// Only the holder can release it: a process that has already lost the lease
|
||||
// must not be able to take it from whoever has it now, which is exactly what
|
||||
// an unconditional delete during a slow shutdown would do.
|
||||
func (s *SQLite) ReleaseLeadership(ctx context.Context, role, owner string) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM leadership WHERE role = ? AND owner = ?`, role, owner)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// leadershipSuite is the same assertions against whichever store it is given,
|
||||
// because the whole point of the lease is that it behaves identically on both
|
||||
// — and the one that matters for HA is the one it was never tested on.
|
||||
func leadershipSuite(t *testing.T, open func(*testing.T) *SQLite) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("the first to ask gets it", func(t *testing.T) {
|
||||
s := open(t)
|
||||
lease, err := s.AcquireLeadership(ctx, "control", "a", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if lease.Owner != "a" {
|
||||
t.Errorf("owner = %q", lease.Owner)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("a second is refused and told who holds it", func(t *testing.T) {
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
held, err := s.AcquireLeadership(ctx, "control", "b", time.Minute)
|
||||
if !errors.Is(err, ErrNotLeader) {
|
||||
t.Fatalf("a second control plane got %v", err)
|
||||
}
|
||||
if held.Owner != "a" {
|
||||
t.Errorf("refused without naming the holder: %+v", held)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("the holder renews without losing when it started", func(t *testing.T) {
|
||||
s := open(t)
|
||||
first, err := s.AcquireLeadership(ctx, "control", "a", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
again, err := s.AcquireLeadership(ctx, "control", "a", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("the holder could not renew: %v", err)
|
||||
}
|
||||
if !again.ExpiresAt.After(first.ExpiresAt) {
|
||||
t.Error("renewing did not move the expiry")
|
||||
}
|
||||
if !again.AcquiredAt.Equal(first.AcquiredAt) {
|
||||
t.Error("renewing reset when this leader took over")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("a lapsed lease is taken", func(t *testing.T) {
|
||||
// The reason a lease and not a lock: a leader that was powered off
|
||||
// cannot release anything, and nobody should have to be paged to make
|
||||
// a standby take over.
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", 20*time.Millisecond); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
lease, err := s.AcquireLeadership(ctx, "control", "b", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("a lapsed lease was not takeable: %v", err)
|
||||
}
|
||||
if lease.Owner != "b" {
|
||||
t.Errorf("owner after takeover = %q", lease.Owner)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("releasing hands over at once", func(t *testing.T) {
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", time.Hour); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.ReleaseLeadership(ctx, "control", "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "b", time.Minute); err != nil {
|
||||
t.Errorf("a released lease was not free: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("a former leader cannot release the current one", func(t *testing.T) {
|
||||
// A process that lost the lease and is shutting down slowly must not
|
||||
// take it from whoever has it now — which is exactly what an
|
||||
// unconditional delete would do.
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", 20*time.Millisecond); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "b", time.Hour); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.ReleaseLeadership(ctx, "control", "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
held, err := s.Leadership(ctx, "control")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if held.Owner != "b" {
|
||||
t.Errorf("the old leader released the new one's lease: %+v", held)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("only one of many wins", func(t *testing.T) {
|
||||
// The case the whole table exists for. Two control planes starting
|
||||
// together is the ordinary way this gets tested in production.
|
||||
s := open(t)
|
||||
const racers = 8
|
||||
var wg sync.WaitGroup
|
||||
won := make([]bool, racers)
|
||||
for i := range racers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := s.AcquireLeadership(ctx, "control", string(rune('a'+i)), time.Minute)
|
||||
won[i] = err == nil
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
n := 0
|
||||
for _, w := range won {
|
||||
if w {
|
||||
n++
|
||||
}
|
||||
}
|
||||
if n != 1 {
|
||||
t.Errorf("%d of %d processes believed they were leader", n, racers)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestLeadershipOnSQLite(t *testing.T) {
|
||||
leadershipSuite(t, func(t *testing.T) *SQLite { return openDB(t) })
|
||||
}
|
||||
|
||||
func TestLeadershipOnPostgres(t *testing.T) {
|
||||
leadershipSuite(t, openPG)
|
||||
}
|
||||
@@ -138,14 +138,14 @@ func (s *SQLite) PutNote(ctx context.Context, projectID, agentID string, n Note)
|
||||
return ErrMemoryNotEnabled
|
||||
}
|
||||
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
q, tx, err := s.begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
res, err := tx.ExecContext(ctx,
|
||||
res, err := q.ExecContext(ctx,
|
||||
`UPDATE memory_notes SET text = ?, tags = ?, updated_at = ?
|
||||
WHERE project_id = ? AND agent_id = ? AND key = ?`,
|
||||
n.Text, tags, now, projectID, agentID, n.Key)
|
||||
@@ -154,7 +154,7 @@ func (s *SQLite) PutNote(ctx context.Context, projectID, agentID string, n Note)
|
||||
}
|
||||
if updated, _ := res.RowsAffected(); updated == 0 {
|
||||
var count int
|
||||
if err := tx.QueryRowContext(ctx,
|
||||
if err := q.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM memory_notes WHERE project_id = ? AND agent_id = ?`,
|
||||
projectID, agentID).Scan(&count); err != nil {
|
||||
return err
|
||||
@@ -162,7 +162,7 @@ func (s *SQLite) PutNote(ctx context.Context, projectID, agentID string, n Note)
|
||||
if count >= MaxNotesPerAgent {
|
||||
return ErrMemoryFull
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`INSERT INTO memory_notes (project_id, agent_id, key, text, tags, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)`,
|
||||
projectID, agentID, n.Key, n.Text, tags, now, now); err != nil {
|
||||
@@ -223,8 +223,20 @@ func (s *SQLite) SearchNotes(ctx context.Context, projectID, agentID, query, tag
|
||||
args := []any{projectID, agentID}
|
||||
|
||||
if match := ftsQuery(query); match != "" {
|
||||
q += ` AND id IN (SELECT rowid FROM memory_notes_fts WHERE memory_notes_fts MATCH ?)`
|
||||
args = append(args, match)
|
||||
// The one query in this store that is not the same sentence in both
|
||||
// dialects. SQLite searches an FTS5 table kept in step by triggers;
|
||||
// PostgreSQL matches a tsvector derived from the row. Writing it out
|
||||
// twice is honest about that — the alternative was a shared spelling
|
||||
// that means something subtly different in each.
|
||||
if s.dialect == PostgresDialect {
|
||||
q += ` AND search @@ plainto_tsquery('simple', ?)`
|
||||
// plainto_tsquery takes words, not FTS5 syntax, so the raw query
|
||||
// goes in rather than the quoted form SQLite needs.
|
||||
args = append(args, query)
|
||||
} else {
|
||||
q += ` AND id IN (SELECT rowid FROM memory_notes_fts WHERE memory_notes_fts MATCH ?)`
|
||||
args = append(args, match)
|
||||
}
|
||||
}
|
||||
if tag != "" {
|
||||
t, err := NormalizeTag(tag)
|
||||
|
||||
@@ -10,7 +10,7 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
//go:embed migrations/sqlite/*.sql
|
||||
//go:embed migrations/sqlite/*.sql migrations/postgres/*.sql
|
||||
var migrationFS embed.FS
|
||||
|
||||
// migrate applies every unapplied migration in order, each in its own
|
||||
@@ -25,8 +25,17 @@ var migrationFS embed.FS
|
||||
// promise this could not keep. Same version numbers, same meaning, different
|
||||
// files.
|
||||
func migrate(ctx context.Context, db *sql.DB, dialect string) error {
|
||||
// BIGINT, not INTEGER. SQLite's INTEGER holds 64 bits; PostgreSQL's is
|
||||
// int4, and a millisecond timestamp does not fit in it — which the first
|
||||
// migration discovers by failing to record that it ran. This table is
|
||||
// created here rather than in a migration file, which is exactly how it
|
||||
// escaped the port of the others.
|
||||
stamp := "BIGINT"
|
||||
if dialect == string(SQLiteDialect) {
|
||||
stamp = "INTEGER"
|
||||
}
|
||||
if _, err := db.ExecContext(ctx,
|
||||
`CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at INTEGER NOT NULL)`,
|
||||
`CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at `+stamp+` NOT NULL)`,
|
||||
); err != nil {
|
||||
return fmt.Errorf("creating the migration table: %w", err)
|
||||
}
|
||||
@@ -63,7 +72,7 @@ func migrate(ctx context.Context, db *sql.DB, dialect string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := applyOne(ctx, db, name, string(body)); err != nil {
|
||||
if err := applyOne(ctx, db, dialect, name, string(body)); err != nil {
|
||||
return fmt.Errorf("migration %s: %w", name, err)
|
||||
}
|
||||
}
|
||||
@@ -86,7 +95,7 @@ func migrationNames(dialect string) ([]string, error) {
|
||||
return names, nil
|
||||
}
|
||||
|
||||
func applyOne(ctx context.Context, db *sql.DB, name, body string) error {
|
||||
func applyOne(ctx context.Context, db *sql.DB, dialect, name, body string) error {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -97,7 +106,7 @@ func applyOne(ctx context.Context, db *sql.DB, name, body string) error {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)`,
|
||||
Dialect(dialect).rebind(`INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)`),
|
||||
name, nowMillis(),
|
||||
); err != nil {
|
||||
return err
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- The tenancy tree. Three levels, each owning the next.
|
||||
CREATE TABLE organizations (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE teams (
|
||||
id TEXT PRIMARY KEY,
|
||||
org_id TEXT NOT NULL REFERENCES organizations(id),
|
||||
name TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX teams_by_org ON teams(org_id);
|
||||
|
||||
CREATE TABLE projects (
|
||||
id TEXT PRIMARY KEY,
|
||||
team_id TEXT NOT NULL REFERENCES teams(id),
|
||||
name TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX projects_by_team ON projects(team_id);
|
||||
|
||||
-- Only sha256(key) is stored, so a database leak hands over no working
|
||||
-- credential. SHA-256 rather than a slow password hash on purpose: the input is
|
||||
-- 256 random bits with no dictionary to attack, so the cost of argon2 would buy
|
||||
-- nothing and charge tens of milliseconds on every request (plan §7.9.1).
|
||||
CREATE TABLE api_keys (
|
||||
id TEXT PRIMARY KEY,
|
||||
key_hash BYTEA NOT NULL UNIQUE,
|
||||
key_prefix TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL REFERENCES projects(id),
|
||||
band TEXT,
|
||||
routing_profile TEXT,
|
||||
created_at BIGINT NOT NULL,
|
||||
expires_at BIGINT,
|
||||
-- Revocation is a soft delete. Usage rows reference the key, and an audit
|
||||
-- trail that loses its subject is not an audit trail.
|
||||
revoked_at BIGINT,
|
||||
last_used_at BIGINT
|
||||
);
|
||||
CREATE INDEX api_keys_by_project ON api_keys(project_id);
|
||||
|
||||
-- Limits hang off any scope rather than being duplicated onto each one. A NULL
|
||||
-- column means "no limit at this level", which is not the same as 0.
|
||||
CREATE TABLE limits (
|
||||
scope_type TEXT NOT NULL,
|
||||
scope_id TEXT NOT NULL,
|
||||
rpm BIGINT,
|
||||
tpm BIGINT,
|
||||
-- Micro-units, not a float: a budget is accumulated by repeated addition, and
|
||||
-- repeated addition of floats drifts.
|
||||
budget_micros BIGINT,
|
||||
budget_period TEXT,
|
||||
allowed_models TEXT,
|
||||
PRIMARY KEY (scope_type, scope_id)
|
||||
);
|
||||
|
||||
-- One row per request, written in batches off the request path.
|
||||
CREATE TABLE usage (
|
||||
ts BIGINT NOT NULL,
|
||||
key_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
served_model TEXT NOT NULL,
|
||||
prompt_tokens BIGINT NOT NULL,
|
||||
completion_tokens BIGINT NOT NULL,
|
||||
cost_micros BIGINT NOT NULL,
|
||||
status BIGINT NOT NULL,
|
||||
latency_ms BIGINT NOT NULL
|
||||
);
|
||||
-- Budget restore reads by key within a period, which is this index exactly.
|
||||
CREATE INDEX usage_by_key_ts ON usage(key_id, ts);
|
||||
CREATE INDEX usage_by_ts ON usage(ts);
|
||||
@@ -0,0 +1,29 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- A key has to reach the person it was issued for, and a key pasted into a chat
|
||||
-- window stays in that chat window forever. A claim is a one-time link instead:
|
||||
-- the operator sends the link, the person opens it once, and the link dies.
|
||||
--
|
||||
-- The key is not stored here in the clear. It is sealed with a key derived from
|
||||
-- the claim token, and the token is never written down — only its hash is. So a
|
||||
-- database leak yields two useless hashes and a blob nothing in the database can
|
||||
-- open, which keeps the property the api_keys table has.
|
||||
CREATE TABLE key_claims (
|
||||
token_hash BYTEA PRIMARY KEY,
|
||||
key_id TEXT NOT NULL REFERENCES api_keys(id),
|
||||
-- The sealed key. Emptied the moment it is claimed, so a claimed link leaves
|
||||
-- nothing behind even before the row is purged.
|
||||
sealed BYTEA NOT NULL,
|
||||
created_at BIGINT NOT NULL,
|
||||
expires_at BIGINT NOT NULL,
|
||||
claimed_at BIGINT
|
||||
);
|
||||
CREATE INDEX key_claims_by_key ON key_claims(key_id);
|
||||
CREATE INDEX key_claims_by_expiry ON key_claims(expires_at);
|
||||
@@ -0,0 +1,56 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Accounts, so a person can sign up, be approved, and then issue their own keys
|
||||
-- without an operator running a command for each one.
|
||||
CREATE TABLE users (
|
||||
id TEXT PRIMARY KEY,
|
||||
-- Lowercased at the boundary. Two accounts differing only in case would be
|
||||
-- two accounts to the database and one to every human who looks at them.
|
||||
email TEXT NOT NULL UNIQUE,
|
||||
name TEXT NOT NULL,
|
||||
-- argon2id, encoded with its own parameters so the cost can be raised later
|
||||
-- without invalidating existing hashes. Slow hashing is right here and wrong
|
||||
-- for API keys: this input is chosen by a human and has a dictionary.
|
||||
password_hash TEXT NOT NULL,
|
||||
-- 'pending' until an operator approves, then 'active'. 'suspended' keeps the
|
||||
-- account and stops it working.
|
||||
state TEXT NOT NULL,
|
||||
-- Assigned at approval: which project this person's keys belong to, and
|
||||
-- therefore which limits they inherit.
|
||||
project_id TEXT REFERENCES projects(id),
|
||||
created_at BIGINT NOT NULL,
|
||||
approved_at BIGINT,
|
||||
approved_by TEXT,
|
||||
last_login_at BIGINT
|
||||
);
|
||||
CREATE INDEX users_by_state ON users(state);
|
||||
|
||||
-- Server-side sessions rather than signed cookies, so signing out and
|
||||
-- suspending an account take effect immediately instead of whenever a token
|
||||
-- happens to expire.
|
||||
CREATE TABLE sessions (
|
||||
token_hash BYTEA PRIMARY KEY,
|
||||
user_id TEXT NOT NULL REFERENCES users(id),
|
||||
-- Per session, submitted with every form. A cookie alone authenticates the
|
||||
-- browser, not the page that asked, which is what makes a cross-site POST
|
||||
-- work.
|
||||
csrf_hash BYTEA NOT NULL,
|
||||
created_at BIGINT NOT NULL,
|
||||
expires_at BIGINT NOT NULL,
|
||||
last_seen BIGINT NOT NULL
|
||||
);
|
||||
CREATE INDEX sessions_by_user ON sessions(user_id);
|
||||
CREATE INDEX sessions_by_expiry ON sessions(expires_at);
|
||||
|
||||
-- A key issued through the profile page belongs to the person who made it, so
|
||||
-- they can see and revoke their own and nobody else's. NULL means a key an
|
||||
-- operator issued directly.
|
||||
ALTER TABLE api_keys ADD COLUMN user_id TEXT REFERENCES users(id);
|
||||
CREATE INDEX api_keys_by_user ON api_keys(user_id);
|
||||
@@ -0,0 +1,20 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Who may see the cluster.
|
||||
--
|
||||
-- Until now the answer was "whoever has a shell on the control plane", because
|
||||
-- the dashboard lived on a loopback-bound port. Serving it over the web needs
|
||||
-- an actual answer, and this is the smallest one that is not a lie: a flag on
|
||||
-- an account that already has a password, a session and a CSRF token.
|
||||
--
|
||||
-- Deliberately not a role system. One bit is enough for "can see the cluster"
|
||||
-- and impossible to get subtly wrong; roles arrive when there is a second
|
||||
-- distinction worth drawing.
|
||||
ALTER TABLE users ADD COLUMN is_admin BIGINT NOT NULL DEFAULT 0;
|
||||
@@ -0,0 +1,34 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Models added while the cluster is running.
|
||||
--
|
||||
-- The catalog came from a YAML file read once at startup, so adding a model
|
||||
-- meant editing a file on the control plane and restarting it: a settling
|
||||
-- window, adoption churn, and on one occasion two minutes of outage. That is a
|
||||
-- reasonable way to manage a fixed production set and a bad one for a platform
|
||||
-- whose promise is that people can try things.
|
||||
--
|
||||
-- These live alongside the ones in configuration rather than replacing them.
|
||||
-- Config-as-code stays declarative — a model in YAML is in the catalog on every
|
||||
-- start, whatever the database says — and this table holds the ones added
|
||||
-- afterwards. A model added here may not take a name or alias that
|
||||
-- configuration already uses, which the catalog's own constructor enforces for
|
||||
-- free, because it refuses a duplicate whatever the source.
|
||||
--
|
||||
-- The definition is stored as JSON rather than as columns. A profile is a
|
||||
-- nested, evolving shape, and every field added to it would otherwise be a
|
||||
-- migration; the catalog already knows how to validate one, and a row that
|
||||
-- fails validation is reported and skipped rather than blocking startup.
|
||||
CREATE TABLE models (
|
||||
name TEXT PRIMARY KEY,
|
||||
definition TEXT NOT NULL,
|
||||
created_at TIMESTAMP NOT NULL,
|
||||
created_by TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
@@ -0,0 +1,57 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Server-side conversations, so a thread can continue on a different model
|
||||
-- than it started on.
|
||||
--
|
||||
-- Off unless a project turns it on. Until this migration ran, llmgw stored no
|
||||
-- message content anywhere — usage rows carry token counts and costs, never
|
||||
-- text — and that stays true for every project that does not ask. An on-prem
|
||||
-- gateway keeping nothing anyone typed is much easier to put in front of a
|
||||
-- team than one keeping everything, so storing is a choice a project makes
|
||||
-- rather than a default it discovers.
|
||||
|
||||
-- retention_hours is the switch and the limit at once. NULL means this project
|
||||
-- stores nothing, which is every project that existed before this and every
|
||||
-- one created since without asking.
|
||||
--
|
||||
-- One column rather than a boolean and a duration: a thread store with no
|
||||
-- expiry is a transcript of everything a team ever typed, kept forever by
|
||||
-- omission. Turning it on therefore requires saying for how long.
|
||||
ALTER TABLE projects ADD COLUMN thread_retention_hours BIGINT;
|
||||
|
||||
CREATE TABLE threads (
|
||||
id TEXT PRIMARY KEY,
|
||||
-- Scoped to a project, and checked on every read. Thread ids are opaque, but
|
||||
-- an unguessable id is a secret and secrets leak into logs, shell history
|
||||
-- and bug reports; the scope is what makes a leaked id useless to someone
|
||||
-- outside the project.
|
||||
project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
||||
created_at BIGINT NOT NULL,
|
||||
-- Touched on every append, so retention measures idleness rather than age:
|
||||
-- expiring a conversation someone is still having would be a strange way to
|
||||
-- keep a promise about storage.
|
||||
updated_at BIGINT NOT NULL
|
||||
);
|
||||
CREATE INDEX threads_by_project ON threads(project_id);
|
||||
CREATE INDEX threads_by_updated ON threads(updated_at);
|
||||
|
||||
CREATE TABLE thread_messages (
|
||||
thread_id TEXT NOT NULL REFERENCES threads(id) ON DELETE CASCADE,
|
||||
-- Position rather than a timestamp, because order is the point and two
|
||||
-- messages can share a millisecond.
|
||||
seq BIGINT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
-- Which model produced an assistant turn. A thread that moved between models
|
||||
-- otherwise reads as one voice, and when the answers disagree that is the
|
||||
-- first thing anyone wants to know.
|
||||
model TEXT,
|
||||
PRIMARY KEY (thread_id, seq)
|
||||
);
|
||||
@@ -0,0 +1,41 @@
|
||||
-- Ported from the SQLite migration of the same number, and the only one that
|
||||
-- is not a mechanical translation.
|
||||
--
|
||||
-- SQLite gets full-text search from an FTS5 virtual table with an external
|
||||
-- content table and three triggers keeping it in step. PostgreSQL derives a
|
||||
-- tsvector from the row itself, so there is nothing to keep in step and no
|
||||
-- trigger to forget: the column cannot disagree with the row it is computed
|
||||
-- from.
|
||||
--
|
||||
-- The search query differs too, and that difference lives in the Go code
|
||||
-- rather than being hidden here — a query that means "match these words" is
|
||||
-- spelled MATCH in one and @@ in the other, and pretending otherwise would
|
||||
-- cost more than admitting it.
|
||||
|
||||
ALTER TABLE projects ADD COLUMN memory_retention_hours BIGINT;
|
||||
|
||||
CREATE TABLE memory_notes (
|
||||
id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
|
||||
-- Scoped to a project and checked in SQL on every read, like threads.
|
||||
project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
||||
-- A namespace inside the project, so two agents sharing a key do not read
|
||||
-- each other's notes.
|
||||
agent_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
-- Normalised to ",a,b," so a tag matches whole, never as part of another.
|
||||
tags TEXT NOT NULL DEFAULT '',
|
||||
created_at BIGINT NOT NULL,
|
||||
-- Touched on every write, so retention measures how long a note has gone
|
||||
-- unrefreshed rather than how old it is.
|
||||
updated_at BIGINT NOT NULL,
|
||||
-- 'simple' rather than a language configuration: these are an agent's own
|
||||
-- notes, in whatever language it chose, and stemming English rules over
|
||||
-- Turkish text would quietly stop finding things.
|
||||
search tsvector GENERATED ALWAYS AS (
|
||||
to_tsvector('simple', coalesce(key, '') || ' ' || coalesce(text, '') || ' ' || coalesce(tags, ''))
|
||||
) STORED,
|
||||
UNIQUE (project_id, agent_id, key)
|
||||
);
|
||||
CREATE INDEX memory_notes_by_updated ON memory_notes(updated_at);
|
||||
CREATE INDEX memory_notes_search ON memory_notes USING GIN (search);
|
||||
@@ -0,0 +1,29 @@
|
||||
-- Ported from the SQLite migration of the same number. BIGINT because
|
||||
-- PostgreSQL's INTEGER is 32 bits and these are milliseconds since the epoch.
|
||||
|
||||
-- A lease on being the control plane that places work.
|
||||
--
|
||||
-- Two control planes against one store both reconcile: both see a model
|
||||
-- wanting two replicas and none running, and both place. The result is twice
|
||||
-- the replicas, each holding GPUs the other's inventory believes are free —
|
||||
-- the same failure a control plane restart used to produce before adoption,
|
||||
-- and harder to see because nothing crashed.
|
||||
--
|
||||
-- A lease rather than a lock held by the connection. PostgreSQL has advisory
|
||||
-- locks and SQLite has nothing like them, and this store keeps one set of
|
||||
-- queries for both; a row with an expiry works everywhere and survives the
|
||||
-- leader being killed rather than asked to stop.
|
||||
--
|
||||
-- One row, enforced by the primary key: there is one thing to be leader of.
|
||||
CREATE TABLE leadership (
|
||||
role TEXT PRIMARY KEY,
|
||||
-- Who holds it. An opaque id per process, so a process cannot be fooled by
|
||||
-- its own stale lease after a restart on the same host.
|
||||
owner TEXT NOT NULL,
|
||||
-- When the lease lapses. The holder renews well before this; anyone else
|
||||
-- may take it once the clock has passed, which is what makes a leader that
|
||||
-- was powered off recoverable without anyone being paged.
|
||||
expires_at BIGINT NOT NULL,
|
||||
-- For an operator asking who has been in charge and since when.
|
||||
acquired_at BIGINT NOT NULL
|
||||
);
|
||||
@@ -0,0 +1,26 @@
|
||||
-- A lease on being the control plane that places work.
|
||||
--
|
||||
-- Two control planes against one store both reconcile: both see a model
|
||||
-- wanting two replicas and none running, and both place. The result is twice
|
||||
-- the replicas, each holding GPUs the other's inventory believes are free —
|
||||
-- the same failure a control plane restart used to produce before adoption,
|
||||
-- and harder to see because nothing crashed.
|
||||
--
|
||||
-- A lease rather than a lock held by the connection. PostgreSQL has advisory
|
||||
-- locks and SQLite has nothing like them, and this store keeps one set of
|
||||
-- queries for both; a row with an expiry works everywhere and survives the
|
||||
-- leader being killed rather than asked to stop.
|
||||
--
|
||||
-- One row, enforced by the primary key: there is one thing to be leader of.
|
||||
CREATE TABLE leadership (
|
||||
role TEXT PRIMARY KEY,
|
||||
-- Who holds it. An opaque id per process, so a process cannot be fooled by
|
||||
-- its own stale lease after a restart on the same host.
|
||||
owner TEXT NOT NULL,
|
||||
-- When the lease lapses. The holder renews well before this; anyone else
|
||||
-- may take it once the clock has passed, which is what makes a leader that
|
||||
-- was powered off recoverable without anyone being paged.
|
||||
expires_at INTEGER NOT NULL,
|
||||
-- For an operator asking who has been in charge and since when.
|
||||
acquired_at INTEGER NOT NULL
|
||||
);
|
||||
@@ -0,0 +1,61 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Open picks a backend from what the address looks like.
|
||||
//
|
||||
// A path is SQLite and a postgres:// URL is PostgreSQL, so an existing
|
||||
// configuration keeps working untouched and moving to PostgreSQL is one line
|
||||
// changed rather than a new setting to discover. A gateway on one machine has
|
||||
// no reason to run a database server, and a gateway that has outgrown that
|
||||
// should not have to be reconfigured twice.
|
||||
//
|
||||
// An address of the form @/path/to/file reads the real address out of that
|
||||
// file, which is how a DSN carrying a password stays out of a configuration
|
||||
// file that gets copied, diffed and pasted into issues.
|
||||
func Open(ctx context.Context, address string) (*SQLite, error) {
|
||||
address, err := resolveAddress(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if isPostgres(address) {
|
||||
return OpenPostgres(ctx, address)
|
||||
}
|
||||
return OpenSQLite(ctx, address)
|
||||
}
|
||||
|
||||
// isPostgres reports whether an address names a PostgreSQL server.
|
||||
//
|
||||
// By scheme, because that is the only part of a DSN that is reliably there.
|
||||
// A bare host:port would be ambiguous with a relative file path, so it is not
|
||||
// accepted: an address this cannot classify should fail loudly rather than
|
||||
// quietly create a SQLite file named after someone's database server.
|
||||
func isPostgres(address string) bool {
|
||||
return strings.HasPrefix(address, "postgres://") ||
|
||||
strings.HasPrefix(address, "postgresql://")
|
||||
}
|
||||
|
||||
// resolveAddress expands an @file indirection.
|
||||
//
|
||||
// The file is read whole and trimmed, so a trailing newline — which every
|
||||
// editor and every `echo` adds — does not become part of a password.
|
||||
func resolveAddress(address string) (string, error) {
|
||||
if !strings.HasPrefix(address, "@") {
|
||||
return address, nil
|
||||
}
|
||||
path := strings.TrimPrefix(address, "@")
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("reading the store address from %s: %w", path, err)
|
||||
}
|
||||
out := strings.TrimSpace(string(b))
|
||||
if out == "" {
|
||||
return "", fmt.Errorf("%s is empty, so there is no store address in it", path)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAPathIsStillSQLite(t *testing.T) {
|
||||
// Every existing configuration names a file, and every one of them has to
|
||||
// keep working without being touched.
|
||||
for _, addr := range []string{
|
||||
"/var/lib/llmgw/llmgw.db", "llmgw.db", "./data/llmgw.db",
|
||||
} {
|
||||
if isPostgres(addr) {
|
||||
t.Errorf("%q was read as a PostgreSQL address", addr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPostgresURLIsRecognised(t *testing.T) {
|
||||
for _, addr := range []string{
|
||||
"postgres://u:p@h:5432/db", "postgresql://u@h/db?sslmode=disable",
|
||||
} {
|
||||
if !isPostgres(addr) {
|
||||
t.Errorf("%q was not read as a PostgreSQL address", addr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAddressFileKeepsThePasswordOutOfTheConfig(t *testing.T) {
|
||||
// A DSN carries a password, and a configuration file gets copied, diffed
|
||||
// and pasted into issues. The indirection is what lets the secret live
|
||||
// somewhere with its own permissions.
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "dsn")
|
||||
// With the trailing newline every editor and every echo adds, which must
|
||||
// not become part of the password.
|
||||
if err := os.WriteFile(p, []byte("postgres://u:secret@h/db\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got, err := resolveAddress("@" + p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != "postgres://u:secret@h/db" {
|
||||
t.Errorf("resolved to %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingOrEmptyAddressFileSaysSo(t *testing.T) {
|
||||
// Rather than falling back to creating a SQLite file named "@/etc/...",
|
||||
// which is what a silent failure here would produce.
|
||||
if _, err := resolveAddress("@/nonexistent/dsn"); err == nil {
|
||||
t.Error("a missing address file was accepted")
|
||||
}
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "empty")
|
||||
if err := os.WriteFile(p, []byte(" \n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err := resolveAddress("@" + p)
|
||||
if err == nil || !strings.Contains(err.Error(), "empty") {
|
||||
t.Errorf("an empty address file gave %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenStillCreatesASQLiteFile(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "llmgw.db")
|
||||
s, err := Open(context.Background(), p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if s.dialect != SQLiteDialect {
|
||||
t.Errorf("a path opened as %q", s.dialect)
|
||||
}
|
||||
if _, err := os.Stat(p); err != nil {
|
||||
t.Errorf("no database file was created: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
// The pure-Go PostgreSQL driver. Pure-Go is not a preference here: the
|
||||
// whole binary is built with CGO_ENABLED=0 so that it runs on a node whose
|
||||
// libc is older than the builder's, and a cgo driver would take that away
|
||||
// from every deployment, including the ones still on SQLite.
|
||||
_ "github.com/jackc/pgx/v5/stdlib"
|
||||
)
|
||||
|
||||
// OpenPostgres opens the store against PostgreSQL.
|
||||
//
|
||||
// The same queries run here as on SQLite; only the placeholders differ, and
|
||||
// the dialect handles that. What differs operationally is concurrency: SQLite
|
||||
// is held to one connection because more would only produce lock contention
|
||||
// that reads as a deadlock, while PostgreSQL is the reason someone moves to
|
||||
// it and wants a pool.
|
||||
//
|
||||
// The DSN carries a password, so it is taken as a string the caller read from
|
||||
// somewhere protected rather than as a path this package knows how to find —
|
||||
// and it is never logged, which is why errors below name the database and not
|
||||
// the DSN.
|
||||
func OpenPostgres(ctx context.Context, dsn string) (*SQLite, error) {
|
||||
if dsn == "" {
|
||||
return nil, fmt.Errorf("store: no PostgreSQL address configured")
|
||||
}
|
||||
db, err := sql.Open("pgx", dsn)
|
||||
if err != nil {
|
||||
// Deliberately not wrapped with the DSN: it holds a password, and an
|
||||
// error string ends up in logs, tickets and screenshots.
|
||||
return nil, fmt.Errorf("opening the PostgreSQL store: %w", err)
|
||||
}
|
||||
|
||||
// Sized for a control plane rather than for a fleet of application
|
||||
// servers. The request path barely touches the store — usage is batched,
|
||||
// budgets are counted in memory, authentication is cached — so the pool
|
||||
// exists for conversation turns and reporting queries, not for inference.
|
||||
db.SetMaxOpenConns(16)
|
||||
db.SetMaxIdleConns(4)
|
||||
// Long enough that a quiet cluster keeps its connections, short enough
|
||||
// that a restarted database is not answered with dead ones for an hour.
|
||||
db.SetConnMaxLifetime(30 * time.Minute)
|
||||
db.SetConnMaxIdleTime(5 * time.Minute)
|
||||
|
||||
if err := db.PingContext(ctx); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("reaching the PostgreSQL store: %w", err)
|
||||
}
|
||||
if err := migrate(ctx, db, string(PostgresDialect)); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &SQLite{
|
||||
raw: db,
|
||||
dialect: PostgresDialect,
|
||||
db: bound{c: PostgresDialect, to: db},
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// openPG opens the test database named by LLMGW_TEST_POSTGRES, or skips.
|
||||
//
|
||||
// Env-gated rather than spun up here, because a store test that needs a
|
||||
// container is a test nobody runs. Pointed at a database, it runs the same
|
||||
// assertions the SQLite tests make — which is the only thing that can catch
|
||||
// the two backends drifting apart.
|
||||
//
|
||||
// Every table is emptied first. The alternative is a test database that
|
||||
// accumulates rows until an assertion about counts starts failing for reasons
|
||||
// that have nothing to do with the change being tested.
|
||||
func openPG(t *testing.T) *SQLite {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("LLMGW_TEST_POSTGRES")
|
||||
if dsn == "" {
|
||||
t.Skip("set LLMGW_TEST_POSTGRES to a database DSN to run these")
|
||||
}
|
||||
requireTestDatabase(t, dsn)
|
||||
|
||||
s, err := OpenPostgres(context.Background(), dsn)
|
||||
if err != nil {
|
||||
t.Fatalf("opening the test database: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s.Close() })
|
||||
|
||||
// CASCADE because the tables reference each other, and TRUNCATE rather
|
||||
// than DELETE so identity columns restart too.
|
||||
_, err = s.db.ExecContext(context.Background(),
|
||||
// leadership included, and it is not in copyOrder: a subtest that
|
||||
// inherits the previous one's lease is refused for reasons that have
|
||||
// nothing to do with what it is testing. It passed on SQLite, where
|
||||
// every subtest gets a fresh file, and failed only here.
|
||||
`TRUNCATE thread_messages, threads, memory_notes, usage, api_keys,
|
||||
key_claims, limits, models, sessions, users, projects, teams,
|
||||
organizations, leadership RESTART IDENTITY CASCADE`)
|
||||
if err != nil {
|
||||
t.Fatalf("clearing the test database: %v", err)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// requireTestDatabase refuses to run against a database that is not obviously
|
||||
// a test one.
|
||||
//
|
||||
// These tests truncate every table before each one. Pointed at a real
|
||||
// deployment that is a silent, complete data loss — every key, every user,
|
||||
// every conversation — and it is not a hypothetical: it happened here, through
|
||||
// an SSH tunnel to the production database that had been opened minutes
|
||||
// earlier for a perfectly good reason. The DSN was correct for what it was
|
||||
// opened for and catastrophic for what it was then used for.
|
||||
//
|
||||
// The name has to say so. A guard that could be satisfied by accident would
|
||||
// not have stopped it.
|
||||
func requireTestDatabase(t *testing.T, dsn string) {
|
||||
t.Helper()
|
||||
u, err := url.Parse(dsn)
|
||||
if err != nil {
|
||||
t.Fatalf("LLMGW_TEST_POSTGRES is not a URL: %v", err)
|
||||
}
|
||||
name := strings.TrimPrefix(u.Path, "/")
|
||||
if !strings.Contains(strings.ToLower(name), "test") {
|
||||
t.Fatalf("refusing to run destructive tests against database %q:\n"+
|
||||
" these TRUNCATE every table before each test.\n"+
|
||||
" Point LLMGW_TEST_POSTGRES at a database whose name contains \"test\".", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresCarriesTheWholeTenancyTree(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
|
||||
proj := tree(t, s)
|
||||
secret, k := mustKey(t, s, "k1", proj)
|
||||
_ = secret
|
||||
|
||||
got, err := s.KeyByHash(ctx, k.Hash)
|
||||
if err != nil {
|
||||
t.Fatalf("reading the key back: %v", err)
|
||||
}
|
||||
if got.ID != k.ID || got.ProjectID != proj {
|
||||
t.Errorf("key came back as %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresKeepsUsageAndSpend(t *testing.T) {
|
||||
// The accounting path, which is the one that decides whether a budget
|
||||
// means anything after a restart.
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
proj := tree(t, s)
|
||||
_, k := mustKey(t, s, "k1", proj)
|
||||
|
||||
now := time.Now()
|
||||
err := s.WriteUsage(ctx, []Usage{
|
||||
{At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
|
||||
PromptTokens: 10, CompletionTokens: 20, CostMicros: 500, Status: 200, LatencyMs: 42},
|
||||
{At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
|
||||
PromptTokens: 1, CompletionTokens: 2, CostMicros: 250, Status: 200, LatencyMs: 7},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("writing usage: %v", err)
|
||||
}
|
||||
|
||||
spend, err := s.SpendSince(ctx, now.Add(-time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if spend[k.ID] != 750 {
|
||||
t.Errorf("spend = %d, want 750", spend[k.ID])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresKeepsThreadsInOrderAndScoped(t *testing.T) {
|
||||
// The newest tables, and the ones with the most SQL per operation —
|
||||
// a transaction, a MAX(seq), an insert per turn.
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
proj := tree(t, s)
|
||||
if err := s.SetThreadRetention(ctx, proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for _, m := range []Message{
|
||||
{Role: "user", Content: "first"},
|
||||
{Role: "assistant", Content: "answer", Model: "small"},
|
||||
{Role: "user", Content: "second"},
|
||||
} {
|
||||
if err := s.AppendThread(ctx, proj, "th1", []Message{m}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got, err := s.Thread(ctx, proj, "th1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != 3 || got[0].Content != "first" || got[2].Content != "second" {
|
||||
t.Fatalf("thread came back as %+v", got)
|
||||
}
|
||||
if got[1].Model != "small" {
|
||||
t.Errorf("the model that answered was lost: %+v", got[1])
|
||||
}
|
||||
|
||||
// And the scope holds here exactly as it does on SQLite.
|
||||
if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "team1", Name: "Other"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.Thread(ctx, "proj2", "th1"); !errors.Is(err, ErrNoThread) {
|
||||
t.Errorf("reading across projects returned %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresEnforcesLimitsTheSameWay(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
proj := tree(t, s)
|
||||
_, k := mustKey(t, s, "k1", proj)
|
||||
|
||||
rpm := 7
|
||||
budget := int64(9000)
|
||||
if err := s.PutLimits(ctx, Limits{
|
||||
Scope: ScopeProject, ScopeID: proj, RPM: &rpm,
|
||||
Budget: &budget, BudgetPeriod: PeriodMonth,
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
eff, err := s.Resolve(ctx, k.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if eff.RPM == nil || *eff.RPM != 7 {
|
||||
t.Errorf("rpm did not resolve down the tree: %+v", eff.RPM)
|
||||
}
|
||||
if eff.Budget == nil || *eff.Budget != 9000 {
|
||||
t.Errorf("budget did not resolve down the tree: %+v", eff.Budget)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresUpsertsRatherThanDuplicating(t *testing.T) {
|
||||
// ON CONFLICT is the one non-trivial construct shared by both dialects.
|
||||
// If it behaved differently here, applying a configuration twice would
|
||||
// start failing instead of being a no-op.
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
tree(t, s)
|
||||
tree(t, s)
|
||||
|
||||
orgs, err := s.Orgs(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orgs) != 1 {
|
||||
t.Errorf("applying the tree twice produced %d orgs", len(orgs))
|
||||
}
|
||||
}
|
||||
@@ -17,7 +17,14 @@ import (
|
||||
|
||||
// SQLite is the single-instance store.
|
||||
type SQLite struct {
|
||||
db *sql.DB
|
||||
// raw is for the things that are not queries: transactions, pings, close.
|
||||
raw *sql.DB
|
||||
dialect Dialect
|
||||
// db is the same handle with placeholders rewritten for the dialect. Every
|
||||
// query goes through it, which is the point — a call site that reached the
|
||||
// driver directly would work on SQLite and fail on PostgreSQL, and only on
|
||||
// the path nobody exercised before deploying.
|
||||
db bound
|
||||
|
||||
// mu guards the default limits, which are set once at wiring but read on
|
||||
// every authentication.
|
||||
@@ -63,11 +70,11 @@ func OpenSQLite(ctx context.Context, path string) (*SQLite, error) {
|
||||
return nil, fmt.Errorf("securing %s: %w", f, err)
|
||||
}
|
||||
}
|
||||
return &SQLite{db: db}, nil
|
||||
return &SQLite{raw: db, dialect: SQLiteDialect, db: bound{c: SQLiteDialect, to: db}}, nil
|
||||
}
|
||||
|
||||
// Close releases the database.
|
||||
func (s *SQLite) Close() error { return s.db.Close() }
|
||||
func (s *SQLite) Close() error { return s.raw.Close() }
|
||||
|
||||
// --- tenancy ---
|
||||
|
||||
@@ -193,7 +200,7 @@ func (s *SQLite) TouchKeys(ctx context.Context, seen map[string]time.Time) error
|
||||
if len(seen) == 0 {
|
||||
return nil
|
||||
}
|
||||
return s.inTx(ctx, func(tx *sql.Tx) error {
|
||||
return s.inTx(ctx, func(tx conn) error {
|
||||
stmt, err := tx.PrepareContext(ctx,
|
||||
`UPDATE api_keys SET last_used_at = ? WHERE id = ? AND (last_used_at IS NULL OR last_used_at < ?)`)
|
||||
if err != nil {
|
||||
@@ -378,7 +385,7 @@ func (s *SQLite) WriteUsage(ctx context.Context, records []Usage) error {
|
||||
if len(records) == 0 {
|
||||
return nil
|
||||
}
|
||||
return s.inTx(ctx, func(tx *sql.Tx) error {
|
||||
return s.inTx(ctx, func(tx conn) error {
|
||||
stmt, err := tx.PrepareContext(ctx,
|
||||
`INSERT INTO usage (ts, key_id, project_id, model, served_model,
|
||||
prompt_tokens, completion_tokens, cost_micros, status, latency_ms)
|
||||
@@ -421,13 +428,26 @@ func (s *SQLite) SpendSince(ctx context.Context, since time.Time) (map[string]in
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
func (s *SQLite) inTx(ctx context.Context, fn func(*sql.Tx) error) error {
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
// begin starts a transaction and returns it twice: bound for queries, raw for
|
||||
// commit and rollback. The bound half is what callers must use, so a statement
|
||||
// inside a transaction is rewritten exactly as one outside it is.
|
||||
func (s *SQLite) begin(ctx context.Context) (conn, *sql.Tx, error) {
|
||||
tx, err := s.raw.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return bound{c: s.dialect, to: tx}, tx, nil
|
||||
}
|
||||
|
||||
// inTx runs fn in a transaction, with the same placeholder rewriting the
|
||||
// outside-transaction path gets.
|
||||
func (s *SQLite) inTx(ctx context.Context, fn func(conn) error) error {
|
||||
tx, err := s.raw.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
if err := fn(tx); err != nil {
|
||||
if err := fn(bound{c: s.dialect, to: tx}); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
@@ -453,7 +473,7 @@ func scanKey(sc scanner) (Key, error) {
|
||||
return k, nil
|
||||
}
|
||||
|
||||
func queryAll[T any](ctx context.Context, db *sql.DB, query string, args []any,
|
||||
func queryAll[T any](ctx context.Context, db conn, query string, args []any,
|
||||
scan func(*sql.Rows) (T, error),
|
||||
) ([]T, error) {
|
||||
rows, err := db.QueryContext(ctx, query, args...)
|
||||
|
||||
@@ -85,7 +85,7 @@ func TestTheKeyItselfIsNeverStored(t *testing.T) {
|
||||
project := tree(t, s)
|
||||
secret, _ := mustKey(t, s, "key1", project)
|
||||
|
||||
rows, err := s.db.Query(`SELECT id, key_hash, key_prefix, name FROM api_keys`)
|
||||
rows, err := s.db.QueryContext(context.Background(), `SELECT id, key_hash, key_prefix, name FROM api_keys`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -137,7 +137,7 @@ func (s *SQLite) AppendThread(ctx context.Context, projectID, threadID string, m
|
||||
return ErrThreadsNotEnabled
|
||||
}
|
||||
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
q, tx, err := s.begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -145,10 +145,10 @@ func (s *SQLite) AppendThread(ctx context.Context, projectID, threadID string, m
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
var owner string
|
||||
err = tx.QueryRowContext(ctx, `SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner)
|
||||
err = q.QueryRowContext(ctx, `SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner)
|
||||
switch {
|
||||
case errors.Is(err, sql.ErrNoRows):
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`INSERT INTO threads (id, project_id, created_at, updated_at) VALUES (?, ?, ?, ?)`,
|
||||
threadID, projectID, now, now); err != nil {
|
||||
return err
|
||||
@@ -163,7 +163,7 @@ func (s *SQLite) AppendThread(ctx context.Context, projectID, threadID string, m
|
||||
}
|
||||
|
||||
var next int64
|
||||
if err := tx.QueryRowContext(ctx,
|
||||
if err := q.QueryRowContext(ctx,
|
||||
`SELECT COALESCE(MAX(seq), -1) + 1 FROM thread_messages WHERE thread_id = ?`,
|
||||
threadID).Scan(&next); err != nil {
|
||||
return err
|
||||
@@ -173,14 +173,14 @@ func (s *SQLite) AppendThread(ctx context.Context, projectID, threadID string, m
|
||||
if m.Model != "" {
|
||||
model = m.Model
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`INSERT INTO thread_messages (thread_id, seq, role, content, model) VALUES (?, ?, ?, ?, ?)`,
|
||||
threadID, next+int64(i), m.Role, m.Content, model); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// Touched on every append, so retention measures idleness rather than age.
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`UPDATE threads SET updated_at = ? WHERE id = ?`, now, threadID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user