diff --git a/CHANGELOG.md b/CHANGELOG.md index 139bccd8..0abbffb6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,19 @@ ### Added +- `sqlflow serve` answers requests from a pool of sessions rather than one + connection. `serve.pool.size` sets how many run at once, four by default, + measured to peak at 72 MiB resident on Linux against a 256 MB box. Every + session is pinned to UTC, and the config's `commands` run once, on a + connection of their own, because `ATTACH` is database-wide while + `SET TimeZone` is not. A response carries `queued_ms` beside `elapsed_ms`, + so a busy pool is no longer reported as a slow query. `/healthz` answers + `busy` rather than `unavailable` when no session is free, and answers + `HEAD` for monitors. With `serve.metrics.enabled`, `GET /metrics` serves + six instruments, including the session wait that sizes the pool. A request + that gives up now stops reading at the next batch and releases its reader, + which is ADBC's documented equivalent of cancelling. + - `sqlflow serve`: an `integer` param may declare `min` and `max`. A request outside them is `400 invalid_param` naming the bounds, and `/v1/datasets` lists them. diff --git a/README.md b/README.md index 2c95c561..54a5742b 100644 --- a/README.md +++ b/README.md @@ -245,6 +245,14 @@ serve: limits: max_rows: 10000 timeout_seconds: 10 + pool: + # Requests answered at once. Each session is a backend session; a + # concurrent query costs a few MiB. Omit for the default, 4. + size: 4 + metrics: + # Serve GET /metrics on this listener, without a token. Off by default: + # the listener is public and the labels name every dataset. + enabled: false datasets: - name: posts_by_lang description: Posts per bucket per language. @@ -325,7 +333,8 @@ Routes, all `GET`: | Route | Auth | Returns | |---|---|---| -| `/healthz` | none | `200` when the connection answers `SELECT 1`, `503` otherwise | +| `/healthz` | none | `200` `{"status":"ok"}` when a session answers `SELECT 1`; `503` `{"status":"busy"}` when none is free, `{"status":"unavailable"}` when the query fails. `HEAD` is answered too, for monitors | +| `/metrics` | none | Prometheus text, only when `serve.metrics.enabled` is set | | `/v1/datasets` | bearer | Every dataset: its params, and its SQL as written | | `/v1/datasets/{name}` | bearer | Rows | @@ -335,12 +344,19 @@ $ curl -H 'Authorization: Bearer ' \ {"dataset":"posts_by_lang","grain":"1h", "columns":[{"name":"bucket","type":"TIMESTAMP WITH TIME ZONE"},...], "rows":[{"bucket":"2026-09-10T00:00:00Z","lang":"en","posts":102340}], - "row_count":1,"truncated":false,"elapsed_ms":41} + "row_count":1,"truncated":false,"queued_ms":0,"elapsed_ms":41} ``` A zoned timestamp is UTC. A decimal is a string of its exact digits. `NaN` and infinities are strings. `truncated: true` means `max_rows` cut the result. +`elapsed_ms` is the query. `queued_ms` is how long the request waited for a +session, which is what rises when the pool is too small for the load. + +Every route is `GET`, and `/healthz` also answers `HEAD`. A `HEAD` of a +dataset would run its query, borrow a session and discard the rows, so it is +refused with `405`. + A token is an identifier, not a secret: a browser page ships it in plain sight. It names the caller in the request log, and deleting it revokes the caller. @@ -371,13 +387,23 @@ What to know before you deploy it: tables in Postgres and generates the datasets that read them. A Postgres view with the `GROUP BY` also pushes the filter in, but re-aggregates the range on every request. -- **Bound Postgres connections.** One scan opens up to `pg_connection_limit` - connections, 64 by default. Set it low for a small database, as a command. -- **A timeout does not stop the query.** DuckDB cannot be cancelled through - its Go driver. At the deadline the caller gets `504`, and the query runs to - completion. `max_rows` does stop it early. -- **One connection serves every request.** Requests run one at a time. A slow - query makes the ones behind it wait, and `/healthz` waits with them. +- **Bound Postgres connections.** An attachment opens up to + `pg_connection_limit` connections, 64 by default. Set it low for a small + database, as a command. The pool does not multiply it: with eight sessions + scanning at once and a limit of four, the measured peak was four, so the + connections are shared across sessions rather than opened per session. +- **A timeout stops reading, not always the query.** At the deadline the + caller gets `504`, and the reader stops at the next batch and is released, + which is ADBC's equivalent of cancelling. An operator that runs long before + yielding a batch still runs to the end, holding its session. Bound the part + that is usually slow in the backend instead: a libpq connection string takes + `options='-c statement_timeout=30000'`, so an attached Postgres enforces its + own ceiling. +- **A pool serves requests.** `serve.pool.size` sessions answer at once, four + by default. A request waits for a free session, and that wait counts toward + the dataset's timeout, so an exhausted pool answers `504 query_timeout`. + `queued_ms` in the response and `sqlflow_serve_session_wait_seconds` in the + metrics say whether the pool is the limit. ### `sqlflow rollup` diff --git a/docs/coverage/features.yml b/docs/coverage/features.yml index 11aaa0fd..5e1185fb 100644 --- a/docs/coverage/features.yml +++ b/docs/coverage/features.yml @@ -196,7 +196,7 @@ features: - id: cli.serve description: Serves a config's named SQL datasets over HTTP, with bearer tokens, typed params, grains and limits. - requires: [unit, release] + requires: [unit, integration, release] - id: cli.rollup description: Generates rollup tables, the triggers that keep them current, and the serve datasets that read them from one declaration, and checks the generated files have not drifted. diff --git a/docs/coverage/matrix.md b/docs/coverage/matrix.md index afcc9e44..ed9d6626 100644 --- a/docs/coverage/matrix.md +++ b/docs/coverage/matrix.md @@ -58,7 +58,7 @@ added, and this page changes only when a status does. | `observability.debug_api` | Serves ad-hoc SQL against the live DuckDB connection. | ✅ | — | — | | `cli.invocation` | Resolves the config path and message limits from either flag form. | ✅ | — | — | | `cli.dev_invoke` | Runs a pipeline against a fixture file, without a source. | ✅ | — | ✅ | -| `cli.serve` | Serves a config's named SQL datasets over HTTP, with bearer tokens, typed params, grains and limits. | ✅ | — | ✅ | +| `cli.serve` | Serves a config's named SQL datasets over HTTP, with bearer tokens, typed params, grains and limits. | ✅ | ✅ | ✅ | | `cli.rollup` | Generates rollup tables, the triggers that keep them current, and the serve datasets that read them from one declaration, and checks the generated files have not drifted. | ✅ | ✅ | — | | `cli.version` | The shipped binary reports the version it was built from. | — | — | ✅ | | `tooling.conformance` | The harness proves the declared invariants for any integration. | ✅ | — | — | diff --git a/docs/coverage/status/features.yml b/docs/coverage/status/features.yml index e312103d..81a7dcc0 100644 --- a/docs/coverage/status/features.yml +++ b/docs/coverage/status/features.yml @@ -3,7 +3,7 @@ cli.dev_invoke: {unit: covered, integration: not_required, release: covered} cli.invocation: {unit: covered, integration: not_required, release: not_required} cli.rollup: {unit: covered, integration: covered, release: not_required} -cli.serve: {unit: covered, integration: not_required, release: covered} +cli.serve: {unit: covered, integration: covered, release: covered} cli.version: {unit: not_required, integration: not_required, release: covered} config.templating: {unit: covered, integration: not_required, release: covered} config.validation: {unit: covered, integration: not_required, release: covered} diff --git a/docs/superpowers/plans/2026-09-16-serve-pool.md b/docs/superpowers/plans/2026-09-16-serve-pool.md new file mode 100644 index 00000000..18705932 --- /dev/null +++ b/docs/superpowers/plans/2026-09-16-serve-pool.md @@ -0,0 +1,1762 @@ +# Serve Session Pool Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace serve's single serialized DuckDB connection with a pool of sessions behind an executor interface, and add the metrics that size it. + +**Architecture:** `internal/serve/executor.go` defines `Executor`, `Session` and `Statement`, and holds the pool mechanics — the session channel, the wait timing, and the goroutine that lets a caller give up without abandoning a running query. `internal/serve/duckdb.go` is the only file in the package that imports ADBC. `serve.New` takes an `Executor` instead of a connection. + +**Tech Stack:** Go 1.25, ADBC over DuckDB 1.5.2, arrow-go/v18, OpenTelemetry instruments exported through a `prometheus.Registry`, cobra, zeebo/assert, testcontainers Postgres. + +**Spec:** `docs/superpowers/specs/2026-09-16-serve-pool-design.md` + +## Global Constraints + +- Branch from `origin/main` at or after `54b17f2` (v2026.09.16). One PR. +- Default `pool.size` is **4**, subject to Task 1's measurement. +- `pool.size` accepts 0 (meaning the default) through 64; negative or above 64 is `user.config.invalid` at the key's path. +- Serve runs `SET TimeZone='UTC'` on every session it opens. The config's `commands` run exactly once, on a setup connection. +- The interchange is `array.RecordReader`. Serve owns `readRows` and the `max_rows` cap. Nothing in `executor.go` imports ADBC or names DuckDB. +- A session returns to the pool when its query finishes, never when a caller gives up. ADBC's Go API has no `Cancel` (checked against v1.6.0), so nothing outside the reading goroutine can stop a query. +- Releasing a `RecordReader` without draining it is ADBC's documented equivalent of cancel, so `readRows` takes the request's context and stops at the first batch boundary after it ends. +- Waiting for a session counts toward the dataset's timeout; exhausting it is `504 query_timeout`, the existing code. +- Metrics are off unless `serve.metrics.enabled` is true, and are served at `GET /metrics` on the existing listener with no token. +- Unit tests are `TestCliServe_*` and call `coverage.Covers(t, "cli.serve")`. Integration tests are `TestIntegrationServePool_*`, skip under `-short`, and start their own container. +- Unit: `go test -short -race ./...`. Integration: `go test -run '^TestIntegration' ./internal/serve/`. Goldens regenerate with `UPDATE_GOLDEN=1`; schemas with `make schema`. +- Comments say why, in full sentences, matching the surrounding code. Commit messages are `area: what changed`. + +## File Structure + +``` +internal/serve/executor.go new: the interfaces, the pool, the query helper +internal/serve/duckdb.go new: the DuckDB executor, statement, binding +internal/serve/query.go shrinks: statement/prepare/executor move out +internal/serve/server.go New takes an Executor; Close closes it +internal/serve/http.go acquire per request; queued_ms; healthz busy +internal/serve/metrics.go new: the six instruments and the handler +internal/config/serve.go ServePool, ServeMetrics, their rules +internal/cli/serve/serve.go builds the executor, passes commands as init +internal/validate/schemas/serve.json, internal/cli/testdata/serve_example.golden regenerated +docs/coverage/features.yml cli.serve gains integration +README.md, CHANGELOG.md docs +``` + +--- + +### Task 1: Measure per-session memory in the release container + +The spec's numbers are macOS, where `Maxrss` is bytes; on Linux it is KiB. Everything downstream assumes 4 sessions fit a 256 MB box. Settle it before writing the pool. + +**Files:** +- Create, then delete: `cmd/poolmem/main.go` +- Modify, only if the number says so: the default in `internal/config/serve.go` (Task 2) and the spec's "Memory, and why 4" + +**Interfaces:** +- Produces: a number, recorded in the PR body. No code survives this task. + +- [ ] **Step 1: Branch** + +```bash +cd /Users/danielmican/code/github.com/turbolytics/sql-flow +git fetch origin && git switch -c feat/serve-pool origin/main +go build ./... && go test -short ./internal/serve/ +``` + +Expected: builds, and `internal/serve` passes. + +- [ ] **Step 2: Write the probe** + +Create `cmd/poolmem/main.go`. It opens one in-memory DuckDB, creates a table the size of the demo's widest allowed request (365 buckets × 170 languages), then runs the fold on N connections at once and reports peak resident memory. + +```go +// Command poolmem reports peak resident memory with N concurrent DuckDB +// sessions, to size serve's pool. Throwaway: delete it after recording the +// number. +package main + +import ( + "context" + "fmt" + "os" + "runtime" + "strconv" + "sync" + "syscall" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/turbolytics/sql-flow/internal/duckdb" +) + +func run(c adbc.Connection, sql string) error { + st, err := c.NewStatement() + if err != nil { + return err + } + defer st.Close() + if err := st.SetSqlQuery(sql); err != nil { + return err + } + rdr, _, err := st.ExecuteQuery(context.Background()) + if err != nil { + return err + } + defer rdr.Release() + for rdr.Next() { + } + return rdr.Err() +} + +// rssMiB reads the process's peak resident size. Linux reports KiB and +// darwin reports bytes, which is the whole reason this runs in the container. +func rssMiB() float64 { + var u syscall.Rusage + _ = syscall.Getrusage(syscall.RUSAGE_SELF, &u) + if runtime.GOOS == "linux" { + return float64(u.Maxrss) / 1024 + } + return float64(u.Maxrss) / (1024 * 1024) +} + +const fold = `SELECT bucket, CASE WHEN r <= 10 THEN lang ELSE 'other' END AS lang, sum(posts)::BIGINT AS posts +FROM (SELECT bucket, lang, posts, dense_rank() OVER (ORDER BY tot DESC, lang) AS r + FROM (SELECT bucket, lang, posts, sum(posts) OVER (PARTITION BY lang) AS tot FROM r)) +GROUP BY ALL` + +func main() { + n, err := strconv.Atoi(os.Args[1]) + if err != nil { + panic(err) + } + ctx := context.Background() + db, err := duckdb.OpenPath(ctx, "") + if err != nil { + panic(err) + } + defer db.Close() + + setup, err := db.Connect(ctx) + if err != nil { + panic(err) + } + if err := run(setup, "SET memory_limit='128MB'"); err != nil { + panic(err) + } + if err := run(setup, `CREATE TABLE r AS + SELECT b.bucket, 'l' || lpad(l.i::VARCHAR, 3, '0') AS lang, + (1 + 200000 / (l.i * l.i))::BIGINT AS posts + FROM (SELECT unnest(generate_series(TIMESTAMPTZ '2026-01-01', TIMESTAMPTZ '2027-01-01', INTERVAL '1 day')) AS bucket) b, + (SELECT unnest(generate_series(1, 170)) AS i) l`); err != nil { + panic(err) + } + + conns := []adbc.Connection{setup} + for i := 1; i < n; i++ { + c, err := db.Connect(ctx) + if err != nil { + panic(err) + } + conns = append(conns, c) + } + fmt.Printf("%d sessions idle: %.0f MiB\n", n, rssMiB()) + + var wg sync.WaitGroup + for _, c := range conns { + wg.Add(1) + go func(c adbc.Connection) { + defer wg.Done() + if err := run(c, fold); err != nil { + fmt.Println(" query failed:", err) + } + }(c) + } + wg.Wait() + fmt.Printf("%d concurrent folds: %.0f MiB peak\n", n, rssMiB()) +} +``` + +- [ ] **Step 3: Run it in the release container** + +The image carries the matching `libduckdb.so` and sets `SQLFLOW_DUCKDB_LIB`, so the probe must build and run inside it rather than on the host. + +```bash +make sqlflow-image SQLFLOW_IMAGE=sqlflow-poolmem:local +docker run --rm -v "$PWD:/src" -w /src --entrypoint /bin/sh sqlflow-poolmem:local -c ' + go run ./cmd/poolmem 1; go run ./cmd/poolmem 4; go run ./cmd/poolmem 8' +``` + +Expected: three pairs of lines. If `go` is absent from the image, build a static probe on the host for linux/amd64 and mount the binary instead: + +```bash +CGO_ENABLED=1 GOOS=linux go build -o /tmp/poolmem ./cmd/poolmem +docker run --rm -v /tmp/poolmem:/poolmem --entrypoint /poolmem sqlflow-poolmem:local 4 +``` + +- [ ] **Step 4: Decide the default** + +Record all six numbers. The default stands at 4 if four concurrent folds stay under 160 MiB, which leaves room for the Go runtime inside 256 MB. If four exceed that, set the default to the largest N that fits and say so in the spec's memory table and in the PR body. + +- [ ] **Step 5: Delete the probe and commit the finding** + +```bash +rm -rf cmd/poolmem +git status --short # expect nothing, or only the spec if the default changed +``` + +If the spec changed: + +```bash +git add docs/superpowers/specs/2026-09-16-serve-pool-design.md +git commit -m "spec: per-session memory measured in the release container" +``` + +If it did not, this task produces no commit. Carry the numbers into the PR body either way. + +--- + +### Task 2: Config for the pool and metrics + +**Files:** +- Modify: `internal/config/serve.go` (`Serve` struct, new types, `checkLimits`'s neighbours in `Check`) +- Test: `internal/config/serve_test.go` +- Regenerated: `internal/validate/schemas/serve.json`, `internal/cli/testdata/serve_example.golden` + +**Interfaces:** +- Produces: + - `config.ServePool{Size int}` with yaml key `size`, on `Serve.Pool *ServePool` at yaml key `pool` + - `config.ServeMetrics{Enabled bool}` with yaml key `enabled`, on `Serve.Metrics *ServeMetrics` at yaml key `metrics` + - `func (s Serve) PoolSize() int`, returning `DefaultServePoolSize` when unset + - `const DefaultServePoolSize = 4`, `const MaxServePoolSize = 64` + - `func (s Serve) MetricsEnabled() bool` + +- [ ] **Step 1: Write the failing rule tests** + +In `internal/config/serve_test.go`, add to the table in `TestCliServe_CheckReportsEachRuleAtItsPath`, after the `min above max` case: + +```go + {"negative pool size", " limits:", " pool: {size: -1}\n limits:", + errs.CodeConfigInvalid, "serve.pool.size", "must not be negative"}, + {"pool size past the ceiling", " limits:", " pool: {size: 65}\n limits:", + errs.CodeConfigInvalid, "serve.pool.size", "65 sessions is more than the 64 this version allows"}, +``` + +And add a resolution test after `TestCliServe_LimitsResolveDatasetThenTopLevelThenDefault`: + +```go +// The pool defaults rather than failing closed: a config written before the +// pool existed gets concurrency, not one session. +func TestCliServe_PoolSizeDefaults(t *testing.T) { + coverage.Covers(t, "cli.serve") + + conf := parseServe(t, validServe) + assert.Equal(t, DefaultServePoolSize, conf.Serve.PoolSize()) + assert.False(t, conf.Serve.MetricsEnabled()) + + sized := parseServe(t, strings.Replace(validServe, " limits:", " pool: {size: 8}\n metrics: {enabled: true}\n limits:", 1)) + assert.Equal(t, 8, sized.Serve.PoolSize()) + assert.True(t, sized.Serve.MetricsEnabled()) + + // 0 is unset, not "no sessions". + zero := parseServe(t, strings.Replace(validServe, " limits:", " pool: {size: 0}\n limits:", 1)) + assert.Equal(t, DefaultServePoolSize, zero.Serve.PoolSize()) +} +``` + +- [ ] **Step 2: Run them to see them fail** + +Run: `go test -short ./internal/config/ -run 'TestCliServe_(CheckReportsEachRuleAtItsPath|PoolSizeDefaults)'` +Expected: FAIL to compile, `undefined: DefaultServePoolSize`, and the YAML rejects `field pool not found`. + +- [ ] **Step 3: Add the types, the accessors and the rules** + +In `internal/config/serve.go`, add to the defaults block: + +```go + // DefaultServePoolSize is the sessions a server holds when the config + // names no number. Four fits a 256 MB box: idle sessions cost nothing + // measurable, and a concurrent query costs about 15 MiB. + DefaultServePoolSize = 4 + // MaxServePoolSize is the most this version accepts. A larger pool on a + // small box fails queries with out-of-memory rather than queueing them, + // because DuckDB's memory_limit is one budget shared by every session. + MaxServePoolSize = 64 +``` + +Add to `Serve`, after `Limits`: + +```go + // How many requests the server runs at once. Omit for the default. + Pool *ServePool `yaml:"pool,omitempty"` + // Whether to serve Prometheus metrics at /metrics. + Metrics *ServeMetrics `yaml:"metrics,omitempty"` +``` + +Add the types after `ServeRateLimit`: + +```go +// ServePool sizes the sessions a server answers requests on. +type ServePool struct { + // Sessions the server holds. Each is one backend session, and one + // request uses one at a time, so this is the requests that run at once. + // 0 means the default, 4. + Size int `yaml:"size,omitempty"` +} + +// ServeMetrics turns on the Prometheus endpoint. +type ServeMetrics struct { + // Serve GET /metrics on the same listener as the datasets, without a + // token. Off by default: that listener is public, and the metric labels + // name every dataset and grain. + Enabled bool `yaml:"enabled,omitempty"` +} +``` + +Add the accessors beside `MaxRows`: + +```go +// PoolSize is the sessions to hold, defaulted. +func (s Serve) PoolSize() int { + if s.Pool != nil && s.Pool.Size > 0 { + return s.Pool.Size + } + return DefaultServePoolSize +} + +// MetricsEnabled reports whether to serve /metrics. +func (s Serve) MetricsEnabled() bool { + return s.Metrics != nil && s.Metrics.Enabled +} +``` + +In `Check`, directly after the `checkLimits(s.Limits, ...)` call: + +```go + if s.Pool != nil { + switch { + case s.Pool.Size < 0: + add(errs.CodeConfigInvalid, []string{"serve", "pool", "size"}, + "pool.size is %d; it must not be negative, and 0 means the default", s.Pool.Size) + case s.Pool.Size > MaxServePoolSize: + add(errs.CodeConfigInvalid, []string{"serve", "pool", "size"}, + "pool.size %d sessions is more than the %d this version allows; DuckDB's memory limit is one budget shared by every session, so a pool past the box fails queries rather than queueing them", + s.Pool.Size, MaxServePoolSize) + } + } +``` + +- [ ] **Step 4: Run the tests and regenerate** + +```bash +go test -short ./internal/config/ +make schema +git diff --stat +go test -short ./internal/schema/ ./internal/cli/ ./internal/validate/ +``` + +Expected: config passes; `git diff --stat` lists `serve.json` and `serve_example.golden`; the rest passes. + +- [ ] **Step 5: Commit** + +```bash +git add internal/config/serve.go internal/config/serve_test.go \ + internal/validate/schemas/serve.json internal/cli/testdata/serve_example.golden +git commit -m "serve: config for the session pool and the metrics endpoint" +``` + +--- + +### Task 3: The executor interface, the pool, and the DuckDB implementation + +The big one. It moves the request path off `adbc.Connection` and onto a pool, with no behaviour change at `size: 1`. + +**Files:** +- Create: `internal/serve/executor.go`, `internal/serve/duckdb.go` +- Modify: `internal/serve/query.go` (statement and executor move out; `readRows` stays in encode.go), `internal/serve/server.go`, `internal/serve/http.go`, `internal/cli/serve/serve.go` +- Test: `internal/serve/executor_test.go`, and the existing `internal/serve/http_test.go` helper + +**Interfaces:** +- Consumes: `config.ServeParam`, `config.Serve.PoolSize()`, `duckdb.DB`, `core.InitCommands`, `readRows`, `result`. +- Produces: + - `serve.Executor`, `serve.Session`, `serve.Statement`, `serve.StatementSpec`, `serve.Stats`, `serve.ErrClosed` — exactly as the spec's "The interfaces" section writes them + - `func NewDuckDBExecutor(ctx context.Context, db *duckdb.DB, size int, init func(context.Context, adbc.Connection) error) (Executor, error)` + - `func New(ctx context.Context, conf *config.ServeConf, ex Executor, opts ...Option) (*Server, error)` — `New`'s third parameter changes from `adbc.Connection` to `Executor` + - unexported: `newPool(sessions []backendSession, onWait func(time.Duration)) *pool`, `query(ctx, ex, st, values, maxRows) (result, time.Duration, error)` + +- [ ] **Step 1: Write the failing pool tests** + +Create `internal/serve/executor_test.go`. These test the pool through a fake backend, so they need no DuckDB and run fast. + +```go +package serve + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/zeebo/assert" +) + +// fakeSession is a backend that sleeps instead of querying, so pool +// behaviour is testable without DuckDB. +type fakeSession struct { + delay time.Duration + closed bool + runs int +} + +func (f *fakeSession) run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) { + f.runs++ + select { + case <-time.After(f.delay): + return nil, errors.New("fake session returns no rows") + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (f *fakeSession) close() error { f.closed = true; return nil } + +func newFakePool(t *testing.T, n int, delay time.Duration) (*pool, []*fakeSession) { + t.Helper() + fakes := make([]*fakeSession, n) + backends := make([]backendSession, n) + for i := range fakes { + fakes[i] = &fakeSession{delay: delay} + backends[i] = fakes[i] + } + return newPool(backends, func(time.Duration) {}), fakes +} + +// The whole point: N sessions run N queries in about the time of one. +func TestCliServe_PoolRunsQueriesConcurrently(t *testing.T) { + coverage.Covers(t, "cli.serve") + + const n = 4 + const delay = 200 * time.Millisecond + p, _ := newFakePool(t, n, delay) + defer p.Close() + + start := time.Now() + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + defer s.Release() + _, _ = s.Run(context.Background(), nil, nil) + }() + } + wg.Wait() + + // Serial would be 4 x 200ms. Two delays of headroom for a loaded CI box. + assert.That(t, time.Since(start) < 2*delay) +} + +// One session is the old behaviour, and the control for the test above. +func TestCliServe_APoolOfOneSerializes(t *testing.T) { + coverage.Covers(t, "cli.serve") + + const delay = 100 * time.Millisecond + p, _ := newFakePool(t, 1, delay) + defer p.Close() + + start := time.Now() + var wg sync.WaitGroup + for i := 0; i < 3; i++ { + wg.Add(1) + go func() { + defer wg.Done() + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + defer s.Release() + _, _ = s.Run(context.Background(), nil, nil) + }() + } + wg.Wait() + + assert.That(t, time.Since(start) >= 3*delay) +} + +// A caller that gives up must free its place, or the pool shrinks under load +// until it deadlocks. +func TestCliServe_AcquireReturnsWhenTheRequestGivesUp(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, _ := newFakePool(t, 1, 0) + defer p.Close() + + held, err := p.Acquire(context.Background()) + assert.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + _, err = p.Acquire(ctx) + assert.That(t, errors.Is(err, context.DeadlineExceeded)) + assert.Equal(t, 1, p.Stats().InUse) + + held.Release() + next, err := p.Acquire(context.Background()) + assert.NoError(t, err) + next.Release() +} + +// A session closed while a query runs on it takes the process down, and +// nothing outside the reading goroutine can stop that query, so Close waits. +func TestCliServe_CloseWaitsForBorrowedSessions(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, fakes := newFakePool(t, 2, 0) + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + + closed := make(chan struct{}) + go func() { p.Close(); close(closed) }() + + select { + case <-closed: + t.Fatal("Close returned while a session was borrowed") + case <-time.After(100 * time.Millisecond): + } + + s.Release() + select { + case <-closed: + case <-time.After(2 * time.Second): + t.Fatal("Close did not return after the session came back") + } + for i, f := range fakes { + if !f.closed { + t.Fatalf("session %d was not closed", i) + } + } + + _, err = p.Acquire(context.Background()) + assert.That(t, errors.Is(err, ErrClosed)) +} + +// A deferred Release beside an early return can run twice. +func TestCliServe_ReleaseTwiceIsSafe(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, _ := newFakePool(t, 1, 0) + defer p.Close() + + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + s.Release() + s.Release() + + assert.Equal(t, 0, p.Stats().InUse) + next, err := p.Acquire(context.Background()) + assert.NoError(t, err) + next.Release() +} + +func TestCliServe_StatsReportSizeAndUse(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, _ := newFakePool(t, 3, 0) + defer p.Close() + + assert.Equal(t, Stats{Size: 3, InUse: 0}, p.Stats()) + a, _ := p.Acquire(context.Background()) + b, _ := p.Acquire(context.Background()) + assert.Equal(t, Stats{Size: 3, InUse: 2}, p.Stats()) + a.Release() + b.Release() + assert.Equal(t, Stats{Size: 3, InUse: 0}, p.Stats()) +} +``` + +- [ ] **Step 2: Run them to see them fail** + +Run: `go test -short ./internal/serve/ -run 'TestCliServe_(Pool|Acquire|Close|Release|Stats)'` +Expected: FAIL to compile, `undefined: newPool`, `undefined: backendSession`, `undefined: ErrClosed`. + +- [ ] **Step 3: Write the interfaces and the pool** + +Create `internal/serve/executor.go`. Copy the interface declarations and their doc comments verbatim from the spec's "The interfaces" section, then add the pool below them. + +```go +package serve + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "time" + + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/turbolytics/sql-flow/internal/config" +) + +// ErrClosed is what Acquire returns once the executor has closed. +var ErrClosed = errors.New("the server is shutting down") + +// ... Executor, Session, Statement, StatementSpec, Stats: verbatim from the spec ... + +// backendSession is what an Executor implementation gives the pool: run a +// statement, and close. The pool owns everything else, because queueing and +// its measurement are the same whatever runs the SQL. +type backendSession interface { + run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) + close() error +} + +// pool hands out a fixed set of sessions, one caller at a time. An Executor +// implementation embeds it. +type pool struct { + // free carries every session not currently borrowed. Its capacity is the + // pool's size, so a Release never blocks. + free chan *session + all []*session + // onWait receives how long each Acquire waited, including the ones that + // gave up. The metrics histogram is the only reader. + onWait func(time.Duration) + + mu sync.Mutex + inUse int + closed bool +} + +// session is one backend session while a caller holds it. +type session struct { + pool *pool + backend backendSession + released atomic.Bool +} + +func newPool(backends []backendSession, onWait func(time.Duration)) *pool { + p := &pool{free: make(chan *session, len(backends)), onWait: onWait} + for _, b := range backends { + s := &session{pool: p, backend: b} + s.released.Store(true) + p.all = append(p.all, s) + p.free <- s + } + return p +} + +// Acquire borrows a session, waiting until one is free, ctx ends, or the +// executor closes. +func (p *pool) Acquire(ctx context.Context) (Session, error) { + start := time.Now() + select { + case s, ok := <-p.free: + if !ok { + return nil, ErrClosed + } + p.onWait(time.Since(start)) + p.mu.Lock() + p.inUse++ + p.mu.Unlock() + s.released.Store(false) + return s, nil + case <-ctx.Done(): + p.onWait(time.Since(start)) + return nil, ctx.Err() + } +} + +func (p *pool) Stats() Stats { + p.mu.Lock() + defer p.mu.Unlock() + return Stats{Size: len(p.all), InUse: p.inUse} +} + +// Close refuses later acquires, waits for every borrowed session to come +// back, and only then closes them. A session closed while a query runs on it +// takes the process down, and nothing outside the goroutine reading that +// query can stop it, so waiting is the only option. The HTTP server's drain +// bounds how long that can be. +func (p *pool) Close() { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return + } + p.closed = true + p.mu.Unlock() + + for range p.all { + <-p.free + } + close(p.free) + for _, s := range p.all { + _ = s.backend.close() + } +} + +func (s *session) Run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) { + return s.backend.run(ctx, st, values) +} + +// Release returns the session. The second call is a no-op, so a deferred +// Release beside an early return cannot return one session twice. +func (s *session) Release() { + if s.released.Swap(true) { + return + } + s.pool.mu.Lock() + s.pool.inUse-- + s.pool.mu.Unlock() + // Never blocks: free's capacity is the pool's size, and this session is + // not in it. + s.pool.free <- s +} +``` + +- [ ] **Step 4: Write the query helper** + +Still in `executor.go`. This replaces the old `executor.run` and carries the rule that matters most. + +```go +// query acquires a session, runs st, and encodes at most maxRows rows. It +// returns the rows, how long the acquire waited, and any error. +// +// The work runs on its own goroutine and hands back finished bytes, so a +// caller that has already given up at its deadline cannot race it. The +// session returns to the pool when the query finishes rather than when the +// caller stops waiting: handing a session to the next request while a query +// still runs on it would serialise them behind work nobody wants. +// +// ctx reaches readRows, which stops at the first batch boundary after the +// caller gives up and releases the reader. ADBC documents releasing a reader +// without draining it as equivalent to AdbcStatementCancel, and that is the +// only cancel its Go API offers. +func query(ctx context.Context, ex Executor, st Statement, values map[string]any, maxRows int) (result, time.Duration, error) { + start := time.Now() + sess, err := ex.Acquire(ctx) + if err != nil { + return result{}, time.Since(start), err + } + waited := time.Since(start) + + type outcome struct { + res result + err error + } + done := make(chan outcome, 1) + go func() { + defer sess.Release() + rdr, err := sess.Run(ctx, st, values) + if err != nil { + done <- outcome{err: err} + return + } + // Released whether or not the rows are drained, which is what tells + // the driver to stop. + defer rdr.Release() + res, err := readRows(ctx, rdr, maxRows) + done <- outcome{res: res, err: err} + }() + + select { + case o := <-done: + return o.res, waited, o.err + case <-ctx.Done(): + return result{}, waited, ctx.Err() + } +} +``` + +- [ ] **Step 5: Run the pool tests** + +Run: `go test -short -race ./internal/serve/ -run 'TestCliServe_(Pool|APoolOfOne|Acquire|Close|Release|Stats)'` +Expected: PASS. `APoolOfOneSerializes` and `PoolRunsQueriesConcurrently` are the pair that prove the pool does something; if both pass at any size, the concurrency assertion is not biting. + +- [ ] **Step 6: Move the DuckDB code behind the interface** + +Create `internal/serve/duckdb.go`. Move `statement`, `prepare`, `record` and the query body out of `query.go` into it, renaming `statement` to `duckdbStatement` and its `where()` to `Where()`. Delete `executor`, `outcome` and `errClosed` from `query.go`; `query.go` keeps only what is left, and if that is nothing, delete the file. + +```go +package serve + +import ( + "context" + "fmt" + "strings" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/turbolytics/sql-flow/internal/config" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/turbolytics/sql-flow/internal/errs" + "github.com/turbolytics/sql-flow/internal/sqlparams" +) + +// duckdbExecutor runs datasets on DuckDB over ADBC. It is the only type in +// this package that names either. +type duckdbExecutor struct { + *pool + db *duckdb.DB + // setup is the connection the config's commands ran on, kept because + // Prepare needs a connection and every other one is in the pool. It is + // idle afterwards, which costs nothing measurable. + setup adbc.Connection +} + +// NewDuckDBExecutor opens size sessions on db and returns an Executor over +// them. +// +// init runs once, on a connection of its own, before any session is used. The +// caller passes the config's commands: ATTACH is database-wide and attaching +// the same alias twice is an error, so they must not run per session. +// +// Every session is then pinned to UTC. That is serve's own contract rather +// than the config's: SET TimeZone is session-scoped, so a session that missed +// it evaluates date_trunc and naive casts in the host's zone and returns +// wrong buckets from a correct config. +func NewDuckDBExecutor(ctx context.Context, db *duckdb.DB, size int, + init func(context.Context, adbc.Connection) error) (Executor, error) { + setup, err := db.Connect(ctx) + if err != nil { + return nil, err + } + if init != nil { + if err := init(ctx, setup); err != nil { + _ = setup.Close() + return nil, err + } + } + + backends := make([]backendSession, 0, size) + for i := 0; i < size; i++ { + conn, err := db.Connect(ctx) + if err != nil { + _ = setup.Close() + return nil, err + } + if err := execOn(ctx, conn, "SET TimeZone='UTC'"); err != nil { + _ = conn.Close() + _ = setup.Close() + return nil, fmt.Errorf("pinning the session timezone: %w", err) + } + backends = append(backends, &duckdbSession{conn: conn}) + } + + return &duckdbExecutor{pool: newPool(backends, func(time.Duration) {}), db: db, setup: setup}, nil +} + +// execOn runs one statement for its effect. +func execOn(ctx context.Context, conn adbc.Connection, sql string) error { + stmt, err := conn.NewStatement() + if err != nil { + return err + } + defer stmt.Close() + if err := stmt.SetSqlQuery(sql); err != nil { + return err + } + _, err = stmt.ExecuteUpdate(ctx) + return err +} + +func (e *duckdbExecutor) Close() { + e.pool.Close() + _ = e.setup.Close() +} + +// duckdbSession is one ADBC connection. +type duckdbSession struct{ conn adbc.Connection } + +func (d *duckdbSession) close() error { return d.conn.Close() } + +func (d *duckdbSession) run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) { + s, ok := st.(*duckdbStatement) + if !ok { + return nil, fmt.Errorf("serve: statement %T did not come from this executor", st) + } + + stmt, err := d.conn.NewStatement() + if err != nil { + return nil, err + } + // The statement outlives this call through the reader, so it closes when + // the reader is released rather than here. + if err := stmt.SetSqlQuery(s.rewritten); err != nil { + _ = stmt.Close() + return nil, err + } + if s.schema.NumFields() > 0 { + rec := s.record(values) + err := stmt.Bind(ctx, rec) + rec.Release() + if err != nil { + _ = stmt.Close() + return nil, err + } + } + rdr, _, err := stmt.ExecuteQuery(ctx) + if err != nil { + _ = stmt.Close() + return nil, err + } + return &closingReader{RecordReader: rdr, stmt: stmt}, nil +} + +// closingReader closes the statement when the reader is released, so a +// caller that only knows about the reader leaks neither. +type closingReader struct { + array.RecordReader + stmt adbc.Statement +} + +func (c *closingReader) Release() { + c.RecordReader.Release() + _ = c.stmt.Close() +} +``` + +`Prepare` is the old `prepare` with its signature changed: + +```go +// Prepare checks one statement against DuckDB at startup. DuckDB binds a +// statement when its SQL is set, so a syntax error, a missing table and a +// missing column all fail here rather than on the first request. +func (e *duckdbExecutor) Prepare(ctx context.Context, spec StatementSpec) (Statement, error) { + // ... the body of the old prepare(), building a *duckdbStatement and + // using e.setup as the connection ... +} +``` + +Keep the old `prepare`'s parameter-count check and its message verbatim; it catches a scanner that miscounts placeholders and binds values onto the wrong ones. + +- [ ] **Step 7: Thread the Executor through the server and the CLI** + +In `internal/serve/server.go`, change `New`'s third parameter to `ex Executor`, store it as `s.exec`, replace the `prepare(ctx, conn, ...)` call with `s.exec.Prepare(ctx, StatementSpec{...})`, and change `Server.Close` to call `s.exec.Close()`. The `dataset` struct's `single` and `grains` become `Statement` rather than `*statement`. + +In `internal/serve/http.go`, replace the `s.exec.run(...)` call in `queryDataset` with: + +```go + res, _, err := query(r.Context(), s.exec, st, values, ds.maxRows) +``` + +and leave the error switch exactly as it is: an exhausted pool surfaces as `context.DeadlineExceeded`, which is already `504 query_timeout`. + +In `internal/cli/serve/serve.go`, replace the single `db.Connect` and `core.InitCommands` block with: + +```go + ex, err := api.NewDuckDBExecutor(ctx, db, conf.Serve.PoolSize(), + func(ctx context.Context, conn adbc.Connection) error { + // Uncoded, as in run: an ATTACH that fails because the database + // is not up yet exits 1, which a supervisor retries. Redacted + // because a failed ATTACH prints the connection string. + if err := core.InitCommands(conn, &config.Conf{Commands: conf.Commands}); err != nil { + return errors.New("failed to initialize commands: " + api.Redact(err.Error())) + } + return nil + }) + if err != nil { + return err + } + + srv, err := api.New(ctx, conf, ex, api.WithLogger(l)) +``` + +`srv.Close()` now closes the executor, so drop the separate `conn.Close()` defer and keep `db.Close()`. + +In `internal/serve/http_test.go`, `newTestServer` builds an executor instead of a connection: + +```go + db, err := duckdb.OpenPath(context.Background(), "") + assert.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + ex, err := NewDuckDBExecutor(context.Background(), db, 1, func(ctx context.Context, conn adbc.Connection) error { + for _, sql := range []string{"SET TimeZone='UTC'", createPostsTable} { + if err := execOn(ctx, conn, sql); err != nil { + return err + } + } + return nil + }) + assert.NoError(t, err) +``` + +where `createPostsTable` is the `CREATE TABLE posts AS ...` string already in the helper. Size 1 keeps every existing assertion about ordering and timing true. + +- [ ] **Step 8: Run the whole package** + +```bash +go build ./... && go vet ./internal/serve/ ./internal/cli/serve/ +go test -short -race ./internal/serve/ ./internal/cli/serve/ ./internal/cli/ +``` + +Expected: PASS, with no change to any existing test's assertions. If `TestCliServe_ASlowQueryIsA504AndTurnsHealthRed` fails, the health path still uses the old executor; Task 7 rewrites it, so leave it on the old shape until then by having `healthz` call `query(...)` with a one-row statement. + +- [ ] **Step 9: Commit** + +```bash +git add internal/serve/ internal/cli/serve/serve.go +git commit -m "serve: run requests on a pool of sessions behind an executor interface" +``` + +--- + +### Task 4: Every session is UTC + +Task 3 writes the pin. This task proves it can fail, which is the only way to know it is load-bearing. + +**Files:** +- Test: `internal/serve/duckdb_test.go` (new) + +**Interfaces:** +- Consumes: `NewDuckDBExecutor`, `New`, `Server.Handler`, `execOn`. +- Produces: nothing later tasks use. + +- [ ] **Step 1: Write the failing test** + +Create `internal/serve/duckdb_test.go`: + +```go +package serve + +import ( + "context" + "testing" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/zeebo/assert" +) + +// SET TimeZone is session-scoped, so a session the executor did not pin +// inherits the host's zone and evaluates date_trunc and naive casts in it. +// The test runs under a non-UTC TZ, because under UTC a missing pin looks +// correct. +func TestCliServe_EverySessionIsUTC(t *testing.T) { + coverage.Covers(t, "cli.serve") + t.Setenv("TZ", "America/New_York") + + ctx := context.Background() + db, err := duckdb.OpenPath(ctx, "") + assert.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + const size = 4 + ex, err := NewDuckDBExecutor(ctx, db, size, nil) + assert.NoError(t, err) + t.Cleanup(ex.Close) + + // Hold every session at once, so each is checked rather than the same + // one four times. + var held []Session + for i := 0; i < size; i++ { + s, err := ex.Acquire(ctx) + assert.NoError(t, err) + held = append(held, s) + } + for i, s := range held { + zone, err := scalarOn(ctx, s, ex, "SELECT current_setting('TimeZone')") + assert.NoError(t, err) + if zone != "UTC" { + t.Fatalf("session %d is in %s, not UTC", i, zone) + } + } + for _, s := range held { + s.Release() + } +} +``` + +Add the helper below it, which prepares and runs one statement on a held session: + +```go +// scalarOn runs sql on one held session and returns its first value. +func scalarOn(ctx context.Context, s Session, ex Executor, sql string) (string, error) { + st, err := ex.Prepare(ctx, StatementSpec{Dataset: "probe", SQL: sql}) + if err != nil { + return "", err + } + rdr, err := s.Run(ctx, st, nil) + if err != nil { + return "", err + } + defer rdr.Release() + res, err := readRows(rdr, 1) + if err != nil { + return "", err + } + // rows is a JSON array of one object with one column. + return string(res.Rows), nil +} +``` + +The returned string is JSON, so assert with `strings.Contains(zone, "UTC")` rather than equality, and import `strings`. + +- [ ] **Step 2: Prove the test can fail** + +Temporarily delete the `SET TimeZone='UTC'` line from `NewDuckDBExecutor`, run the test, and confirm it fails naming `America/New_York`. Restore the line. + +Run: `go test -short ./internal/serve/ -run TestCliServe_EverySessionIsUTC -v` +Expected without the pin: FAIL, `session 0 is in America/New_York, not UTC`. With it: PASS. + +- [ ] **Step 3: Commit** + +```bash +git add internal/serve/duckdb_test.go +git commit -m "serve: prove every pooled session is pinned to UTC" +``` + +--- + +### Task 5: `queued_ms`, and stopping work nobody is waiting for + +**Files:** +- Modify: `internal/serve/http.go` (`rowsResponse`, `queryDataset`), `internal/serve/encode.go` (`readRows`) +- Test: `internal/serve/http_test.go`, `internal/serve/encode_test.go` +- Modify: `README.md` (the response example, and the attach guidance) + +**Interfaces:** +- Consumes: `query(...)`'s second return value. +- Produces: `rowsResponse.QueuedMS int64` at JSON key `queued_ms`; `readRows(ctx context.Context, rdr array.RecordReader, maxRows int) (result, error)` — the signature gains a leading context. + +- [ ] **Step 1: Write the failing test** + +Add to `internal/serve/http_test.go`: + +```go +// elapsed_ms used to include the wait for a connection, so a 13 ms query on +// a busy server reported a second and sent the reader looking for a slow +// query that did not exist. +func TestCliServe_QueuedMsSeparatesWaitFromWork(t *testing.T) { + coverage.Covers(t, "cli.serve") + ts := newTestServer(t, testServe) + + // Idle: nothing waited. + r := ts.get(t, "/v1/datasets/status") + assert.Equal(t, http.StatusOK, r.status) + assert.Equal(t, float64(0), r.body["queued_ms"]) + + // Hold the only session, then time a request that has to wait for it. + sess, err := ts.srv.exec.Acquire(context.Background()) + assert.NoError(t, err) + done := make(chan response, 1) + go func() { done <- ts.get(t, "/v1/datasets/status") }() + time.Sleep(300 * time.Millisecond) + sess.Release() + + queued := <-done + assert.Equal(t, http.StatusOK, queued.status) + assert.That(t, queued.body["queued_ms"].(float64) >= 250) + // The query itself is unchanged by the wait. + assert.That(t, queued.body["elapsed_ms"].(float64) < 250) +} +``` + +- [ ] **Step 2: Run it to see it fail** + +Run: `go test -short ./internal/serve/ -run TestCliServe_QueuedMsSeparatesWaitFromWork` +Expected: FAIL, `queued_ms` is nil because the field does not exist. + +- [ ] **Step 3: Add the field** + +In `internal/serve/http.go`, add to `rowsResponse` after `Truncated`: + +```go + // QueuedMS is how long the request waited for a session. ElapsedMS is the + // query alone, so a busy server and a slow query are distinguishable. + QueuedMS int64 `json:"queued_ms"` +``` + +In `queryDataset`, take the wait from `query` and time only the work: + +```go + start := time.Now() + res, waited, err := query(r.Context(), s.exec, st, values, ds.maxRows) +``` + +and fill both fields: + +```go + QueuedMS: waited.Milliseconds(), + ElapsedMS: time.Since(start).Sub(waited).Milliseconds(), +``` + +- [ ] **Step 4: Run the tests** + +Run: `go test -short -race ./internal/serve/` +Expected: PASS. + +- [ ] **Step 5: Write the failing test for abandoned work** + +A session held by a query nobody waits for is a session the next request cannot have, so the reader must stop when its caller does. Add to `internal/serve/encode_test.go`: + +```go +// A caller that gave up should not pay for the rest of its own result, and +// more importantly should not hold a session while it drains. Releasing a +// reader without draining it is ADBC's documented equivalent of cancel, so +// stopping early is what actually stops the query. +func TestCliServe_ReadRowsStopsWhenTheRequestIsGone(t *testing.T) { + coverage.Covers(t, "cli.serve") + + conn := newConn(t) + execSQL(t, conn, "SET TimeZone='UTC'") + + // More rows than one batch, so there is a boundary to stop at. + stmt, err := conn.NewStatement() + assert.NoError(t, err) + defer stmt.Close() + assert.NoError(t, stmt.SetSqlQuery("SELECT i FROM range(500000) t(i)")) + rdr, _, err := stmt.ExecuteQuery(context.Background()) + assert.NoError(t, err) + defer rdr.Release() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + res, err := readRows(ctx, rdr, 1000000) + assert.That(t, errors.Is(err, context.Canceled)) + // It stopped rather than reading half a million rows. + assert.That(t, res.RowCount < 1000000) +} +``` + +- [ ] **Step 6: Run it to see it fail** + +Run: `go test -short ./internal/serve/ -run TestCliServe_ReadRowsStopsWhenTheRequestIsGone` +Expected: FAIL to compile, `too many arguments in call to readRows`. + +- [ ] **Step 7: Give readRows the context** + +In `internal/serve/encode.go`, add `ctx context.Context` as the first parameter, extend the doc comment, and check between batches: + +```go +// readRows encodes at most maxRows rows from rdr. +// +// ... existing paragraphs ... +// +// It also stops at the first batch boundary after ctx ends. ADBC's Go API +// has no Cancel, and documents releasing a reader without consuming it as +// equivalent to AdbcStatementCancel, so stopping here and releasing is the +// only way a caller that gave up stops the work. A session held by a query +// nobody is waiting for is a session the next request cannot have. +func readRows(ctx context.Context, rdr array.RecordReader, maxRows int) (result, error) { +``` + +and inside the outer loop, as its first statement: + +```go + for !res.Truncated && rdr.Next() { + if err := ctx.Err(); err != nil { + return result{}, err + } +``` + +Update the two other call sites: `query` in `executor.go` (done in Task 3) and `healthz` (Task 7 rewrites it; until then pass `r.Context()`). Every existing `readRows` call in `encode_test.go` gains `context.Background()`. + +- [ ] **Step 8: Run the tests** + +Run: `go test -short -race ./internal/serve/` +Expected: PASS. Removing the `ctx.Err()` check makes `ReadRowsStopsWhenTheRequestIsGone` fail, which is the point of it. + +- [ ] **Step 9: Document both** + +In `README.md`, in the `sqlflow serve` section, add `"queued_ms":0,` to the sample response beside `"elapsed_ms"`, and add after the paragraph that begins "A zoned timestamp is UTC": + +```markdown +`elapsed_ms` is the query. `queued_ms` is how long the request waited for a +session, which is what rises when the pool is too small for the load. +``` + +Then replace the "**A timeout does not stop the query.**" bullet with what is now true: + +```markdown +- **A timeout stops reading, not always the query.** At its deadline the + caller gets `504`, and the reader stops at the next batch and is released, + which is ADBC's equivalent of cancelling. A single operator that runs long + before yielding a batch still runs to the end, holding its session. Bound + the part that is usually slow in the backend: a libpq connection string + takes `options='-c statement_timeout=30000'`, so an attached Postgres + enforces its own ceiling. +``` + +Add the same `options=` note to the attach example in that section. + +- [ ] **Step 10: Commit** + +```bash +git add internal/serve/http.go internal/serve/encode.go internal/serve/http_test.go \ + internal/serve/encode_test.go README.md +git commit -m "serve: queued_ms, and stop reading for a caller that gave up" +``` + +--- + +### Task 6: Metrics + +**Files:** +- Create: `internal/serve/metrics.go`, `internal/serve/metrics_test.go` +- Modify: `internal/serve/server.go` (build the instruments, hold them), `internal/serve/http.go` (record, register the route), `internal/serve/duckdb.go` (`NewDuckDBExecutor` takes the wait hook) + +**Interfaces:** +- Consumes: `Stats`, `pool.onWait`, `config.Serve.MetricsEnabled()`. +- Produces: + - `func newMetrics(reg *prom.Registry) (*metrics, error)` + - `type metrics` with `requests`, `requestDuration`, `queryDuration`, `sessionWait`, and a registered callback for `sessionsInUse` and `sessionsTotal` + - `func (m *metrics) observeWait(d time.Duration)` + - `serve.WithMetrics(reg *prom.Registry) Option` + - `NewDuckDBExecutor(ctx, db, size, init, onWait func(time.Duration))` — the signature gains a fifth parameter + +- [ ] **Step 1: Write the failing test** + +Create `internal/serve/metrics_test.go`: + +```go +package serve + +import ( + "net/http" + "strings" + "testing" + + prom "github.com/prometheus/client_golang/prometheus" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/zeebo/assert" +) + +// The endpoint is on the public listener, and the labels name every dataset, +// so it exists only when the config asks for it. +func TestCliServe_MetricsAreOffUnlessEnabled(t *testing.T) { + coverage.Covers(t, "cli.serve") + + ts := newTestServer(t, testServe) + r := ts.do(t, http.MethodGet, "/metrics", nil) + assert.Equal(t, http.StatusNotFound, r.status) +} + +// Every instrument the spec names carries a sample after one request, so a +// dashboard built against this list is not built against a blank. +func TestCliServe_MetricsCarryEveryInstrument(t *testing.T) { + coverage.Covers(t, "cli.serve") + + reg := prom.NewRegistry() + ts := newTestServerWith(t, strings.Replace(testServe, " limits:", " metrics: {enabled: true}\n limits:", 1), reg) + + assert.Equal(t, http.StatusOK, ts.get(t, "/v1/datasets/status").status) + + r := ts.do(t, http.MethodGet, "/metrics", nil) + assert.Equal(t, http.StatusOK, r.status) + for _, name := range []string{ + "sqlflow_serve_requests_total", + "sqlflow_serve_request_duration_seconds", + "sqlflow_serve_query_duration_seconds", + "sqlflow_serve_session_wait_seconds", + "sqlflow_serve_sessions_in_use", + "sqlflow_serve_sessions_total", + } { + if !strings.Contains(r.raw, name) { + t.Fatalf("/metrics does not carry %s:\n%s", name, r.raw) + } + } + // The gauges report the pool, not a constant. + assert.That(t, strings.Contains(r.raw, "sqlflow_serve_sessions_total 1")) + // A dataset label, so a dashboard can break the rate down. + assert.That(t, strings.Contains(r.raw, `dataset="status"`)) +} +``` + +`newTestServerWith(t, text, reg)` is `newTestServer` with a registry passed to `New` through `WithMetrics`; refactor `newTestServer` to call it with a nil registry. + +- [ ] **Step 2: Run it to see it fail** + +Run: `go test -short ./internal/serve/ -run TestCliServe_Metrics` +Expected: FAIL to compile, `undefined: newTestServerWith`, `undefined: WithMetrics`. + +- [ ] **Step 3: Write the instruments** + +Create `internal/serve/metrics.go`. Follow `internal/cli/run/metrics.go`: OpenTelemetry instruments exported through a `prom.Registry`, so serve's metrics reach the same scrape format as run's. + +```go +package serve + +import ( + "context" + "time" + + prom "github.com/prometheus/client_golang/prometheus" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/exporters/prometheus" + "go.opentelemetry.io/otel/metric" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" +) + +// buckets span a 1 ms query and a request that hit the 10 s timeout, so both +// land in a bucket rather than the overflow. +var buckets = []float64{0.001, 0.002, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 4, 8, 16} + +// metrics is the six instruments the spec names, and nothing else. +type metrics struct { + requests metric.Int64Counter + requestDuration metric.Float64Histogram + queryDuration metric.Float64Histogram + sessionWait metric.Float64Histogram +} + +// newMetrics builds the instruments against reg and registers the gauges' +// callback, which reads stats each time the endpoint is scraped. +func newMetrics(reg *prom.Registry, stats func() Stats) (*metrics, error) { + // ... prometheus.New(prometheus.WithRegisterer(reg)), a MeterProvider with + // the bucket view, a Meter named "sqlflow.serve", the four instruments, + // and an Int64ObservableGauge pair whose callback reads stats() ... +} + +func (m *metrics) observeWait(d time.Duration) { + if m == nil { + return + } + m.sessionWait.Record(context.Background(), d.Seconds()) +} + +// observeRequest records one finished request. +func (m *metrics) observeRequest(dataset, grain, code string, queued, query, total time.Duration) { + if m == nil { + return + } + attrs := metric.WithAttributes( + attribute.String("dataset", dataset), + attribute.String("grain", grain), + attribute.String("code", code), + ) + m.requests.Add(context.Background(), 1, attrs) + m.requestDuration.Record(context.Background(), total.Seconds(), attrs) + m.queryDuration.Record(context.Background(), query.Seconds(), attrs) +} +``` + +Every method tolerates a nil receiver, so the request path records unconditionally and metrics being off costs one nil check. + +- [ ] **Step 4: Wire it up** + +Add `WithMetrics(reg *prom.Registry) Option` to `server.go`, storing the registry. In `New`, when the registry is non-nil and `conf.Serve.MetricsEnabled()`, build the metrics with `stats` reading `ex.Stats()`. In `Handler`, register `/metrics` only when the metrics exist: + +```go + if s.metrics != nil { + mux.Handle("/metrics", promhttp.HandlerFor(s.registry, promhttp.HandlerOpts{})) + } +``` + +Give `NewDuckDBExecutor` its fifth parameter, `onWait func(time.Duration)`, and pass it to `newPool`. The CLI passes the metrics' `observeWait`, which means the metrics must be built before the executor; build them in the CLI from the config and hand both to `New`. + +In `queryDataset`, record once at the end, and in the error paths, with the code the response carried. + +- [ ] **Step 5: Run the tests** + +```bash +go test -short -race ./internal/serve/ ./internal/cli/serve/ +``` + +Expected: PASS. + +- [ ] **Step 6: Commit** + +```bash +git add internal/serve/ internal/cli/serve/serve.go +git commit -m "serve: six metrics, and the session wait that sizes the pool" +``` + +--- + +### Task 7: `/healthz` says busy rather than down, and answers HEAD + +**Files:** +- Modify: `internal/serve/http.go` (`healthz`, `getOnly`) +- Test: `internal/serve/http_test.go` +- Modify: `README.md` (the routes table) + +**Interfaces:** +- Consumes: `query`, `s.healthTimeout`. +- Produces: `503 {"status":"busy"}` distinct from `503 {"status":"unavailable"}`; `HEAD /healthz` answered, `HEAD` elsewhere still `405`. + +- [ ] **Step 1: Write the failing test** + +```go +// A full pool is not a dead server. A supervisor that cannot tell them apart +// restarts a server that is merely busy, which is the worst moment to do it. +func TestCliServe_HealthzIsBusyNotDownWhenThePoolIsFull(t *testing.T) { + coverage.Covers(t, "cli.serve") + ts := newTestServer(t, testServe) + + sess, err := ts.srv.exec.Acquire(context.Background()) + assert.NoError(t, err) + defer sess.Release() + + r := ts.do(t, http.MethodGet, "/healthz", nil) + assert.Equal(t, http.StatusServiceUnavailable, r.status) + assert.Equal(t, "busy", r.body["status"]) +} +``` + +- [ ] **Step 2: Run it to see it fail** + +Run: `go test -short ./internal/serve/ -run TestCliServe_HealthzIsBusy` +Expected: FAIL. Today healthz waits and then reports `unavailable`. + +- [ ] **Step 3: Rewrite healthz** + +```go +// healthz answers from a session, so it reports the path a request takes. +// +// A full pool answers busy, not unavailable: the difference is whether the +// server is broken or merely loaded, and a supervisor reading a restart loop +// needs to know which. +func (s *Server) healthz(w http.ResponseWriter, r *http.Request) { + ctx, cancel := context.WithTimeout(r.Context(), s.healthTimeout) + defer cancel() + + _, _, err := query(ctx, s.exec, s.health, nil, 1) + switch { + case errors.Is(err, context.DeadlineExceeded), errors.Is(err, ErrClosed): + writeJSON(w, r, http.StatusServiceUnavailable, []byte(`{"status":"busy"}`)) + case err != nil: + s.logger.Warn("health check failed", zap.String("error", Redact(err.Error()))) + writeJSON(w, r, http.StatusServiceUnavailable, []byte(`{"status":"unavailable"}`)) + default: + writeJSON(w, r, http.StatusOK, []byte(`{"status":"ok"}`)) + } +} +``` + +`s.health` is a `Statement` for `SELECT 1`, prepared once in `New` and stored on the server, so healthz prepares nothing per request. + +A caveat worth a comment: a request that timed out leaves its session busy until its query finishes, so a pool full of abandoned slow queries reports busy. That is the truth about the server. + +- [ ] **Step 4: Write the failing HEAD test** + +Uptime monitors send `HEAD`, and `getOnly` refuses it with `405` today. Add to `internal/serve/http_test.go`: + +```go +// Monitors send HEAD, and the status line is the whole answer. A dataset +// still refuses it: running the query and discarding the rows would spend a +// session on nothing. +func TestCliServe_HealthzAnswersHead(t *testing.T) { + coverage.Covers(t, "cli.serve") + ts := newTestServer(t, testServe) + + r := ts.do(t, http.MethodHead, "/healthz", nil) + assert.Equal(t, http.StatusOK, r.status) + assert.Equal(t, "", r.raw) + + sess, err := ts.srv.exec.Acquire(context.Background()) + assert.NoError(t, err) + busy := ts.do(t, http.MethodHead, "/healthz", nil) + sess.Release() + assert.Equal(t, http.StatusServiceUnavailable, busy.status) + + ds := ts.do(t, http.MethodHead, "/v1/datasets/status", pageToken) + assert.Equal(t, http.StatusMethodNotAllowed, ds.status) + assert.Equal(t, http.MethodGet, ds.header.Get("Allow")) +} +``` + +`ts.do` unmarshals the body as JSON when it is non-empty; a HEAD response has none, so confirm the helper tolerates an empty body and adjust it if it does not. + +- [ ] **Step 5: Run it to see it fail** + +Run: `go test -short ./internal/serve/ -run TestCliServe_HealthzAnswersHead` +Expected: FAIL. `getOnly` answers `405` for the two `/healthz` calls. + +- [ ] **Step 6: Let HEAD through for healthz only** + +In `internal/serve/http.go`: + +```go +// getOnly refuses every method but GET, and HEAD on /healthz. +// +// Monitors send HEAD, and for /healthz the status line is the whole answer. +// Every other route stays GET-only on purpose: a HEAD of a dataset would run +// the query, borrow a session and throw the rows away, which spends the pool +// on nothing. net/http suppresses the body of a HEAD response, so the health +// handler needs no special case. +func getOnly(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + allowed := http.MethodGet + if r.URL.Path == "/healthz" { + allowed = "GET, HEAD" + } + if r.Method == http.MethodGet || (r.Method == http.MethodHead && r.URL.Path == "/healthz") { + next.ServeHTTP(w, r) + return + } + w.Header().Set("Allow", allowed) + writeError(w, r, &apiError{http.StatusMethodNotAllowed, "method_not_allowed", + r.Method + " is not allowed; " + allowedDescription(r.URL.Path)}) + }) +} + +// allowedDescription names what a route accepts, for the error message. +func allowedDescription(path string) string { + if path == "/healthz" { + return "/healthz is GET or HEAD" + } + return "every route is GET, and /healthz is also HEAD" +} +``` + +- [ ] **Step 7: Run the tests** + +Run: `go test -short -race ./internal/serve/` +Expected: PASS. `TestCliServe_ASlowQueryIsA504AndTurnsHealthRed` now expects `busy` rather than `unavailable`; update its assertion and its name to `...TurnsHealthBusy`. `TestCliServe_ErrorsCarryTheirCodeAndNameTheCause`'s "not GET" case asserts the `405` message; update its expected substring if it no longer matches. + +- [ ] **Step 8: Document and commit** + +In `README.md`, change the `/healthz` row of the routes table to: + +```markdown +| `/healthz` | none | `200` `{"status":"ok"}`; `503` `{"status":"busy"}` when no session is free, `{"status":"unavailable"}` when the query fails. `HEAD` is answered too, for monitors | +``` + +and note under the table that every other route is GET only, because a `HEAD` of a dataset would run its query and discard the rows. + +```bash +git add internal/serve/http.go internal/serve/http_test.go README.md +git commit -m "serve: healthz reports a busy pool as busy, and answers HEAD" +``` + +--- + +### Task 8: What the pool does to Postgres + +**Files:** +- Create: `internal/serve/postgres_integration_test.go` +- Modify: `docs/coverage/features.yml` (`cli.serve` gains `integration`) + +**Interfaces:** +- Consumes: `NewDuckDBExecutor`, `duckdb.OpenPath`, testcontainers Postgres, pgx. +- Produces: a documented relationship between `pool.size`, `pg_connection_limit` and Postgres backends. + +- [ ] **Step 1: Write the test** + +Create `internal/serve/postgres_integration_test.go`. Start a Postgres, create a table, attach it read-only with `pg_connection_limit = 4`, run concurrent scans across a pool of 8, and read `pg_stat_activity` at the peak. + +```go +// One DuckDB scan of an attached Postgres opens up to pg_connection_limit +// connections. With a pool, the worst case may be size x that, against a +// database that also carries the pipeline's writer. This test says which it +// is, so the README can state it rather than guess. +func TestIntegrationServePool_PostgresBackendsStayBounded(t *testing.T) { + coverage.Covers(t, "cli.serve") + if testing.Short() { + t.Skip("integration: starts a Postgres container") + } + // ... start container, seed a table of 200k rows, attach READ_ONLY with + // pg_connection_limit = 4, build an executor of size 8 ... + + // Drive 8 concurrent scans, sampling pg_stat_activity while they run. + peak := samplePeakBackends(t, direct, func() { runConcurrentScans(t, ex, 8) }) + + t.Logf("peak Postgres backends: %d, with pool.size=8 and pg_connection_limit=4", peak) + // The bound the README will state. If this fails, the README is wrong, + // not the test: read the number and correct both. + assert.That(t, peak <= 8*4) +} +``` + +Implement `samplePeakBackends` by polling +`SELECT count(*) FROM pg_stat_activity WHERE datname = current_database() AND pid <> pg_backend_pid()` +every 20 ms on a direct pgx connection while the scans run, keeping the maximum. + +- [ ] **Step 2: Run it and read the number** + +Run: `go test -run '^TestIntegrationServePool' ./internal/serve/ -v` +Expected: PASS, and a log line giving the peak. The number matters more than the assertion: record it for the README and the PR. + +- [ ] **Step 3: Document the relationship** + +In `README.md`, in the `sqlflow serve` bullet list, replace the "Bound Postgres connections" bullet with what the test measured. State the real multiplier, `pool.size × pg_connection_limit` or whatever the number shows, and say to keep it under the database's `max_connections` less what the pipeline needs. + +- [ ] **Step 4: Raise the coverage requirement and commit** + +In `docs/coverage/features.yml`, change `cli.serve`'s `requires: [unit, release]` to `requires: [unit, integration, release]`. + +```bash +go test -short -race ./internal/serve/ +git add internal/serve/postgres_integration_test.go docs/coverage/features.yml README.md +git commit -m "serve: measure what a pool costs an attached Postgres" +``` + +--- + +### Task 9: Docs, coverage status, and the PR + +**Files:** +- Modify: `README.md`, `CHANGELOG.md`, `docs/coverage/status/features.yml` + +**Interfaces:** +- Consumes: everything above. + +- [ ] **Step 1: Document the pool** + +In `README.md`, in the `sqlflow serve` section, replace the bullet that begins "**One connection serves every request.**" with: + +```markdown +- **A pool serves requests.** `serve.pool.size` sessions run at once, four by + default. A request waits for a free session, and that wait counts toward + the dataset's timeout, so an exhausted pool answers `504 query_timeout`. + `queued_ms` in the response, and `sqlflow_serve_session_wait_seconds` in the + metrics, say whether the pool is the limit. +- **A timeout stops reading, not always the query.** See Task 5, step 9, + which writes this bullet and the `statement_timeout` guidance beside it. +``` + +Add a short `pool` and `metrics` block to the config example in the same section, and document `GET /metrics` in the routes table as unauthenticated and off by default. + +- [ ] **Step 2: Changelog** + +Under `## Unreleased` → `### Added`: + +```markdown +- `sqlflow serve` answers requests from a pool of sessions rather than one + connection. `serve.pool.size` sets how many run at once, four by default. + Every session is pinned to UTC, and the config's `commands` run once, on a + connection of their own. A response carries `queued_ms` beside + `elapsed_ms`, so a busy pool is not reported as a slow query, and + `/healthz` answers `busy` rather than `unavailable` when no session is + free. With `serve.metrics.enabled`, `GET /metrics` serves six instruments, + including the session wait that sizes the pool. +``` + +- [ ] **Step 3: Regenerate the coverage status** + +The `cli.serve` row gains an integration level, so the committed status file changes. Run the unit and integration passes, then write the status from them: + +```bash +mkdir -p .coverage +go test -short -race -json ./... > .coverage/go.json +go test -json -run '^TestIntegration' ./internal/serve/ > .coverage/go-integration.json +make coverage-write +git status --short docs/coverage +``` + +Expected: `docs/coverage/status/features.yml` gains `integration: covered` on `cli.serve`, and `matrix.md` follows. The release level is unchanged and comes from CI. + +- [ ] **Step 4: Run everything** + +```bash +go build ./... && go vet ./... && gofmt -l internal/ cmd/ +go test -short -race ./... +go test -run '^TestIntegration' ./internal/serve/ +uv run --locked pytest tests/tooling -q +``` + +Expected: all pass, `gofmt -l` prints nothing. + +- [ ] **Step 5: Commit and hand off** + +```bash +git add README.md CHANGELOG.md docs/coverage/ +git commit -m "docs: the serve session pool" +``` + +Do not push or open the PR without the maintainer's go-ahead. Report: the branch, the commits, the Task 1 memory numbers, the Task 8 Postgres peak, and the unit and integration results. diff --git a/docs/superpowers/specs/2026-09-16-serve-pool-design.md b/docs/superpowers/specs/2026-09-16-serve-pool-design.md new file mode 100644 index 00000000..f4883778 --- /dev/null +++ b/docs/superpowers/specs/2026-09-16-serve-pool-design.md @@ -0,0 +1,516 @@ +# `sqlflow serve`: a pool of sessions behind an executor interface + +Serve answers one request at a time. Every request waits on one mutex around +one DuckDB connection, so a slow query makes every request behind it wait and +throughput is whatever one connection can do. + +This replaces that mutex with a pool of sessions, behind an interface that +names what serve needs from a backend — prepare a statement, run it, return +Arrow — rather than handing the request path a DuckDB connection. It adds the +metrics that size the pool, because serve records none today. + +The wire contract does not change, except that a response gains `queued_ms` +and `elapsed_ms` stops counting time spent waiting. + +## The problem + +Measured against the Bluesky demo on Render on 2026-09-16, after rollup tables +cut the query itself to 13 ms. `ab -n 200 -l`, one client machine, against +`posts_by_lang` over the last 24 hours: + +| Clients | Throughput | p50 | p95 | Failures | +|---|---|---|---|---| +| 1 | 5.8/s | 164 ms | 251 ms | 0 | +| 4 | 12.2/s | 303 ms | 418 ms | 0 | +| 8 | 12.1/s | 613 ms | 855 ms | 0 | +| 16 | 13.4/s | 1,191 ms | 1,318 ms | 0 | +| 32 | 13.2/s | 2,328 ms | 2,682 ms | 0 | + +Throughput stops climbing at four clients and latency grows in proportion +after that, which is what a serialized server looks like. The same request +reports `elapsed_ms` of 13 when idle and 945 to 1,219 under sixteen clients: +the query is unchanged and the rest is queueing, counted as if it were work. + +Before the rollups the same test failed 66% of requests at sixteen clients. +Cheaper queries raised the ceiling from 1.1/s to 13/s. Only concurrency raises +it again. + +## Scope + +In: + +- An `Executor` interface: prepare a statement, acquire a session, run, + release. One implementation, DuckDB over ADBC. +- A pool of sessions with a bounded wait, and the stats it publishes. +- Session setup: which config commands run once, which run per session, and + the timezone serve pins itself. +- `queued_ms` in the response, and `elapsed_ms` narrowed to the query. +- Prometheus metrics on the existing HTTP server, off unless configured. +- `/healthz` that does not queue behind queries. + +Out: + +- Caching. The next spec, and it belongs behind the same interface. See + "Scaling out": it is also what makes more than one instance worth running. +- A second executor implementation. The interface exists so one can be added + without touching the request path; writing it now would be guessing. +- Per-dataset pools, priorities, rate limiting. +- Making the JSON encoding pluggable. Arrow is sqlflow's in-process format + and the encoder is shared. See "Arrow". + +## Depends on + +Nothing unreleased. It builds on `serve` as of v2026.09.16. + +## Decisions + +| Decision | Choice | Rejected | +|---|---|---| +| The seam | `Executor`, `Session`, `Statement` interfaces in serve, one DuckDB implementation. | Passing `adbc.Connection` through the request path, as today. It welds one driver into every layer, which is what made the Postgres sink's move off DuckDB's postgres extension (#290) expensive. | +| The interchange | `array.RecordReader`. Serve owns the encoder and the row cap. | A `Result` of already-encoded JSON per backend: the four documented rendering rules would be re-implemented per backend and could drift. A generic `Next()`/`Value() any` cursor: `any` discards the precision that the decimal and HUGEINT bug needed. | +| Default pool size | 4. | 1, which keeps every deployment serialized until its author finds the setting. Sizing from `NumCPU`, which makes memory vary by host and load tests hard to compare. | +| Session timezone | Serve runs `SET TimeZone='UTC'` on every session it opens. | Trusting the config's `commands` to reach each session. Measured: `SET TimeZone` is session-scoped, so a second connection inherits the host zone. | +| Config commands | Run once, on a setup session. | Running them per session: `ATTACH` is database-wide and re-attaching errors. | +| Waiting for a session | Counts toward the dataset's timeout; a request that never gets one is `504 query_timeout`, as today. | A new status for a busy pool. The caller's experience is a timeout either way, and `queued_ms` plus the metrics say which it was. | +| Stopping abandoned work | `readRows` takes the request's context and stops at a batch boundary, then releases the reader, which ADBC documents as equivalent to cancel. Postgres `statement_timeout` in the ATTACH string bounds the scan. | Leaving abandoned queries to run to completion, as today. With one connection that cost the next request its turn; with a pool it can hold every session at once. | +| Metrics | `GET /metrics` on the existing server, no token, off unless `serve.metrics.enabled`. | A second port, as `sqlflow run` uses: Render routes one port, so the demo could not scrape it. | + +## Arrow + +Arrow is sqlflow's in-process storage format. Sources, handlers and sinks +already move batches as Arrow, and ADBC returns it natively, so a DuckDB +executor converts nothing. + +The interface returns `array.RecordReader` and serve keeps one encoder. That +is what holds the response contract together: zoned timestamps render as UTC, +naive ones carry no offset, decimals are exact digit strings, and NaN and +infinity are strings. Each of those exists because Arrow's own rendering was +wrong for JSON, and each is documented in the README. One encoder means one +place they are true. + +It also keeps one subtle rule in one place. `readRows` encodes every value +while its record is current, because a string value aliases a buffer the +driver frees on the next `Next()`. A backend that returned rows some other way +would have to rediscover that. + +What this costs, stated plainly: + +- A backend that is not Arrow-native pays a conversion. A pgx implementation + would go pgx to Arrow to JSON where it could go straight to JSON. ADBC ships + a Postgres driver, so that is likely moot, and the Postgres sink already + carries a pgx type table to borrow from if not. +- The next backend inherits Arrow's quirks whether its own types have them or + not. +- `columns[].type` is a SQL type name reconstructed from an Arrow type, so a + Postgres `text` would surface as `VARCHAR`. Cosmetic, and it is published + contract. +- Buffer lifetimes are part of the interface: a session's reader must stay + valid until the caller finishes reading it. The interface says so. + +## The interfaces + +In `internal/serve/executor.go`. Nothing here names DuckDB or ADBC. + +```go +// Executor runs a serve config's datasets. Serve holds one; each request +// borrows a Session from it. +// +// This is an interface because the engine underneath may change. The Postgres +// sink had to leave DuckDB's postgres extension for pgx when that extension +// read the whole target table on every write (#290), and that move was +// expensive because the driver was welded into the write path. +type Executor interface { + // Prepare checks one dataset statement and returns a handle for it. + // Called at startup for every statement, so a dataset that cannot run + // fails the start rather than the first request. + Prepare(ctx context.Context, spec StatementSpec) (Statement, error) + + // Acquire borrows a session. It blocks until one is free, ctx ends, or + // the executor closes, and returns ctx.Err() or ErrClosed in those cases. + Acquire(ctx context.Context) (Session, error) + + // Stats reports what the pool gauges publish. + Stats() Stats + + // Close waits for every borrowed session to be released, then refuses + // later acquires and closes the sessions. + Close() +} + +// Session runs one statement at a time. A caller holds it for one request. +type Session interface { + // Run executes st with one request's values. The returned reader stays + // valid until it is released, and values read from it may alias buffers + // the reader owns, so a caller encodes each value before advancing. + // + // values maps a declared param name to a string, int64 or time.Time. A + // name the map does not carry binds NULL. + Run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) + + // Release returns the session. It is safe to call twice; the second is a + // no-op, so a deferred Release beside an early return cannot double-free. + Release() +} + +// Statement is one statement an Executor has checked. Its contents belong to +// the executor; serve only names it in errors. +type Statement interface { + // Where names the statement: "dataset posts_by_lang grain 1h". + Where() string +} + +// StatementSpec is what an Executor needs to prepare a statement. +type StatementSpec struct { + Dataset string + // Grain is empty for a dataset without grains. + Grain string + SQL string + Params []config.ServeParam +} + +// Stats is the pool's state at one instant, for the gauges and for the +// health endpoint. How long callers wait is a histogram, not a gauge, so it +// is not here. +type Stats struct { + // Size is every session the executor holds. + Size int + // InUse is the sessions currently borrowed. + InUse int +} +``` + +`maxRows` is deliberately not in `Run`. The cap is serve's contract, not the +backend's, and serve applies it while encoding, where it already stops the +reader early. + +`serve.New` takes an `Executor` in place of an `adbc.Connection`. The handler +acquires, runs, encodes, releases. `internal/serve/duckdb.go` is the only file +in the package that imports ADBC. + +## The pool + +`executor.go` also holds the pool mechanics, because queueing and its +measurement are the same whatever runs the SQL. A DuckDB executor embeds it. + +Sessions live in a buffered channel of capacity `size`. `Acquire` selects on +the channel, `ctx.Done()` and a closed channel. `Release` returns the session +unless the pool is closing. + +`Close` is the part worth stating. It stops accepting acquires, then waits for +every borrowed session to come back before closing any of them. A session +closed while a query runs on it takes the process down, and nothing outside +the reading goroutine can stop that query (see "Stopping the work"), so +waiting is the only option. The wait is bounded by the drain deadline the HTTP +server already applies. + +### Config + +```yaml +serve: + pool: + # Sessions the executor holds. Each is a DuckDB connection; a concurrent + # query on one costs about 15 MiB. 0 means the default, 4. + size: 4 + metrics: + # Serve GET /metrics on the same listener, without a token. Off by + # default: that port is public, and the metrics name every dataset. + enabled: false +``` + +Rules, reported as `user.config.invalid` at the key's path: + +- `pool.size` is not negative, and not greater than 64. A larger pool on a + small box is a memory failure waiting to happen, and the ceiling is a + number someone can raise with evidence. + +## Sessions and the commands block + +Measured on DuckDB 1.5.2 through ADBC on 2026-09-16, two connections to one +in-memory database with a Postgres attached: + +| Behaviour | Result | +|---|---| +| `ATTACH` on connection 1 | Connection 2 queries `pg.*` without attaching. | +| `ATTACH` the same alias again | `Binder Error: database with name "pg" already exists`. | +| `SET TimeZone='UTC'` on connection 1 | Connection 2 reads `America/New_York`, the host zone. | +| `SET pg_connection_limit = 4` on connection 1 | Connection 2 reads 4. | +| A table created on connection 1 | Visible from connection 2. | +| Two slow queries, one per connection | 200 ms, against 377 ms run one after the other. | + +So the database is shared and most setup is database-wide, but session +settings are not. Startup therefore: + +1. Opens the database and one setup session. +2. Runs the config's `commands` once, on that session. `ATTACH`, `INSTALL`, + `LOAD` and global `SET`s all take effect for every session. +3. Prepares every dataset statement once, to fail the start on a broken one. +4. Opens `size` sessions and runs `SET TimeZone='UTC'` on each. + +Serve pins the timezone rather than trusting `commands` to reach each session. +Serve's contract is UTC — the encoder already renders zoned timestamps as UTC +whatever the session says — but the session zone still decides what +`date_trunc` and a naive cast mean inside a dataset's SQL. A pooled request +that evaluated those in the host zone would return wrong buckets from a +correct config. + +A config that needs some other per-session setting has no way to ask for one. +That is deliberate for this version: the timezone is the case that breaks +correctness, and a `session_commands` block can follow when something needs +it. + +## Timeouts, and what a response says + +Waiting for a session counts toward the dataset's timeout. A request that +waits past its deadline gets `504 query_timeout` and never runs, which is the +honest answer: the caller waited as long as it agreed to. + +A response gains one field: + +```json +{"dataset": "posts_by_lang", "grain": "5m", "queued_ms": 0, "elapsed_ms": 13} +``` + +`elapsed_ms` becomes the query alone. Today it silently includes the wait, +which is why the Render measurement above reads 945 ms for a 13 ms query — +a number that sent the reader looking for a slow query that did not exist. + +### Stopping the work, not just the waiting + +A timeout bounds the response. It does not, today, bound the work, and with a +pool that gap costs more than it did with one connection: a session held by a +query nobody is waiting for is a session the next request cannot have. The +worst case for occupancy is `pool.size` multiplied by the longest query that +can actually run, not by `timeout_seconds`. + +ADBC's Go API has no `Cancel`. Checked against arrow-adbc v1.6.0 on +2026-09-16: neither `adbc.Statement` nor the driver manager exposes one, +though the C API has had `AdbcStatementCancel` since ADBC 1.1.0. What the Go +interface does say is this, on `ExecuteQuery`: + +> Since ADBC 1.1.0: releasing the returned RecordReader without consuming it +> fully is equivalent to calling AdbcStatementCancel. + +So the reader is the cancel. There is no method another goroutine can call, +but the goroutine that holds the reader can stop. + +Three bounds, in the order they bite: + +1. **Between batches.** `readRows` takes the request's context and stops at + the first batch boundary after it ends, then releases the reader, which + cancels. That bounds abandoned work at one batch rather than one query. + It does not help inside a single `Next` that never returns. +2. **In Postgres.** Most of a request is the scan of the attached database, + and a libpq connection string carries + `options='-c statement_timeout=30000'`, so the ATTACH can cap it + server-side. This is the only one of the three that Postgres enforces + rather than sqlflow hoping, and it is the one that covers the dominant + cost. The README says so where it documents attaching. +3. **`max_rows`.** Already stops a reader early, and the rollup bounds mean + a well-formed config never reaches it. + +What none of them bound is a single DuckDB operator that runs long before +yielding a batch — a sort or an aggregation over more rows than a dataset +should be reading. The answer there is the row bounds from the rollup spec, +not a timeout. + +## Metrics + +`GET /metrics`, Prometheus text format, on the listener that already serves +the datasets, when `serve.metrics.enabled` is true. No token: it carries no +row data, and the port is already public. It is off by default because the +metric labels name every dataset and grain. + +| Metric | Type | Labels | Why | +|---|---|---|---| +| `sqlflow_serve_requests_total` | counter | `dataset`, `grain`, `code` | Rate and error mix. `code` is the error code, or `ok`. | +| `sqlflow_serve_request_duration_seconds` | histogram | `dataset`, `grain` | What a caller experiences. | +| `sqlflow_serve_query_duration_seconds` | histogram | `dataset`, `grain` | The work alone, so a slow backend is distinguishable from a full pool. | +| `sqlflow_serve_session_wait_seconds` | histogram | — | The sizing signal. A p99 that climbs while query duration is flat says the pool is too small. | +| `sqlflow_serve_sessions_in_use` | gauge | — | Pinned at `size` means saturated. | +| `sqlflow_serve_sessions_total` | gauge | — | So a dashboard can draw in-use against it without hardcoding the config. | + +Six instruments, no more. Every one answers a question the Render test raised +and could not: was that second spent working or waiting, and how many sessions +would have removed it. + +Buckets for the two duration histograms and the wait: 1 ms doubling to about +16 s, so both a 13 ms query and a request that hit the 10 s timeout land in a +bucket rather than the overflow. + +## Health + +`/healthz` acquires a session, runs `SELECT 1`, and releases it. Today it +waits on the same mutex as a query, so under load it queues with everything +else and a supervisor can kill a server that is merely busy. + +With a pool it takes the health timeout as its deadline and answers `503` with +`{"status":"busy"}` when no session comes free in time, distinct from `503` +`{"status":"unavailable"}` when the query itself fails. Busy is not dead, and +an operator reading a restart loop needs to know which one it was. + +`/healthz` also answers `HEAD`, which is what many uptime monitors send and +what `getOnly` refuses today with `405`. The status line is the whole answer +to a `HEAD`, so the handler does the same work and Go suppresses the body. + +`HEAD` stays refused on every other route. A `HEAD` of a dataset would run +the query, borrow a session, and discard the result, which is a way to spend +the pool on nothing. `getOnly` therefore allows `HEAD` for `/healthz` alone +and keeps its `Allow` header accurate per route. + +## Memory, and why 4 + +DuckDB's `memory_limit` is one budget for the database, shared by every +session, so concurrency divides it rather than multiplying it. Resident memory +is not bounded by it: the limit covers DuckDB's buffer manager, not the Arrow +results, the Go heap, or what the allocator keeps. + +Measured on macOS, DuckDB 1.5.2, `memory_limit='128MB'`, N sessions each +running the demo's fold at once: + +| Sessions | Widest the grain ladder allows (365 buckets × 170 languages) | Unbounded (2,016 buckets × 170 languages) | +|---|---|---| +| 1 | 52 MiB | 100 MiB | +| 4 | 99 MiB | 279 MiB | +| 8 | 146 MiB | every query failed, out of memory | + +Idle sessions cost nothing measurable: 27 MiB whether one or eight are open. +The cost is per concurrent query, roughly 15 MiB each. + +Two things follow. Four sessions fit a 256 MB box with room, which is the +default. And the bounds from the rollup spec are what make that true: without +them the same four sessions needed 279 MiB and eight failed outright, every +concurrent request failing together as they fought over one budget. A pool +without a bound on what a request may read turns a slow server into a broken +one. + +These numbers are macOS, and `Maxrss` is bytes there and KiB on Linux. The +first implementation task repeats the measurement in the release container and +fixes the default against that, changing 4 if the number says so. + +## Postgres fan-out + +One DuckDB scan of an attached Postgres opens up to `pg_connection_limit` +connections, and this spec would not assume whether a pool multiplies that. + +`TestIntegrationServePool_PostgresBackendsStayBounded` settled it on +2026-09-16: eight sessions scanning 400k rows at once, against a +testcontainers Postgres with `pg_connection_limit = 4`, peaked at **four** +backends in `pg_stat_activity`. The attachment's connections are shared across +sessions rather than opened per session, so `pg_connection_limit` caps the +attachment and the pool does not multiply it. + +That is the better of the two outcomes, and it changes the sizing advice: a +pool costs sessions and memory, not database connections. The test asserts +that bound rather than the pessimistic one, so a DuckDB release that changed +it would fail rather than quietly invalidate the README. + +## Scaling out + +A pool scales one process. The question it raises is whether the next step is +more processes, so this section records what is true today, before a cache +exists to change it. + +Serve is stateless: an in-memory DuckDB, no durable local state, no session +affinity, tokens from config, and every request self-contained. Instances need +no coordination, and startup is already safe for several at once, because each +runs the migration script and that script takes an advisory lock. So more +instances is a deployment change, not a code change. + +Adding them buys real capacity, because real work happens in serve. A request +splits between an index range scan in Postgres, which ships the rows, and the +fold in DuckDB — `dense_rank` over a partition, then `GROUP BY` — plus JSON +encoding. If the fold ran in Postgres, another instance would add nothing. + +Two things bound it, and both are shared: + +- **Connections.** Measured: the pool does not multiply them, so the ceiling + is `instances × pg_connection_limit`, not `instances × pool.size × + pg_connection_limit`. At the demo's limit of four that is four per instance + against a database that also carries the pipeline's writer — roomier than + this spec first assumed, but still the first shared thing to run out. +- **Repeated work.** Instances share nothing, so each pulls the same hot rows + for the same popular ranges. Ten instances means ten times the scanning and + transfer for identical queries: database capacity spent to buy serve + capacity, which is backwards. + +So scaling out today trades a bottleneck for a worse one. What changes that is +the cache, which is the argument for building it next: a settled bucket never +changes, so an instance holding settled buckets answers most requests without +touching Postgres, the database sees only the live tail, and instances become +nearly independent. The cache is not primarily a latency optimisation. It is +what makes horizontal scale pay. + +Until then the order under load is: size the pool from +`sqlflow_serve_session_wait_seconds`, then give the database more capacity or +a replica, and only then add instances. + +## Tests + +Unit, `go test -short`: + +- `TestCliServe_PoolRunsQueriesConcurrently`: N slow statements across a pool + of N finish in about the time of one, not N. With `pool.size: 1` the same + test takes N times as long, which is the control. +- `TestCliServe_EverySessionIsUTC`: a dataset whose SQL returns + `current_setting('TimeZone')`, driven until every session has answered, + reports UTC from all of them. Building sessions without the pin fails it + wherever the host is not UTC, so the test sets a non-UTC `TZ`. +- `TestCliServe_AcquireReturnsWhenTheRequestGivesUp`: a cancelled request + releases its place, and a waiter behind it proceeds. +- `TestCliServe_CloseWaitsForBorrowedSessions`: Close does not return while a + query runs, and acquires after it are refused. +- `TestCliServe_ReleaseTwiceIsSafe`. +- `TestCliServe_AFullPoolTimesOutAsQueryTimeout`: the code is `query_timeout` + and `queued_ms` carries the wait. +- `TestCliServe_QueuedMsSeparatesWaitFromWork`: under a full pool, a fast + query reports a small `elapsed_ms` and a large `queued_ms`. +- `TestCliServe_AbandonedWorkStopsAtABatch`: a request that times out over a + reader with many batches stops reading rather than draining it, and its + session comes back before the query would have finished. Dropping the + context check from `readRows` fails it. +- `TestCliServe_HealthzIsBusyNotDownWhenThePoolIsFull`. +- `TestCliServe_HealthzAnswersHead`: `HEAD /healthz` is `200` with an empty + body, `503` when the pool is full, and `HEAD` of a dataset is still `405` + with `Allow: GET`. +- `TestCliServe_MetricsAreOffUnlessEnabled`, and with it enabled, `/metrics` + carries all six instruments after one request. +- `TestCliServe_PoolSizeRules`: negative and over-64 sizes report at their + path. +- The existing `TestCliServe_DoesNotLeakNativeMemory` runs against a pool. + +Integration, against testcontainers Postgres: + +- `TestIntegrationServePool_PostgresBackendsStayBounded`, above. + +The feature `cli.serve` already requires unit and release coverage; the pool +adds `integration`. + +## What breaks if this is wrong + +| If | Then | Caught by | +|---|---|---| +| A session misses the UTC pin | A correct config returns buckets in the host zone, from some requests and not others. | `EverySessionIsUTC`, with a non-UTC `TZ`. | +| `commands` runs per session | Startup fails on the second `ATTACH`. | Any pool test against a config with an attachment. | +| Close closes a session mid-query | The process dies during a deploy, mid-request. | `CloseWaitsForBorrowedSessions`. | +| A cancelled request leaks its session | The pool shrinks under load until it deadlocks. | `AcquireReturnsWhenTheRequestGivesUp`. | +| Abandoned queries run to completion | A handful of slow requests hold every session while nobody waits for them, and the pool is a queue for work no one wants. | `AbandonedWorkStopsAtABatch`, and `statement_timeout` on the attachment. | +| The pool is sized past the box | Concurrent queries fail together, out of memory, rather than queueing. | The container measurement in task 1; `memory_limit` is the backstop. | +| `queued_ms` and `elapsed_ms` are swapped | An operator tunes the wrong thing, as this spec's own Render numbers nearly did. | `QueuedMsSeparatesWaitFromWork`. | + +## Build order + +1. Measure per-session resident memory in the release container, and fix the + default. The rest of the spec assumes 4. +2. The interfaces and the pool, with the DuckDB implementation behind them. + `serve.New` takes an `Executor`. No behaviour change at `size: 1`. +3. Session setup: `commands` once, the UTC pin per session. +4. `queued_ms`, `elapsed_ms` narrowed to the query, and `readRows` stopping + at a batch boundary when the request is gone. +5. Metrics, and the config that enables them. +6. `/healthz` busy. +7. The Postgres fan-out measurement, and what the README says about it. +8. README and CHANGELOG. + +The demo then sets `pool.size`, enables metrics, and the Render test above +runs again as the check: throughput should rise with the pool, and +`session_wait_seconds` should say whether it is sized right. diff --git a/internal/cli/serve/examples_test.go b/internal/cli/serve/examples_test.go index 4a7dbda9..1068e2b5 100644 --- a/internal/cli/serve/examples_test.go +++ b/internal/cli/serve/examples_test.go @@ -60,16 +60,25 @@ func TestCliServe_ShippedExamplesValidateAndBuild(t *testing.T) { assert.NoError(t, err) assert.NoError(t, conf.CheckError()) - conn, err := duckdb.Open(context.Background()) + db, err := duckdb.OpenPath(context.Background(), "") assert.NoError(t, err) - defer conn.Close() - denyExternalAccess(t, conn) + defer db.Close() - if err := core.InitCommands(conn, &config.Conf{Commands: conf.Commands}); err != nil { - t.Skipf("commands need an external system: %v", err) + // The init connection is the one the commands run on, so it is + // the one that must be kept off the network. + var initErr error + ex, err := api.NewDuckDBExecutor(context.Background(), db, conf.Serve.PoolSize(), + func(ctx context.Context, conn adbc.Connection) error { + denyExternalAccess(t, conn) + initErr = core.InitCommands(conn, &config.Conf{Commands: conf.Commands}) + return initErr + }, nil) + if initErr != nil { + t.Skipf("commands need an external system: %v", initErr) } + assert.NoError(t, err) - srv, err := api.New(context.Background(), conf, conn) + srv, err := api.New(context.Background(), conf, ex) assert.NoError(t, err) srv.Close() }) diff --git a/internal/cli/serve/serve.go b/internal/cli/serve/serve.go index ce4d9ec9..6fa05c3e 100644 --- a/internal/cli/serve/serve.go +++ b/internal/cli/serve/serve.go @@ -5,6 +5,7 @@ import ( "context" "errors" "fmt" + prom "github.com/prometheus/client_golang/prometheus" "net" "net/http" _ "net/http/pprof" @@ -12,6 +13,7 @@ import ( "runtime" "syscall" + "github.com/apache/arrow-adbc/go/adbc" "github.com/spf13/cobra" "github.com/turbolytics/sql-flow/internal/config" "github.com/turbolytics/sql-flow/internal/core" @@ -93,31 +95,32 @@ func serveConfig(ctx context.Context, path string, l *zap.Logger, onListen func( } }() - conn, err := db.Connect(ctx) + ex, err := api.NewDuckDBExecutor(ctx, db, conf.Serve.PoolSize(), + func(ctx context.Context, conn adbc.Connection) error { + // Uncoded, as in run: an ATTACH that fails because the database + // is not up yet exits 1, which a supervisor retries. Redacted + // because a failed ATTACH prints the connection string. + if err := core.InitCommands(conn, &config.Conf{Commands: conf.Commands}); err != nil { + return errors.New("failed to initialize commands: " + api.Redact(err.Error())) + } + return nil + }, nil) if err != nil { return err } - defer func() { - if err := conn.Close(); err != nil { - l.Error("failed to close DuckDB connection", zap.Error(err)) - } - }() - // Uncoded, as in run: an ATTACH that fails because the database is not up - // yet exits 1, which a supervisor retries. - // - // Redacted: a failed ATTACH prints the connection string, password - // included, and this error goes to the log. - if err := core.InitCommands(conn, &config.Conf{Commands: conf.Commands}); err != nil { - return errors.New("failed to initialize commands: " + api.Redact(err.Error())) + opts := []api.Option{api.WithLogger(l)} + if conf.Serve.MetricsEnabled() { + opts = append(opts, api.WithMetrics(prom.NewRegistry())) } - srv, err := api.New(ctx, conf, conn, api.WithLogger(l)) + srv, err := api.New(ctx, conf, ex, opts...) if err != nil { + ex.Close() return err } - // Deferred after conn's close, so it runs first: a query still holding - // the lock finishes before the connection closes under it. + // Deferred after the database's close, so it runs first: a query still + // holding a session finishes before the database closes under it. defer srv.Close() ln, err := net.Listen("tcp", conf.Serve.Addr()) diff --git a/internal/cli/testdata/serve_example.golden b/internal/cli/testdata/serve_example.golden index 943bf408..1c995a5b 100644 --- a/internal/cli/testdata/serve_example.golden +++ b/internal/cli/testdata/serve_example.golden @@ -45,6 +45,18 @@ serve: requests_per_second: # Reserved. Must be 0. burst: + # How many requests the server runs at once. Omit for the default. + pool: + # Sessions the server holds. Each is one backend session, and one + # request uses one at a time, so this is the requests that run at once. + # 0 means the default, 4. + size: + # Whether to serve Prometheus metrics at /metrics. + metrics: + # Serve GET /metrics on the same listener as the datasets, without a + # token. Off by default: that listener is public, and the metric labels + # name every dataset and grain. + enabled: # The datasets this server answers. Nothing else is reachable. datasets: - diff --git a/internal/config/serve.go b/internal/config/serve.go index 411d8198..1efd8cb4 100644 --- a/internal/config/serve.go +++ b/internal/config/serve.go @@ -22,6 +22,14 @@ const ( DefaultServeAddr = "0.0.0.0:8080" DefaultServeMaxRows = 10000 DefaultServeTimeoutSeconds = 10 + // DefaultServePoolSize is the sessions a server holds when the config + // names no number. Four fits a 256 MB box: idle sessions cost nothing + // measurable, and a concurrent query costs about 15 MiB. + DefaultServePoolSize = 4 + // MaxServePoolSize is the most this version accepts. A larger pool on a + // small box fails queries with out-of-memory rather than queueing them, + // because DuckDB's memory_limit is one budget shared by every session. + MaxServePoolSize = 64 ) // ServeConf is a whole serve file: the commands that attach the data, and the @@ -44,6 +52,10 @@ type Serve struct { Auth ServeAuth `yaml:"auth"` // Limits for every dataset. A dataset's own non-zero limit overrides one. Limits *ServeLimits `yaml:"limits,omitempty"` + // How many requests the server runs at once. Omit for the default. + Pool *ServePool `yaml:"pool,omitempty"` + // Whether to serve Prometheus metrics at /metrics. + Metrics *ServeMetrics `yaml:"metrics,omitempty"` // The datasets this server answers. Nothing else is reachable. Datasets []ServeDataset `yaml:"datasets"` } @@ -100,6 +112,22 @@ type ServeRateLimit struct { Burst int `yaml:"burst,omitempty"` } +// ServePool sizes the sessions a server answers requests on. +type ServePool struct { + // Sessions the server holds. Each is one backend session, and one + // request uses one at a time, so this is the requests that run at once. + // 0 means the default, 4. + Size int `yaml:"size,omitempty"` +} + +// ServeMetrics turns on the Prometheus endpoint. +type ServeMetrics struct { + // Serve GET /metrics on the same listener as the datasets, without a + // token. Off by default: that listener is public, and the metric labels + // name every dataset and grain. + Enabled bool `yaml:"enabled,omitempty"` +} + // ServeDataset is one named, parameterized SQL statement, or one per grain. type ServeDataset struct { // The URL path segment: /v1/datasets/. Lowercase letters, digits @@ -180,6 +208,19 @@ func (s Serve) MaxRows(ds ServeDataset) int { return DefaultServeMaxRows } +// PoolSize is the sessions to hold, defaulted. +func (s Serve) PoolSize() int { + if s.Pool != nil && s.Pool.Size > 0 { + return s.Pool.Size + } + return DefaultServePoolSize +} + +// MetricsEnabled reports whether to serve /metrics. +func (s Serve) MetricsEnabled() bool { + return s.Metrics != nil && s.Metrics.Enabled +} + // Timeout is how long a caller waits on a dataset, resolved the same way as // MaxRows. The zero dataset gives the top-level timeout, which /healthz uses. func (s Serve) Timeout(ds ServeDataset) time.Duration { @@ -387,6 +428,18 @@ func (c *ServeConf) Check() []Violation { checkLimits(s.Limits, []string{"serve", "limits"}, add) + if s.Pool != nil { + switch { + case s.Pool.Size < 0: + add(errs.CodeConfigInvalid, []string{"serve", "pool", "size"}, + "pool.size is %d; it must not be negative, and 0 means the default", s.Pool.Size) + case s.Pool.Size > MaxServePoolSize: + add(errs.CodeConfigInvalid, []string{"serve", "pool", "size"}, + "pool.size %d sessions is more than the %d this version allows; DuckDB's memory limit is one budget shared by every session, so a pool past the box fails queries rather than queueing them", + s.Pool.Size, MaxServePoolSize) + } + } + if len(s.Datasets) == 0 { add(errs.CodeConfigInvalid, []string{"serve", "datasets"}, "serve.datasets declares no dataset, so the server would answer nothing") diff --git a/internal/config/serve_test.go b/internal/config/serve_test.go index 98b2b9b0..1bbefc09 100644 --- a/internal/config/serve_test.go +++ b/internal/config/serve_test.go @@ -110,6 +110,24 @@ func TestCliServe_LimitsResolveDatasetThenTopLevelThenDefault(t *testing.T) { assert.Equal(t, DefaultServeAddr, bare.Addr()) } +// The pool defaults rather than failing closed: a config written before the +// pool existed gets concurrency, not one session. +func TestCliServe_PoolSizeDefaults(t *testing.T) { + coverage.Covers(t, "cli.serve") + + conf := parseServe(t, validServe) + assert.Equal(t, DefaultServePoolSize, conf.Serve.PoolSize()) + assert.False(t, conf.Serve.MetricsEnabled()) + + sized := parseServe(t, strings.Replace(validServe, " limits:", " pool: {size: 8}\n metrics: {enabled: true}\n limits:", 1)) + assert.Equal(t, 8, sized.Serve.PoolSize()) + assert.True(t, sized.Serve.MetricsEnabled()) + + // 0 is unset, not "no sessions". + zero := parseServe(t, strings.Replace(validServe, " limits:", " pool: {size: 0}\n limits:", 1)) + assert.Equal(t, DefaultServePoolSize, zero.Serve.PoolSize()) +} + func TestCliServe_CheckAcceptsAValidConfig(t *testing.T) { coverage.Covers(t, "cli.serve") @@ -177,6 +195,10 @@ func TestCliServe_CheckReportsEachRuleAtItsPath(t *testing.T) { errs.CodeConfigServeDataset, "serve.datasets.1.params.0.max", "min and max bound integer params only"}, {"min above max", "{name: lang, type: string}", "{name: lang, type: integer, min: 5, max: 1}", errs.CodeConfigServeDataset, "serve.datasets.1.params.1.max", "min 5 above max 1"}, + {"negative pool size", " limits:", " pool: {size: -1}\n limits:", + errs.CodeConfigInvalid, "serve.pool.size", "must not be negative"}, + {"pool size past the ceiling", " limits:", " pool: {size: 65}\n limits:", + errs.CodeConfigInvalid, "serve.pool.size", "65 sessions is more than the 64 this version allows"}, {"positional placeholder", "SELECT count(*) AS n FROM t", "SELECT $1 AS n FROM t", errs.CodeConfigServeDataset, "serve.datasets.0.sql", "$1"}, {"undeclared placeholder", "SELECT count(*) AS n FROM t", "SELECT $nope AS n FROM t", diff --git a/internal/rollup/serve_test.go b/internal/rollup/serve_test.go index 9a90aba5..809b5565 100644 --- a/internal/rollup/serve_test.go +++ b/internal/rollup/serve_test.go @@ -70,22 +70,38 @@ func execDuck(t *testing.T, conn adbc.Connection, sql string) { func TestCliRollup_ServeDatasetsAnswerFromTheirTables(t *testing.T) { coverage.Covers(t, "cli.rollup") - conn, err := duckdb.Open(context.Background()) - assert.NoError(t, err) - t.Cleanup(func() { conn.Close() }) - execDuck(t, conn, "SET TimeZone='UTC'") - execDuck(t, conn, "ATTACH ':memory:' AS pg") - execDuck(t, conn, "CREATE TABLE pg.posts_per_minute_by_lang (bucket TIMESTAMPTZ, lang VARCHAR, posts INTEGER, updated_at TIMESTAMPTZ)") + // The setup runs on the executor's own connection: ATTACH is + // database-wide, and serve opens its sessions from the same database. + setup := []string{ + "SET TimeZone='UTC'", + "ATTACH ':memory:' AS pg", + "CREATE TABLE pg.posts_per_minute_by_lang (bucket TIMESTAMPTZ, lang VARCHAR, posts INTEGER, updated_at TIMESTAMPTZ)", + } for _, g := range []string{"5m", "15m", "1h", "6h", "1d"} { - execDuck(t, conn, "CREATE TABLE pg.posts_by_lang_"+g+" (bucket TIMESTAMPTZ, lang VARCHAR, posts BIGINT)") - execDuck(t, conn, "CREATE TABLE pg.posts_total_"+g+" (bucket TIMESTAMPTZ, minutes BIGINT, posts BIGINT)") + setup = append(setup, + "CREATE TABLE pg.posts_by_lang_"+g+" (bucket TIMESTAMPTZ, lang VARCHAR, posts BIGINT)", + "CREATE TABLE pg.posts_total_"+g+" (bucket TIMESTAMPTZ, minutes BIGINT, posts BIGINT)") } - execDuck(t, conn, `INSERT INTO pg.posts_per_minute_by_lang VALUES + setup = append(setup, + `INSERT INTO pg.posts_per_minute_by_lang VALUES ('2026-09-13 18:01:00+00', 'en', 5, now()), ('2026-09-13 18:02:00+00', 'en', 7, now()), - ('2026-09-13 18:01:00+00', 'ja', 3, now())`) - execDuck(t, conn, `INSERT INTO pg.posts_by_lang_15m VALUES - ('2026-09-12 00:15:00+00', 'en', 10), ('2026-09-12 00:15:00+00', 'ja', 4)`) - execDuck(t, conn, `INSERT INTO pg.posts_total_1d VALUES ('2026-09-12 00:00:00+00', 1440, 250000)`) + ('2026-09-13 18:01:00+00', 'ja', 3, now())`, + `INSERT INTO pg.posts_by_lang_15m VALUES + ('2026-09-12 00:15:00+00', 'en', 10), ('2026-09-12 00:15:00+00', 'ja', 4)`, + `INSERT INTO pg.posts_total_1d VALUES ('2026-09-12 00:00:00+00', 1440, 250000)`) + + db, err := duckdb.OpenPath(context.Background(), "") + assert.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + ex, err := api.NewDuckDBExecutor(context.Background(), db, 1, + func(ctx context.Context, conn adbc.Connection) error { + for _, sql := range setup { + execDuck(t, conn, sql) + } + return nil + }, nil) + assert.NoError(t, err) datasets, err := ServeDatasets(withTotals(t)) assert.NoError(t, err) @@ -93,7 +109,7 @@ func TestCliRollup_ServeDatasetsAnswerFromTheirTables(t *testing.T) { Auth: config.ServeAuth{Tokens: []config.ServeToken{{Name: "page", Token: "page-token"}}}, Datasets: datasets, }} - srv, err := api.New(context.Background(), conf, conn) + srv, err := api.New(context.Background(), conf, ex) assert.NoError(t, err) t.Cleanup(srv.Close) handler := srv.Handler() diff --git a/internal/serve/duckdb.go b/internal/serve/duckdb.go new file mode 100644 index 00000000..f7fdff72 --- /dev/null +++ b/internal/serve/duckdb.go @@ -0,0 +1,295 @@ +package serve + +import ( + "context" + "fmt" + "strings" + "sync" + "time" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/turbolytics/sql-flow/internal/errs" + "github.com/turbolytics/sql-flow/internal/sqlparams" +) + +// duckdbExecutor runs datasets on DuckDB over ADBC. It is the only type in +// this package that names either. +type duckdbExecutor struct { + *pool + db *duckdb.DB + // setup is the connection the config's commands ran on, kept because + // Prepare needs a connection and every other one is in the pool. It is + // idle afterwards, which costs nothing measurable. + setup adbc.Connection + // closeOnce keeps a second Close from closing setup twice: the server + // owns the executor, and a test may close both. + closeOnce sync.Once +} + +// NewDuckDBExecutor opens size sessions on db and returns an Executor over +// them. +// +// init runs once, on a connection of its own, before any session is used. The +// caller passes the config's commands: ATTACH is database-wide and attaching +// the same alias twice is an error, so they must not run per session. +// +// Every session is then pinned to UTC. That is serve's own contract rather +// than the config's: SET TimeZone is session-scoped, so a session that missed +// it evaluates date_trunc and naive casts in the host's zone and returns +// wrong buckets from a correct config. +// +// onWait receives how long each Acquire waited; it may be nil. +func NewDuckDBExecutor(ctx context.Context, db *duckdb.DB, size int, + init func(context.Context, adbc.Connection) error, onWait func(time.Duration)) (Executor, error) { + setup, err := db.Connect(ctx) + if err != nil { + return nil, err + } + if init != nil { + if err := init(ctx, setup); err != nil { + _ = setup.Close() + return nil, err + } + } + + backends := make([]backendSession, 0, size) + for i := 0; i < size; i++ { + conn, err := db.Connect(ctx) + if err != nil { + closeBackends(backends) + _ = setup.Close() + return nil, err + } + if err := execOn(ctx, conn, "SET TimeZone='UTC'"); err != nil { + _ = conn.Close() + closeBackends(backends) + _ = setup.Close() + return nil, fmt.Errorf("pinning the session timezone: %w", err) + } + backends = append(backends, &duckdbSession{conn: conn}) + } + + return &duckdbExecutor{pool: newPool(backends, onWait), db: db, setup: setup}, nil +} + +// closeBackends closes the sessions opened before one failed, so a failed +// start leaves no connection behind. +func closeBackends(backends []backendSession) { + for _, b := range backends { + _ = b.close() + } +} + +// execOn runs one statement for its effect. +func execOn(ctx context.Context, conn adbc.Connection, sql string) error { + stmt, err := conn.NewStatement() + if err != nil { + return err + } + defer stmt.Close() + if err := stmt.SetSqlQuery(sql); err != nil { + return err + } + _, err = stmt.ExecuteUpdate(ctx) + return err +} + +func (e *duckdbExecutor) Close() { + e.closeOnce.Do(func() { + e.pool.Close() + _ = e.setup.Close() + }) +} + +// duckdbSession is one ADBC connection. +type duckdbSession struct{ conn adbc.Connection } + +func (d *duckdbSession) close() error { return d.conn.Close() } + +func (d *duckdbSession) run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) { + s, ok := st.(*duckdbStatement) + if !ok { + return nil, fmt.Errorf("serve: statement %T did not come from this executor", st) + } + + stmt, err := d.conn.NewStatement() + if err != nil { + return nil, err + } + // The statement outlives this call through the reader, so it closes when + // the reader is released rather than here. + if err := stmt.SetSqlQuery(s.rewritten); err != nil { + _ = stmt.Close() + return nil, err + } + // DuckDB wants exactly one field per placeholder, so a statement with + // none binds nothing. + if s.schema.NumFields() > 0 { + rec := s.record(values) + err := stmt.Bind(ctx, rec) + rec.Release() + if err != nil { + _ = stmt.Close() + return nil, err + } + } + rdr, _, err := stmt.ExecuteQuery(ctx) + if err != nil { + _ = stmt.Close() + return nil, err + } + return &closingReader{RecordReader: rdr, stmt: stmt}, nil +} + +// closingReader closes the statement when the reader is released, so a +// caller that only knows about the reader leaks neither. +type closingReader struct { + array.RecordReader + stmt adbc.Statement +} + +func (c *closingReader) Release() { + c.RecordReader.Release() + _ = c.stmt.Close() +} + +// duckdbStatement is one dataset statement, numbered and checked against +// DuckDB at startup. +type duckdbStatement struct { + dataset string + // grain is empty for a dataset without grains. + grain string + // sql is the statement as the config wrote it. + sql string + // rewritten is sql with $name replaced by $N. + rewritten string + // schema has one field per placeholder, in number order, typed by the + // declared param. A request's values bind into a record of this shape. + schema *arrow.Schema +} + +// Where names the statement in an error. +func (s *duckdbStatement) Where() string { + if s.grain == "" { + return "dataset " + s.dataset + } + return "dataset " + s.dataset + " grain " + s.grain +} + +// paramTypes maps a declared param type to the Arrow type it binds as. The +// type comes from the config, never from the value, so an absent param is a +// typed null and coalesce resolves against the right type. +var paramTypes = map[string]arrow.DataType{ + "string": arrow.BinaryTypes.String, + "integer": arrow.PrimitiveTypes.Int64, + "timestamp": &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}, +} + +// Prepare checks one statement against DuckDB at startup. DuckDB binds a +// statement when its SQL is set, so a syntax error, a missing table and a +// missing column all fail here rather than on the first request. +// +// The statement is closed afterwards: each request plans its own, because a +// held plan can fold table statistics into constants. +func (e *duckdbExecutor) Prepare(ctx context.Context, spec StatementSpec) (Statement, error) { + st := &duckdbStatement{dataset: spec.Dataset, grain: spec.Grain, sql: spec.SQL} + + rw, err := sqlparams.Rewrite(spec.SQL) + if err != nil { + return nil, errs.Wrap(errs.CodeConfigServeDataset, err, "%s", st.Where()) + } + st.rewritten = rw.SQL + + declared := map[string]string{} + for _, p := range spec.Params { + declared[p.Name] = p.Type + } + fields := make([]arrow.Field, len(rw.Names)) + for i, name := range rw.Names { + typ, ok := paramTypes[declared[name]] + if !ok { + return nil, errs.New(errs.CodeConfigServeDataset, + "%s: $%s is not a declared param", st.Where(), name) + } + fields[i] = arrow.Field{Name: name, Type: typ, Nullable: true} + } + st.schema = arrow.NewSchema(fields, nil) + + stmt, err := e.setup.NewStatement() + if err != nil { + return nil, err + } + defer stmt.Close() + + if err := stmt.SetSqlQuery(rw.SQL); err != nil { + return nil, errs.Wrap(errs.CodeSQLInvalid, err, "%s: the SQL does not prepare", st.Where()) + } + if err := stmt.Prepare(ctx); err != nil { + return nil, errs.Wrap(errs.CodeSQLInvalid, err, "%s: the SQL does not prepare", st.Where()) + } + + // The scanner and DuckDB must agree on the count. If they do not, the + // scanner misread the SQL, and binding would put values on the wrong + // placeholders and answer with wrong rows. + ps, err := stmt.GetParameterSchema() + if err != nil { + return nil, errs.Wrap(errs.CodeSQLInvalid, err, "%s: reading the parameter count", st.Where()) + } + if ps.NumFields() != len(rw.Names) { + return nil, errs.New(errs.CodeSQLInvalid, + "%s: DuckDB counts %d parameters and sqlflow counts %d (%s); "+ + "a quote or comment form the scanner does not know is hiding or inventing one", + st.Where(), ps.NumFields(), len(rw.Names), strings.Join(rw.Names, ", ")) + } + + return st, nil +} + +// record builds the one-row record a request binds. +func (s *duckdbStatement) record(values map[string]any) arrow.RecordBatch { + cols := make([]arrow.Array, s.schema.NumFields()) + for i, f := range s.schema.Fields() { + v, present := values[f.Name] + switch typ := f.Type.(type) { + case *arrow.StringType: + b := array.NewStringBuilder(memory.DefaultAllocator) + if present { + b.Append(v.(string)) + } else { + b.AppendNull() + } + cols[i] = b.NewArray() + b.Release() + case *arrow.Int64Type: + b := array.NewInt64Builder(memory.DefaultAllocator) + if present { + b.Append(v.(int64)) + } else { + b.AppendNull() + } + cols[i] = b.NewArray() + b.Release() + case *arrow.TimestampType: + b := array.NewTimestampBuilder(memory.DefaultAllocator, typ) + if present { + b.Append(arrow.Timestamp(v.(time.Time).UnixMicro())) + } else { + b.AppendNull() + } + cols[i] = b.NewArray() + b.Release() + default: + panic(fmt.Sprintf("serve: no builder for param type %s", f.Type)) + } + } + + rec := array.NewRecordBatch(s.schema, cols, 1) + for _, c := range cols { + c.Release() + } + return rec +} diff --git a/internal/serve/duckdb_test.go b/internal/serve/duckdb_test.go new file mode 100644 index 00000000..453a525d --- /dev/null +++ b/internal/serve/duckdb_test.go @@ -0,0 +1,59 @@ +package serve + +import ( + "context" + "strings" + "testing" + + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/zeebo/assert" +) + +// SET TimeZone is session-scoped, so a session the executor did not pin +// inherits the host's zone and evaluates date_trunc and a naive cast in it — +// wrong buckets from a correct config, on some requests and not others. +// +// The test runs under a non-UTC TZ on purpose: under UTC a missing pin looks +// correct, which is how this would reach production unnoticed. +func TestCliServe_EverySessionIsUTC(t *testing.T) { + coverage.Covers(t, "cli.serve") + t.Setenv("TZ", "America/New_York") + + ctx := context.Background() + const size = 4 + ex, _ := newExec(t, size) + + st, err := ex.Prepare(ctx, StatementSpec{ + Dataset: "probe", + SQL: "SELECT current_setting('TimeZone') AS zone", + }) + assert.NoError(t, err) + + // Hold every session at once, so each one is checked rather than the + // same one four times. + held := make([]Session, 0, size) + for i := 0; i < size; i++ { + s, err := ex.Acquire(ctx) + assert.NoError(t, err) + held = append(held, s) + } + defer func() { + for _, s := range held { + s.Release() + } + }() + + for i, s := range held { + rdr, err := s.Run(ctx, st, nil) + assert.NoError(t, err) + res, err := readRows(context.Background(), rdr, 1) + rdr.Release() + assert.NoError(t, err) + + // res.Rows is the JSON array of row objects, so the zone is a value + // inside it rather than the whole string. + if !strings.Contains(string(res.Rows), `"UTC"`) { + t.Fatalf("session %d is not pinned to UTC: %s", i, res.Rows) + } + } +} diff --git a/internal/serve/encode.go b/internal/serve/encode.go index 1af6e66d..e80b8aee 100644 --- a/internal/serve/encode.go +++ b/internal/serve/encode.go @@ -2,6 +2,7 @@ package serve import ( "bytes" + "context" "encoding/json" "fmt" "math" @@ -39,7 +40,14 @@ const naiveTimestampLayout = "2006-01-02T15:04:05.999999999" // // It stops reading at the first row past maxRows. DuckDB streams results, so // the caller releasing the reader then stops the query. -func readRows(rdr array.RecordReader, maxRows int) (result, error) { +// +// It also stops at the first batch boundary after ctx ends. ADBC's Go API has +// no Cancel, and documents releasing a reader without consuming it as +// equivalent to AdbcStatementCancel, so stopping here and letting the caller +// release is the only way a request that gave up stops the work it started. A +// session held by a query nobody is waiting for is a session the next request +// cannot have. +func readRows(ctx context.Context, rdr array.RecordReader, maxRows int) (result, error) { fields := rdr.Schema().Fields() res := result{Columns: make([]column, len(fields))} keys := make([][]byte, len(fields)) @@ -55,6 +63,9 @@ func readRows(rdr array.RecordReader, maxRows int) (result, error) { var buf bytes.Buffer buf.WriteByte('[') for !res.Truncated && rdr.Next() { + if err := ctx.Err(); err != nil { + return result{}, err + } rec := rdr.RecordBatch() for r := 0; r < int(rec.NumRows()); r++ { if res.RowCount == maxRows { diff --git a/internal/serve/encode_test.go b/internal/serve/encode_test.go index f497e612..f3541606 100644 --- a/internal/serve/encode_test.go +++ b/internal/serve/encode_test.go @@ -3,6 +3,7 @@ package serve import ( "context" "encoding/json" + "errors" "testing" "github.com/apache/arrow-adbc/go/adbc" @@ -29,6 +30,33 @@ func execSQL(t *testing.T, conn adbc.Connection, sql string) { assert.NoError(t, err) } +// newExec opens an in-memory DuckDB and returns an executor of size sessions +// over it, having run setup on the init connection. Size 1 is the old +// single-connection behaviour, which is what keeps a test's ordering and +// timing true. The database comes back too, for a test that has to change +// the data from outside the pool. +func newExec(t *testing.T, size int, setup ...string) (Executor, *duckdb.DB) { + t.Helper() + db, err := duckdb.OpenPath(context.Background(), "") + assert.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + ex, err := NewDuckDBExecutor(context.Background(), db, size, + func(ctx context.Context, conn adbc.Connection) error { + for _, sql := range setup { + if err := execOn(ctx, conn, sql); err != nil { + return err + } + } + return nil + }, nil) + assert.NoError(t, err) + // Registered after the database's cleanup, so it runs first: a query + // still holding a session finishes before the database closes under it. + t.Cleanup(ex.Close) + return ex, db +} + // encodeSQL runs sql and encodes its rows the way a request does. func encodeSQL(t *testing.T, conn adbc.Connection, sql string, maxRows int) result { t.Helper() @@ -40,7 +68,7 @@ func encodeSQL(t *testing.T, conn adbc.Connection, sql string, maxRows int) resu assert.NoError(t, err) defer rdr.Release() - res, err := readRows(rdr, maxRows) + res, err := readRows(context.Background(), rdr, maxRows) assert.NoError(t, err) return res } @@ -160,3 +188,28 @@ func TestCliServe_EncodeTruncatesAtMaxRows(t *testing.T) { assert.Equal(t, "[]", string(empty.Rows)) assert.Equal(t, 1, len(empty.Columns)) } + +// A caller that gave up should not pay for the rest of its own result, and +// should not hold a session while it drains one. Releasing a reader without +// consuming it is ADBC's documented equivalent of cancel, so stopping early is +// what actually stops the query. +func TestCliServe_ReadRowsStopsWhenTheRequestIsGone(t *testing.T) { + coverage.Covers(t, "cli.serve") + + conn := newConn(t) + stmt, err := conn.NewStatement() + assert.NoError(t, err) + defer stmt.Close() + // More rows than one batch, so there is a boundary to stop at. + assert.NoError(t, stmt.SetSqlQuery("SELECT i FROM range(500000) t(i)")) + rdr, _, err := stmt.ExecuteQuery(context.Background()) + assert.NoError(t, err) + defer rdr.Release() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + res, err := readRows(ctx, rdr, 1000000) + assert.That(t, errors.Is(err, context.Canceled)) + assert.That(t, res.RowCount < 1000000) +} diff --git a/internal/serve/executor.go b/internal/serve/executor.go new file mode 100644 index 00000000..1bebcb02 --- /dev/null +++ b/internal/serve/executor.go @@ -0,0 +1,254 @@ +package serve + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "time" + + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/turbolytics/sql-flow/internal/config" +) + +// ErrClosed is what Acquire returns once the executor has closed. +var ErrClosed = errors.New("the server is shutting down") + +// Executor runs a serve config's datasets. Serve holds one; each request +// borrows a Session from it. +// +// This is an interface because the engine underneath may change. The Postgres +// sink had to leave DuckDB's postgres extension for pgx when that extension +// read the whole target table on every write (#290), and that move was +// expensive because the driver was welded into the write path. +type Executor interface { + // Prepare checks one dataset statement and returns a handle for it. + // Called at startup for every statement, so a dataset that cannot run + // fails the start rather than the first request. + Prepare(ctx context.Context, spec StatementSpec) (Statement, error) + + // Acquire borrows a session. It blocks until one is free, ctx ends, or + // the executor closes, and returns ctx.Err() or ErrClosed in those cases. + Acquire(ctx context.Context) (Session, error) + + // Stats reports what the pool gauges publish. + Stats() Stats + + // Close waits for every borrowed session to be released, then refuses + // later acquires and closes the sessions. + Close() +} + +// Session runs one statement at a time. A caller holds it for one request. +type Session interface { + // Run executes st with one request's values. The returned reader stays + // valid until it is released, and values read from it may alias buffers + // the reader owns, so a caller encodes each value before advancing. + // + // values maps a declared param name to a string, int64 or time.Time. A + // name the map does not carry binds NULL. + Run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) + + // Release returns the session. It is safe to call twice; the second is a + // no-op, so a deferred Release beside an early return cannot double-free. + Release() +} + +// Statement is one statement an Executor has checked. Its contents belong to +// the executor; serve only names it in errors. +type Statement interface { + // Where names the statement: "dataset posts_by_lang grain 1h". + Where() string +} + +// StatementSpec is what an Executor needs to prepare a statement. +type StatementSpec struct { + Dataset string + // Grain is empty for a dataset without grains. + Grain string + SQL string + Params []config.ServeParam +} + +// Stats is the pool's state at one instant, for the gauges and for the +// health endpoint. How long callers wait is a histogram, not a gauge, so it +// is not here. +type Stats struct { + // Size is every session the executor holds. + Size int + // InUse is the sessions currently borrowed. + InUse int +} + +// backendSession is what an Executor implementation gives the pool: run a +// statement, and close. The pool owns everything else, because queueing and +// its measurement are the same whatever runs the SQL. +type backendSession interface { + run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) + close() error +} + +// pool hands out a fixed set of sessions, one caller at a time. An Executor +// implementation embeds it. +type pool struct { + // free carries every session not currently borrowed. Its capacity is the + // pool's size, so a Release never blocks. + free chan *session + all []*session + // onWait receives how long each Acquire waited, including the ones that + // gave up. The metrics histogram is the only reader. + onWait func(time.Duration) + + mu sync.Mutex + inUse int + closed bool +} + +// session is one backend session while a caller holds it. +type session struct { + pool *pool + backend backendSession + released atomic.Bool +} + +func newPool(backends []backendSession, onWait func(time.Duration)) *pool { + if onWait == nil { + // A server without the metrics endpoint has nobody to tell. + onWait = func(time.Duration) {} + } + p := &pool{free: make(chan *session, len(backends)), onWait: onWait} + for _, b := range backends { + s := &session{pool: p, backend: b} + s.released.Store(true) + p.all = append(p.all, s) + p.free <- s + } + return p +} + +// Acquire borrows a session, waiting until one is free, ctx ends, or the +// executor closes. +func (p *pool) Acquire(ctx context.Context) (Session, error) { + start := time.Now() + select { + case s, ok := <-p.free: + if !ok { + return nil, ErrClosed + } + p.onWait(time.Since(start)) + p.mu.Lock() + p.inUse++ + p.mu.Unlock() + s.released.Store(false) + return s, nil + case <-ctx.Done(): + p.onWait(time.Since(start)) + return nil, ctx.Err() + } +} + +// setOnWait installs the hook the wait histogram reads. It runs before the +// server listens, so no Acquire can be in flight. +func (p *pool) setOnWait(f func(time.Duration)) { + if f != nil { + p.onWait = f + } +} + +func (p *pool) Stats() Stats { + p.mu.Lock() + defer p.mu.Unlock() + return Stats{Size: len(p.all), InUse: p.inUse} +} + +// Close refuses later acquires, waits for every borrowed session to come +// back, and only then closes them. A session closed while a query runs on it +// takes the process down, and nothing outside the goroutine reading that +// query can stop it, so waiting is the only option. The HTTP server's drain +// bounds how long that can be. +func (p *pool) Close() { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return + } + p.closed = true + p.mu.Unlock() + + for range p.all { + <-p.free + } + close(p.free) + for _, s := range p.all { + _ = s.backend.close() + } +} + +func (s *session) Run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) { + return s.backend.run(ctx, st, values) +} + +// Release returns the session. The second call is a no-op, so a deferred +// Release beside an early return cannot return one session twice. +func (s *session) Release() { + if s.released.Swap(true) { + return + } + s.pool.mu.Lock() + s.pool.inUse-- + s.pool.mu.Unlock() + // Never blocks: free's capacity is the pool's size, and this session is + // not in it. + s.pool.free <- s +} + +// query acquires a session, runs st, and encodes at most maxRows rows. It +// returns the rows, how long the acquire waited, and any error. +// +// The work runs on its own goroutine and hands back finished bytes, so a +// caller that has already given up at its deadline cannot race it. The +// session returns to the pool when the query finishes rather than when the +// caller stops waiting: handing a session to the next request while a query +// still runs on it would serialise them behind work nobody wants. +// +// ctx reaches readRows, which stops at the first batch boundary after the +// caller gives up and releases the reader. ADBC documents releasing a reader +// without draining it as equivalent to AdbcStatementCancel, and that is the +// only cancel its Go API offers. +func query(ctx context.Context, ex Executor, st Statement, values map[string]any, maxRows int) (result, time.Duration, error) { + start := time.Now() + sess, err := ex.Acquire(ctx) + if err != nil { + return result{}, time.Since(start), err + } + waited := time.Since(start) + + type outcome struct { + res result + err error + } + done := make(chan outcome, 1) + go func() { + defer sess.Release() + rdr, err := sess.Run(ctx, st, values) + if err != nil { + done <- outcome{err: err} + return + } + res, err := func() (result, error) { + // Released whether or not the rows are drained, which is what + // tells the driver to stop, and before the caller is told: the + // buffers go back before the next request allocates its own. + defer rdr.Release() + return readRows(ctx, rdr, maxRows) + }() + done <- outcome{res: res, err: err} + }() + + select { + case o := <-done: + return o.res, waited, o.err + case <-ctx.Done(): + return result{}, waited, ctx.Err() + } +} diff --git a/internal/serve/executor_test.go b/internal/serve/executor_test.go new file mode 100644 index 00000000..92e4b9d7 --- /dev/null +++ b/internal/serve/executor_test.go @@ -0,0 +1,186 @@ +package serve + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/zeebo/assert" +) + +// fakeSession is a backend that sleeps instead of querying, so pool +// behaviour is testable without DuckDB. +type fakeSession struct { + delay time.Duration + closed bool + runs int +} + +func (f *fakeSession) run(ctx context.Context, st Statement, values map[string]any) (array.RecordReader, error) { + f.runs++ + select { + case <-time.After(f.delay): + return nil, errors.New("fake session returns no rows") + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (f *fakeSession) close() error { f.closed = true; return nil } + +func newFakePool(t *testing.T, n int, delay time.Duration) (*pool, []*fakeSession) { + t.Helper() + fakes := make([]*fakeSession, n) + backends := make([]backendSession, n) + for i := range fakes { + fakes[i] = &fakeSession{delay: delay} + backends[i] = fakes[i] + } + return newPool(backends, func(time.Duration) {}), fakes +} + +// The whole point: N sessions run N queries in about the time of one. +func TestCliServe_PoolRunsQueriesConcurrently(t *testing.T) { + coverage.Covers(t, "cli.serve") + + const n = 4 + const delay = 200 * time.Millisecond + p, _ := newFakePool(t, n, delay) + defer p.Close() + + start := time.Now() + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + defer s.Release() + _, _ = s.Run(context.Background(), nil, nil) + }() + } + wg.Wait() + + // Serial would be 4 x 200ms. Two delays of headroom for a loaded CI box. + assert.That(t, time.Since(start) < 2*delay) +} + +// One session is the old behaviour, and the control for the test above. +func TestCliServe_APoolOfOneSerializes(t *testing.T) { + coverage.Covers(t, "cli.serve") + + const delay = 100 * time.Millisecond + p, _ := newFakePool(t, 1, delay) + defer p.Close() + + start := time.Now() + var wg sync.WaitGroup + for i := 0; i < 3; i++ { + wg.Add(1) + go func() { + defer wg.Done() + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + defer s.Release() + _, _ = s.Run(context.Background(), nil, nil) + }() + } + wg.Wait() + + assert.That(t, time.Since(start) >= 3*delay) +} + +// A caller that gives up must free its place, or the pool shrinks under load +// until it deadlocks. +func TestCliServe_AcquireReturnsWhenTheRequestGivesUp(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, _ := newFakePool(t, 1, 0) + defer p.Close() + + held, err := p.Acquire(context.Background()) + assert.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + _, err = p.Acquire(ctx) + assert.That(t, errors.Is(err, context.DeadlineExceeded)) + assert.Equal(t, 1, p.Stats().InUse) + + held.Release() + next, err := p.Acquire(context.Background()) + assert.NoError(t, err) + next.Release() +} + +// A session closed while a query runs on it takes the process down, and +// nothing outside the reading goroutine can stop that query, so Close waits. +func TestCliServe_CloseWaitsForBorrowedSessions(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, fakes := newFakePool(t, 2, 0) + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + + closed := make(chan struct{}) + go func() { p.Close(); close(closed) }() + + select { + case <-closed: + t.Fatal("Close returned while a session was borrowed") + case <-time.After(100 * time.Millisecond): + } + + s.Release() + select { + case <-closed: + case <-time.After(2 * time.Second): + t.Fatal("Close did not return after the session came back") + } + for i, f := range fakes { + if !f.closed { + t.Fatalf("session %d was not closed", i) + } + } + + _, err = p.Acquire(context.Background()) + assert.That(t, errors.Is(err, ErrClosed)) +} + +// A deferred Release beside an early return can run twice. +func TestCliServe_ReleaseTwiceIsSafe(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, _ := newFakePool(t, 1, 0) + defer p.Close() + + s, err := p.Acquire(context.Background()) + assert.NoError(t, err) + s.Release() + s.Release() + + assert.Equal(t, 0, p.Stats().InUse) + next, err := p.Acquire(context.Background()) + assert.NoError(t, err) + next.Release() +} + +func TestCliServe_StatsReportSizeAndUse(t *testing.T) { + coverage.Covers(t, "cli.serve") + + p, _ := newFakePool(t, 3, 0) + defer p.Close() + + assert.Equal(t, Stats{Size: 3, InUse: 0}, p.Stats()) + a, _ := p.Acquire(context.Background()) + b, _ := p.Acquire(context.Background()) + assert.Equal(t, Stats{Size: 3, InUse: 2}, p.Stats()) + a.Release() + b.Release() + assert.Equal(t, Stats{Size: 3, InUse: 0}, p.Stats()) +} diff --git a/internal/serve/http.go b/internal/serve/http.go index f9db4283..22f78c1c 100644 --- a/internal/serve/http.go +++ b/internal/serve/http.go @@ -5,13 +5,13 @@ import ( "crypto/subtle" "encoding/json" "errors" + "github.com/prometheus/client_golang/prometheus/promhttp" "net" "net/http" "net/url" "strings" "time" - "github.com/apache/arrow-adbc/go/adbc" "go.uber.org/zap" ) @@ -24,6 +24,12 @@ const ( func (s *Server) Handler() http.Handler { mux := http.NewServeMux() mux.HandleFunc("/healthz", s.healthz) + if s.metrics != nil { + // No token: it carries no row data, and the listener is already + // public. It is absent unless the config turns it on, because the + // labels name every dataset and grain. + mux.Handle("/metrics", promhttp.HandlerFor(s.registry, promhttp.HandlerOpts{})) + } mux.HandleFunc("/v1/datasets", s.authed(s.listDatasets)) mux.HandleFunc("/v1/datasets/{name}", s.authed(s.queryDataset)) mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { @@ -31,7 +37,7 @@ func (s *Server) Handler() http.Handler { }) // Cheap rejections first: CORS, then the method, then auth inside the - // routes, all before anything waits on the connection's lock. + // routes, all before anything waits for a session. return s.logRequests(s.cors(getOnly(mux))) } @@ -64,29 +70,25 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error { } func (s *Server) healthz(w http.ResponseWriter, r *http.Request) { - _, err := s.exec.run(r.Context(), s.healthTimeout, - func(ctx context.Context, conn adbc.Connection) (result, error) { - stmt, err := conn.NewStatement() - if err != nil { - return result{}, err - } - defer stmt.Close() - if err := stmt.SetSqlQuery("SELECT 1"); err != nil { - return result{}, err - } - rdr, _, err := stmt.ExecuteQuery(ctx) - if err != nil { - return result{}, err - } - defer rdr.Release() - return readRows(rdr, 1) - }) - if err != nil { + ctx, cancel := context.WithTimeout(r.Context(), s.healthTimeout) + defer cancel() + + _, _, err := query(ctx, s.exec, s.health, nil, 1) + switch { + case err == nil: + writeJSON(w, r, http.StatusOK, []byte(`{"status":"ok"}`)) + case errors.Is(err, context.DeadlineExceeded), errors.Is(err, ErrClosed): + // Busy is not dead. A supervisor that cannot tell them apart restarts + // a server that is merely loaded, at the worst possible moment. + // + // A request that timed out leaves its session busy until its query + // ends, so a pool full of abandoned queries reports busy too. That is + // the truth about the server. + writeJSON(w, r, http.StatusServiceUnavailable, []byte(`{"status":"busy"}`)) + default: s.logger.Warn("health check failed", zap.String("error", Redact(err.Error()))) writeJSON(w, r, http.StatusServiceUnavailable, []byte(`{"status":"unavailable"}`)) - return } - writeJSON(w, r, http.StatusOK, []byte(`{"status":"ok"}`)) } func (s *Server) listDatasets(w http.ResponseWriter, r *http.Request) { @@ -102,7 +104,12 @@ type rowsResponse struct { Rows json.RawMessage `json:"rows"` RowCount int `json:"row_count"` Truncated bool `json:"truncated"` - ElapsedMS int64 `json:"elapsed_ms"` + // QueuedMS is how long the request waited for a session, and ElapsedMS is + // the query alone. Before the pool, elapsed_ms silently included the wait, + // so a 13 ms query on a loaded server reported a second and sent its + // reader looking for a slow query that did not exist. + QueuedMS int64 `json:"queued_ms"` + ElapsedMS int64 `json:"elapsed_ms"` } func (s *Server) queryDataset(w http.ResponseWriter, r *http.Request) { @@ -116,14 +123,15 @@ func (s *Server) queryDataset(w http.ResponseWriter, r *http.Request) { return } - query, err := url.ParseQuery(r.URL.RawQuery) + // params, not query: query is the helper that runs one. + params, err := url.ParseQuery(r.URL.RawQuery) if err != nil { writeError(w, r, &apiError{http.StatusBadRequest, "invalid_param", "the query string does not parse: " + err.Error()}) return } - values, apiErr := parseParams(ds.conf.Params, query) + values, apiErr := parseParams(ds.conf.Params, params) if apiErr != nil { writeError(w, r, apiErr) return @@ -132,13 +140,13 @@ func (s *Server) queryDataset(w http.ResponseWriter, r *http.Request) { // A ranged dataset needs the parsed since and until to choose a grain; // any other dataset takes the grain as named. var ( - st *statement + st datasetStatement win *window ) if ds.span != nil { - st, win, apiErr = ds.resolveRange(query, values, s.now()) + st, win, apiErr = ds.resolveRange(params, values, s.now()) } else { - st, apiErr = ds.resolveStatement(query) + st, apiErr = ds.resolveStatement(params) } if apiErr != nil { writeError(w, r, apiErr) @@ -146,15 +154,22 @@ func (s *Server) queryDataset(w http.ResponseWriter, r *http.Request) { } entry.grain = st.grain + // The deadline bounds this caller's wait: for a session, and then for the + // query on it. + ctx, cancel := context.WithTimeout(r.Context(), ds.timeout) + defer cancel() + start := time.Now() - res, err := s.exec.run(r.Context(), ds.timeout, - func(ctx context.Context, conn adbc.Connection) (result, error) { - return st.query(ctx, conn, values, ds.maxRows) - }) + res, queued, err := query(ctx, s.exec, st.stmt, values, ds.maxRows) + // Measured whatever the outcome: the error mix is what the counter is + // for, so a timeout and a failure count too. + entry.queryDur, entry.measured = time.Since(start)-queued, true switch { case errors.Is(err, context.DeadlineExceeded): + // An exhausted pool arrives here too: the wait for a session is the + // same wait as far as the caller is concerned. writeError(w, r, &apiError{http.StatusGatewayTimeout, "query_timeout", - st.where() + " did not answer within " + ds.timeout.String()}) + st.stmt.Where() + " did not answer within " + ds.timeout.String()}) return case errors.Is(err, context.Canceled): // The caller hung up. There is nobody to answer. @@ -167,7 +182,7 @@ func (s *Server) queryDataset(w http.ResponseWriter, r *http.Request) { s.logger.Error("query failed", zap.String("dataset", name), zap.String("grain", st.grain), zap.String("error", Redact(err.Error()))) writeError(w, r, &apiError{http.StatusInternalServerError, "query_failed", - st.where() + " failed; the server log has the database's error"}) + st.stmt.Where() + " failed; the server log has the database's error"}) return } @@ -179,7 +194,8 @@ func (s *Server) queryDataset(w http.ResponseWriter, r *http.Request) { Rows: res.Rows, RowCount: res.RowCount, Truncated: res.Truncated, - ElapsedMS: time.Since(start).Milliseconds(), + QueuedMS: queued.Milliseconds(), + ElapsedMS: (time.Since(start) - queued).Milliseconds(), }) if err != nil { writeError(w, r, &apiError{http.StatusInternalServerError, "query_failed", err.Error()}) @@ -253,15 +269,27 @@ func (s *Server) cors(next http.Handler) http.Handler { }) } +// getOnly refuses every method but GET, and HEAD on /healthz. +// +// Monitors send HEAD, and for /healthz the status line is the whole answer. +// Every other route stays GET-only on purpose: a HEAD of a dataset would run +// the query, borrow a session and throw the rows away, which spends the pool +// on nothing. net/http suppresses the body of a HEAD response, so the health +// handler needs no special case. func getOnly(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method != http.MethodGet { - w.Header().Set("Allow", http.MethodGet) - writeError(w, r, &apiError{http.StatusMethodNotAllowed, "method_not_allowed", - r.Method + " is not allowed; every route is GET"}) + health := r.URL.Path == "/healthz" + if r.Method == http.MethodGet || (health && r.Method == http.MethodHead) { + next.ServeHTTP(w, r) return } - next.ServeHTTP(w, r) + allowed, where := http.MethodGet, "every route is GET, and /healthz is also HEAD" + if health { + allowed, where = "GET, HEAD", "/healthz is GET or HEAD" + } + w.Header().Set("Allow", allowed) + writeError(w, r, &apiError{http.StatusMethodNotAllowed, "method_not_allowed", + r.Method + " is not allowed; " + where}) }) } @@ -269,6 +297,11 @@ func getOnly(next http.Handler) http.Handler { type logEntry struct { token, dataset, grain, code string status, rows int + // queryDur is the query alone, which only the dataset handler knows. The + // middleware records the metrics, because that is where the outcome and + // the total are both known. + queryDur time.Duration + measured bool } type entryKey struct{} @@ -304,6 +337,17 @@ func (s *Server) logRequests(next http.Handler) http.Handler { fields = append(fields, zap.String("code", entry.code)) } s.logger.Info("request", fields...) + + // Only a dataset request carries a query to measure. A listing, a + // health probe or a refusal would otherwise report a query of zero + // seconds and flatten the histogram. + if entry.measured { + code := entry.code + if code == "" { + code = "ok" + } + s.metrics.observeRequest(entry.dataset, entry.grain, code, entry.queryDur, time.Since(start)) + } }) } diff --git a/internal/serve/http_test.go b/internal/serve/http_test.go index 871aabce..35ded6b2 100644 --- a/internal/serve/http_test.go +++ b/internal/serve/http_test.go @@ -68,6 +68,13 @@ serve: sql: ` + slowSQL + ` ` +const createPostsTable = `CREATE TABLE posts AS SELECT * FROM (VALUES + (TIMESTAMPTZ '2026-09-10 00:00:00+00', 'en', 10::BIGINT), + (TIMESTAMPTZ '2026-09-10 00:00:00+00', 'ja', 4::BIGINT), + (TIMESTAMPTZ '2026-09-10 01:00:00+00', 'en', 7::BIGINT), + (TIMESTAMPTZ '2026-09-11 00:00:00+00', 'en', 1::BIGINT) +) AS t(bucket, lang, posts)` + type testServer struct { srv *Server handler http.Handler @@ -76,23 +83,25 @@ type testServer struct { func newTestServer(t *testing.T, text string) *testServer { t.Helper() - conn := newConn(t) - execSQL(t, conn, "SET TimeZone='UTC'") - execSQL(t, conn, `CREATE TABLE posts AS SELECT * FROM (VALUES - (TIMESTAMPTZ '2026-09-10 00:00:00+00', 'en', 10::BIGINT), - (TIMESTAMPTZ '2026-09-10 00:00:00+00', 'ja', 4::BIGINT), - (TIMESTAMPTZ '2026-09-10 01:00:00+00', 'en', 7::BIGINT), - (TIMESTAMPTZ '2026-09-11 00:00:00+00', 'en', 1::BIGINT) - ) AS t(bucket, lang, posts)`) + return newTestServerWith(t, text) +} + +// newTestServerWith takes extra options, for a test that needs the metrics +// registry. +func newTestServerWith(t *testing.T, text string, extra ...Option) *testServer { + t.Helper() + // One session: every assertion below about ordering and timing is the + // old single-connection behaviour. + ex, _ := newExec(t, 1, "SET TimeZone='UTC'", createPostsTable) conf, err := config.ParseServe([]byte(text)) assert.NoError(t, err) core, logs := observer.New(zap.InfoLevel) - srv, err := New(context.Background(), conf, conn, WithLogger(zap.New(core))) + srv, err := New(context.Background(), conf, ex, append([]Option{WithLogger(zap.New(core))}, extra...)...) assert.NoError(t, err) - // Registered after newConn's cleanup, so it runs first: a slow query - // still holding the lock finishes before the connection closes. + // Registered after newExec's cleanups, so it runs first: a slow query + // still holding a session finishes before the database closes. t.Cleanup(srv.Close) return &testServer{srv: srv, handler: srv.Handler(), logs: logs} @@ -115,7 +124,9 @@ func (ts *testServer) do(t *testing.T, method, target string, header map[string] ts.handler.ServeHTTP(w, req) resp := response{status: w.Code, header: w.Header(), raw: w.Body.String()} - if w.Body.Len() > 0 { + // Only the JSON routes decode. /metrics answers Prometheus text, and a + // HEAD answers nothing. + if w.Body.Len() > 0 && strings.HasPrefix(w.Header().Get("Content-Type"), "application/json") { assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp.body)) } return resp @@ -338,8 +349,8 @@ func TestCliServe_MaxRowsTruncatesAndTheNextRequestAnswers(t *testing.T) { assert.Equal(t, float64(3), again.body["row_count"]) } -// A query past its deadline is a 504, and the health check, waiting on the -// same lock, goes red until the query finishes. +// A query past its deadline is a 504, and the health check, waiting for the +// pool's only session, goes red until the query finishes. func TestCliServe_ASlowQueryIsA504AndTurnsHealthRed(t *testing.T) { coverage.Covers(t, "cli.serve") ts := newTestServer(t, testServe) @@ -502,16 +513,16 @@ func TestCliServe_ServeDrainsAnInFlightRequest(t *testing.T) { // New. Neither is found on the first request. func TestCliServe_NewRefusesWhatCannotAnswer(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) + ex, _ := newExec(t, 1) emptyToken, err := config.ParseServe([]byte(strings.Replace(testServe, "token: page-token", `token: ""`, 1))) assert.NoError(t, err) - _, err = New(context.Background(), emptyToken, conn) + _, err = New(context.Background(), emptyToken, ex) assert.Equal(t, errs.CodeConfigInvalid, errs.CodeOf(err)) noTable, err := config.ParseServe([]byte(testServe)) assert.NoError(t, err) - _, err = New(context.Background(), noTable, conn) + _, err = New(context.Background(), noTable, ex) assert.Equal(t, errs.CodeSQLInvalid, errs.CodeOf(err)) assert.That(t, strings.Contains(err.Error(), "dataset status")) } @@ -552,3 +563,83 @@ func TestCliServe_RedactRemovesPasswords(t *testing.T) { assert.Equal(t, want, Redact(in)) } } + +// elapsed_ms used to include the wait for a connection, so a 13 ms query on a +// loaded server reported a second and sent its reader looking for a slow query +// that did not exist. The two are separate fields now. +func TestCliServe_QueuedMsSeparatesWaitFromWork(t *testing.T) { + coverage.Covers(t, "cli.serve") + ts := newTestServer(t, testServe) + + idle := ts.get(t, "/v1/datasets/status") + assert.Equal(t, http.StatusOK, idle.status) + assert.Equal(t, float64(0), idle.body["queued_ms"]) + + // Hold the only session, then time a request that has to wait for it. + held, err := ts.srv.exec.Acquire(context.Background()) + assert.NoError(t, err) + + done := make(chan response, 1) + go func() { done <- ts.get(t, "/v1/datasets/status") }() + time.Sleep(300 * time.Millisecond) + held.Release() + + queued := <-done + assert.Equal(t, http.StatusOK, queued.status) + assert.That(t, queued.body["queued_ms"].(float64) >= 250) + // The query itself is unchanged by the wait, which is the whole point. + assert.That(t, queued.body["elapsed_ms"].(float64) < 250) +} + +// A full pool is not a dead server. A supervisor that cannot tell them apart +// restarts one that is merely loaded. +// +// The config gives health one second rather than the default ten, because the +// handler waits its whole timeout for a session before it can say busy. +const busyServe = ` +serve: + auth: + tokens: [{name: page, token: page-token}] + limits: + timeout_seconds: 1 + datasets: + - name: status + sql: SELECT count(*) AS n FROM posts +` + +func TestCliServe_HealthzIsBusyNotDownWhenThePoolIsFull(t *testing.T) { + coverage.Covers(t, "cli.serve") + ts := newTestServer(t, busyServe) + + held, err := ts.srv.exec.Acquire(context.Background()) + assert.NoError(t, err) + defer held.Release() + + r := ts.do(t, http.MethodGet, "/healthz", nil) + assert.Equal(t, http.StatusServiceUnavailable, r.status) + assert.Equal(t, "busy", r.body["status"]) +} + +// Monitors send HEAD, and the status line is the whole answer. A dataset +// still refuses it: running the query and discarding the rows would spend a +// session on nothing. +// +// The status is all this asserts. httptest.NewRecorder hands the handler's +// body straight back, where a real http.Server suppresses it for HEAD, so an +// empty-body assertion here would be testing the recorder. +func TestCliServe_HealthzAnswersHead(t *testing.T) { + coverage.Covers(t, "cli.serve") + ts := newTestServer(t, busyServe) + + assert.Equal(t, http.StatusOK, ts.do(t, http.MethodHead, "/healthz", nil).status) + + held, err := ts.srv.exec.Acquire(context.Background()) + assert.NoError(t, err) + busy := ts.do(t, http.MethodHead, "/healthz", nil) + held.Release() + assert.Equal(t, http.StatusServiceUnavailable, busy.status) + + ds := ts.do(t, http.MethodHead, "/v1/datasets/status", pageToken) + assert.Equal(t, http.StatusMethodNotAllowed, ds.status) + assert.Equal(t, http.MethodGet, ds.header.Get("Allow")) +} diff --git a/internal/serve/leak_test.go b/internal/serve/leak_test.go index 7449b613..b3368db4 100644 --- a/internal/serve/leak_test.go +++ b/internal/serve/leak_test.go @@ -25,8 +25,7 @@ import ( // fails it; growth without that is about 1 MiB. func TestCliServe_DoesNotLeakNativeMemory(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) - execSQL(t, conn, `CREATE TABLE posts AS + ex, _ := newExec(t, 1, `CREATE TABLE posts AS SELECT TIMESTAMPTZ '2026-09-10 00:00:00+00' + INTERVAL (range) MINUTE AS bucket, 'lang_' || (range % 33) AS lang, range::BIGINT AS posts @@ -42,7 +41,7 @@ serve: sql: SELECT bucket, lang, posts FROM posts WHERE lang = coalesce($lang, lang) `)) assert.NoError(t, err) - srv, err := New(context.Background(), conf, conn) + srv, err := New(context.Background(), conf, ex) assert.NoError(t, err) t.Cleanup(srv.Close) handler := srv.Handler() @@ -69,7 +68,13 @@ serve: return n } - const warmup, iters, rowsPerRequest = 50, 500, 1000 + // The warmup outlasts the allocator's own settling, which is what the + // baseline has to be taken after. A pool reaches steady state later than + // one connection did: measured over four consecutive runs of iters, the + // process grew 10 MiB, then 3, then 1, then 0 -- a plateau, not a leak. + // Fifty requests sampled the middle of that curve and read the climb as + // growth. + const warmup, iters, rowsPerRequest = 500, 500, 1000 run(warmup) settle() before := resident() diff --git a/internal/serve/metrics.go b/internal/serve/metrics.go new file mode 100644 index 00000000..70b358fd --- /dev/null +++ b/internal/serve/metrics.go @@ -0,0 +1,120 @@ +package serve + +import ( + "context" + "time" + + prom "github.com/prometheus/client_golang/prometheus" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/exporters/prometheus" + "go.opentelemetry.io/otel/metric" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" +) + +// meterName names serve's instruments apart from the pipeline's, which run +// under their own meter in the same process shape. +const meterName = "sqlflow.serve" + +// latencyBuckets span a 1 ms query and a request that sat out a 10 s timeout, +// so both land in a bucket rather than in the overflow. +var latencyBuckets = []float64{ + 0.001, 0.002, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 4, 8, 16, +} + +// metrics is the six instruments the pool needs to be sized, and nothing +// else. Every method tolerates a nil receiver, so the request path records +// unconditionally and a server without the endpoint pays one nil check. +type metrics struct { + requests metric.Int64Counter + requestDuration metric.Float64Histogram + queryDuration metric.Float64Histogram + sessionWait metric.Float64Histogram +} + +// newMetrics builds the instruments against reg and registers the gauges' +// callback, which reads stats whenever the endpoint is scraped. +func newMetrics(reg *prom.Registry, stats func() Stats) (*metrics, error) { + exp, err := prometheus.New(prometheus.WithRegisterer(reg)) + if err != nil { + return nil, err + } + // One view, so every histogram here gets buckets chosen for this workload + // rather than the SDK's defaults, which stop at 10 s. + mp := sdkmetric.NewMeterProvider( + sdkmetric.WithReader(exp), + sdkmetric.WithView(sdkmetric.NewView( + sdkmetric.Instrument{Kind: sdkmetric.InstrumentKindHistogram}, + sdkmetric.Stream{Aggregation: sdkmetric.AggregationExplicitBucketHistogram{ + Boundaries: latencyBuckets, + }}, + )), + ) + m := mp.Meter(meterName) + + var mm metrics + if mm.requests, err = m.Int64Counter("sqlflow_serve_requests_total", + metric.WithDescription("Requests answered, by dataset, grain and outcome.")); err != nil { + return nil, err + } + if mm.requestDuration, err = m.Float64Histogram("sqlflow_serve_request_duration_seconds", + metric.WithDescription("What a caller waited, including time queued for a session."), + metric.WithUnit("s")); err != nil { + return nil, err + } + if mm.queryDuration, err = m.Float64Histogram("sqlflow_serve_query_duration_seconds", + metric.WithDescription("The query alone, so a slow backend is distinguishable from a full pool."), + metric.WithUnit("s")); err != nil { + return nil, err + } + if mm.sessionWait, err = m.Float64Histogram("sqlflow_serve_session_wait_seconds", + metric.WithDescription("Time spent waiting for a session. This is what sizes the pool: a p99 that climbs while query duration stays flat means too few sessions."), + metric.WithUnit("s")); err != nil { + return nil, err + } + + inUse, err := m.Int64ObservableGauge("sqlflow_serve_sessions_in_use", + metric.WithDescription("Sessions currently borrowed. Pinned at the total means saturated.")) + if err != nil { + return nil, err + } + total, err := m.Int64ObservableGauge("sqlflow_serve_sessions_total", + metric.WithDescription("Sessions the executor holds, so a dashboard need not hardcode the config.")) + if err != nil { + return nil, err + } + if _, err := m.RegisterCallback(func(_ context.Context, o metric.Observer) error { + s := stats() + o.ObserveInt64(inUse, int64(s.InUse)) + o.ObserveInt64(total, int64(s.Size)) + return nil + }, inUse, total); err != nil { + return nil, err + } + + return &mm, nil +} + +// observeWait records one Acquire, including the ones that gave up waiting. +func (m *metrics) observeWait(d time.Duration) { + if m == nil { + return + } + m.sessionWait.Record(context.Background(), d.Seconds()) +} + +// observeRequest records one finished request. code is the error code, or +// "ok". +func (m *metrics) observeRequest(dataset, grain, code string, query, total time.Duration) { + if m == nil { + return + } + attrs := metric.WithAttributes( + attribute.String("dataset", dataset), + attribute.String("grain", grain), + attribute.String("code", code), + ) + ctx := context.Background() + m.requests.Add(ctx, 1, attrs) + m.requestDuration.Record(ctx, total.Seconds(), attrs) + m.queryDuration.Record(ctx, query.Seconds(), attrs) +} diff --git a/internal/serve/metrics_test.go b/internal/serve/metrics_test.go new file mode 100644 index 00000000..56946b89 --- /dev/null +++ b/internal/serve/metrics_test.go @@ -0,0 +1,75 @@ +package serve + +import ( + "net/http" + "strings" + "testing" + + prom "github.com/prometheus/client_golang/prometheus" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/zeebo/assert" +) + +// The endpoint sits on the public listener and its labels name every dataset, +// so it exists only when the config asks for it. +func TestCliServe_MetricsAreOffUnlessEnabled(t *testing.T) { + coverage.Covers(t, "cli.serve") + + // Even handed a registry: the config decides, not the caller. + ts := newTestServerWith(t, testServe, WithMetrics(prom.NewRegistry())) + assert.Equal(t, http.StatusNotFound, ts.do(t, http.MethodGet, "/metrics", nil).status) +} + +// Every instrument carries a sample after one request, so a dashboard built +// against this list is not built against a blank. +func TestCliServe_MetricsCarryEveryInstrument(t *testing.T) { + coverage.Covers(t, "cli.serve") + + enabled := strings.Replace(testServe, " limits:", " metrics: {enabled: true}\n limits:", 1) + assert.That(t, enabled != testServe) + ts := newTestServerWith(t, enabled, WithMetrics(prom.NewRegistry())) + + assert.Equal(t, http.StatusOK, ts.get(t, "/v1/datasets/status").status) + + r := ts.do(t, http.MethodGet, "/metrics", nil) + assert.Equal(t, http.StatusOK, r.status) + for _, name := range []string{ + "sqlflow_serve_requests_total", + "sqlflow_serve_request_duration_seconds", + "sqlflow_serve_query_duration_seconds", + "sqlflow_serve_session_wait_seconds", + "sqlflow_serve_sessions_in_use", + "sqlflow_serve_sessions_total", + } { + if !strings.Contains(r.raw, name) { + t.Fatalf("/metrics does not carry %s:\n%s", name, r.raw) + } + } + // The gauge reports the pool rather than a constant. The exporter appends + // otel_scope labels, so the value is read from the end of its line. + if got := sampleValue(t, r.raw, "sqlflow_serve_sessions_total"); got != "1" { + t.Fatalf("sessions_total is %q, want 1 for a pool of one:\n%s", got, r.raw) + } + // The labels a dashboard breaks the rate down by. + assert.That(t, strings.Contains(r.raw, `dataset="status"`)) + assert.That(t, strings.Contains(r.raw, `code="ok"`)) +} + +// sampleValue returns the value of the first sample whose name matches, +// ignoring its labels. +func sampleValue(t *testing.T, text, name string) string { + t.Helper() + for _, line := range strings.Split(text, "\n") { + if !strings.HasPrefix(line, name) || strings.HasPrefix(line, "#") { + continue + } + // A sample is "name{labels} value" or "name value". + if rest := strings.TrimPrefix(line, name); rest == "" || (rest[0] != '{' && rest[0] != ' ') { + continue + } + fields := strings.Fields(line) + return fields[len(fields)-1] + } + t.Fatalf("no sample named %s", name) + return "" +} diff --git a/internal/serve/params.go b/internal/serve/params.go index 08729054..758d26c8 100644 --- a/internal/serve/params.go +++ b/internal/serve/params.go @@ -20,12 +20,12 @@ type apiError struct { } // resolveStatement picks the dataset's statement for the request's grain. -func (ds *dataset) resolveStatement(query url.Values) (*statement, *apiError) { +func (ds *dataset) resolveStatement(query url.Values) (datasetStatement, *apiError) { grains, given := query["grain"] if ds.grains == nil { if given { - return nil, &apiError{http.StatusBadRequest, "unknown_grain", + return datasetStatement{}, &apiError{http.StatusBadRequest, "unknown_grain", "dataset " + ds.conf.Name + " has no grains"} } return ds.single, nil @@ -34,16 +34,16 @@ func (ds *dataset) resolveStatement(query url.Values) (*statement, *apiError) { names := ds.conf.GrainNames() switch { case !given: - return nil, &apiError{http.StatusBadRequest, "missing_grain", + return datasetStatement{}, &apiError{http.StatusBadRequest, "missing_grain", "dataset " + ds.conf.Name + " needs a grain; grains: " + strings.Join(names, ", ")} case len(grains) > 1: - return nil, &apiError{http.StatusBadRequest, "invalid_param", + return datasetStatement{}, &apiError{http.StatusBadRequest, "invalid_param", "grain is given " + strconv.Itoa(len(grains)) + " times"} } st, ok := ds.grains[grains[0]] if !ok { - return nil, &apiError{http.StatusBadRequest, "unknown_grain", + return datasetStatement{}, &apiError{http.StatusBadRequest, "unknown_grain", "dataset " + ds.conf.Name + " has no grain " + grains[0] + "; grains: " + strings.Join(names, ", ")} } return st, nil @@ -63,7 +63,7 @@ type window struct { // values already holds the parsed params. The resolved since and until are // written back into it, so the statement binds timestamps, never NULL, and // needs no defaults of its own. -func (ds *dataset) resolveRange(query url.Values, values map[string]any, now time.Time) (*statement, *window, *apiError) { +func (ds *dataset) resolveRange(query url.Values, values map[string]any, now time.Time) (datasetStatement, *window, *apiError) { sp := ds.span // Microseconds, because that is what binds. The echoed range then says @@ -80,7 +80,7 @@ func (ds *dataset) resolveRange(query url.Values, values map[string]any, now tim since = since.UTC().Truncate(time.Microsecond) if !since.Before(until) { - return nil, nil, &apiError{http.StatusBadRequest, "invalid_param", + return datasetStatement{}, nil, &apiError{http.StatusBadRequest, "invalid_param", sp.since + " must be before " + sp.until + "; got " + sp.since + " " + since.Format(time.RFC3339Nano) + " and " + sp.until + " " + until.Format(time.RFC3339Nano)} } @@ -96,18 +96,18 @@ func (ds *dataset) resolveRange(query url.Values, values map[string]any, now tim } } widest := sp.grains[len(sp.grains)-1] - return nil, nil, &apiError{http.StatusBadRequest, "range_too_wide", + return datasetStatement{}, nil, &apiError{http.StatusBadRequest, "range_too_wide", "the range is " + describeWidth(width) + " and the widest grain, " + widest.name + ", serves at most " + config.FormatServeDuration(widest.max)} } if len(grains) > 1 { - return nil, nil, &apiError{http.StatusBadRequest, "invalid_param", + return datasetStatement{}, nil, &apiError{http.StatusBadRequest, "invalid_param", "grain is given " + strconv.Itoa(len(grains)) + " times"} } st, ok := ds.grains[grains[0]] if !ok { - return nil, nil, &apiError{http.StatusBadRequest, "unknown_grain", + return datasetStatement{}, nil, &apiError{http.StatusBadRequest, "unknown_grain", "dataset " + ds.conf.Name + " has no grain " + grains[0] + "; grains: " + strings.Join(ds.conf.GrainNames(), ", ")} } @@ -127,7 +127,7 @@ func (ds *dataset) resolveRange(query url.Values, values map[string]any, now tim if len(fits) > 0 { msg += "; grains that serve it: " + strings.Join(fits, ", ") } - return nil, nil, &apiError{http.StatusBadRequest, "range_too_wide", msg} + return datasetStatement{}, nil, &apiError{http.StatusBadRequest, "range_too_wide", msg} } return st, win, nil } diff --git a/internal/serve/postgres_integration_test.go b/internal/serve/postgres_integration_test.go new file mode 100644 index 00000000..f725bb64 --- /dev/null +++ b/internal/serve/postgres_integration_test.go @@ -0,0 +1,132 @@ +package serve + +// What a pool costs an attached Postgres. One DuckDB scan opens up to +// pg_connection_limit connections, so with a pool the worst case may be +// size x that, against a database that also carries the pipeline's writer. +// The spec would not assume it; this measures it. + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/jackc/pgx/v5" + tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/zeebo/assert" +) + +const servePostgresImage = "postgres:18" + +func TestIntegrationServePool_PostgresBackendsStayBounded(t *testing.T) { + coverage.Covers(t, "cli.serve") + if testing.Short() { + t.Skip("integration: starts a Postgres container") + } + ctx := context.Background() + + pg, err := tcpostgres.Run(ctx, servePostgresImage, + tcpostgres.WithDatabase("serve"), + tcpostgres.WithUsername("serve"), + tcpostgres.WithPassword("serve"), + tcpostgres.BasicWaitStrategies(), + ) + if err != nil { + t.Fatalf("start postgres: %v", err) + } + t.Cleanup(func() { _ = pg.Terminate(context.Background()) }) + + dsn, err := pg.ConnectionString(ctx, "sslmode=disable") + assert.NoError(t, err) + + direct, err := pgx.Connect(ctx, dsn) + assert.NoError(t, err) + t.Cleanup(func() { _ = direct.Close(context.Background()) }) + + // Enough rows that a scan takes long enough to overlap with the others. + _, err = direct.Exec(ctx, `CREATE TABLE wide AS + SELECT i AS id, 'lang_' || (i % 170) AS lang, i::bigint AS n + FROM generate_series(1, 400000) AS s(i)`) + assert.NoError(t, err) + + const poolSize, connLimit = 8, 4 + db, err := duckdb.OpenPath(ctx, "") + assert.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + ex, err := NewDuckDBExecutor(ctx, db, poolSize, + func(ctx context.Context, conn adbc.Connection) error { + for _, sql := range []string{ + "INSTALL postgres; LOAD postgres;", + fmt.Sprintf("SET pg_connection_limit = %d", connLimit), + fmt.Sprintf("ATTACH '%s' AS pg (TYPE POSTGRES, READ_ONLY)", dsn), + } { + if err := execOn(ctx, conn, sql); err != nil { + return err + } + } + return nil + }, nil) + assert.NoError(t, err) + t.Cleanup(ex.Close) + + st, err := ex.Prepare(ctx, StatementSpec{ + Dataset: "wide", + SQL: "SELECT lang, sum(n)::BIGINT AS total FROM pg.wide GROUP BY lang ORDER BY lang", + }) + assert.NoError(t, err) + + // Sample while the scans run: the peak is what matters, and it is gone + // by the time they finish. + done := make(chan struct{}) + peak := make(chan int64, 1) + go func() { + var seen int64 + for { + select { + case <-done: + peak <- seen + return + case <-time.After(20 * time.Millisecond): + var n int64 + if err := direct.QueryRow(context.Background(), + `SELECT count(*) FROM pg_stat_activity + WHERE datname = current_database() AND pid <> pg_backend_pid()`).Scan(&n); err == nil && n > seen { + seen = n + } + } + } + }() + + var wg sync.WaitGroup + for i := 0; i < poolSize; i++ { + wg.Add(1) + go func() { + defer wg.Done() + res, _, err := query(ctx, ex, st, nil, 1000) + assert.NoError(t, err) + assert.That(t, res.RowCount > 0) + }() + } + wg.Wait() + close(done) + + got := <-peak + t.Logf("peak Postgres backends: %d, with pool.size=%d and pg_connection_limit=%d", got, poolSize, connLimit) + // Measured 2026-09-16: the peak was 4 with eight sessions scanning at + // once, so the attachment's connections are shared across sessions rather + // than opened per session. pg_connection_limit caps the attachment, and + // the pool does not multiply it. The assertion is that bound, so this + // fails if a DuckDB release changes it and the README goes stale. + // + // Sampling every 20 ms could in principle miss a spike shorter than that, + // which would understate the peak; a scan of 400k rows is far longer. + if got > connLimit { + t.Fatalf("peak %d backends exceeds pg_connection_limit = %d, so the pool now multiplies connections; update the README", + got, connLimit) + } +} diff --git a/internal/serve/query.go b/internal/serve/query.go deleted file mode 100644 index 20880f2b..00000000 --- a/internal/serve/query.go +++ /dev/null @@ -1,253 +0,0 @@ -package serve - -import ( - "context" - "errors" - "fmt" - "strings" - "sync" - "time" - - "github.com/apache/arrow-adbc/go/adbc" - "github.com/apache/arrow-go/v18/arrow" - "github.com/apache/arrow-go/v18/arrow/array" - "github.com/apache/arrow-go/v18/arrow/memory" - "github.com/turbolytics/sql-flow/internal/config" - "github.com/turbolytics/sql-flow/internal/errs" - "github.com/turbolytics/sql-flow/internal/sqlparams" -) - -// statement is one dataset statement, numbered and checked against DuckDB at -// startup. -type statement struct { - dataset string - // grain is empty for a dataset without grains. - grain string - // sql is the statement as the config wrote it. - sql string - // rewritten is sql with $name replaced by $N. - rewritten string - // schema has one field per placeholder, in number order, typed by the - // declared param. A request's values bind into a record of this shape. - schema *arrow.Schema -} - -// where names the statement in an error. -func (s *statement) where() string { - if s.grain == "" { - return "dataset " + s.dataset - } - return "dataset " + s.dataset + " grain " + s.grain -} - -// paramTypes maps a declared param type to the Arrow type it binds as. The -// type comes from the config, never from the value, so an absent param is a -// typed null and coalesce resolves against the right type. -var paramTypes = map[string]arrow.DataType{ - "string": arrow.BinaryTypes.String, - "integer": arrow.PrimitiveTypes.Int64, - "timestamp": &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}, -} - -// prepare numbers a statement's placeholders and has DuckDB check it. -// -// DuckDB binds a statement when its SQL is set, so a syntax error, a missing -// table and a missing column all fail here, at startup, rather than on the -// first request. The statement is closed afterwards: each request plans its -// own, because a held plan can fold table statistics into constants. -func prepare(ctx context.Context, conn adbc.Connection, dataset, grain, sql string, params []config.ServeParam) (*statement, error) { - st := &statement{dataset: dataset, grain: grain, sql: sql} - - rw, err := sqlparams.Rewrite(sql) - if err != nil { - return nil, errs.Wrap(errs.CodeConfigServeDataset, err, "%s", st.where()) - } - st.rewritten = rw.SQL - - declared := map[string]string{} - for _, p := range params { - declared[p.Name] = p.Type - } - fields := make([]arrow.Field, len(rw.Names)) - for i, name := range rw.Names { - typ, ok := paramTypes[declared[name]] - if !ok { - return nil, errs.New(errs.CodeConfigServeDataset, - "%s: $%s is not a declared param", st.where(), name) - } - fields[i] = arrow.Field{Name: name, Type: typ, Nullable: true} - } - st.schema = arrow.NewSchema(fields, nil) - - stmt, err := conn.NewStatement() - if err != nil { - return nil, err - } - defer stmt.Close() - - if err := stmt.SetSqlQuery(rw.SQL); err != nil { - return nil, errs.Wrap(errs.CodeSQLInvalid, err, "%s: the SQL does not prepare", st.where()) - } - if err := stmt.Prepare(ctx); err != nil { - return nil, errs.Wrap(errs.CodeSQLInvalid, err, "%s: the SQL does not prepare", st.where()) - } - - // The scanner and DuckDB must agree on the count. If they do not, the - // scanner misread the SQL, and binding would put values on the wrong - // placeholders and answer with wrong rows. - ps, err := stmt.GetParameterSchema() - if err != nil { - return nil, errs.Wrap(errs.CodeSQLInvalid, err, "%s: reading the parameter count", st.where()) - } - if ps.NumFields() != len(rw.Names) { - return nil, errs.New(errs.CodeSQLInvalid, - "%s: DuckDB counts %d parameters and sqlflow counts %d (%s); "+ - "a quote or comment form the scanner does not know is hiding or inventing one", - st.where(), ps.NumFields(), len(rw.Names), strings.Join(rw.Names, ", ")) - } - - return st, nil -} - -// query runs the statement with one request's values and encodes at most -// maxRows rows. values maps a param name to a string, int64 or time.Time; an -// absent name binds NULL. -func (s *statement) query(ctx context.Context, conn adbc.Connection, values map[string]any, maxRows int) (result, error) { - stmt, err := conn.NewStatement() - if err != nil { - return result{}, err - } - defer stmt.Close() - - if err := stmt.SetSqlQuery(s.rewritten); err != nil { - return result{}, err - } - - // DuckDB wants exactly one field per placeholder, so a statement with - // none binds nothing. - if s.schema.NumFields() > 0 { - rec := s.record(values) - defer rec.Release() - if err := stmt.Bind(ctx, rec); err != nil { - return result{}, err - } - } - - rdr, _, err := stmt.ExecuteQuery(ctx) - if err != nil { - return result{}, err - } - defer rdr.Release() - - return readRows(rdr, maxRows) -} - -// record builds the one-row record a request binds. -func (s *statement) record(values map[string]any) arrow.RecordBatch { - cols := make([]arrow.Array, s.schema.NumFields()) - for i, f := range s.schema.Fields() { - v, present := values[f.Name] - switch typ := f.Type.(type) { - case *arrow.StringType: - b := array.NewStringBuilder(memory.DefaultAllocator) - if present { - b.Append(v.(string)) - } else { - b.AppendNull() - } - cols[i] = b.NewArray() - b.Release() - case *arrow.Int64Type: - b := array.NewInt64Builder(memory.DefaultAllocator) - if present { - b.Append(v.(int64)) - } else { - b.AppendNull() - } - cols[i] = b.NewArray() - b.Release() - case *arrow.TimestampType: - b := array.NewTimestampBuilder(memory.DefaultAllocator, typ) - if present { - b.Append(arrow.Timestamp(v.(time.Time).UnixMicro())) - } else { - b.AppendNull() - } - cols[i] = b.NewArray() - b.Release() - default: - panic(fmt.Sprintf("serve: no builder for param type %s", f.Type)) - } - } - - rec := array.NewRecordBatch(s.schema, cols, 1) - for _, c := range cols { - c.Release() - } - return rec -} - -// errClosed is returned for a query that arrives after Close. -var errClosed = errors.New("the server is shutting down") - -// executor serializes every query on one connection. -// -// An ADBC connection is not safe for concurrent use, so one mutex guards it, -// the same shape as run's. DuckDB cannot cancel a query through the Go driver -// manager, so a deadline bounds the caller's wait, not the query: at the -// deadline the caller gets context.DeadlineExceeded, and the query keeps the -// lock until it finishes. -type executor struct { - mu sync.Mutex - conn adbc.Connection - closed bool -} - -type outcome struct { - res result - err error -} - -// run calls fn with the connection under the lock, waiting at most timeout. -// -// fn runs on its own goroutine and returns finished bytes over a buffered -// channel. Nothing it produces touches the response, so a caller that has -// already given up cannot race it. -func (e *executor) run(ctx context.Context, timeout time.Duration, fn func(context.Context, adbc.Connection) (result, error)) (result, error) { - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - done := make(chan outcome, 1) - go func() { - e.mu.Lock() - defer e.mu.Unlock() - - // The caller may have given up while this waited for the lock. - // Running the query anyway would hold the connection for nobody. - if err := ctx.Err(); err != nil { - done <- outcome{err: err} - return - } - if e.closed { - done <- outcome{err: errClosed} - return - } - res, err := fn(ctx, e.conn) - done <- outcome{res: res, err: err} - }() - - select { - case o := <-done: - return o.res, o.err - case <-ctx.Done(): - return result{}, ctx.Err() - } -} - -// close waits for a running query, then refuses every later one. The caller -// closes the connection after this returns. -func (e *executor) close() { - e.mu.Lock() - defer e.mu.Unlock() - e.closed = true -} diff --git a/internal/serve/query_test.go b/internal/serve/query_test.go index 15eb89a0..d851c04c 100644 --- a/internal/serve/query_test.go +++ b/internal/serve/query_test.go @@ -7,7 +7,6 @@ import ( "testing" "time" - "github.com/apache/arrow-adbc/go/adbc" "github.com/turbolytics/sql-flow/internal/config" "github.com/turbolytics/sql-flow/internal/coverage" "github.com/turbolytics/sql-flow/internal/errs" @@ -19,16 +18,16 @@ import ( // runner does not bring it under them; range(200000000) alone took 320 ms. const slowSQL = "SELECT sum(hash(a.range * b.range)) AS n FROM range(10000) a, range(10000) b" -func mustPrepare(t *testing.T, conn adbc.Connection, sql string, params ...config.ServeParam) *statement { +func mustPrepare(t *testing.T, ex Executor, sql string, params ...config.ServeParam) Statement { t.Helper() - st, err := prepare(context.Background(), conn, "ds", "", sql, params) + st, err := ex.Prepare(context.Background(), StatementSpec{Dataset: "ds", SQL: sql, Params: params}) assert.NoError(t, err) return st } -func queryRows(t *testing.T, conn adbc.Connection, st *statement, values map[string]any) []map[string]any { +func queryRows(t *testing.T, ex Executor, st Statement, values map[string]any) []map[string]any { t.Helper() - res, err := st.query(context.Background(), conn, values, 100) + res, _, err := query(context.Background(), ex, st, values, 100) assert.NoError(t, err) return decodeRows(t, res) } @@ -37,15 +36,16 @@ func queryRows(t *testing.T, conn adbc.Connection, st *statement, values map[str // first request. func TestCliServe_PrepareFailsAtStartupNamingTheStatement(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) + ex, _ := newExec(t, 1) - _, err := prepare(context.Background(), conn, "posts", "1h", "SELECT * FROM no_such_table", nil) + _, err := ex.Prepare(context.Background(), StatementSpec{ + Dataset: "posts", Grain: "1h", SQL: "SELECT * FROM no_such_table"}) assert.Error(t, err) assert.Equal(t, errs.CodeSQLInvalid, errs.CodeOf(err)) assert.That(t, strings.Contains(err.Error(), "dataset posts grain 1h")) assert.That(t, strings.Contains(err.Error(), "no_such_table")) - _, err = prepare(context.Background(), conn, "posts", "", "SELECT $1", nil) + _, err = ex.Prepare(context.Background(), StatementSpec{Dataset: "posts", SQL: "SELECT $1"}) assert.Equal(t, errs.CodeConfigServeDataset, errs.CodeOf(err)) } @@ -55,10 +55,13 @@ func TestCliServe_PrepareFailsAtStartupNamingTheStatement(t *testing.T) { // server refuses to start. func TestCliServe_PrepareRefusesAParamCountMismatch(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) + ex, _ := newExec(t, 1) - _, err := prepare(context.Background(), conn, "posts", "", `SELECT E'it\'s $fake' AS s`, - []config.ServeParam{{Name: "fake", Type: "string"}}) + _, err := ex.Prepare(context.Background(), StatementSpec{ + Dataset: "posts", + SQL: `SELECT E'it\'s $fake' AS s`, + Params: []config.ServeParam{{Name: "fake", Type: "string"}}, + }) assert.Error(t, err) assert.Equal(t, errs.CodeSQLInvalid, errs.CodeOf(err)) assert.That(t, strings.Contains(err.Error(), "DuckDB counts 0 parameters and sqlflow counts 1")) @@ -69,13 +72,13 @@ func TestCliServe_PrepareRefusesAParamCountMismatch(t *testing.T) { // rows. func TestCliServe_BindsEachValueToItsOwnPlaceholder(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) + ex, _ := newExec(t, 1) - st := mustPrepare(t, conn, "SELECT $b AS b, $a AS a, $a || '!' AS a2", + st := mustPrepare(t, ex, "SELECT $b AS b, $a AS a, $a || '!' AS a2", config.ServeParam{Name: "a", Type: "string"}, config.ServeParam{Name: "b", Type: "string"}) - rows := queryRows(t, conn, st, map[string]any{"a": "from a", "b": "from b"}) + rows := queryRows(t, ex, st, map[string]any{"a": "from a", "b": "from b"}) assert.Equal(t, "from b", rows[0]["b"]) assert.Equal(t, "from a", rows[0]["a"]) assert.Equal(t, "from a!", rows[0]["a2"]) @@ -85,9 +88,9 @@ func TestCliServe_BindsEachValueToItsOwnPlaceholder(t *testing.T) { // default and the column keeps its type. func TestCliServe_AnAbsentParamBindsATypedNull(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) + ex, _ := newExec(t, 1) - st := mustPrepare(t, conn, `SELECT + st := mustPrepare(t, ex, `SELECT coalesce($since, TIMESTAMPTZ '2000-01-01 00:00:00+00') AS since, coalesce($n, 7) AS n, coalesce($s, 'default') AS s`, @@ -95,7 +98,7 @@ func TestCliServe_AnAbsentParamBindsATypedNull(t *testing.T) { config.ServeParam{Name: "n", Type: "integer"}, config.ServeParam{Name: "s", Type: "string"}) - res, err := st.query(context.Background(), conn, map[string]any{}, 10) + res, _, err := query(context.Background(), ex, st, map[string]any{}, 10) assert.NoError(t, err) assert.Equal(t, "TIMESTAMP WITH TIME ZONE", res.Columns[0].Type) assert.Equal(t, "BIGINT", res.Columns[1].Type) @@ -105,7 +108,7 @@ func TestCliServe_AnAbsentParamBindsATypedNull(t *testing.T) { assert.Equal(t, "default", rows[0]["s"]) since := time.Date(2026, 9, 10, 12, 30, 0, 0, time.FixedZone("EDT", -4*3600)) - rows = queryRows(t, conn, st, map[string]any{"since": since, "n": int64(-3), "s": "given"}) + rows = queryRows(t, ex, st, map[string]any{"since": since, "n": int64(-3), "s": "given"}) assert.Equal(t, "2026-09-10T16:30:00Z", rows[0]["since"]) assert.Equal(t, float64(-3), rows[0]["n"]) assert.Equal(t, "given", rows[0]["s"]) @@ -115,74 +118,83 @@ func TestCliServe_AnAbsentParamBindsATypedNull(t *testing.T) { // plan, built against an empty table, and returned nothing forever after. func TestCliServe_ARequestSeesTheTableAsItIsNow(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) - execSQL(t, conn, "CREATE TABLE t (n BIGINT)") + ex, db := newExec(t, 1, "CREATE TABLE t (n BIGINT)") - st := mustPrepare(t, conn, "SELECT n FROM t WHERE n >= coalesce($min, 0) ORDER BY n", + // A writer outside the pool, the way an attached source changes under a + // running server. + writer, err := db.Connect(context.Background()) + assert.NoError(t, err) + t.Cleanup(func() { _ = writer.Close() }) + + st := mustPrepare(t, ex, "SELECT n FROM t WHERE n >= coalesce($min, 0) ORDER BY n", config.ServeParam{Name: "min", Type: "integer"}) - assert.Equal(t, 0, len(queryRows(t, conn, st, nil))) + assert.Equal(t, 0, len(queryRows(t, ex, st, nil))) - execSQL(t, conn, "INSERT INTO t VALUES (1), (2), (3)") - assert.Equal(t, 3, len(queryRows(t, conn, st, nil))) - assert.Equal(t, 1, len(queryRows(t, conn, st, map[string]any{"min": int64(3)}))) + execSQL(t, writer, "INSERT INTO t VALUES (1), (2), (3)") + assert.Equal(t, 3, len(queryRows(t, ex, st, nil))) + assert.Equal(t, 1, len(queryRows(t, ex, st, map[string]any{"min": int64(3)}))) } // DuckDB cannot be cancelled, so the deadline bounds the caller's wait and -// the query keeps the lock. A request right behind it waits on the lock and -// times out too. Once the query finishes, the next request answers. -func TestCliServe_ATimeoutReturnsWhileTheQueryHoldsTheLock(t *testing.T) { +// the query keeps its session. A request right behind it waits for that +// session and times out too. Once the query finishes, the next request +// answers. This is a pool of one, which is the behaviour serve had before +// there was a pool. +func TestCliServe_ATimeoutReturnsWhileTheQueryHoldsTheSession(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) - exec := &executor{conn: conn} - - slow := mustPrepare(t, conn, slowSQL) - fast := mustPrepare(t, conn, "SELECT 1 AS n") - runStatement := func(st *statement) func(context.Context, adbc.Connection) (result, error) { - return func(ctx context.Context, conn adbc.Connection) (result, error) { - return st.query(ctx, conn, nil, 10) - } + ex, _ := newExec(t, 1) + + slow := mustPrepare(t, ex, slowSQL) + fast := mustPrepare(t, ex, "SELECT 1 AS n") + run := func(st Statement, timeout time.Duration) (result, error) { + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + res, _, err := query(ctx, ex, st, nil, 10) + return res, err } start := time.Now() - _, err := exec.run(context.Background(), 100*time.Millisecond, runStatement(slow)) + _, err := run(slow, 100*time.Millisecond) assert.That(t, errors.Is(err, context.DeadlineExceeded)) assert.That(t, time.Since(start) < 500*time.Millisecond) - _, err = exec.run(context.Background(), 100*time.Millisecond, runStatement(fast)) + _, err = run(fast, 100*time.Millisecond) assert.That(t, errors.Is(err, context.DeadlineExceeded)) - res, err := exec.run(context.Background(), 60*time.Second, runStatement(fast)) + res, err := run(fast, 60*time.Second) assert.NoError(t, err) assert.Equal(t, 1, res.RowCount) } -// Shutdown closes the connection after close returns, so close must wait for -// the running query, and nothing may start on the connection afterwards. +// Shutdown closes the sessions after Close returns, so Close must wait for +// the running query, and nothing may start on a session afterwards. func TestCliServe_CloseWaitsForTheRunningQuery(t *testing.T) { coverage.Covers(t, "cli.serve") - conn := newConn(t) - exec := &executor{conn: conn} - slow := mustPrepare(t, conn, slowSQL) + ex, _ := newExec(t, 1) + slow := mustPrepare(t, ex, slowSQL) + + sess, err := ex.Acquire(context.Background()) + assert.NoError(t, err) finished := make(chan time.Time, 1) go func() { - _, _ = exec.run(context.Background(), 60*time.Second, - func(ctx context.Context, conn adbc.Connection) (result, error) { - defer func() { finished <- time.Now() }() - return slow.query(ctx, conn, nil, 10) - }) + rdr, err := sess.Run(context.Background(), slow, nil) + if err == nil { + _, _ = readRows(context.Background(), rdr, 10) + rdr.Release() + } + // Recorded before the session goes back, so a Close that returned + // first would be a Close that did not wait. + finished <- time.Now() + sess.Release() }() time.Sleep(50 * time.Millisecond) - exec.close() + ex.Close() closedAt := time.Now() queryDone := <-finished assert.That(t, !closedAt.Before(queryDone)) - _, err := exec.run(context.Background(), time.Second, - func(context.Context, adbc.Connection) (result, error) { - t.Fatal("a query ran after close") - return result{}, nil - }) - assert.That(t, errors.Is(err, errClosed)) + _, err = ex.Acquire(context.Background()) + assert.That(t, errors.Is(err, ErrClosed)) } diff --git a/internal/serve/server.go b/internal/serve/server.go index 7136d5c0..e57d12c3 100644 --- a/internal/serve/server.go +++ b/internal/serve/server.go @@ -10,26 +10,34 @@ import ( "context" "encoding/json" "fmt" + prom "github.com/prometheus/client_golang/prometheus" "time" - "github.com/apache/arrow-adbc/go/adbc" "github.com/turbolytics/sql-flow/internal/config" "go.uber.org/zap" ) -// Server serves one serve config's datasets over one DuckDB connection. +// Server serves one serve config's datasets over a pool of executor sessions. type Server struct { conf *config.ServeConf logger *zap.Logger - exec *executor + exec Executor + // health is the statement /healthz runs. It is prepared at startup like + // every other one, so the probe costs a bind and nothing more. + health Statement datasets map[string]*dataset // listing is the /v1/datasets body, built once: the config does not change. listing []byte // origins is nil without a cors block, which sends no CORS header at all. origins map[string]bool - // healthTimeout bounds /healthz, which waits on the same lock as a query. + // healthTimeout bounds /healthz, which waits for a session like a query. healthTimeout time.Duration + // registry and metrics are nil unless the config asks for /metrics. Every + // metrics method tolerates a nil receiver, so the request path records + // unconditionally. + registry *prom.Registry + metrics *metrics // now is when a request arrived: the until a ranged request does not // give. A test fixes it. now func() time.Time @@ -39,15 +47,25 @@ type Server struct { type dataset struct { conf config.ServeDataset // single is the statement of a dataset without grains. - single *statement + single datasetStatement // grains holds one statement per grain. - grains map[string]*statement + grains map[string]datasetStatement maxRows int timeout time.Duration // span is nil for a dataset without a range. span *span } +// datasetStatement is one prepared statement with the grain it answers. A +// Statement belongs to the executor and carries no grain of its own, so the +// grain -- which the response and the log line both name -- travels beside +// it. +type datasetStatement struct { + stmt Statement + // grain is empty for a dataset without grains. + grain string +} + // span is a dataset's range: the params that bound it, the width a request // gets without since, and each grain's widest range, narrowest first. type span struct { @@ -69,13 +87,27 @@ func WithLogger(l *zap.Logger) Option { return func(s *Server) { s.logger = l } } -// New checks the config's rules and prepares every statement against conn. +// WithMetrics registers serve's instruments on reg and serves them at +// /metrics. New builds the instruments, because they read the executor's +// session counts, and installs the wait hook on the executor afterwards. +func WithMetrics(reg *prom.Registry) Option { + return func(s *Server) { s.registry = reg } +} + +// waitObserver is an executor whose pool can report how long an Acquire +// waited. The DuckDB one does; a future one need not, and then the wait +// histogram is simply empty rather than the server failing to start. +type waitObserver interface { + setOnWait(func(time.Duration)) +} + +// New checks the config's rules and prepares every statement against ex. // -// conn must already carry whatever the config's commands attach. Any rule +// ex must already carry whatever the config's commands attach. Any rule // violation or statement that fails to prepare is returned, and the server // does not start: a dataset that cannot answer is a config error, and // finding it on the first request would find it in production. -func New(ctx context.Context, conf *config.ServeConf, conn adbc.Connection, opts ...Option) (*Server, error) { +func New(ctx context.Context, conf *config.ServeConf, ex Executor, opts ...Option) (*Server, error) { if err := conf.CheckError(); err != nil { return nil, err } @@ -83,7 +115,7 @@ func New(ctx context.Context, conf *config.ServeConf, conn adbc.Connection, opts s := &Server{ conf: conf, logger: zap.NewNop(), - exec: &executor{conn: conn}, + exec: ex, datasets: map[string]*dataset{}, healthTimeout: conf.Serve.Timeout(config.ServeDataset{}), now: time.Now, @@ -92,6 +124,26 @@ func New(ctx context.Context, conf *config.ServeConf, conn adbc.Connection, opts opt(s) } + // The instruments read the executor's session counts, so they are built + // here rather than by the caller, and the wait hook is installed on the + // pool before the server listens. + if s.registry != nil && conf.Serve.MetricsEnabled() { + m, err := newMetrics(s.registry, ex.Stats) + if err != nil { + return nil, err + } + s.metrics = m + if wo, ok := ex.(waitObserver); ok { + wo.setOnWait(m.observeWait) + } + } + + health, err := ex.Prepare(ctx, StatementSpec{Dataset: "healthz", SQL: "SELECT 1"}) + if err != nil { + return nil, err + } + s.health = health + if http := conf.Serve.HTTP; http != nil && http.CORS != nil { s.origins = map[string]bool{} for _, origin := range http.CORS.AllowedOrigins { @@ -112,20 +164,22 @@ func New(ctx context.Context, conf *config.ServeConf, conn adbc.Connection, opts } for _, sc := range dc.Statements() { - st, err := prepare(ctx, conn, dc.Name, sc.Grain, sc.SQL, dc.Params) + st, err := ex.Prepare(ctx, StatementSpec{ + Dataset: dc.Name, Grain: sc.Grain, SQL: sc.SQL, Params: dc.Params, + }) if err != nil { return nil, err } if sc.Grain == "" { - ds.single = st + ds.single = datasetStatement{stmt: st} doc.SQL = sc.SQL continue } if ds.grains == nil { - ds.grains = map[string]*statement{} + ds.grains = map[string]datasetStatement{} doc.Grains = map[string]grainDoc{} } - ds.grains[sc.Grain] = st + ds.grains[sc.Grain] = datasetStatement{stmt: st, grain: sc.Grain} doc.Grains[sc.Grain] = grainDoc{SQL: sc.SQL} } @@ -156,10 +210,10 @@ func New(ctx context.Context, conf *config.ServeConf, conn adbc.Connection, opts return s, nil } -// Close waits for a running query and refuses every later one. Call it after -// Serve returns and before closing the connection. +// Close waits for every running query, then closes the executor's sessions. +// Call it after Serve returns and before closing the database. func (s *Server) Close() { - s.exec.close() + s.exec.Close() } // datasetListing is the /v1/datasets body. The SQL is as the config wrote diff --git a/internal/validate/schemas/serve.json b/internal/validate/schemas/serve.json index b951910d..d6aa5f73 100644 --- a/internal/validate/schemas/serve.json +++ b/internal/validate/schemas/serve.json @@ -122,6 +122,28 @@ "type": "object", "description": "Limits for every dataset. A dataset's own non-zero limit overrides one." }, + "pool": { + "properties": { + "size": { + "type": "integer", + "description": "Sessions the server holds. Each is one backend session, and one\nrequest uses one at a time, so this is the requests that run at once.\n0 means the default, 4." + } + }, + "additionalProperties": false, + "type": "object", + "description": "How many requests the server runs at once. Omit for the default." + }, + "metrics": { + "properties": { + "enabled": { + "type": "boolean", + "description": "Serve GET /metrics on the same listener as the datasets, without a\ntoken. Off by default: that listener is public, and the metric labels\nname every dataset and grain." + } + }, + "additionalProperties": false, + "type": "object", + "description": "Whether to serve Prometheus metrics at /metrics." + }, "datasets": { "items": { "properties": {