Allow mid-session SET of duckgres.worker_ttl

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,

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=<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.

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

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

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)
+	}
+}

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)

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

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

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

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) != "" {

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
+		},
+	}
+}