From 89e62ff78061e98b89069c5452b022e6a90b05fb Mon Sep 17 00:00:00 2001 From: James Greenhill Date: Mon, 10 Aug 2026 23:29:09 +0000 Subject: [PATCH] Allow mid-session SET of duckgres.worker_ttl duckgres.worker_ttl (how long a worker stays hot-idle/warm after its last session ends) was connect-time only: clients that cannot set libpq startup options had no way to extend warm retention, so their worker was reaped at the default TTL whenever the connection went idle. Add SET/SHOW/RESET support on all three protocol paths (simple, batch, extended), mirroring duckgres.s3_cache: - the transpiler intercepts and validates the GUC (22023 on a bad duration, never forwarded to DuckDB, never intercepted inside a multi-statement batch) - the connection layer applies it through a control-plane hook that updates the bound worker's pool-side hot-idle TTL, gated on DUCKGRES_K8S_ALLOW_CLIENT_WORKER_PROFILE and clamped to DUCKGRES_K8S_WORKER_MAX_TTL exactly like the startup option - session state only flips after the apply succeeds, so SHOW never reports a TTL the worker won't park with; SET acquires a worker on lazily-activated connections and the override is re-applied across exploratory-tier worker switches - standalone/process backends get session-state-only SET/SHOW (they have no per-worker hot-idle TTL) --- CLAUDE.md | 13 +- README.md | 26 ++ controlplane/control.go | 8 + controlplane/k8s_pool_lifecycle.go | 30 ++ controlplane/k8s_pool_worker_ttl_test.go | 79 ++++ controlplane/multitenant.go | 16 +- controlplane/org_reserved_pool.go | 11 + controlplane/session_mgr.go | 43 +++ controlplane/worker_pool.go | 16 + controlplane/worker_profile.go | 14 + controlplane/worker_ttl.go | 70 ++++ controlplane/worker_ttl_test.go | 223 +++++++++++ docs/design/worker-ttl-pool.md | 6 + server/conn.go | 47 +++ server/conn_extended_query.go | 41 +++ server/conn_query_exec.go | 25 ++ server/conn_tier.go | 55 ++- server/conn_worker_ttl.go | 168 +++++++++ server/exports.go | 12 + server/worker_ttl_test.go | 448 +++++++++++++++++++++++ transpiler/config.go | 16 + transpiler/transform/setshow.go | 137 +++++-- transpiler/transform/transform.go | 14 + transpiler/transpiler.go | 11 +- transpiler/worker_ttl_test.go | 154 ++++++++ 25 files changed, 1615 insertions(+), 68 deletions(-) create mode 100644 controlplane/k8s_pool_worker_ttl_test.go create mode 100644 controlplane/worker_ttl.go create mode 100644 controlplane/worker_ttl_test.go create mode 100644 server/conn_worker_ttl.go create mode 100644 server/worker_ttl_test.go create mode 100644 transpiler/worker_ttl_test.go diff --git a/CLAUDE.md b/CLAUDE.md index e85ea0e4..eb04ee01 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -377,13 +377,16 @@ connect) and **escalation** (small → standard, one-way). Env-only knobs: answers itself — `SET`/`SHOW duckgres.query_source`, ignored SETs, no-ops, `pg_stat_activity`, the empty query — MUST NOT acquire. That is the point of the whole feature; adding an acquire to an engine-free path silently deletes - the benefit. The `duckgres.s3_cache` GUC is the exception, on all three - protocol paths, because unlike the other duckgres GUCs it is WORKER state (it - rebuilds the worker's `ducklake_s3` secret): **`SET` always acquires** — the - swap needs a worker to apply to — and **`SHOW` acquires only when + the benefit. The `duckgres.s3_cache` and `duckgres.worker_ttl` GUCs are the + exception, on all three protocol paths, because unlike the other duckgres + GUCs they are WORKER state (the s3_cache secret swap; the worker_ttl + pool-side hot-idle TTL): **`SET` always acquires** — the apply needs a + worker to land on — and **`SHOW duckgres.s3_cache` acquires only when `c.hasPendingS3Cache`**, i.e. when a connect-time option has not been applied yet and answering first would report a transport the session is about to - leave (see the s3_cache section below). + leave (see the s3_cache section below). `SHOW duckgres.worker_ttl` never + acquires: until a worker exists, the connect-time baseline is the truthful + answer (there is no pending worker-side TTL state). - **A pinning FIRST statement acquires the standard profile directly** (one acquire, `pinned=true` → `MarkConnectionPinned`), never small-then-escalate. - **The pin set is the state boundary, and every member is load-bearing:** DML, diff --git a/README.md b/README.md index 958124e6..f1fbc5c9 100644 --- a/README.md +++ b/README.md @@ -411,6 +411,32 @@ Requests must be positive and no greater than the configured maximum. Leaving the maximum unset disables client overrides, and clients cannot request an unlimited timeout because idle sessions retain worker capacity. +### Per-session worker TTL + +On the remote/K8s backend, a worker whose last session ends is parked +`hot_idle` (warm, quickly reusable by the same org) and retired once its TTL +expires — 1 minute by default. Clients can override that TTL per connection, +either at connect time or mid-session: + +```bash +PGOPTIONS='-c duckgres.worker_ttl=20m' psql "host= dbname=ducklake sslmode=require" +``` + +```sql +SET duckgres.worker_ttl = '20m'; -- Go duration; 0s retires the worker at session end +SHOW duckgres.worker_ttl; -- the TTL this session's worker will park with +RESET duckgres.worker_ttl; -- back to the connect-time value +``` + +The mid-session form exists for clients that cannot set startup options; it +takes effect on the bound worker immediately and governs the park when the +session ends. Both forms are gated on +`DUCKGRES_K8S_ALLOW_CLIENT_WORKER_PROFILE` (a mid-session `SET` is rejected +with 22023 when the gate is off) and clamped to +`DUCKGRES_K8S_WORKER_MAX_TTL`. The TTL is stamped with whole-minute precision, +so sub-minute overrides round down. On the standalone/process backends there is +no hot-idle TTL to override; `SET`/`SHOW` are accepted as session state only. + ### PostHog Logging Duckgres can optionally export structured logs to [PostHog Logs](https://posthog.com/docs/logs) via the OpenTelemetry Protocol (OTLP). Logs are always written to stderr regardless of this setting. diff --git a/controlplane/control.go b/controlplane/control.go index 67ce0e4f..e20b85d7 100644 --- a/controlplane/control.go +++ b/controlplane/control.go @@ -1524,6 +1524,14 @@ func (cp *ControlPlane) handleConnection(conn net.Conn) { } cc := server.NewClientConn(cp.srv, tlsConn, reader, writer, username, orgID, database, applicationName, sessionExec, pid, secretKey, workerID, workerPod) server.SetConnectionIdleTimeout(cc, clientIdleTimeout) + // Install the mid-session `SET duckgres.worker_ttl` capability + // (remote/k8s only — the process backend and standalone have no per-worker + // hot-idle TTL to override, so there SET/SHOW are session-state-only). + // Apply/Current resolve the session's worker by pid at call time, so they + // follow the worker across lazy activation and tier escalation. + if cp.isRemoteBackend { + server.SetConnectionWorkerTTLControl(cc, cp.workerTTLControlFor(sessions, pid, initialProfile, clog)) + } // Stamp the PostHog team id (config-snapshot read, no I/O) so this // connection's product-analytics events carry a PostHog-native key. Same // resolution the compute meter uses: the connecting user's team, else the diff --git a/controlplane/k8s_pool_lifecycle.go b/controlplane/k8s_pool_lifecycle.go index 0cfb5ca5..39a0b003 100644 --- a/controlplane/k8s_pool_lifecycle.go +++ b/controlplane/k8s_pool_lifecycle.go @@ -139,6 +139,36 @@ func (p *K8sWorkerPool) RetireIfDrainingAndEmpty(id int, origin LifecycleOrigin) go p.retireWorkerPod(id, w) } +// SetWorkerTTL overrides the TTL stamped on the worker's record when it next +// parks hot → hot_idle — the pool-side half of `SET duckgres.worker_ttl`. +// In-memory only is sufficient: workerRecordFor re-reads profile.TTL on every +// persist, the park write (commitHotIdleLocked) is what the expiry queries +// consult, and ttl_minutes on a HOT row is never used for reaping. Returns +// false when the worker is gone. +func (p *K8sWorkerPool) SetWorkerTTL(id int, ttl time.Duration) bool { + p.mu.Lock() + defer p.mu.Unlock() + w, ok := p.workers[id] + if !ok { + return false + } + w.profile.TTL = ttl + return true +} + +// WorkerTTL reports the TTL the worker would park with if its last session +// ended now (0 = the deployment default applies at reap time). ok=false when +// the worker is gone. +func (p *K8sWorkerPool) WorkerTTL(id int) (time.Duration, bool) { + p.mu.RLock() + defer p.mu.RUnlock() + w, ok := p.workers[id] + if !ok { + return 0, false + } + return w.profile.TTL, true +} + // TransitionToHotIdleIfNoSessions decrements the worker's active session count // and transitions a hot worker to hot_idle when its last session ends. The worker // keeps its org assignment and DuckLake attachment so it can be quickly reclaimed diff --git a/controlplane/k8s_pool_worker_ttl_test.go b/controlplane/k8s_pool_worker_ttl_test.go new file mode 100644 index 00000000..8d66feea --- /dev/null +++ b/controlplane/k8s_pool_worker_ttl_test.go @@ -0,0 +1,79 @@ +//go:build kubernetes + +package controlplane + +import ( + "testing" + "time" +) + +// TestK8sPoolSetWorkerTTL asserts the pool-side override behind +// SET duckgres.worker_ttl: the worker's profile TTL is updated in place (under +// the pool lock) and unknown workers are reported. +func TestK8sPoolSetWorkerTTL(t *testing.T) { + pool, _ := newTestK8sPool(t, 5) + worker := &ManagedWorker{ + ID: 5, + activeSessions: 1, + profile: WorkerProfile{CPU: "8", Memory: "16Gi", TTL: time.Minute}, + done: make(chan struct{}), + } + pool.workers[worker.ID] = worker + + if !pool.SetWorkerTTL(5, 20*time.Minute) { + t.Fatal("SetWorkerTTL(5) = false, want true") + } + ttl, ok := pool.WorkerTTL(5) + if !ok || ttl != 20*time.Minute { + t.Fatalf("WorkerTTL(5) = %s, %v; want 20m, true", ttl, ok) + } + + if pool.SetWorkerTTL(99, time.Minute) { + t.Fatal("SetWorkerTTL(99) = true, want false (unknown worker)") + } + if _, ok := pool.WorkerTTL(99); ok { + t.Fatal("WorkerTTL(99) ok=true, want false (unknown worker)") + } +} + +// TestK8sPoolSetWorkerTTLPersistsAtPark asserts the override actually governs +// reaping: when the worker's last session ends, the hot_idle record the +// reapers read carries the OVERRIDDEN ttl_minutes, not the connect-time one. +func TestK8sPoolSetWorkerTTLPersistsAtPark(t *testing.T) { + pool, _ := newTestK8sPool(t, 5) + store := &captureRuntimeWorkerStore{} + pool.runtimeStore = store + worker := &ManagedWorker{ + ID: 5, + activeSessions: 1, + profile: WorkerProfile{CPU: "8", Memory: "16Gi", TTL: time.Minute}, + done: make(chan struct{}), + } + if err := worker.SetSharedState(SharedWorkerState{ + Lifecycle: WorkerLifecycleHot, + Assignment: &WorkerAssignment{OrgID: "analytics"}, + }); err != nil { + t.Fatalf("SetSharedState: %v", err) + } + pool.workers[worker.ID] = worker + + if !pool.SetWorkerTTL(5, 20*time.Minute) { + t.Fatal("SetWorkerTTL(5) = false, want true") + } + if !pool.TransitionToHotIdleIfNoSessions(worker.ID) { + t.Fatal("expected the worker to park to hot_idle") + } + + var hotIdle *int + for i := range store.records { + if store.records[i].State == "hot_idle" { + hotIdle = &store.records[i].TTLMinutes + } + } + if hotIdle == nil { + t.Fatal("no hot_idle record persisted at park") + } + if *hotIdle != 20 { + t.Fatalf("parked record ttl_minutes = %d, want 20 (the override)", *hotIdle) + } +} diff --git a/controlplane/multitenant.go b/controlplane/multitenant.go index d4bd2fe6..23686932 100644 --- a/controlplane/multitenant.go +++ b/controlplane/multitenant.go @@ -99,19 +99,9 @@ func (a *orgRouterAdapter) MetadataProxySessions() *metadataProxySessionRegistry return a.metadataSessions } -// effectiveDefaultWorkerTTL resolves the janitor's hot-idle retention: the -// operator default TTL (DUCKGRES_K8S_WORKER_DEFAULT_TTL → -// K8sConfig.WorkerDefaultTTL) when set, otherwise the single built-in -// defaultWorkerTTL (1m — the same fallback sized-but-no-ttl requests get at -// profile resolution, so there is exactly ONE default TTL however a worker -// came to have no explicit one). The full per-request precedence is: -// client GUC > org default > deployment default TTL > built-in 1m. -func effectiveDefaultWorkerTTL(configured time.Duration) time.Duration { - if configured > 0 { - return configured - } - return defaultWorkerTTL -} +// effectiveDefaultWorkerTTL moved to worker_profile.go (untagged) so the +// duckgres.worker_ttl session-GUC hook can resolve the same default in every +// build flavor. func (a *orgRouterAdapter) StackForOrg(orgID string) (WorkerPool, *SessionManager, *MemoryRebalancer, bool) { stack, ok := a.router.StackForOrg(orgID) diff --git a/controlplane/org_reserved_pool.go b/controlplane/org_reserved_pool.go index bb37d3c5..acfd75a8 100644 --- a/controlplane/org_reserved_pool.go +++ b/controlplane/org_reserved_pool.go @@ -43,6 +43,17 @@ func NewOrgReservedPool(shared *K8sWorkerPool, orgID string, maxWorkers int, ima return pool } +// SetWorkerTTL delegates the duckgres.worker_ttl session-GUC override to the +// shared pool (the org's workers live in shared.workers). +func (p *OrgReservedPool) SetWorkerTTL(id int, ttl time.Duration) bool { + return p.shared.SetWorkerTTL(id, ttl) +} + +// WorkerTTL delegates to the shared pool; see K8sWorkerPool.WorkerTTL. +func (p *OrgReservedPool) WorkerTTL(id int) (time.Duration, bool) { + return p.shared.WorkerTTL(id) +} + func (p *OrgReservedPool) AcquireWorker(ctx context.Context, profile *WorkerProfile) (worker *ManagedWorker, err error) { // End-to-end acquire latency (fast path included), observed via a named- // return defer so every exit — reuse, capacity error, ctx cancel — lands in diff --git a/controlplane/session_mgr.go b/controlplane/session_mgr.go index 104cfda4..931c668e 100644 --- a/controlplane/session_mgr.go +++ b/controlplane/session_mgr.go @@ -945,6 +945,49 @@ func (sm *SessionManager) WorkerIDForPID(pid int32) int { return -1 } +// SetWorkerTTLForPID overrides the hot-idle TTL of the worker bound to pid's +// session — the apply half of the `duckgres.worker_ttl` session GUC. It is a +// no-op success when the pool has no per-worker TTL (the process backend: +// its idle reaping is pool-global, and the GUC hook is only installed for the +// remote backend anyway). An error means the override did NOT take effect +// (the session or its worker is gone), so the caller must not update its +// session state. +func (sm *SessionManager) SetWorkerTTLForPID(pid int32, ttl time.Duration) error { + sm.mu.RLock() + s, ok := sm.sessions[pid] + sm.mu.RUnlock() + if !ok { + return fmt.Errorf("no session for pid %d", pid) + } + pool, ok := sm.pool.(workerTTLPool) + if !ok { + return nil + } + if !pool.SetWorkerTTL(s.WorkerID, ttl) { + return fmt.Errorf("worker %d is no longer in the pool", s.WorkerID) + } + return nil +} + +// WorkerTTLForPID reports the hot-idle TTL the worker bound to pid's session +// would park with right now — the SHOW half of the `duckgres.worker_ttl` +// session GUC. ok=false when there is no session, no worker, or the pool has +// no per-worker TTL. A zero TTL with ok=true means the deployment default +// applies at reap time (the caller resolves it). +func (sm *SessionManager) WorkerTTLForPID(pid int32) (time.Duration, bool) { + sm.mu.RLock() + s, ok := sm.sessions[pid] + sm.mu.RUnlock() + if !ok { + return 0, false + } + pool, ok := sm.pool.(workerTTLPool) + if !ok { + return 0, false + } + return pool.WorkerTTL(s.WorkerID) +} + // SessionForWorker returns the session bound to the given cluster-unique worker // id, or ok=false if this stack has none. One session per worker, so the first // (only) pid in the worker's index is authoritative. Used by the admin diff --git a/controlplane/worker_pool.go b/controlplane/worker_pool.go index c1f8eb13..ea8b8a42 100644 --- a/controlplane/worker_pool.go +++ b/controlplane/worker_pool.go @@ -89,6 +89,22 @@ type WorkerPool interface { ShutdownAll() } +// workerTTLPool is the optional WorkerPool capability behind the +// `duckgres.worker_ttl` session GUC: reading and overriding the hot-idle TTL +// a worker parks with when its last session ends. K8sWorkerPool implements it +// (and OrgReservedPool delegates to the shared pool); the process-backend +// FlightWorkerPool deliberately does not — its idle reaping is pool-global +// (WorkerIdleTimeout), not per-worker, so there is nothing to override. +type workerTTLPool interface { + // SetWorkerTTL overrides the hot-idle TTL stamped on the worker's record + // when it next parks hot → hot_idle. false when the worker is gone. + SetWorkerTTL(id int, ttl time.Duration) bool + // WorkerTTL reports the worker's current hot-idle TTL (0 = the + // deployment default applies at reap time). ok=false when the worker is + // gone. + WorkerTTL(id int) (time.Duration, bool) +} + // K8sWorkerPoolConfig holds the configuration for creating a K8sWorkerPool. type K8sWorkerPoolConfig struct { Namespace string diff --git a/controlplane/worker_profile.go b/controlplane/worker_profile.go index a7356c06..e79d1936 100644 --- a/controlplane/worker_profile.go +++ b/controlplane/worker_profile.go @@ -242,6 +242,20 @@ func firstNonEmpty(a, b string) string { return b } +// effectiveDefaultWorkerTTL resolves the hot-idle retention floor: the +// operator default TTL (DUCKGRES_K8S_WORKER_DEFAULT_TTL → +// K8sConfig.WorkerDefaultTTL) when set, otherwise the single built-in +// defaultWorkerTTL (1m — the same fallback sized-but-no-ttl requests get at +// profile resolution, so there is exactly ONE default TTL however a worker +// came to have no explicit one). The full per-request precedence is: +// client GUC > org default > deployment default TTL > built-in 1m. +func effectiveDefaultWorkerTTL(configured time.Duration) time.Duration { + if configured > 0 { + return configured + } + return defaultWorkerTTL +} + func requestedWorkerVCPUs(profile *WorkerProfile, workerCPURequest string) (int, error) { cpu := strings.TrimSpace(workerCPURequest) if profile != nil && strings.TrimSpace(profile.CPU) != "" { diff --git a/controlplane/worker_ttl.go b/controlplane/worker_ttl.go new file mode 100644 index 00000000..3ebe9034 --- /dev/null +++ b/controlplane/worker_ttl.go @@ -0,0 +1,70 @@ +package controlplane + +import ( + "context" + "log/slog" + "time" + + "github.com/posthog/duckgres/server" + "github.com/posthog/duckgres/transpiler/transform" +) + +// errWorkerTTLOverrideDisabled rejects a mid-session `SET duckgres.worker_ttl` +// when the deployment does not trust client-supplied worker settings +// (DUCKGRES_K8S_ALLOW_CLIENT_WORKER_PROFILE off) — the same trust boundary the +// duckgres.worker_* startup options have, except a startup option is silently +// ignored there while a SET must not pretend it took effect. +var errWorkerTTLOverrideDisabled = &transform.CodedError{ + Code: "22023", // invalid_parameter_value + Message: "duckgres.worker_ttl overrides are not enabled on this server", +} + +// sessionWorkerTTLBaseline is the connect-time TTL a session's SHOW falls back +// to and RESET restores: the session profile's TTL when one was resolved +// (startup GUC / org default / exploratory tier), else the deployment default +// TTL, else the built-in 1m — the same chain resolveWorkerProfile applies. +func sessionWorkerTTLBaseline(profile *WorkerProfile, k K8sConfig) time.Duration { + if profile != nil && profile.TTL > 0 { + return profile.TTL + } + return effectiveDefaultWorkerTTL(k.WorkerDefaultTTL) +} + +// workerTTLControlFor builds the per-connection server.WorkerTTLControl behind +// the mid-session `duckgres.worker_ttl` GUC. Apply updates the bound worker's +// pool-side hot-idle TTL (gated on AllowClientWorkerProfile, clamped to +// WorkerMaxTTL — both exactly like the startup option); Current lets SHOW +// report the TTL the bound worker would actually park with (a reused hot-idle +// worker can carry a previous request's TTL, which beats the baseline). +func (cp *ControlPlane) workerTTLControlFor(sessions *SessionManager, pid int32, initialProfile *WorkerProfile, clog *slog.Logger) *server.WorkerTTLControl { + return &server.WorkerTTLControl{ + Baseline: sessionWorkerTTLBaseline(initialProfile, cp.cfg.K8s), + Apply: func(_ context.Context, ttl time.Duration) (time.Duration, error) { + if !cp.cfg.K8s.AllowClientWorkerProfile { + return 0, errWorkerTTLOverrideDisabled + } + applied := ttl + if max := cp.cfg.K8s.WorkerMaxTTL; max > 0 && ttl > max { + clog.Warn("Clamped duckgres.worker_ttl override to the deployment maximum.", + "requested", ttl.String(), "max", max.String()) + applied = max + } + if err := sessions.SetWorkerTTLForPID(pid, applied); err != nil { + return 0, err + } + return applied, nil + }, + Current: func() (time.Duration, bool) { + ttl, ok := sessions.WorkerTTLForPID(pid) + if !ok { + return 0, false + } + if ttl <= 0 { + // A default-profile worker carries TTL 0 = "the deployment + // default applies at reap time"; resolve it for SHOW. + ttl = effectiveDefaultWorkerTTL(cp.cfg.K8s.WorkerDefaultTTL) + } + return ttl, true + }, + } +} diff --git a/controlplane/worker_ttl_test.go b/controlplane/worker_ttl_test.go new file mode 100644 index 00000000..87e3c2f1 --- /dev/null +++ b/controlplane/worker_ttl_test.go @@ -0,0 +1,223 @@ +package controlplane + +import ( + "context" + "errors" + "log/slog" + "testing" + "time" +) + +// ttlFakePool is a WorkerPool that also implements the worker-TTL capability, +// recording SetWorkerTTL calls. +type ttlFakePool struct { + workers map[int]*ManagedWorker + setCalls []int + setTTLs []time.Duration +} + +func (p *ttlFakePool) AcquireWorker(context.Context, *WorkerProfile) (*ManagedWorker, error) { + return nil, errors.New("not implemented") +} + +func (p *ttlFakePool) ReleaseWorker(int) {} + +func (p *ttlFakePool) RetireWorker(int) {} + +func (p *ttlFakePool) RetireWorkerIfNoSessions(int) bool { return false } + +func (p *ttlFakePool) Worker(id int) (*ManagedWorker, bool) { + w, ok := p.workers[id] + return w, ok +} + +func (p *ttlFakePool) SpawnMinWorkers(int) error { return nil } + +func (p *ttlFakePool) HealthCheckLoop(context.Context, time.Duration, WorkerCrashHandler, ProgressHandler) { +} + +func (p *ttlFakePool) SetMaxWorkers(int) {} + +func (p *ttlFakePool) ShutdownAll() {} + +func (p *ttlFakePool) SetWorkerTTL(id int, ttl time.Duration) bool { + if _, ok := p.workers[id]; !ok { + return false + } + p.setCalls = append(p.setCalls, id) + p.setTTLs = append(p.setTTLs, ttl) + return true +} + +func (p *ttlFakePool) WorkerTTL(id int) (time.Duration, bool) { + w, ok := p.workers[id] + if !ok { + return 0, false + } + return w.profile.TTL, true +} + +// ttlLessFakePool is a WorkerPool WITHOUT the worker-TTL capability (the +// process-backend shape). +type ttlLessFakePool struct{} + +func (p *ttlLessFakePool) AcquireWorker(context.Context, *WorkerProfile) (*ManagedWorker, error) { + return nil, errors.New("not implemented") +} + +func (p *ttlLessFakePool) ReleaseWorker(int) {} + +func (p *ttlLessFakePool) RetireWorker(int) {} + +func (p *ttlLessFakePool) RetireWorkerIfNoSessions(int) bool { return false } + +func (p *ttlLessFakePool) Worker(int) (*ManagedWorker, bool) { return nil, false } + +func (p *ttlLessFakePool) SpawnMinWorkers(int) error { return nil } + +func (p *ttlLessFakePool) HealthCheckLoop(context.Context, time.Duration, WorkerCrashHandler, ProgressHandler) { +} + +func (p *ttlLessFakePool) SetMaxWorkers(int) {} + +func (p *ttlLessFakePool) ShutdownAll() {} + +// TestSessionManagerSetWorkerTTLForPID asserts the session→worker routing of a +// duckgres.worker_ttl override: the pool is asked to stamp the TTL on the +// worker bound to the session's pid. +func TestSessionManagerSetWorkerTTLForPID(t *testing.T) { + pool := &ttlFakePool{workers: map[int]*ManagedWorker{5: {ID: 5}}} + sm := NewSessionManager(pool, nil) + sm.sessions[1001] = &ManagedSession{WorkerID: 5} + + if err := sm.SetWorkerTTLForPID(1001, 20*time.Minute); err != nil { + t.Fatalf("SetWorkerTTLForPID: %v", err) + } + if len(pool.setCalls) != 1 || pool.setCalls[0] != 5 || pool.setTTLs[0] != 20*time.Minute { + t.Fatalf("pool calls = %v/%v, want worker 5 with 20m", pool.setCalls, pool.setTTLs) + } + + // No session for the pid: an error, and the pool is untouched. + if err := sm.SetWorkerTTLForPID(1002, time.Minute); err == nil { + t.Fatal("SetWorkerTTLForPID with unknown pid: nil error, want failure") + } + if len(pool.setCalls) != 1 { + t.Fatalf("unknown pid reached the pool: calls = %v", pool.setCalls) + } + + // Session exists but the worker is gone (raced with retirement): an error. + sm.sessions[1003] = &ManagedSession{WorkerID: 99} + if err := sm.SetWorkerTTLForPID(1003, time.Minute); err == nil { + t.Fatal("SetWorkerTTLForPID with missing worker: nil error, want failure") + } +} + +// TestSessionManagerSetWorkerTTLPoolWithoutCapability asserts the process +// backend shape: a pool without the TTL capability makes the apply a no-op +// success (the hook is only installed for the remote backend anyway, so this +// is defensive). +func TestSessionManagerSetWorkerTTLPoolWithoutCapability(t *testing.T) { + sm := NewSessionManager(&ttlLessFakePool{}, nil) + sm.sessions[1001] = &ManagedSession{WorkerID: 5} + if err := sm.SetWorkerTTLForPID(1001, 20*time.Minute); err != nil { + t.Fatalf("SetWorkerTTLForPID on a TTL-less pool: %v, want no-op success", err) + } + if _, ok := sm.WorkerTTLForPID(1001); ok { + t.Fatal("WorkerTTLForPID on a TTL-less pool: ok=true, want false") + } +} + +// TestSessionManagerWorkerTTLForPID asserts the SHOW-facing read path reports +// the bound worker's current pool-side TTL. +func TestSessionManagerWorkerTTLForPID(t *testing.T) { + pool := &ttlFakePool{workers: map[int]*ManagedWorker{ + 5: {ID: 5, profile: WorkerProfile{CPU: "8", Memory: "16Gi", TTL: 7 * time.Minute}}, + }} + sm := NewSessionManager(pool, nil) + sm.sessions[1001] = &ManagedSession{WorkerID: 5} + + ttl, ok := sm.WorkerTTLForPID(1001) + if !ok || ttl != 7*time.Minute { + t.Fatalf("WorkerTTLForPID = %s, %v; want 7m, true", ttl, ok) + } + if _, ok := sm.WorkerTTLForPID(1002); ok { + t.Fatal("WorkerTTLForPID with unknown pid: ok=true, want false") + } +} + +// TestWorkerTTLControlForGateDisabled asserts the mid-session override honors +// the same trust boundary as the duckgres.worker_* startup options: with +// AllowClientWorkerProfile off the apply is rejected with 22023 and never +// reaches the pool. +func TestWorkerTTLControlForGateDisabled(t *testing.T) { + pool := &ttlFakePool{workers: map[int]*ManagedWorker{5: {ID: 5}}} + sm := NewSessionManager(pool, nil) + sm.sessions[1001] = &ManagedSession{WorkerID: 5} + + cp := &ControlPlane{} + cp.cfg.K8s.AllowClientWorkerProfile = false + ctl := cp.workerTTLControlFor(sm, 1001, nil, slog.Default()) + + if _, err := ctl.Apply(context.Background(), 20*time.Minute); err == nil { + t.Fatal("Apply with the gate off: nil error, want 22023 rejection") + } else { + var coded interface{ SQLState() string } + if !errors.As(err, &coded) || coded.SQLState() != "22023" { + t.Fatalf("Apply error = %v, want SQLSTATE 22023", err) + } + } + if len(pool.setCalls) != 0 { + t.Fatalf("gated apply reached the pool: calls = %v", pool.setCalls) + } +} + +// TestWorkerTTLControlForClamps asserts the apply honors the deployment's +// WorkerMaxTTL ceiling exactly like the startup option: the value stamped on +// the worker AND the value reported back to the session are the clamped one. +func TestWorkerTTLControlForClamps(t *testing.T) { + pool := &ttlFakePool{workers: map[int]*ManagedWorker{5: {ID: 5}}} + sm := NewSessionManager(pool, nil) + sm.sessions[1001] = &ManagedSession{WorkerID: 5} + + cp := &ControlPlane{} + cp.cfg.K8s.AllowClientWorkerProfile = true + cp.cfg.K8s.WorkerMaxTTL = time.Hour + ctl := cp.workerTTLControlFor(sm, 1001, nil, slog.Default()) + + applied, err := ctl.Apply(context.Background(), 24*time.Hour) + if err != nil { + t.Fatalf("Apply: %v", err) + } + if applied != time.Hour { + t.Fatalf("Apply returned %s, want the clamped 1h", applied) + } + if len(pool.setTTLs) != 1 || pool.setTTLs[0] != time.Hour { + t.Fatalf("pool TTLs = %v, want [1h]", pool.setTTLs) + } + + // Within the ceiling: applied as-is. + if _, err := ctl.Apply(context.Background(), 20*time.Minute); err != nil { + t.Fatalf("Apply(20m): %v", err) + } + if pool.setTTLs[1] != 20*time.Minute { + t.Fatalf("pool TTLs = %v, want [1h 20m]", pool.setTTLs) + } +} + +// TestSessionWorkerTTLBaseline pins the connect-time baseline SHOW falls back +// to: the session profile's TTL when one was resolved (sized / org default / +// exploratory), else the deployment default TTL, else the built-in 1m. +func TestSessionWorkerTTLBaseline(t *testing.T) { + var k K8sConfig + if got := sessionWorkerTTLBaseline(nil, k); got != defaultWorkerTTL { + t.Fatalf("baseline(nil profile) = %s, want built-in %s", got, defaultWorkerTTL) + } + k.WorkerDefaultTTL = 70 * time.Minute + if got := sessionWorkerTTLBaseline(nil, k); got != 70*time.Minute { + t.Fatalf("baseline(nil profile, deployment default) = %s, want 70m", got) + } + p := &WorkerProfile{CPU: "8", Memory: "16Gi", TTL: 48 * time.Hour} + if got := sessionWorkerTTLBaseline(p, k); got != 48*time.Hour { + t.Fatalf("baseline(concrete profile) = %s, want 48h", got) + } +} diff --git a/docs/design/worker-ttl-pool.md b/docs/design/worker-ttl-pool.md index 94d9e74b..d19bb944 100644 --- a/docs/design/worker-ttl-pool.md +++ b/docs/design/worker-ttl-pool.md @@ -50,6 +50,12 @@ options=-c duckgres.worker_cpu=8 -c duckgres.worker_memory=16Gi -c duckgres.work to `[min,max]` and ttl to `[0,maxTTL]` per deployment (out-of-range → clamp + warn). Gate off → every request uses the defaults. +`duckgres.worker_ttl` is additionally settable mid-session +(`SET duckgres.worker_ttl = '20m'` / `SHOW` / `RESET`) for clients that cannot +set startup options: the override updates the bound worker's pool-side TTL +(gated on `AllowClientWorkerProfile`, clamped to `WorkerMaxTTL`, like the +startup option) and is re-applied across exploratory-tier worker switches. + TTL resolution, per request (the same chain whether the request is sized or not — there is exactly ONE default TTL however a worker comes to have no explicit one): diff --git a/server/conn.go b/server/conn.go index 098b4a2e..23ba69a0 100644 --- a/server/conn.go +++ b/server/conn.go @@ -86,6 +86,8 @@ type preparedStmt struct { querySourceShow bool // True if this is SHOW duckgres.query_source (answered from session state) s3CacheSet *string // non-nil: SET duckgres.s3_cache; pointed-to value to apply to session s3CacheShow bool // True if this is SHOW duckgres.s3_cache (answered from session state) + workerTTLSet *string // non-nil: SET duckgres.worker_ttl; pointed-to value to apply to session + workerTTLShow bool // True if this is SHOW duckgres.worker_ttl (answered from session state) described bool // True if Describe(S) was called on this statement statements []string // Multi-statement rewrite (e.g., writable CTE) cleanupStatements []string // Cleanup statements for multi-statement (DROP temp tables, COMMIT) @@ -295,6 +297,17 @@ type clientConn struct { pendingS3Cache string hasPendingS3Cache bool + // workerTTLOverride holds the `duckgres.worker_ttl` session GUC state + // (another duckgres-namespaced custom parameter, NOT forwarded to + // DuckDB): non-nil when this session overrode how long its worker stays + // hot-idle after the session ends. Only flipped by applyWorkerTTLSetting + // AFTER the control plane stamped the value on the bound worker, so SHOW + // never reports a TTL the worker won't park with. workerTTLCtl is the + // control-plane capability behind the GUC (nil outside the remote/k8s + // backend, where SET/SHOW are session-state-only). See conn_worker_ttl.go. + workerTTLOverride *time.Duration + workerTTLCtl *WorkerTTLControl + // fatalErr parks a connection-terminating error raised inside an // extended-query handler. Those handlers are void (the protocol reports // their failures as ErrorResponse + skip-until-Sync), but a failed tier @@ -1682,6 +1695,40 @@ func (c *clientConn) handleQuery(body []byte) (retErr error) { return nil } + // Handle the duckgres.worker_ttl custom GUC (SET / SHOW). Intercepted + // session-side; applied by overriding the bound worker's pool-side + // hot-idle TTL, never forwarded to DuckDB. A failed apply fails the SET + // so the session state never diverges from the TTL the worker will + // actually park with. + if result.WorkerTTLSet != nil { + // Lazy activation: like duckgres.s3_cache this is WORKER state (the + // control plane's pool record for the bound worker), so it must have + // a worker to apply to. Not a pinning statement, so it acquires the + // exploratory worker. + if err := c.activateForStatement(query, false); err != nil { + return err + } + if err := c.applyWorkerTTLSetting(*result.WorkerTTLSet); err != nil { + c.sendError("ERROR", workerTTLApplyErrorSQLState(err), err.Error()) + } else { + _ = c.writeCommandComplete("SET") + } + _ = c.writeReadyForQuery(c.txStatus) + _ = c.flushWriter() + return nil + } + if result.WorkerTTLShow { + // No lazy activation: the connect-time baseline (or the built-in + // default) is the truthful answer until a worker exists — there is no + // pending worker-side state the way a connect-time s3_cache=off is. + _ = c.sendRowDescription([]string{WorkerTTLGUCName}, []ColumnTyper{staticColumnType("VARCHAR")}) + _ = c.sendDataRowWithFormats([]interface{}{c.workerTTLValue()}, nil, nil) + _ = c.writeCommandComplete("SHOW") + _ = c.writeReadyForQuery(c.txStatus) + _ = c.flushWriter() + return nil + } + // Handle ignored SET parameters if result.IsIgnoredSet { c.logger().Debug("Ignoring PostgreSQL-specific SET.", "query", query) diff --git a/server/conn_extended_query.go b/server/conn_extended_query.go index dae2f68e..a292e3c5 100644 --- a/server/conn_extended_query.go +++ b/server/conn_extended_query.go @@ -221,6 +221,8 @@ func (c *clientConn) handleParse(body []byte) { querySourceShow: result.QuerySourceShow, // SHOW duckgres.query_source s3CacheSet: result.S3CacheSet, // SET duckgres.s3_cache (custom GUC) s3CacheShow: result.S3CacheShow, // SHOW duckgres.s3_cache + workerTTLSet: result.WorkerTTLSet, // SET duckgres.worker_ttl (custom GUC) + workerTTLShow: result.WorkerTTLShow, // SHOW duckgres.worker_ttl statements: result.Statements, // Multi-statement rewrite (writable CTE) cleanupStatements: result.CleanupStatements, // Cleanup statements pinsWorker: pinsWorker, // Exploratory tier: escalate before Describe/Execute @@ -325,6 +327,17 @@ func (c *clientConn) handleDescribe(body []byte) { return } + // duckgres.worker_ttl custom GUC: same shape as query_source above. + if ps.workerTTLSet != nil { + _ = wire.WriteNoData(c.writer) + return + } + if ps.workerTTLShow { + _ = c.sendRowDescription([]string{WorkerTTLGUCName}, []ColumnTyper{staticColumnType("VARCHAR")}) + ps.described = true + return + } + // For queries that return results, we need to send RowDescription // For other queries, send NoData returnsResults := queryReturnsResults(ps.query) @@ -739,6 +752,34 @@ func (c *clientConn) handleExecute(body []byte) { return } + // duckgres.worker_ttl custom GUC (SET / SHOW): intercepted session-side, + // applied via the bound worker's pool-side hot-idle TTL override. + // Determined by the transpiler during Parse. A failed apply errors the + // Execute so the session state never diverges from the TTL the worker + // will actually park with. + if p.stmt.workerTTLSet != nil { + // Lazy activation: the override needs a worker to apply to (see the + // matching site in handleQuery). Not pinning, so the exploratory tier + // is enough. + if err := c.activateForStatement(p.stmt.query, false); err != nil { + return + } + if err := c.applyWorkerTTLSetting(*p.stmt.workerTTLSet); err != nil { + c.sendError("ERROR", workerTTLApplyErrorSQLState(err), err.Error()) + return + } + _ = c.writeCommandComplete("SET") + return + } + if p.stmt.workerTTLShow { + if !p.described { + _ = c.sendRowDescription([]string{WorkerTTLGUCName}, []ColumnTyper{staticColumnType("VARCHAR")}) + } + _ = c.sendDataRowWithFormats([]interface{}{c.workerTTLValue()}, p.resultFormats, nil) + _ = c.writeCommandComplete("SHOW") + return + } + // Check if this is a PostgreSQL-specific SET command that should be ignored // (determined by transpiler during Parse) if p.stmt.isIgnoredSet { diff --git a/server/conn_query_exec.go b/server/conn_query_exec.go index f55c9d4f..5b1e2cf4 100644 --- a/server/conn_query_exec.go +++ b/server/conn_query_exec.go @@ -702,6 +702,31 @@ func (c *clientConn) executeSingleStatement(query string) (errSent bool, fatalEr return false, nil } + // duckgres.worker_ttl custom GUC (SET / SHOW): intercepted session-side + // and applied via the bound worker's pool-side hot-idle TTL override. A + // failed apply aborts the rest of the batch — later statements may depend + // on the requested warm-retention state. + if result.WorkerTTLSet != nil { + // Lazy activation: the override needs a worker to apply to (see the + // matching site in handleQuery). Not pinning, so the exploratory tier + // is enough. + if err := c.activateForStatement(query, false); err != nil { + return false, err + } + if err := c.applyWorkerTTLSetting(*result.WorkerTTLSet); err != nil { + c.sendError("ERROR", workerTTLApplyErrorSQLState(err), err.Error()) + return true, nil + } + _ = c.writeCommandComplete("SET") + return false, nil + } + if result.WorkerTTLShow { + _ = c.sendRowDescription([]string{WorkerTTLGUCName}, []ColumnTyper{staticColumnType("VARCHAR")}) + _ = c.sendDataRowWithFormats([]interface{}{c.workerTTLValue()}, nil, nil) + _ = c.writeCommandComplete("SHOW") + return false, nil + } + if result.IsIgnoredSet { _ = c.writeCommandComplete("SET") return false, nil diff --git a/server/conn_tier.go b/server/conn_tier.go index d5930601..b96cf03b 100644 --- a/server/conn_tier.go +++ b/server/conn_tier.go @@ -174,6 +174,22 @@ var exploratoryEscalationsTotal = promauto.NewCounterVec(prometheus.CounterOpts{ // post-swap `duckgres.s3_cache` re-apply failed. See failS3CacheReapply. var errS3CacheReapplyFailed = errors.New("s3_cache re-apply after worker switch failed") +// errWorkerTTLReapplyFailed is the duckgres.worker_ttl twin of +// errS3CacheReapplyFailed: the escalation succeeded, only the post-swap TTL +// re-apply failed — also statement-scoped, never connection-fatal. +var errWorkerTTLReapplyFailed = errors.New("worker_ttl re-apply after worker switch failed") + +// workerTTLReapplyError tags a worker_ttl re-apply failure with +// errWorkerTTLReapplyFailed WITHOUT the sentinel's text landing in the +// client-visible message (mirrors s3CacheReapplyError). +type workerTTLReapplyError struct{ err error } + +func (e *workerTTLReapplyError) Error() string { return e.err.Error() } + +func (e *workerTTLReapplyError) Unwrap() error { return e.err } + +func (e *workerTTLReapplyError) Is(target error) bool { return target == errWorkerTTLReapplyFailed } + // s3CacheReapplyError tags a re-apply failure with errS3CacheReapplyFailed // WITHOUT the sentinel's text landing in the client-visible message — the // wrapped error already names the GUC and the worker switch. @@ -229,6 +245,15 @@ func (c *clientConn) escalateWorker(ctx context.Context, reason string) error { c.s3CacheOff = false return &s3CacheReapplyError{err: err} } + // Same re-assertion for a duckgres.worker_ttl override: it is pool-side + // per-worker state, and the new worker's profile carries the TTL resolved + // at connect time. On failure the session override is cleared (the new + // worker parks with its own baseline, so session state must match) and + // the STATEMENT fails; the escalation stands, as above. + if err := c.reapplyWorkerTTLAfterWorkerSwitch(ctx); err != nil { + c.workerTTLOverride = nil + return &workerTTLReapplyError{err: err} + } return nil } @@ -322,11 +347,12 @@ func (c *clientConn) failWorkerEscalation(query string, escErr error, clientMess // failEscalation is the single entry point for an escalateWorker error: it // routes the connection-fatal shape (the switcher failed, previous session // gone) to failWorkerEscalation, and the statement-scoped shape (the swap -// succeeded, only the s3_cache re-apply failed) to failS3CacheReapply. Every -// call site uses it so the two can never be confused at one of them. +// succeeded, only a post-swap GUC re-apply failed) to +// failReapplyAfterWorkerSwitch. Every call site uses it so the two can never +// be confused at one of them. func (c *clientConn) failEscalation(query string, escErr error, clientMessage string) error { - if errors.Is(escErr, errS3CacheReapplyFailed) { - return c.failS3CacheReapply(query, escErr) + if errors.Is(escErr, errS3CacheReapplyFailed) || errors.Is(escErr, errWorkerTTLReapplyFailed) { + return c.failReapplyAfterWorkerSwitch(query, escErr) } return c.failWorkerEscalation(query, escErr, clientMessage) } @@ -338,24 +364,23 @@ func (c *clientConn) failEscalation(query string, escErr error, clientMessage st // client has its ErrorResponse and, on the simple protocol, its ReadyForQuery. var errStatementAborted = errors.New("statement aborted") -// failS3CacheReapply reports a post-escalation `duckgres.s3_cache` re-apply -// failure. Unlike a failed escalation this is NOT connection-fatal: the -// escalation succeeded, so the connection has a healthy session on the standard -// worker and the pin stands (deliberately not rolled back — the swap really -// happened). Only the transport could not be re-asserted, so: +// failReapplyAfterWorkerSwitch reports a post-escalation session-GUC +// re-apply failure (duckgres.s3_cache or duckgres.worker_ttl). Unlike a +// failed escalation this is NOT connection-fatal: the escalation succeeded, +// so the connection has a healthy session on the standard worker and the pin +// stands (deliberately not rolled back — the swap really happened). Only the +// worker-side state could not be re-asserted, so: // // - the statement fails with a normal ERROR (XX000) naming the re-apply, -// rather than a benchmark quietly continuing through the cache; -// - the session flag was already reset by escalateWorker to the worker's -// ACTUAL transport (proxied — a fresh session always starts on the cache -// proxy), so SHOW stays truthful and the client can retry -// `SET duckgres.s3_cache = off`; +// rather than the connection quietly continuing with divergent state; +// - the session state was already reset by escalateWorker to the worker's +// ACTUAL state, so SHOW stays truthful and the client can retry the SET; // - the connection stays alive. // // ReadyForQuery is written only on the simple protocol; inside an // extended-query message Sync owns it, and writing one here would desync the // client's response accounting. -func (c *clientConn) failS3CacheReapply(query string, err error) error { +func (c *clientConn) failReapplyAfterWorkerSwitch(query string, err error) error { c.logQueryError(query, err) c.sendError("ERROR", "XX000", err.Error()) c.setTxError() diff --git a/server/conn_worker_ttl.go b/server/conn_worker_ttl.go new file mode 100644 index 00000000..d93541ac --- /dev/null +++ b/server/conn_worker_ttl.go @@ -0,0 +1,168 @@ +package server + +import ( + "context" + "errors" + "fmt" + "time" +) + +// WorkerTTLGUCName is the duckgres-namespaced session GUC controlling how long +// the session's worker stays hot-idle (warm, reusable) after its last session +// ends, on the remote/k8s backend: +// +// SET duckgres.worker_ttl = '20m' +// +// It is the mid-session form of the `-c duckgres.worker_ttl=...` startup +// option (controlplane/worker_profile.go), for clients that cannot set startup +// options. Used as the SHOW result column label. +const WorkerTTLGUCName = "duckgres.worker_ttl" + +// workerTTLApplyTimeout bounds the control-plane apply hook. The apply is an +// in-memory pool mutation (no worker RPC), so this is generous. +const workerTTLApplyTimeout = 5 * time.Second + +// defaultWorkerTTLFallback is what SHOW reports when nothing better is known +// (standalone / process backend with no override): the control plane's +// built-in hot-idle TTL default (defaultWorkerTTL, 1m), mirrored here because +// the import direction is controlplane -> server. +const defaultWorkerTTLFallback = time.Minute + +// WorkerTTLControl is the optional control-plane capability behind the +// `duckgres.worker_ttl` session GUC, installed on remote/k8s connections. It +// is how the connection layer reaches the bound worker's pool-side profile — +// the hot-idle TTL lives in the control plane's worker pool, not in the +// worker process, so unlike duckgres.s3_cache this is NOT an executor +// capability. Connections without it (standalone, process backend) get +// session-state-only SET/SHOW, which is correct because those deployments +// have no hot-idle worker TTL to override. +type WorkerTTLControl struct { + // Baseline is the TTL resolved at connect time (startup GUC > org default + // > deployment default > built-in 1m). RESET restores it on the worker, + // and SHOW falls back to it when no worker is bound yet. + Baseline time.Duration + + // Apply overrides the bound worker's hot-idle TTL, returning the value + // actually applied (the control plane may clamp to WorkerMaxTTL). A + // returned *transform.CodedError preserves its SQLSTATE to the client; + // any other error surfaces as XX000. + Apply func(ctx context.Context, ttl time.Duration) (applied time.Duration, err error) + + // Current reports the TTL the bound worker would park with NOW (ok=false + // when no worker is bound — a lazily activated connection before its + // first engine statement). It beats Baseline for SHOW because a reused + // hot-idle worker can carry a previous request's TTL. + Current func() (ttl time.Duration, ok bool) +} + +// effectiveWorkerTTL resolves the value SHOW reports: the session override +// wins, then the bound worker's current TTL, then the connect-time baseline, +// then the built-in default. +func (c *clientConn) effectiveWorkerTTL() time.Duration { + if c.workerTTLOverride != nil { + return *c.workerTTLOverride + } + if c.workerTTLCtl != nil { + if c.workerTTLCtl.Current != nil { + if cur, ok := c.workerTTLCtl.Current(); ok { + return cur + } + } + if c.workerTTLCtl.Baseline > 0 { + return c.workerTTLCtl.Baseline + } + } + return defaultWorkerTTLFallback +} + +// workerTTLValue is the SHOW-facing rendering of the session state. +func (c *clientConn) workerTTLValue() string { + return c.effectiveWorkerTTL().String() +} + +// applyWorkerTTLSetting applies an already-normalized `duckgres.worker_ttl` +// value (a canonical Go duration, or "" = reset to the connect-time baseline) +// to the session. When the effective state changes and the control plane +// installed a WorkerTTLControl, the apply hook runs FIRST and the session +// state is only updated on success — a SET that failed to take effect on the +// worker must error, not leave SHOW claiming a TTL the worker won't park +// with. Callers pass validated values only (transform.NormalizeWorkerTTL, +// rejecting anything else with 22023 before this is reached). +func (c *clientConn) applyWorkerTTLSetting(value string) error { + var override *time.Duration + if value != "" { + d, err := time.ParseDuration(value) + if err != nil { + // Unreachable: every SET path validates via NormalizeWorkerTTL + // first. Defensive so a future caller cannot store unparseable + // state. + return err + } + override = &d + } + // No-op when the effective state does not change (a redundant SET must + // not re-invoke the control plane, mirroring applyS3CacheSetting). + if (override == nil) == (c.workerTTLOverride == nil) && + (override == nil || *override == *c.workerTTLOverride) { + return nil + } + if c.workerTTLCtl != nil && c.workerTTLCtl.Apply != nil { + target := c.workerTTLCtl.Baseline + if override != nil { + target = *override + } + c.ensureConnectionContext() + ctx, cancel := context.WithTimeout(c.ctx, workerTTLApplyTimeout) + defer cancel() + applied, err := c.workerTTLCtl.Apply(ctx, target) + if err != nil { + return fmt.Errorf("failed to apply %s: %w", WorkerTTLGUCName, err) + } + // Store the value the worker ACTUALLY got (the hook may have clamped + // it), so SHOW never reports a TTL the worker won't park with. + if override != nil { + *override = applied + } + c.logger().Info("Set duckgres.worker_ttl.", "ttl", applied.String()) + } + c.workerTTLOverride = override + return nil +} + +// reapplyWorkerTTLAfterWorkerSwitch re-asserts this session's TTL override on +// a freshly acquired worker after a tier escalation. The override is pool-side +// per-worker state, and the escalated worker's profile carries the TTL +// resolved at connect time — without the re-apply the connection's warm +// retention would silently revert the moment it escalated. On failure the +// caller resets the session override (the new worker parks with its own +// baseline TTL, so session state must match). No-op when the session never +// overrode the TTL. +func (c *clientConn) reapplyWorkerTTLAfterWorkerSwitch(ctx context.Context) error { + if c.workerTTLOverride == nil { + return nil + } + if c.workerTTLCtl == nil || c.workerTTLCtl.Apply == nil { + return nil + } + applyCtx, cancel := context.WithTimeout(ctx, workerTTLApplyTimeout) + defer cancel() + applied, err := c.workerTTLCtl.Apply(applyCtx, *c.workerTTLOverride) + if err != nil { + return fmt.Errorf("failed to re-apply %s on the new worker: %w", WorkerTTLGUCName, err) + } + *c.workerTTLOverride = applied + c.logger().Info("Re-applied duckgres.worker_ttl after worker switch.", "ttl", applied.String()) + return nil +} + +// workerTTLApplyErrorSQLState picks the client-facing SQLSTATE for a failed +// SET duckgres.worker_ttl: a coded rejection from the control plane (the +// AllowClientWorkerProfile gate's 22023) keeps its code; anything else is an +// internal apply failure (XX000). +func workerTTLApplyErrorSQLState(err error) string { + var coded interface{ SQLState() string } + if errors.As(err, &coded) { + return coded.SQLState() + } + return "XX000" +} diff --git a/server/exports.go b/server/exports.go index ecc3c0f4..b49666d3 100644 --- a/server/exports.go +++ b/server/exports.go @@ -258,6 +258,18 @@ func SetPendingS3CacheOption(cc *clientConn, raw string) { } } +// SetConnectionWorkerTTLControl installs the control-plane capability behind +// the `duckgres.worker_ttl` session GUC on a remote/k8s connection. Call +// before RunMessageLoop; the hooks run on the message-loop goroutine (the +// same one that handles SET/SHOW and tier escalation), so they are +// single-threaded with statement handling. Connections without it +// (standalone, process backend) get session-state-only SET/SHOW. +func SetConnectionWorkerTTLControl(cc *clientConn, ctl *WorkerTTLControl) { + if cc != nil { + cc.workerTTLCtl = ctl + } +} + // SetConnectionDatabase updates the PostgreSQL-visible database name for a // control-plane connection after the fact. The eager connect path knows the // resolved catalog before it builds the connection; the lazily-activated path diff --git a/server/worker_ttl_test.go b/server/worker_ttl_test.go new file mode 100644 index 00000000..c51c9f37 --- /dev/null +++ b/server/worker_ttl_test.go @@ -0,0 +1,448 @@ +package server + +import ( + "bytes" + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/posthog/duckgres/transpiler/transform" +) + +// recordingWorkerTTLControl is a WorkerTTLControl whose Apply records the +// requested TTLs, so a test can assert what the conn asked the control plane +// to stamp on its worker (and in what order relative to statements). +type recordingWorkerTTLControl struct { + baseline time.Duration + current time.Duration + hasCur bool + applied []time.Duration + err error + // clampTo, when positive, makes Apply return min(requested, clampTo), + // mirroring the control plane's WorkerMaxTTL clamp. + clampTo time.Duration +} + +func (r *recordingWorkerTTLControl) control() *WorkerTTLControl { + return &WorkerTTLControl{ + Baseline: r.baseline, + Apply: func(_ context.Context, ttl time.Duration) (time.Duration, error) { + if r.err != nil { + return 0, r.err + } + r.applied = append(r.applied, ttl) + if r.clampTo > 0 && ttl > r.clampTo { + return r.clampTo, nil + } + return ttl, nil + }, + Current: func() (time.Duration, bool) { + return r.current, r.hasCur + }, + } +} + +func newWorkerTTLConn(exec QueryExecutor, rec *recordingWorkerTTLControl) (*clientConn, *bytes.Buffer) { + c, out := newBufferedConn(exec) + if rec != nil { + c.workerTTLCtl = rec.control() + } + return c, out +} + +// TestWorkerTTLSetAppliesThroughControl asserts the core contract of the SET +// path: `SET duckgres.worker_ttl = '20m'` invokes the control-plane hook with +// 20m BEFORE the session state flips, SHOW then reports the override, and a +// redundant SET to the same value does NOT re-invoke the hook. RESET restores +// the connect-time baseline on the worker. +func TestWorkerTTLSetAppliesThroughControl(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: time.Minute} + c, out := newWorkerTTLConn(&selectOneExecutor{}, rec) + + if got := c.workerTTLValue(); got != "1m0s" { + t.Fatalf("fresh session workerTTLValue() = %q, want %q (baseline)", got, "1m0s") + } + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(SET 20m): %v", err) + } + if len(rec.applied) != 1 || rec.applied[0] != 20*time.Minute { + t.Fatalf("applied after SET 20m = %v, want [20m]", rec.applied) + } + if got := c.workerTTLValue(); got != "20m0s" { + t.Fatalf("workerTTLValue() = %q after SET, want %q", got, "20m0s") + } + + // SHOW reports the session state. + out.Reset() + if err := c.handleQuery([]byte("SHOW duckgres.worker_ttl\x00")); err != nil { + t.Fatalf("handleQuery(SHOW): %v", err) + } + sawValue := false + for _, m := range parseWireMsgs(t, out.Bytes()) { + if m.typ == 'D' && bytes.Contains(m.body, []byte("20m0s")) { + sawValue = true + } + } + if !sawValue { + t.Fatalf("SHOW duckgres.worker_ttl did not report '20m0s'") + } + + // Redundant SET to the same value: no control-plane call. + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(redundant SET): %v", err) + } + if len(rec.applied) != 1 { + t.Fatalf("redundant SET re-invoked the control plane: applied = %v", rec.applied) + } + + // RESET restores the connect-time baseline on the worker. + if err := c.handleQuery([]byte("RESET duckgres.worker_ttl\x00")); err != nil { + t.Fatalf("handleQuery(RESET): %v", err) + } + if len(rec.applied) != 2 || rec.applied[1] != time.Minute { + t.Fatalf("applied after RESET = %v, want [20m 1m]", rec.applied) + } + if got := c.workerTTLValue(); got != "1m0s" { + t.Fatalf("workerTTLValue() = %q after RESET, want baseline %q", got, "1m0s") + } +} + +// TestWorkerTTLSetApplyFailureKeepsState asserts a failed control-plane apply +// fails the SET (ErrorResponse, no CommandComplete) and leaves the session +// state on its previous value — SHOW must never claim a TTL the worker will +// not park with. +func TestWorkerTTLSetApplyFailureKeepsState(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: time.Minute, err: errors.New("pool: boom")} + c, out := newWorkerTTLConn(&selectOneExecutor{}, rec) + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(SET 20m): %v", err) + } + msgs := parseWireMsgs(t, out.Bytes()) + if !errorResponseWith(msgs, "XX000", WorkerTTLGUCName) { + t.Fatalf("failed apply did not surface as XX000 ErrorResponse; msgs=%s", describeMsgs(msgs)) + } + for _, m := range msgs { + if m.typ == 'C' && bytes.Contains(m.body, []byte("SET")) { + t.Fatalf("failed SET still produced CommandComplete(SET); msgs=%s", describeMsgs(msgs)) + } + } + if got := c.workerTTLValue(); got != "1m0s" { + t.Fatalf("failed apply flipped session state: workerTTLValue() = %q, want baseline %q", got, "1m0s") + } +} + +// TestWorkerTTLSetGateRejection asserts a control-plane rejection carrying a +// SQLSTATE (the AllowClientWorkerProfile gate's 22023) surfaces with THAT +// code, not the generic XX000 an apply failure gets. +func TestWorkerTTLSetGateRejection(t *testing.T) { + rec := &recordingWorkerTTLControl{ + baseline: time.Minute, + err: &transform.CodedError{Code: "22023", Message: "duckgres.worker_ttl overrides are not enabled on this server"}, + } + c, out := newWorkerTTLConn(&selectOneExecutor{}, rec) + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(SET 20m): %v", err) + } + msgs := parseWireMsgs(t, out.Bytes()) + if !errorResponseWith(msgs, "22023", "not enabled") { + t.Fatalf("gate rejection did not surface as 22023; msgs=%s", describeMsgs(msgs)) + } + if got := c.workerTTLValue(); got != "1m0s" { + t.Fatalf("rejected SET flipped session state: workerTTLValue() = %q", got) + } +} + +// TestWorkerTTLSetClampReportsClamped asserts that when the control plane +// clamps the requested TTL (WorkerMaxTTL), the session state stores the +// CLAMPED value so SHOW never reports a TTL the worker will not park with. +func TestWorkerTTLSetClampReportsClamped(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: time.Minute, clampTo: time.Hour} + c, _ := newWorkerTTLConn(&selectOneExecutor{}, rec) + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '24h'\x00")); err != nil { + t.Fatalf("handleQuery(SET 24h): %v", err) + } + if len(rec.applied) != 1 || rec.applied[0] != 24*time.Hour { + t.Fatalf("applied after SET 24h = %v, want [24h] (the hook sees the full request)", rec.applied) + } + if got := c.workerTTLValue(); got != "1h0m0s" { + t.Fatalf("workerTTLValue() = %q after a clamped SET, want %q", got, "1h0m0s") + } +} + +// TestWorkerTTLSetWithoutControlIsStateOnly asserts the standalone/in-process +// case: a connection without the control-plane capability gets +// session-state-only SET/SHOW (no error) — those deployments have no hot-idle +// worker TTL to override. +func TestWorkerTTLSetWithoutControlIsStateOnly(t *testing.T) { + c, out := newWorkerTTLConn(&selectOneExecutor{}, nil) + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(SET 20m): %v", err) + } + msgs := parseWireMsgs(t, out.Bytes()) + sawSet := false + for _, m := range msgs { + if m.typ == 'C' && bytes.Contains(m.body, []byte("SET")) { + sawSet = true + } + } + if !sawSet { + t.Fatalf("SET without control did not complete; msgs=%s", describeMsgs(msgs)) + } + if got := c.workerTTLValue(); got != "20m0s" { + t.Fatalf("workerTTLValue() = %q after SET, want %q (state-only)", got, "20m0s") + } +} + +// TestWorkerTTLShowDefaults pins the SHOW fallback order: the session +// override wins, then the bound worker's current TTL, then the connect-time +// baseline, then the built-in default. +func TestWorkerTTLShowDefaults(t *testing.T) { + // No control at all (standalone): the documented built-in default. + c, _ := newWorkerTTLConn(&selectOneExecutor{}, nil) + if got := c.workerTTLValue(); got != "1m0s" { + t.Fatalf("standalone workerTTLValue() = %q, want %q", got, "1m0s") + } + + // A bound worker's CURRENT TTL beats the connect-time baseline (e.g. the + // connection reused a hot-idle worker carrying a previous request's TTL). + rec := &recordingWorkerTTLControl{baseline: 5 * time.Minute, current: 7 * time.Minute, hasCur: true} + c, _ = newWorkerTTLConn(&selectOneExecutor{}, rec) + if got := c.workerTTLValue(); got != "7m0s" { + t.Fatalf("workerTTLValue() = %q, want current %q", got, "7m0s") + } + + // The session override beats both. + if err := c.applyWorkerTTLSetting("20m0s"); err != nil { + t.Fatalf("applyWorkerTTLSetting: %v", err) + } + if got := c.workerTTLValue(); got != "20m0s" { + t.Fatalf("workerTTLValue() = %q with override, want %q", got, "20m0s") + } +} + +// TestWorkerTTLInvalidSetRejected asserts the simple-protocol SET path +// rejects a non-duration value with 22023, does not echo the client input, +// and leaves the session state untouched. +func TestWorkerTTLInvalidSetRejected(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: time.Minute} + c, out := newWorkerTTLConn(&selectOneExecutor{}, rec) + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = 'garbage'\x00")); err != nil { + t.Fatalf("handleQuery: %v", err) + } + msgs := parseWireMsgs(t, out.Bytes()) + if !errorResponseWith(msgs, "22023", "duration") { + t.Fatalf("no 22023 ErrorResponse describing the expected shape; msgs=%s", describeMsgs(msgs)) + } + if errorResponseWith(msgs, "garbage") { + t.Fatalf("rejection echoes the offending value; msgs=%s", describeMsgs(msgs)) + } + if len(rec.applied) != 0 { + t.Fatalf("rejected SET reached the control plane: applied = %v", rec.applied) + } + if got := c.workerTTLValue(); got != "1m0s" { + t.Fatalf("rejected SET flipped session state: workerTTLValue() = %q", got) + } +} + +// TestWorkerTTLMixedBatch asserts the split-batch path: +// `SET duckgres.worker_ttl = '20m'; SELECT 1` applies the GUC (control-plane +// call included) and still runs the trailing SELECT. +func TestWorkerTTLMixedBatch(t *testing.T) { + exec := &selectOneExecutor{} + rec := &recordingWorkerTTLControl{baseline: time.Minute} + c, out := newWorkerTTLConn(exec, rec) + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'; SELECT 1\x00")); err != nil { + t.Fatalf("handleQuery: %v", err) + } + if len(rec.applied) != 1 || rec.applied[0] != 20*time.Minute { + t.Fatalf("applied = %v, want [20m]", rec.applied) + } + if exec.queryCalls != 1 { + t.Fatalf("SELECT did not run: QueryContext called %d times, want 1", exec.queryCalls) + } + msgs := parseWireMsgs(t, out.Bytes()) + var sawSet, sawSelect bool + for _, m := range msgs { + if m.typ == 'C' { + if bytes.Contains(m.body, []byte("SET")) { + sawSet = true + } + if bytes.Contains(m.body, []byte("SELECT")) { + sawSelect = true + } + } + } + if !sawSet || !sawSelect { + t.Fatalf("batch did not complete both statements (SET=%v SELECT=%v); msgs=%s", sawSet, sawSelect, describeMsgs(msgs)) + } +} + +// TestWorkerTTLMixedBatchApplyFailureAborts asserts a failed apply inside a +// batch aborts the remaining statements — they may depend on the requested +// warm-retention state. +func TestWorkerTTLMixedBatchApplyFailureAborts(t *testing.T) { + exec := &selectOneExecutor{} + rec := &recordingWorkerTTLControl{baseline: time.Minute, err: errors.New("boom")} + c, out := newWorkerTTLConn(exec, rec) + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'; SELECT 1\x00")); err != nil { + t.Fatalf("handleQuery: %v", err) + } + if exec.queryCalls != 0 { + t.Fatalf("SELECT ran after the failed SET: QueryContext called %d times, want 0", exec.queryCalls) + } + msgs := parseWireMsgs(t, out.Bytes()) + if !errorResponseWith(msgs, "XX000") { + t.Fatalf("failed batched SET did not surface XX000; msgs=%s", describeMsgs(msgs)) + } +} + +// TestWorkerTTLExtendedParse asserts the extended-protocol path: an invalid +// value is rejected at Parse time; a valid SET parses, applies at Execute +// time through the control-plane hook, and Describe returns NoData. +func TestWorkerTTLExtendedParse(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: time.Minute} + c, out := newWorkerTTLConn(&selectOneExecutor{}, rec) + c.stmts = make(map[string]*preparedStmt) + c.portals = make(map[string]*portal) + + // Invalid value: rejected at Parse, nothing stored. + body := append([]byte("s1\x00SET duckgres.worker_ttl = 'garbage'\x00"), 0, 0) + c.handleParse(body) + _ = c.flushWriter() + msgs := parseWireMsgs(t, out.Bytes()) + if !errorResponseWith(msgs, "22023", "duration") { + t.Fatalf("extended Parse of invalid SET not rejected with 22023; msgs=%s", describeMsgs(msgs)) + } + if _, ok := c.stmts["s1"]; ok { + t.Fatalf("rejected Parse still stored the prepared statement") + } + + // Valid SET: parses with workerTTLSet populated. + out.Reset() + body = append([]byte("s2\x00SET duckgres.worker_ttl = '20m'\x00"), 0, 0) + c.handleParse(body) + _ = c.flushWriter() + st, ok := c.stmts["s2"] + if !ok { + t.Fatalf("valid SET did not parse: msgs=%s", describeMsgs(parseWireMsgs(t, out.Bytes()))) + } + if st.workerTTLSet == nil || *st.workerTTLSet != "20m0s" { + t.Fatalf("prepared stmt workerTTLSet = %v, want 20m0s", st.workerTTLSet) + } + + // Bind + Execute applies through the hook. + out.Reset() + // Bind message: portal name, statement name (NUL-terminated), int16 format + // count (0), int16 param count (0), int16 result-format count (0). + bindBody := append([]byte("\x00s2\x00"), 0, 0, 0, 0, 0, 0) + c.handleBind(bindBody) + c.handleExecute(append([]byte("\x00"), 0, 0, 0, 0)) + _ = c.flushWriter() + if len(rec.applied) != 1 || rec.applied[0] != 20*time.Minute { + t.Fatalf("applied after extended Execute = %v, want [20m]", rec.applied) + } + if got := c.workerTTLValue(); got != "20m0s" { + t.Fatalf("workerTTLValue() = %q after extended SET, want %q", got, "20m0s") + } +} + +// TestWorkerTTLSetActivatesLazyConnection asserts that a SET on a +// not-yet-acquired (exploratory-tier lazy) connection binds a worker FIRST — +// the TTL is worker-side pool state, so applying without a worker would leave +// SHOW reporting a TTL no worker parks with. +func TestWorkerTTLSetActivatesLazyConnection(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: 48 * time.Hour} + c, _ := newWorkerTTLConn(nil, rec) + activations := 0 + c.sessionActivator = func(_ context.Context, pinned bool) (QueryExecutor, int, string, error) { + activations++ + if pinned { + t.Errorf("SET duckgres.worker_ttl activated pinned; it is not a pinning statement") + } + return &selectOneExecutor{}, 7, "worker-7", nil + } + + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(SET 20m): %v", err) + } + if activations != 1 { + t.Fatalf("activations = %d, want 1 (SET must bind a worker before applying)", activations) + } + if len(rec.applied) != 1 || rec.applied[0] != 20*time.Minute { + t.Fatalf("applied = %v, want [20m]", rec.applied) + } + if got := c.workerTTLValue(); got != "20m0s" { + t.Fatalf("workerTTLValue() = %q, want %q", got, "20m0s") + } +} + +// TestWorkerTTLReappliedOnWorkerEscalation asserts that a session carrying a +// TTL override carries it onto the worker it escalates to: the new worker's +// profile carries the connect-time baseline, so without the re-apply the +// connection's warm retention would silently revert at escalation. +func TestWorkerTTLReappliedOnWorkerEscalation(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: 48 * time.Hour} + c, _ := newWorkerTTLConn(&selectOneExecutor{}, rec) + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(SET 20m): %v", err) + } + + c.onExploratoryWorker = true + c.workerSwitcher = func(context.Context, string) (QueryExecutor, int, string, error) { + return &selectOneExecutor{}, 8, "worker-8", nil + } + if err := c.escalateWorker(context.Background(), escalateReasonState); err != nil { + t.Fatalf("escalateWorker: %v", err) + } + if len(rec.applied) != 2 || rec.applied[1] != 20*time.Minute { + t.Fatalf("applied after escalation = %v, want the override re-applied to the new worker", rec.applied) + } + if got := c.workerTTLValue(); got != "20m0s" { + t.Fatalf("workerTTLValue() = %q after escalation, want %q (override preserved)", got, "20m0s") + } +} + +// TestWorkerTTLEscalationReapplyFailureResetsState asserts that when the +// override cannot be re-applied on the new worker, the statement fails AND +// the session state is reset to match the TTL the worker will actually park +// with (the connect-time baseline) — SHOW must never lie. The ESCALATION +// itself succeeded, so the failure is statement-scoped, not connection-fatal. +func TestWorkerTTLEscalationReapplyFailureResetsState(t *testing.T) { + rec := &recordingWorkerTTLControl{baseline: 48 * time.Hour} + c, _ := newWorkerTTLConn(&selectOneExecutor{}, rec) + if err := c.handleQuery([]byte("SET duckgres.worker_ttl = '20m'\x00")); err != nil { + t.Fatalf("handleQuery(SET 20m): %v", err) + } + + rec.err = errors.New("pool: worker gone") + c.onExploratoryWorker = true + c.workerSwitcher = func(context.Context, string) (QueryExecutor, int, string, error) { + return &selectOneExecutor{}, 8, "worker-8", nil + } + err := c.escalateWorker(context.Background(), escalateReasonState) + if err == nil { + t.Fatal("escalateWorker returned nil, want the re-apply failure") + } + if !errors.Is(err, errWorkerTTLReapplyFailed) { + t.Fatalf("error %v is not tagged errWorkerTTLReapplyFailed; callers would terminate the connection", err) + } + if !strings.Contains(err.Error(), "duckgres.worker_ttl") { + t.Fatalf("error %q does not name the GUC", err) + } + if got := c.workerTTLValue(); got != "48h0m0s" { + t.Fatalf("workerTTLValue() = %q after a failed re-apply; state must match the worker's actual TTL %q", got, "48h0m0s") + } + if c.onExploratoryWorker { + t.Fatal("onExploratoryWorker = true after a failed re-apply; the escalation itself SUCCEEDED and must not be rolled back") + } +} diff --git a/transpiler/config.go b/transpiler/config.go index a1b96584..85210317 100644 --- a/transpiler/config.go +++ b/transpiler/config.go @@ -117,6 +117,22 @@ type Result struct { // answers it from session state (defaulting to "on"). S3CacheShow bool + // WorkerTTLSet, when non-nil, indicates a `SET duckgres.worker_ttl = + // ''` on the duckgres-namespaced custom GUC. The connection layer + // applies the pointed-to value to the session's bound worker (overriding + // how long the worker stays hot-idle after the session ends) and + // acknowledges it as "SET" — it is NOT forwarded to DuckDB, which would + // reject the unknown setting. The value is pre-validated and canonical: a + // Go duration string ("20m0s", "0s") or "" (reset to the connect-time + // TTL) — an invalid value never reaches here (it surfaces as Error with + // SQLSTATE 22023 instead). + WorkerTTLSet *string + + // WorkerTTLShow indicates a `SHOW duckgres.worker_ttl`; the connection + // layer answers it from session state (the TTL the session's worker will + // park with). + WorkerTTLShow bool + // Error is set when a transform detects an error that should be returned to the client // (e.g., unrecognized configuration parameter in SHOW command) Error error diff --git a/transpiler/transform/setshow.go b/transpiler/transform/setshow.go index 380f0382..6594e9a9 100644 --- a/transpiler/transform/setshow.go +++ b/transpiler/transform/setshow.go @@ -4,6 +4,7 @@ import ( "fmt" "regexp" "strings" + "time" pg_query "github.com/pganalyze/pg_query_go/v6" ) @@ -99,6 +100,43 @@ func errInvalidS3Cache() *CodedError { } } +// workerTTLParam is the duckgres-namespaced custom GUC that overrides how long +// the session's worker stays hot-idle (warm) after its last session ends +// (remote/k8s backend). Intercepted here and applied by the connection layer +// (which updates the bound worker's pool-side TTL); it is NEVER forwarded to +// DuckDB. See clientConn.applyWorkerTTLSetting in server/. +const workerTTLParam = "duckgres.worker_ttl" + +// NormalizeWorkerTTL validates a client-supplied duckgres.worker_ttl value: a +// non-negative Go duration ("20m", "1h30m", "0s" — zero retires the worker as +// soon as the session ends, matching the startup option). Matching ignores +// surrounding whitespace; the returned value is the canonical duration string. +// Empty is valid and means "reset to default" (the session then parks with the +// TTL resolved at connect time). An invalid value returns a 22023 +// (invalid_parameter_value) CodedError that, like NormalizeQuerySource, does +// not echo the offending client input. +func NormalizeWorkerTTL(raw string) (string, error) { + v := strings.TrimSpace(raw) + if v == "" { + return "", nil + } + d, err := time.ParseDuration(v) + if err != nil || d < 0 { + return "", errInvalidWorkerTTL() + } + return d.String(), nil +} + +// errInvalidWorkerTTL is the SET-time rejection for a bad duckgres.worker_ttl +// value: 22023 invalid_parameter_value, same treatment as duckgres.s3_cache. +func errInvalidWorkerTTL() *CodedError { + return &CodedError{ + Code: "22023", // invalid_parameter_value + Message: fmt.Sprintf("invalid value for %q: must be a non-negative Go duration (e.g. \"20m\")", + workerTTLParam), + } +} + // duckdbShowCommands are DuckDB-specific SHOW commands that should be passed // through to DuckDB rather than treated as PostgreSQL config parameters. var duckdbShowCommands = map[string]bool{ @@ -247,18 +285,18 @@ func NewSetShowTransform() *SetShowTransform { "wal_receiver_timeout": true, // Session settings (silently accept) - "datestyle": true, - "intervalstyle": true, - "standard_conforming_strings": true, - "escape_string_warning": true, - "array_nulls": true, - "backslash_quote": true, - "default_with_oids": true, - "quote_all_identifiers": true, - "sql_inheritance": true, - "transform_null_equals": true, - "lo_compat_privileges": true, - "operator_precedence_warning": true, + "datestyle": true, + "intervalstyle": true, + "standard_conforming_strings": true, + "escape_string_warning": true, + "array_nulls": true, + "backslash_quote": true, + "default_with_oids": true, + "quote_all_identifiers": true, + "sql_inheritance": true, + "transform_null_equals": true, + "lo_compat_privileges": true, + "operator_precedence_warning": true, // Server version settings (commonly queried) "server_version": true, @@ -266,8 +304,8 @@ func NewSetShowTransform() *SetShowTransform { "server_encoding": true, // Timezone (DuckDB has its own timezone setting) - "timezone": true, - "log_timezone": true, + "timezone": true, + "log_timezone": true, "timezone_abbreviations": true, }, PassthroughParams: map[string]bool{ @@ -396,6 +434,35 @@ func (t *SetShowTransform) Transform(tree *pg_query.ParseResult, result *Result) return true, nil } + // duckgres.worker_ttl: same interception contract as + // duckgres.s3_cache above (single-statement only; RESET / + // SET ... TO DEFAULT map to the empty value = the TTL resolved + // at connect time; invalid durations rejected with 22023 via + // result.Error). The connection layer applies the value to the + // session's bound worker — never forwarded to DuckDB. + if paramName == workerTTLParam && !multiStatement { + value := "" + if n.VariableSetStmt.Kind == pg_query.VariableSetKind_VAR_SET_VALUE { + extracted := false + if len(n.VariableSetStmt.Args) == 1 { + if v, ok := searchPathValue(n.VariableSetStmt.Args[0]); ok { + value, extracted = v, true + } + } + if !extracted { + result.Error = errInvalidWorkerTTL() + return true, nil + } + } + norm, err := NormalizeWorkerTTL(value) + if err != nil { + result.Error = err + return true, nil + } + result.WorkerTTLSet = &norm + return true, nil + } + if paramName == "search_path" { if sql, ok := normalizeSearchPathSet(n.VariableSetStmt); ok { result.SQLOverride = sql @@ -462,6 +529,14 @@ func (t *SetShowTransform) Transform(tree *pg_query.ParseResult, result *Result) return true, nil } + // duckgres.worker_ttl: answered from session state by the + // connection layer (the TTL the session's worker will park + // with), not DuckDB. + if paramName == workerTTLParam && !multiStatement { + result.WorkerTTLShow = true + return true, nil + } + // Passthrough params: SHOW → SELECT value FROM duckdb_settings() WHERE name = '...' // (DuckDB's SHOW describes a table, not a setting) if t.PassthroughParams[paramName] { @@ -577,11 +652,11 @@ func searchPathValue(node *pg_query.Node) (string, bool) { func (t *SetShowTransform) getDefaultValue(paramName string) string { defaults := map[string]string{ // Client connection settings - "application_name": "duckgres", - "client_encoding": "UTF8", - "statement_timeout": "0", - "lock_timeout": "0", - "extra_float_digits": "1", + "application_name": "duckgres", + "client_encoding": "UTF8", + "statement_timeout": "0", + "lock_timeout": "0", + "extra_float_digits": "1", "client_min_messages": "notice", // Transaction settings @@ -614,18 +689,18 @@ func (t *SetShowTransform) getDefaultValue(paramName string) string { "timezone": "UTC", // Other commonly queried settings - "max_identifier_length": "63", - "default_tablespace": "", - "temp_tablespaces": "", - "lc_collate": "en_US.UTF-8", - "lc_ctype": "en_US.UTF-8", - "lc_messages": "en_US.UTF-8", - "lc_monetary": "en_US.UTF-8", - "lc_numeric": "en_US.UTF-8", - "lc_time": "en_US.UTF-8", - "integer_datetimes": "on", - "is_superuser": "on", - "session_authorization": "duckdb", + "max_identifier_length": "63", + "default_tablespace": "", + "temp_tablespaces": "", + "lc_collate": "en_US.UTF-8", + "lc_ctype": "en_US.UTF-8", + "lc_messages": "en_US.UTF-8", + "lc_monetary": "en_US.UTF-8", + "lc_numeric": "en_US.UTF-8", + "lc_time": "en_US.UTF-8", + "integer_datetimes": "on", + "is_superuser": "on", + "session_authorization": "duckdb", } if val, ok := defaults[paramName]; ok { return val diff --git a/transpiler/transform/transform.go b/transpiler/transform/transform.go index 263b381f..6c6289c8 100644 --- a/transpiler/transform/transform.go +++ b/transpiler/transform/transform.go @@ -51,6 +51,20 @@ type Result struct { // rather than DuckDB. S3CacheShow bool + // WorkerTTLSet, when non-nil, indicates the statement is + // `SET duckgres.worker_ttl = ''` (a duckgres-namespaced custom GUC + // that must NOT be forwarded to DuckDB). The pointed-to string is the + // validated canonical Go duration ("20m0s", "0s", or "" = reset to the + // connect-time TTL; an invalid value sets Error with 22023 instead). The + // connection layer applies it to the session's bound worker, overriding + // how long the worker stays hot-idle after the session ends. + WorkerTTLSet *string + + // WorkerTTLShow indicates the statement is `SHOW duckgres.worker_ttl`. + // The connection layer answers it from session state (the TTL the + // session's worker will park with) rather than DuckDB. + WorkerTTLShow bool + // Error is set when a transform detects an error that should be returned to the client // (e.g., unrecognized configuration parameter in SHOW command) Error error diff --git a/transpiler/transpiler.go b/transpiler/transpiler.go index 2b8b466c..56a3d44b 100644 --- a/transpiler/transpiler.go +++ b/transpiler/transpiler.go @@ -280,11 +280,12 @@ func (t *Transpiler) transpileWithFlags(sql string, flags TransformFlags) (*Resu } // duckgres-namespaced custom GUCs (duckgres.query_source, - // duckgres.s3_cache): intercepted session-side, never - // deparsed/forwarded to DuckDB. Return early so the connection layer - // can store/apply/echo them. + // duckgres.s3_cache, duckgres.worker_ttl): intercepted session-side, + // never deparsed/forwarded to DuckDB. Return early so the connection + // layer can store/apply/echo them. if transformResult.QuerySourceSet != nil || transformResult.QuerySourceShow || - transformResult.S3CacheSet != nil || transformResult.S3CacheShow { + transformResult.S3CacheSet != nil || transformResult.S3CacheShow || + transformResult.WorkerTTLSet != nil || transformResult.WorkerTTLShow { return &Result{ SQL: sql, ParamCount: transformResult.ParamCount, @@ -292,6 +293,8 @@ func (t *Transpiler) transpileWithFlags(sql string, flags TransformFlags) (*Resu QuerySourceShow: transformResult.QuerySourceShow, S3CacheSet: transformResult.S3CacheSet, S3CacheShow: transformResult.S3CacheShow, + WorkerTTLSet: transformResult.WorkerTTLSet, + WorkerTTLShow: transformResult.WorkerTTLShow, }, nil } diff --git a/transpiler/worker_ttl_test.go b/transpiler/worker_ttl_test.go new file mode 100644 index 00000000..ab3fc6f3 --- /dev/null +++ b/transpiler/worker_ttl_test.go @@ -0,0 +1,154 @@ +package transpiler + +import ( + "errors" + "strings" + "testing" +) + +// TestTranspile_WorkerTTLSet asserts that `SET duckgres.worker_ttl = ...` is +// intercepted as a duckgres-namespaced custom GUC (WorkerTTLSet populated, not +// forwarded to DuckDB) and the value is normalized to its canonical Go +// duration string. +func TestTranspile_WorkerTTLSet(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + {"set minutes", "SET duckgres.worker_ttl = '20m'", "20m0s"}, + {"set hours", "SET duckgres.worker_ttl = '24h'", "24h0m0s"}, + {"compound duration", "SET duckgres.worker_ttl = '1h30m'", "1h30m0s"}, + {"zero disables warm retention", "SET duckgres.worker_ttl = '0s'", "0s"}, + {"set local", "SET LOCAL duckgres.worker_ttl = '20m'", "20m0s"}, + {"case-insensitive name", "SET DUCKGRES.WORKER_TTL = '20m'", "20m0s"}, + {"whitespace trimmed", "SET duckgres.worker_ttl = ' 20m '", "20m0s"}, + {"empty string resets to default", "SET duckgres.worker_ttl = ''", ""}, + {"set to default clears to empty", "SET duckgres.worker_ttl TO DEFAULT", ""}, + {"reset clears to empty", "RESET duckgres.worker_ttl", ""}, + } + + tr := New(DefaultConfig()) + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := tr.Transpile(tt.input) + if err != nil { + t.Fatalf("Transpile(%q) error: %v", tt.input, err) + } + if result.WorkerTTLSet == nil { + t.Fatalf("Transpile(%q): WorkerTTLSet = nil, want non-nil (error=%v)", tt.input, result.Error) + } + if got := *result.WorkerTTLSet; got != tt.want { + t.Errorf("Transpile(%q): WorkerTTLSet = %q, want %q", tt.input, got, tt.want) + } + // Custom GUC must never be forwarded to DuckDB. + if result.WorkerTTLShow { + t.Errorf("Transpile(%q): WorkerTTLShow = true, want false", tt.input) + } + // Must not leak into the sibling GUC interceptions. + if result.S3CacheSet != nil || result.S3CacheShow || result.QuerySourceSet != nil || result.QuerySourceShow { + t.Errorf("Transpile(%q): sibling GUC fields populated, want untouched", tt.input) + } + }) + } +} + +// TestTranspile_WorkerTTLSetInvalidRejected asserts that a SET with a value +// that is not a valid non-negative Go duration surfaces Result.Error with +// SQLSTATE 22023 and does NOT populate WorkerTTLSet. The error message must +// describe the expected shape but must NOT echo the offending value +// (arbitrary client input flowing into logs / the recent-errors ring). +func TestTranspile_WorkerTTLSetInvalidRejected(t *testing.T) { + tr := New(DefaultConfig()) + + longJunk := strings.Repeat("x", 10*1024) + inputs := map[string]string{ + "garbage": "SET duckgres.worker_ttl = 'garbage'", + "missing unit": "SET duckgres.worker_ttl = '20'", + "negative duration": "SET duckgres.worker_ttl = '-5m'", + "10KB string": "SET duckgres.worker_ttl = '" + longJunk + "'", + "integer constant": "SET duckgres.worker_ttl = 2", + "multiple values": "SET duckgres.worker_ttl = '20m', '30m'", + "set local garbage": "SET LOCAL duckgres.worker_ttl = 'garbage'", + } + for name, in := range inputs { + t.Run(name, func(t *testing.T) { + result, err := tr.Transpile(in) + if err != nil { + t.Fatalf("Transpile(%.80q) error: %v", in, err) + } + if result.Error == nil { + got := "" + if result.WorkerTTLSet != nil { + got = *result.WorkerTTLSet + } + t.Fatalf("Transpile(%.80q): Error = nil, want 22023 rejection (WorkerTTLSet=%.80q)", in, got) + } + var coded interface{ SQLState() string } + if !errors.As(result.Error, &coded) || coded.SQLState() != "22023" { + t.Errorf("Transpile(%.80q): Error SQLSTATE = %v, want 22023", in, result.Error) + } + msg := result.Error.Error() + if !strings.Contains(msg, "duration") { + t.Errorf("error message must describe the expected duration shape, got %q", msg) + } + if strings.Contains(msg, "garbage") || strings.Contains(msg, longJunk[:64]) { + t.Errorf("error message must not echo the offending value, got %.120q", msg) + } + if result.WorkerTTLSet != nil { + t.Errorf("Transpile(%.80q): WorkerTTLSet = %q, want nil on rejection", in, *result.WorkerTTLSet) + } + }) + } +} + +// TestTranspile_WorkerTTLShow asserts `SHOW duckgres.worker_ttl` is +// intercepted (answered session-side) rather than treated as an unrecognized +// config parameter or forwarded to DuckDB. +func TestTranspile_WorkerTTLShow(t *testing.T) { + tr := New(DefaultConfig()) + result, err := tr.Transpile("SHOW duckgres.worker_ttl") + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + if !result.WorkerTTLShow { + t.Fatalf("WorkerTTLShow = false, want true (error=%v)", result.Error) + } + if result.Error != nil { + t.Errorf("Error = %v, want nil (must not be treated as unrecognized param)", result.Error) + } + if result.WorkerTTLSet != nil { + t.Errorf("WorkerTTLSet = %v, want nil", result.WorkerTTLSet) + } +} + +// TestTranspile_WorkerTTLMultiStatementNotIntercepted mirrors the sibling GUC +// guard: transpiling a MULTI-statement batch containing a duckgres.worker_ttl +// statement must NOT surface WorkerTTLSet/WorkerTTLShow on the whole-batch +// Result (the early return would swallow every statement after the GUC one). +// The connection layer splits the batch and re-transpiles each statement +// individually — where the single-statement interception then fires. +func TestTranspile_WorkerTTLMultiStatementNotIntercepted(t *testing.T) { + tr := New(DefaultConfig()) + + cases := []string{ + "SET duckgres.worker_ttl = '20m'; SHOW duckgres.worker_ttl", + "SET duckgres.worker_ttl = '20m'; SELECT 1", + "SHOW duckgres.worker_ttl; SELECT 1", + } + for _, in := range cases { + t.Run(in, func(t *testing.T) { + result, err := tr.Transpile(in) + if err != nil { + t.Fatalf("Transpile(%q) error: %v", in, err) + } + if result.WorkerTTLSet != nil { + t.Errorf("Transpile(%q): WorkerTTLSet = %v, want nil for a multi-statement batch (would swallow trailing statements)", in, *result.WorkerTTLSet) + } + if result.WorkerTTLShow { + t.Errorf("Transpile(%q): WorkerTTLShow = true, want false for a multi-statement batch (would swallow trailing statements)", in) + } + }) + } +}