diff --git a/README.md b/README.md
index f7a19a2b..035cdd72 100644
--- a/README.md
+++ b/README.md
@@ -92,6 +92,8 @@ labels, aggregation rules, PromQL examples, and admission metric migration.
| `duckgres_session_admission_queue_depth{org}` | Gauge | Local callers waiting after successful durable enqueue; sum across replicas |
| `duckgres_session_admission_active_vcpus{org}` | Gauge | Requested vCPUs held by local live lease handles; cleanup-pending durable rows are excluded |
| `duckgres_session_admission_limit_vcpus{org}` | Gauge | Config-reconciled effective org cap for active org stacks; zero means unlimited, max across replicas |
+| `duckgres_session_admission_active_memory_bytes{org}` | Gauge | Requested memory bytes held by local live lease handles; cleanup-pending durable rows are excluded |
+| `duckgres_session_admission_limit_memory_bytes{org}` | Gauge | Config-reconciled effective org memory cap for active org stacks; zero means unlimited, max across replicas |
| `duckgres_session_admission_reclaim_pending` | Gauge | Activated cleanup intents awaiting or executing exact database reclamation |
| `duckgres_session_admission_reclaim_attempts_total{outcome}` | Counter | Exact cleanup attempts by `success` or `error` outcome |
| `duckgres_session_admission_reclaim_reservations_in_use` | Gauge | Cleanup-ownership slots held before enqueue, while queued or live, and during pending cleanup |
@@ -231,7 +233,7 @@ column cannot be added to a populated table).
- [Performance Harness](docs/perf-harness-runbook.md): Local smoke and nightly operations for performance testing.
- [Dev Scenario Runner](docs/runbooks/scenario-dev.md): Scheduled and manually dispatched scenario runs against the configured dev environment.
- [Control Plane Rollout](docs/runbooks/control-plane-rollout.md): Zero-downtime deployment process for the control plane itself.
-- [Org Connection Admission](docs/runbooks/org-connection-admission.md): Global vCPU admission, exact cleanup ownership, failure recovery, and operational metrics.
+- [Org Connection Admission](docs/runbooks/org-connection-admission.md): Global CPU/memory admission, exact cleanup ownership, failure recovery, and operational metrics.
- [Managed Warehouse Provisioning Recovery](docs/runbooks/managed-warehouse-provisioning-recovery.md): Diagnose a failed warehouse whose Duckling dependencies were repaired and verify automatic convergence back to ready.
- [Managed Warehouse Deprovision](docs/runbooks/managed-warehouse-deprovision.md): Destructive teardown process for managed warehouse infrastructure and org cleanup.
- [Resharding Operations](docs/runbooks/resharding.md): Runner recovery, durable respawn reset, safety checks, and local verification.
@@ -389,7 +391,7 @@ Run with config file:
| `DUCKGRES_IDLE_TIMEOUT` | Connection idle timeout (e.g., `30m`, `1h`, `-1` to disable) | `24h` |
| `DUCKGRES_CLIENT_IDLE_TIMEOUT_MAX` | Maximum client-requested `duckgres.idle_timeout`; unset disables client overrides | disabled |
| `DUCKGRES_SESSION_INIT_TIMEOUT` | Session startup metadata initialization and catalog probe timeout | `10s` |
-| `DUCKGRES_WORKER_QUEUE_TIMEOUT` | Max time to wait for worker acquisition and per-org/per-user vCPU resource admission; the managed K8s queue TTL uses this value | `60s` |
+| `DUCKGRES_WORKER_QUEUE_TIMEOUT` | Max time to wait for worker acquisition and org memory/vCPU or user vCPU resource admission; the managed K8s queue TTL uses this value | `60s` |
| `DUCKGRES_ADMISSION_RECLAIMER_MAX_RESERVATIONS` | Max queued/live admission identities whose cleanup ownership one control plane may retain; new admissions are rejected before enqueue when full | `4096` |
| `DUCKGRES_HANDOVER_DRAIN_TIMEOUT` | Max time to drain planned shutdowns and upgrades before forcing exit | `24h` in process mode, `15m` in remote K8s mode |
| `DUCKGRES_SNI_ROUTING_MODE` | Multi-tenant managed-hostname routing: `off`, `passthrough`, or `enforce`. Postgres uses the requested dbname first; managed SNI must resolve to the same org, and SNI supplies the database only when dbname is empty. | `off` |
@@ -839,7 +841,7 @@ Built-in rate limiting protects against brute-force authentication attacks:
- **Failed attempt tracking**: Bans IPs after too many failed auth attempts
- **Connection limits**: Limits concurrent connections per IP and, when configured, total concurrent sessions in standalone mode.
-- **K8s multi-tenant resource limits**: Org and user `max_vcpus` bound the sum of active worker pod vCPUs admitted through runtime-store leases. 0 means unlimited.
+- **K8s multi-tenant resource limits**: Org and user `max_vcpus` bound active admitted worker pod vCPUs, while org `max_memory` independently bounds active admitted worker pod memory. Zero means unlimited.
- **Auto-cleanup**: Expired records are automatically cleaned up
```yaml
@@ -1034,8 +1036,8 @@ Managed-warehouse contract notes:
- At most one managed-warehouse row exists per team. The row may be absent before first provisioning or after cleanup, but there is never more than one active warehouse contract for a team.
- Each org has a `data_imports_table_naming_version`. Migration `000034` assigns `legacy_batch_v1` to orgs that already exist and changes the database default to `copy_v1` for orgs created afterward. `GET /api/v1/orgs/:id/teams` returns the org-level value alongside the team rows so every data-import reader and writer derives the same physical table name. Operators can change the policy in the admin console or with `PUT /api/v1/orgs/:id` using `{"data_imports_table_naming_version":"copy_v1"}`. Migrate existing tables before changing an org that has already written data.
- The admin API exposes that contract at `GET /api/v1/teams/:name/warehouse` and `PUT /api/v1/teams/:name/warehouse`. Team list/get responses also include a nested `warehouse` object when present.
-- Org rows support optional `max_vcpus` on `POST /api/v1/orgs` and `PUT /api/v1/orgs/:id`. In K8s multi-tenant mode, this caps the org's active admitted worker pod vCPUs; `0` means unlimited.
-- Orgs automatically created during warehouse provisioning start with `max_vcpus=150`; `0` remains the explicit unlimited sentinel.
+- Org rows support optional `max_vcpus` and `max_memory` on `POST /api/v1/orgs` and `PUT /api/v1/orgs/:id`. In K8s multi-tenant mode these independently cap the org's active admitted worker pod vCPUs and memory. `max_memory` is a Kubernetes quantity such as `240Gi`; empty or `0` means unlimited.
+- Orgs automatically created during warehouse provisioning start with `max_vcpus=150` and unlimited memory. `0` remains the explicit unlimited sentinel.
- User rows support an optional `max_vcpus` field on `POST /api/v1/users` and `PUT /api/v1/orgs/:id/users/:username`. `max_vcpus` limits the user's active admitted worker pod vCPUs in K8s multi-tenant mode; `0` means unlimited.
- `PUT /api/v1/orgs/:id/teams/:team_id/project-reader` creates or rotates the generated SQL login for a PostHog project. The login can read every current and future table in the project's team, data-import, and modeled-data schemas, plus its legacy events/persons relations. Writes, unqualified application relations, external-reader and introspection functions, other projects' schemas, and their catalog metadata are denied by the PostgreSQL query gateway. The plaintext password is returned only by the rotation response.
- The typed sections are `warehouse_database`, `metadata_store`, `s3`, `worker_identity`, and structured secret refs for `warehouse_database_credentials`, `metadata_store_credentials`, `s3_credentials`, and `runtime_config`. In shared worker mode, every non-empty secret ref must store an explicit `namespace`, and it must match `worker_identity.namespace`.
diff --git a/controlplane/admin/api.go b/controlplane/admin/api.go
index 80cef2c2..8f6a8332 100644
--- a/controlplane/admin/api.go
+++ b/controlplane/admin/api.go
@@ -286,6 +286,7 @@ func (s *gormAPIStore) UpdateOrg(name string, updates configstore.Org) (*configs
fields := map[string]interface{}{
"max_workers": updates.MaxWorkers,
"max_vcpus": updates.MaxVCPUs,
+ "max_memory": updates.MaxMemory,
// Org default worker profile: written unconditionally so an explicit
// empty string CLEARS the default (the handler's presence-merge keeps
// omitted fields at their stored values before this runs).
@@ -1043,6 +1044,9 @@ func (h *apiHandler) updateOrg(c *gin.Context) {
if _, ok := fields["max_vcpus"]; ok {
merged.MaxVCPUs = updates.MaxVCPUs
}
+ if _, ok := fields["max_memory"]; ok {
+ merged.MaxMemory = updates.MaxMemory
+ }
// Org default worker profile: present-in-payload wins, including an
// explicit "" which clears the default.
if _, ok := fields["default_worker_cpu"]; ok {
@@ -1075,6 +1079,7 @@ func (h *apiHandler) updateOrg(c *gin.Context) {
addChange("database_name", existing.DatabaseName, merged.DatabaseName)
addChange("max_workers", existing.MaxWorkers, merged.MaxWorkers)
addChange("max_vcpus", existing.MaxVCPUs, merged.MaxVCPUs)
+ addChange("max_memory", orgStr(existing.MaxMemory), orgStr(merged.MaxMemory))
addChange("default_worker_cpu", orgStr(existing.DefaultWorkerCPU), orgStr(merged.DefaultWorkerCPU))
addChange("default_worker_memory", orgStr(existing.DefaultWorkerMemory), orgStr(merged.DefaultWorkerMemory))
addChange("default_worker_ttl", orgStr(existing.DefaultWorkerTTL), orgStr(merged.DefaultWorkerTTL))
@@ -1432,6 +1437,11 @@ func validateOrgMutationPayload(org *configstore.Org) error {
if org.MaxVCPUs < 0 {
return fmt.Errorf("max_vcpus: value %d must be >= 0", org.MaxVCPUs)
}
+ maxMemory, err := configstore.NormalizeOrgMaxMemory(org.MaxMemory)
+ if err != nil {
+ return fmt.Errorf("max_memory: %w", err)
+ }
+ org.MaxMemory = maxMemory
return nil
}
diff --git a/controlplane/admin/api_test.go b/controlplane/admin/api_test.go
index d9144f69..b2aa93b0 100644
--- a/controlplane/admin/api_test.go
+++ b/controlplane/admin/api_test.go
@@ -81,6 +81,7 @@ func (s *fakeAPIStore) UpdateOrg(name string, updates configstore.Org) (*configs
}
org.MaxWorkers = updates.MaxWorkers
org.MaxVCPUs = updates.MaxVCPUs
+ org.MaxMemory = updates.MaxMemory
// Mirrors gormAPIStore: "" = preserve, non-empty renames.
if updates.DatabaseName != "" {
org.DatabaseName = updates.DatabaseName
@@ -2547,6 +2548,124 @@ func TestUpdateOrgRejectsNegativeMaxVCPUs(t *testing.T) {
}
}
+func TestCreateOrgAcceptsMaxMemory(t *testing.T) {
+ store := newFakeAPIStore()
+ router := newTestAPIRouter(store)
+
+ rec := adminJSON(t, router, http.MethodPost, "/api/v1/orgs",
+ `{"name":"analytics","database_name":"analytics","team_id":1,"max_memory":"120Gi"}`)
+ if rec.Code != http.StatusCreated {
+ t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusCreated, rec.Body.String())
+ }
+ if got := store.orgs["analytics"].MaxMemory; got != "120Gi" {
+ t.Fatalf("stored max_memory = %q, want 120Gi", got)
+ }
+ var response configstore.Org
+ if err := json.Unmarshal(rec.Body.Bytes(), &response); err != nil {
+ t.Fatalf("decode response: %v", err)
+ }
+ if response.MaxMemory != "120Gi" {
+ t.Fatalf("response max_memory = %q, want 120Gi", response.MaxMemory)
+ }
+}
+
+func TestUpdateOrgMaxMemoryPersistsAndAudits(t *testing.T) {
+ store := newFakeAPIStore()
+ store.orgs["analytics"] = &configstore.Org{
+ Name: "analytics", MaxWorkers: 2, MaxMemory: "120Gi",
+ }
+
+ gin.SetMode(gin.TestMode)
+ router := gin.New()
+ var detail string
+ router.Use(func(c *gin.Context) {
+ c.Next()
+ detail = c.GetString(ctxAuditDetailKey)
+ })
+ registerAPIWithStore(router.Group("/api/v1"), store, nil, nil)
+
+ rec := adminJSON(t, router, http.MethodPut, "/api/v1/orgs/analytics", `{"max_memory":"240Gi"}`)
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String())
+ }
+ if got := store.orgs["analytics"].MaxMemory; got != "240Gi" {
+ t.Fatalf("stored max_memory = %q, want 240Gi", got)
+ }
+ if got := store.orgs["analytics"].MaxWorkers; got != 2 {
+ t.Fatalf("max_workers = %d, want preserved 2", got)
+ }
+ if !strings.Contains(detail, "max_memory 120Gi → 240Gi") {
+ t.Fatalf("audit detail = %q, want max_memory change", detail)
+ }
+}
+
+func TestUpdateOrgOmittingMaxMemoryPreservesIt(t *testing.T) {
+ store := newFakeAPIStore()
+ store.orgs["analytics"] = &configstore.Org{
+ Name: "analytics", MaxWorkers: 2, MaxMemory: "120Gi",
+ }
+ router := newTestAPIRouter(store)
+
+ rec := adminJSON(t, router, http.MethodPut, "/api/v1/orgs/analytics", `{"max_workers":3}`)
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String())
+ }
+ if got := store.orgs["analytics"].MaxMemory; got != "120Gi" {
+ t.Fatalf("max_memory = %q, want preserved 120Gi", got)
+ }
+}
+
+func TestUpdateOrgMaxMemoryCanBeCleared(t *testing.T) {
+ for _, tc := range []struct {
+ name string
+ body string
+ }{
+ {name: "empty string", body: `{"max_memory":""}`},
+ {name: "null", body: `{"max_memory":null}`},
+ {name: "zero quantity", body: `{"max_memory":"0"}`},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ store := newFakeAPIStore()
+ store.orgs["analytics"] = &configstore.Org{Name: "analytics", MaxMemory: "120Gi"}
+ router := newTestAPIRouter(store)
+
+ rec := adminJSON(t, router, http.MethodPut, "/api/v1/orgs/analytics", tc.body)
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String())
+ }
+ if got := store.orgs["analytics"].MaxMemory; got != "" {
+ t.Fatalf("max_memory = %q, want cleared", got)
+ }
+ })
+ }
+}
+
+func TestOrgMutationRejectsInvalidMaxMemory(t *testing.T) {
+ for _, tc := range []struct {
+ name string
+ body string
+ }{
+ {name: "invalid quantity", body: `{"max_memory":"lots"}`},
+ {name: "negative quantity", body: `{"max_memory":"-1Gi"}`},
+ {name: "decimal overflow", body: `{"max_memory":"10000000000000000000"}`},
+ {name: "binary overflow", body: `{"max_memory":"8Ei"}`},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ store := newFakeAPIStore()
+ store.orgs["analytics"] = &configstore.Org{Name: "analytics", MaxMemory: "120Gi"}
+ router := newTestAPIRouter(store)
+
+ rec := adminJSON(t, router, http.MethodPut, "/api/v1/orgs/analytics", tc.body)
+ if rec.Code != http.StatusBadRequest {
+ t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusBadRequest, rec.Body.String())
+ }
+ if got := store.orgs["analytics"].MaxMemory; got != "120Gi" {
+ t.Fatalf("invalid update changed max_memory to %q", got)
+ }
+ })
+ }
+}
+
func TestUpdateOrgDataImportsTableNamingVersion(t *testing.T) {
for _, tc := range []struct {
name string
diff --git a/controlplane/admin/embed_ui_test.go b/controlplane/admin/embed_ui_test.go
index 9d4bc088..32b2e488 100644
--- a/controlplane/admin/embed_ui_test.go
+++ b/controlplane/admin/embed_ui_test.go
@@ -5,6 +5,7 @@ package admin
import (
"net/http"
"net/http/httptest"
+ "os"
"strings"
"testing"
@@ -49,3 +50,19 @@ func TestRegisterUIServesSPA(t *testing.T) {
})
}
}
+
+func TestModelsUIExplainsAndEditsOrgMaxMemory(t *testing.T) {
+ raw, err := os.ReadFile("static/models.html")
+ if err != nil {
+ t.Fatalf("read models UI: %v", err)
+ }
+ body := string(raw)
+ for _, want := range []string{
+ `max_memory: "Maximum active worker pod memory`,
+ `max_memory: "text"`,
+ } {
+ if !strings.Contains(body, want) {
+ t.Fatalf("models UI missing %q", want)
+ }
+ }
+}
diff --git a/controlplane/admin/monitoring.go b/controlplane/admin/monitoring.go
index c7b2d302..302cc143 100644
--- a/controlplane/admin/monitoring.go
+++ b/controlplane/admin/monitoring.go
@@ -45,6 +45,7 @@ type monitoringWarehouse struct {
type monitoringLimits struct {
MaxWorkers int `json:"max_workers"`
MaxVCPUs int `json:"max_vcpus"`
+ MaxMemoryBytes int64 `json:"max_memory_bytes"`
DefaultWorkerCPU string `json:"default_worker_cpu"`
DefaultWorkerMemory string `json:"default_worker_memory"`
DefaultWorkerTTLSeconds int `json:"default_worker_ttl_seconds"`
@@ -292,6 +293,7 @@ func (h *monitoringHandler) snapshot(c *gin.Context) {
Limits: monitoringLimits{
MaxWorkers: org.MaxWorkers,
MaxVCPUs: org.MaxVCPUs,
+ MaxMemoryBytes: org.MaxMemoryBytes,
DefaultWorkerCPU: configuredDefaults.CPU,
DefaultWorkerMemory: configuredDefaults.Memory,
DefaultWorkerTTLSeconds: int(configuredDefaults.TTL / time.Second),
diff --git a/controlplane/admin/monitoring_test.go b/controlplane/admin/monitoring_test.go
index 6336b799..a5bec602 100644
--- a/controlplane/admin/monitoring_test.go
+++ b/controlplane/admin/monitoring_test.go
@@ -57,6 +57,7 @@ func TestMonitoringSnapshotIsOrgScopedSanitizedAndReportsPartialCoverage(t *test
Name: "org-a",
MaxWorkers: 4,
MaxVCPUs: 8,
+ MaxMemoryBytes: 16 * 1024 * 1024 * 1024,
DefaultWorkerCPU: "1",
DefaultWorkerMemory: "2Gi",
DefaultWorkerTTL: "30m",
@@ -131,6 +132,9 @@ func TestMonitoringSnapshotIsOrgScopedSanitizedAndReportsPartialCoverage(t *test
if got.Limits.DefaultWorkerCPU != "1" || got.Limits.DefaultWorkerMemory != "2Gi" || got.Limits.DefaultWorkerTTLSeconds != 1800 {
t.Fatalf("configured defaults = %+v, want current org defaults", got.Limits)
}
+ if got.Limits.MaxMemoryBytes != 16*1024*1024*1024 {
+ t.Fatalf("max memory bytes = %d, want 16Gi", got.Limits.MaxMemoryBytes)
+ }
if got.Totals.Workers != 2 || got.Totals.AllocatedCPUCores != 2.75 || got.Totals.AllocatedMemoryBytes != 7*1024*1024*1024 {
t.Fatalf("worker totals = %+v, want 2 workers / 2.75 cores / 7Gi", got.Totals)
}
diff --git a/controlplane/admin/static/models.html b/controlplane/admin/static/models.html
index 84359c19..64027f29 100644
--- a/controlplane/admin/static/models.html
+++ b/controlplane/admin/static/models.html
@@ -366,6 +366,7 @@
Session Expired
database_name: "Database name clients connect with (dbname=). Unique across orgs.",
max_workers: "Maximum concurrent worker pods this org may run. 0 means use the deployment default.",
max_vcpus: "Maximum active worker pod vCPUs this org may have admitted. 0 means unlimited.",
+ max_memory: "Maximum active worker pod memory this org may have admitted, as a Kubernetes quantity. Empty means unlimited.",
default_worker_cpu: "Operator-set default worker pod CPU (k8s quantity, e.g. 2). Empty means unset. Applied when a connection does not size itself via duckgres.worker_* options.",
default_worker_memory: "Operator-set default worker pod memory (k8s quantity, e.g. 8Gi). Empty means unset.",
default_worker_ttl: "Operator-set default hot-idle TTL for workers (Go duration string, e.g. 75m). Empty means unset.",
@@ -453,6 +454,7 @@
Session Expired
pid: "Client backend PID associated with the request.",
protocol: "Session interface for the request (postgres or admin).",
requested_vcpus: "Worker pod vCPUs requested by this admission entry.",
+ requested_memory_bytes: "Worker pod memory bytes requested by this admission entry.",
enqueued_at: "When the request joined the FIFO queue; orders fairness.",
expires_at: "When this queue entry expires and is reclaimed.",
granted_at: "When admission was granted, if applicable.",
@@ -468,7 +470,8 @@
Session Expired
pid: "Client backend PID holding the lease.",
protocol: "Session interface for the lease (postgres or admin).",
requested_vcpus: "Worker pod vCPUs charged to the org and user budgets.",
- acquired_at: "When the lease was acquired; counts toward active vCPU resource budgets.",
+ requested_memory_bytes: "Worker pod memory bytes charged to the org budget.",
+ acquired_at: "When the lease was acquired; counts toward active vCPU and memory resource budgets.",
created_at: "When this lease record was created.",
updated_at: "When this lease record was last modified.",
},
@@ -483,8 +486,8 @@
Session Expired
"managed-warehouses": "The config-store source of truth for an org's managed warehouse: image, metadata/object stores, catalogs, identity, and provisioning status.",
"cp-instances": "A runtime coordination record for one control-plane process, with liveness state and heartbeat.",
"worker-records": "A durable runtime record for one worker pod: lifecycle state, owner, profile, TTL, and credential expiry.",
- "org-connection-queue": "A cluster-wide FIFO resource admission request for one org/user session; rows expire quickly and coordinate fairness across CP replicas.",
- "org-connection-leases": "A durable cluster-wide resource admission lease for a live session; active leases count toward org and user vCPU budgets.",
+ "org-connection-queue": "A cluster-wide FIFO CPU and memory admission request for one org/user session; rows expire quickly and coordinate fairness across CP replicas.",
+ "org-connection-leases": "A durable cluster-wide resource admission lease for a live session; active leases count toward org memory and org/user vCPU budgets.",
};
// Per-org aux data computed client-side: the live /models/orgs listing does
@@ -715,6 +718,7 @@
Session Expired
const ORG_EDIT_FIELDS = {
max_workers: "int",
max_vcpus: "int",
+ max_memory: "text",
default_worker_cpu: "text",
default_worker_memory: "text",
default_worker_ttl: "text",
diff --git a/controlplane/admin/ui/src/pages/OrgDetail.test.tsx b/controlplane/admin/ui/src/pages/OrgDetail.test.tsx
index 1ca375b4..42e080d8 100644
--- a/controlplane/admin/ui/src/pages/OrgDetail.test.tsx
+++ b/controlplane/admin/ui/src/pages/OrgDetail.test.tsx
@@ -54,6 +54,7 @@ const ORG: Org = {
hostname_alias: null,
max_workers: 1,
max_vcpus: 2,
+ max_memory: "120Gi",
default_worker_cpu: "2",
default_worker_memory: "8Gi",
default_worker_ttl: "75m",
@@ -215,6 +216,19 @@ describe("Org detail", () => {
expect(orgUpdate).toHaveBeenCalledWith(expect.objectContaining({ database_name: "acme-inc" }));
});
+ it("shows and saves the org memory limit", async () => {
+ const user = userEvent.setup();
+ renderPage(false);
+
+ const input = screen.getByLabelText(/max memory/i);
+ expect(input).toHaveValue("120Gi");
+ await user.clear(input);
+ await user.type(input, "240Gi");
+ await user.click(screen.getByText("Save changes"));
+
+ expect(orgUpdate).toHaveBeenCalledWith(expect.objectContaining({ max_memory: "240Gi" }));
+ });
+
it("keeps every other setting editable for an org whose stored name predates the rule", async () => {
// The premise of the break-glass surface: grandfathered rows like
// "ACME INC" must NOT wedge the whole org-config card behind a forced
diff --git a/controlplane/admin/ui/src/pages/OrgDetail.tsx b/controlplane/admin/ui/src/pages/OrgDetail.tsx
index 889e4aac..b74cab02 100644
--- a/controlplane/admin/ui/src/pages/OrgDetail.tsx
+++ b/controlplane/admin/ui/src/pages/OrgDetail.tsx
@@ -54,6 +54,7 @@ interface FormState {
database_name: string;
max_workers: string;
max_vcpus: string;
+ max_memory: string;
default_worker_cpu: string;
default_worker_memory: string;
default_worker_ttl: string;
@@ -66,6 +67,7 @@ function orgToForm(o: {
database_name: string;
max_workers: number;
max_vcpus: number;
+ max_memory: string;
default_worker_cpu: string;
default_worker_memory: string;
default_worker_ttl: string;
@@ -77,6 +79,7 @@ function orgToForm(o: {
database_name: o.database_name,
max_workers: String(o.max_workers),
max_vcpus: String(o.max_vcpus),
+ max_memory: o.max_memory,
default_worker_cpu: o.default_worker_cpu,
default_worker_memory: o.default_worker_memory,
default_worker_ttl: o.default_worker_ttl,
@@ -149,6 +152,7 @@ export function OrgDetail() {
const body: OrgUpdate = {
max_workers: Number(form.max_workers) || 0,
max_vcpus: Number(form.max_vcpus) || 0,
+ max_memory: form.max_memory,
default_worker_cpu: form.default_worker_cpu,
default_worker_memory: form.default_worker_memory,
default_worker_ttl: form.default_worker_ttl,
@@ -248,6 +252,14 @@ export function OrgDetail() {
onChange={(e) => set("max_vcpus", e.target.value)}
/>
+
+ set("max_memory", e.target.value)}
+ />
+
{v === 0 ? "∞" : fmtInt(v)};
},
},
+ {
+ accessorKey: "max_memory",
+ header: "Max memory",
+ cell: ({ getValue }) => {
+ const v = getValue() as string;
+ return {v === "" ? "∞" : v};
+ },
+ },
{
id: "users",
header: "Users",
@@ -164,7 +172,7 @@ export function Orgs() {
)}
{orgs.isLoading ? (
-
+
) : orgs.isError ? (
orgs.refetch()} />
) : (
diff --git a/controlplane/admin/ui/src/types/api.ts b/controlplane/admin/ui/src/types/api.ts
index 6de8b9be..33fa1c57 100644
--- a/controlplane/admin/ui/src/types/api.ts
+++ b/controlplane/admin/ui/src/types/api.ts
@@ -49,6 +49,8 @@ export interface Org {
teams?: OrgTeam[];
max_workers: number;
max_vcpus: number;
+ // Kubernetes memory quantity; empty means unlimited.
+ max_memory: string;
default_worker_cpu: string;
default_worker_memory: string;
default_worker_ttl: string;
@@ -70,6 +72,7 @@ export interface OrgUpdate {
database_name?: string;
max_workers?: number;
max_vcpus?: number;
+ max_memory?: string;
default_worker_cpu?: string;
default_worker_memory?: string;
default_worker_ttl?: string;
diff --git a/controlplane/admission_metrics.go b/controlplane/admission_metrics.go
index f85b578e..fc7d9969 100644
--- a/controlplane/admission_metrics.go
+++ b/controlplane/admission_metrics.go
@@ -33,6 +33,16 @@ var sessionAdmissionLimitVCPUsGauge = promauto.NewGaugeVec(prometheus.GaugeOpts{
Help: "Effective org session-admission vCPU limit for active org stacks, reconciled from this control-plane process's current config snapshot; zero means unlimited and max by org collapses replica duplication.",
}, []string{"org"})
+var sessionAdmissionActiveMemoryBytesGauge = promauto.NewGaugeVec(prometheus.GaugeOpts{
+ Name: "duckgres_session_admission_active_memory_bytes",
+ Help: "Requested memory bytes held by live local admission lease handles; cleanup-pending durable rows are excluded.",
+}, []string{"org"})
+
+var sessionAdmissionLimitMemoryBytesGauge = promauto.NewGaugeVec(prometheus.GaugeOpts{
+ Name: "duckgres_session_admission_limit_memory_bytes",
+ Help: "Effective org session-admission memory limit in bytes for active org stacks, reconciled from this control-plane process's current config snapshot; zero means unlimited and max by org collapses replica duplication.",
+}, []string{"org"})
+
func observeSessionAdmissionTerminal(org, outcome, reason string, wait time.Duration) {
if wait < 0 {
wait = 0
@@ -56,7 +66,7 @@ func normalizedSessionAdmissionReason(reason string) string {
switch reason {
case "", "none":
return "none"
- case "org_vcpu", "user_vcpu", "org_user_vcpu", "user_ineligible", "resharding", "fifo", "store_error", "mixed":
+ case "org_vcpu", "user_vcpu", "org_user_vcpu", "org_memory", "resource_mixed", "user_ineligible", "resharding", "fifo", "store_error", "mixed":
return reason
default:
return "store_error"
@@ -64,10 +74,11 @@ func normalizedSessionAdmissionReason(reason string) string {
}
type sessionAdmissionReasons struct {
- firstNonVCPU string
- mixedNonVCPU bool
- orgVCPU bool
- userVCPU bool
+ firstNonResource string
+ mixedNonResource bool
+ orgVCPU bool
+ userVCPU bool
+ orgMemory bool
}
func (r *sessionAdmissionReasons) add(reason string) {
@@ -86,17 +97,26 @@ func (r *sessionAdmissionReasons) add(reason string) {
r.orgVCPU = true
r.userVCPU = true
return
+ case "org_memory":
+ r.orgMemory = true
+ return
}
- if r.firstNonVCPU == "" {
- r.firstNonVCPU = reason
+ if r.firstNonResource == "" {
+ r.firstNonResource = reason
return
}
- if r.firstNonVCPU != reason {
- r.mixedNonVCPU = true
+ if r.firstNonResource != reason {
+ r.mixedNonResource = true
}
}
func (r sessionAdmissionReasons) value() string {
+ if r.orgMemory && (r.orgVCPU || r.userVCPU) {
+ return "resource_mixed"
+ }
+ if r.orgMemory {
+ return "org_memory"
+ }
if r.orgVCPU && r.userVCPU {
return "org_user_vcpu"
}
@@ -106,11 +126,11 @@ func (r sessionAdmissionReasons) value() string {
if r.userVCPU {
return "user_vcpu"
}
- if r.mixedNonVCPU {
+ if r.mixedNonResource {
return "mixed"
}
- if r.firstNonVCPU == "" {
+ if r.firstNonResource == "" {
return "none"
}
- return r.firstNonVCPU
+ return r.firstNonResource
}
diff --git a/controlplane/admission_metrics_test.go b/controlplane/admission_metrics_test.go
index e9a9cc13..5074de2c 100644
--- a/controlplane/admission_metrics_test.go
+++ b/controlplane/admission_metrics_test.go
@@ -16,6 +16,8 @@ func TestSessionAdmissionReasonsValue(t *testing.T) {
{name: "user vCPU takes precedence", reasons: []string{"fifo", "user_vcpu", "store_error"}, want: "user_vcpu"},
{name: "both vCPU limits", reasons: []string{"org_vcpu", "fifo", "user_vcpu"}, want: "org_user_vcpu"},
{name: "both vCPU limits in reverse order", reasons: []string{"user_vcpu", "resharding", "org_vcpu"}, want: "org_user_vcpu"},
+ {name: "org memory takes precedence over queue reason", reasons: []string{"fifo", "org_memory"}, want: "org_memory"},
+ {name: "memory and vCPU blockers are resource mixed", reasons: []string{"org_memory", "org_vcpu"}, want: "resource_mixed"},
{name: "unknown reason is bounded", reasons: []string{"future_dynamic_reason"}, want: "store_error"},
}
@@ -31,3 +33,9 @@ func TestSessionAdmissionReasonsValue(t *testing.T) {
})
}
}
+
+func TestNormalizedSessionAdmissionReasonKeepsResourceMixed(t *testing.T) {
+ if got := normalizedSessionAdmissionReason("resource_mixed"); got != "resource_mixed" {
+ t.Fatalf("normalized reason = %q, want resource_mixed", got)
+ }
+}
diff --git a/controlplane/configstore/migrations/000037_add_org_max_memory.sql b/controlplane/configstore/migrations/000037_add_org_max_memory.sql
new file mode 100644
index 00000000..3de8567a
--- /dev/null
+++ b/controlplane/configstore/migrations/000037_add_org_max_memory.sql
@@ -0,0 +1,9 @@
+-- +goose Up
+
+ALTER TABLE duckgres_orgs
+ ADD COLUMN IF NOT EXISTS max_memory VARCHAR(32) NOT NULL DEFAULT '';
+
+-- +goose Down
+
+ALTER TABLE duckgres_orgs
+ DROP COLUMN IF EXISTS max_memory;
diff --git a/controlplane/configstore/models.go b/controlplane/configstore/models.go
index 192c1354..d772aa90 100644
--- a/controlplane/configstore/models.go
+++ b/controlplane/configstore/models.go
@@ -15,6 +15,10 @@ type Org struct {
HostnameAlias *string `gorm:"size:255;uniqueIndex" json:"hostname_alias"`
MaxWorkers int `gorm:"default:0" json:"max_workers"`
MaxVCPUs int `gorm:"column:max_vcpus;default:0" json:"max_vcpus"`
+ // MaxMemory is the org-wide active-session memory ceiling as a Kubernetes
+ // resource quantity. Empty means unlimited; the config snapshot carries
+ // the parsed byte value used by readers that do not need the raw quantity.
+ MaxMemory string `gorm:"column:max_memory;size:32;not null;default:''" json:"max_memory"`
// DefaultWorkerCPU/Memory/TTL are the org's operator-set default worker
// profile: the pod shape (k8s resource quantities, e.g. "2"/"8Gi") and
// hot-idle TTL (Go duration string, e.g. "75m" — stored as a string for
@@ -507,23 +511,25 @@ func (WorkerRecord) TableName() string { return "worker_records" }
// OrgResourceLimits is the current resource-admission ceiling for an org and
// the connecting user. 0 means unlimited for either dimension.
type OrgResourceLimits struct {
- OrgMaxVCPUs int
- UserMaxVCPUs int
+ OrgMaxVCPUs int
+ OrgMaxMemoryBytes int64
+ UserMaxVCPUs int
}
// OrgConnectionQueueEntry is a cluster-wide FIFO admission request for one org
// connection. Rows expire quickly; they coordinate fairness across CP replicas.
type OrgConnectionQueueEntry struct {
- RequestID string `gorm:"primaryKey;size:64" json:"request_id"`
- OrgID string `gorm:"size:255;not null;index:idx_org_connection_queue_pending,priority:1" json:"org_id"`
- Username string `gorm:"size:255;index" json:"username"`
- CPInstanceID string `gorm:"size:255;not null;index" json:"cp_instance_id"`
- PID int32 `gorm:"not null" json:"pid"`
- Protocol string `gorm:"size:32;not null" json:"protocol"`
- RequestedVCPUs int `gorm:"column:requested_vcpus;not null;default:1" json:"requested_vcpus"`
- EnqueuedAt time.Time `gorm:"not null;index:idx_org_connection_queue_pending,priority:2" json:"enqueued_at"`
- ExpiresAt time.Time `gorm:"not null;index" json:"expires_at"`
- GrantedAt *time.Time `gorm:"index" json:"granted_at,omitempty"`
+ RequestID string `gorm:"primaryKey;size:64" json:"request_id"`
+ OrgID string `gorm:"size:255;not null;index:idx_org_connection_queue_pending,priority:1" json:"org_id"`
+ Username string `gorm:"size:255;index" json:"username"`
+ CPInstanceID string `gorm:"size:255;not null;index" json:"cp_instance_id"`
+ PID int32 `gorm:"not null" json:"pid"`
+ Protocol string `gorm:"size:32;not null" json:"protocol"`
+ RequestedVCPUs int `gorm:"column:requested_vcpus;not null;default:1" json:"requested_vcpus"`
+ RequestedMemoryBytes int64 `gorm:"column:requested_memory_bytes;not null;default:0" json:"requested_memory_bytes"`
+ EnqueuedAt time.Time `gorm:"not null;index:idx_org_connection_queue_pending,priority:2" json:"enqueued_at"`
+ ExpiresAt time.Time `gorm:"not null;index" json:"expires_at"`
+ GrantedAt *time.Time `gorm:"index" json:"granted_at,omitempty"`
// canceled_at was dropped: cancellation is a hard DELETE of the row, so the
// column was never set to a non-NULL value.
CreatedAt time.Time `json:"created_at"`
@@ -536,17 +542,18 @@ func (OrgConnectionQueueEntry) TableName() string { return "org_connection_queue
// session. Capacity checks count active leases, ignoring owners whose CP row
// has expired.
type OrgConnectionLease struct {
- LeaseID string `gorm:"primaryKey;size:64" json:"lease_id"`
- RequestID string `gorm:"size:64;not null;uniqueIndex" json:"request_id"`
- OrgID string `gorm:"size:255;not null;index" json:"org_id"`
- Username string `gorm:"size:255;index" json:"username"`
- CPInstanceID string `gorm:"size:255;not null;index" json:"cp_instance_id"`
- PID int32 `gorm:"not null" json:"pid"`
- Protocol string `gorm:"size:32;not null" json:"protocol"`
- RequestedVCPUs int `gorm:"column:requested_vcpus;not null;default:1" json:"requested_vcpus"`
- AcquiredAt time.Time `gorm:"not null" json:"acquired_at"`
- CreatedAt time.Time `json:"created_at"`
- UpdatedAt time.Time `json:"updated_at"`
+ LeaseID string `gorm:"primaryKey;size:64" json:"lease_id"`
+ RequestID string `gorm:"size:64;not null;uniqueIndex" json:"request_id"`
+ OrgID string `gorm:"size:255;not null;index" json:"org_id"`
+ Username string `gorm:"size:255;index" json:"username"`
+ CPInstanceID string `gorm:"size:255;not null;index" json:"cp_instance_id"`
+ PID int32 `gorm:"not null" json:"pid"`
+ Protocol string `gorm:"size:32;not null" json:"protocol"`
+ RequestedVCPUs int `gorm:"column:requested_vcpus;not null;default:1" json:"requested_vcpus"`
+ RequestedMemoryBytes int64 `gorm:"column:requested_memory_bytes;not null;default:0" json:"requested_memory_bytes"`
+ AcquiredAt time.Time `gorm:"not null" json:"acquired_at"`
+ CreatedAt time.Time `json:"created_at"`
+ UpdatedAt time.Time `json:"updated_at"`
}
func (OrgConnectionLease) TableName() string { return "org_connection_leases" }
@@ -563,6 +570,7 @@ type OrgConfig struct {
HostnameAlias string // empty when no alias is configured
MaxWorkers int
MaxVCPUs int
+ MaxMemoryBytes int64 // parsed org max_memory; 0 = unlimited
DefaultWorkerCPU string // org default worker profile: pod cpu quantity ("" = unset)
DefaultWorkerMemory string // org default worker profile: pod memory quantity ("" = unset)
DefaultWorkerTTL string // org default worker profile: hot-idle TTL, Go duration string ("" = unset)
diff --git a/controlplane/configstore/org_connection_metrics.go b/controlplane/configstore/org_connection_metrics.go
index 8aeb03d9..addd6016 100644
--- a/controlplane/configstore/org_connection_metrics.go
+++ b/controlplane/configstore/org_connection_metrics.go
@@ -11,9 +11,11 @@ const (
orgConnectionAdmissionOutcomeGranted = "granted"
orgConnectionAdmissionOutcomeAlreadyGranted = "already_granted"
orgConnectionAdmissionOutcomeBlockedOrgVCPU = "blocked_org_vcpu"
+ orgConnectionAdmissionOutcomeBlockedOrgMemory = "blocked_org_memory"
orgConnectionAdmissionOutcomeBlockedUserVCPU = "blocked_user_vcpu"
orgConnectionAdmissionOutcomeBlockedOrgUserVCPU = "blocked_org_user_vcpu"
orgConnectionAdmissionOutcomeRejectedOrgVCPU = "rejected_org_vcpu"
+ orgConnectionAdmissionOutcomeRejectedOrgMemory = "rejected_org_memory"
orgConnectionAdmissionOutcomeRejectedUserVCPU = "rejected_user_vcpu"
orgConnectionAdmissionOutcomeIneligibleUser = "ineligible_user"
orgConnectionAdmissionOutcomeInactive = "inactive_request"
@@ -62,12 +64,16 @@ func orgConnectionAdmissionEvaluationForOutcome(outcome string) OrgConnectionAdm
return OrgConnectionAdmissionEvaluation{Decision: "already_granted", Reason: "none"}
case orgConnectionAdmissionOutcomeBlockedOrgVCPU:
return OrgConnectionAdmissionEvaluation{Decision: "blocked", Reason: "org_vcpu"}
+ case orgConnectionAdmissionOutcomeBlockedOrgMemory:
+ return OrgConnectionAdmissionEvaluation{Decision: "blocked", Reason: "org_memory"}
case orgConnectionAdmissionOutcomeBlockedUserVCPU:
return OrgConnectionAdmissionEvaluation{Decision: "blocked", Reason: "user_vcpu"}
case orgConnectionAdmissionOutcomeBlockedOrgUserVCPU:
return OrgConnectionAdmissionEvaluation{Decision: "blocked", Reason: "org_user_vcpu"}
case orgConnectionAdmissionOutcomeRejectedOrgVCPU:
return OrgConnectionAdmissionEvaluation{Decision: "rejected", Reason: "org_vcpu"}
+ case orgConnectionAdmissionOutcomeRejectedOrgMemory:
+ return OrgConnectionAdmissionEvaluation{Decision: "rejected", Reason: "org_memory"}
case orgConnectionAdmissionOutcomeRejectedUserVCPU:
return OrgConnectionAdmissionEvaluation{Decision: "rejected", Reason: "user_vcpu"}
case orgConnectionAdmissionOutcomeIneligibleUser:
diff --git a/controlplane/configstore/org_connection_metrics_test.go b/controlplane/configstore/org_connection_metrics_test.go
index 57a6bffc..512773f8 100644
--- a/controlplane/configstore/org_connection_metrics_test.go
+++ b/controlplane/configstore/org_connection_metrics_test.go
@@ -11,9 +11,11 @@ func TestOrgConnectionAdmissionEvaluationForOutcome(t *testing.T) {
{orgConnectionAdmissionOutcomeGranted, "granted_current", "none"},
{orgConnectionAdmissionOutcomeAlreadyGranted, "already_granted", "none"},
{orgConnectionAdmissionOutcomeBlockedOrgVCPU, "blocked", "org_vcpu"},
+ {orgConnectionAdmissionOutcomeBlockedOrgMemory, "blocked", "org_memory"},
{orgConnectionAdmissionOutcomeBlockedUserVCPU, "blocked", "user_vcpu"},
{orgConnectionAdmissionOutcomeBlockedOrgUserVCPU, "blocked", "org_user_vcpu"},
{orgConnectionAdmissionOutcomeRejectedOrgVCPU, "rejected", "org_vcpu"},
+ {orgConnectionAdmissionOutcomeRejectedOrgMemory, "rejected", "org_memory"},
{orgConnectionAdmissionOutcomeRejectedUserVCPU, "rejected", "user_vcpu"},
{orgConnectionAdmissionOutcomeIneligibleUser, "blocked", "user_ineligible"},
{orgConnectionAdmissionOutcomeResharding, "blocked", "resharding"},
diff --git a/controlplane/configstore/org_connections.go b/controlplane/configstore/org_connections.go
index 1e6ea454..da27a903 100644
--- a/controlplane/configstore/org_connections.go
+++ b/controlplane/configstore/org_connections.go
@@ -14,11 +14,13 @@ import (
const (
missingOwnerOrgConnectionLeaseGrace = 5 * time.Minute
legacyOrgConnectionRequestedVCPUs = 1
+ maximumOrgConnectionMemoryBytes = int64(^uint64(0) >> 1)
)
// ErrOrgConnectionAdmissionRejected identifies a request that can never fit
-// under its configured hard vCPU ceiling. Temporary saturation does not wrap
-// this sentinel; those requests stay queued until capacity becomes available.
+// under its configured hard resource ceilings. Temporary saturation does not
+// wrap this sentinel; those requests stay queued until capacity becomes
+// available.
var ErrOrgConnectionAdmissionRejected = errors.New("org connection admission rejected")
// OrgConnectionAdmissionRef is the immutable identity of one connection
@@ -47,22 +49,28 @@ func (r OrgConnectionAdmissionRef) validate() error {
type OrgConnectionAdmissionRejectionReason string
const (
- OrgConnectionAdmissionRejectedOrgVCPU OrgConnectionAdmissionRejectionReason = "org_vcpu"
- OrgConnectionAdmissionRejectedUserVCPU OrgConnectionAdmissionRejectionReason = "user_vcpu"
+ OrgConnectionAdmissionRejectedOrgVCPU OrgConnectionAdmissionRejectionReason = "org_vcpu"
+ OrgConnectionAdmissionRejectedOrgMemory OrgConnectionAdmissionRejectionReason = "org_memory"
+ OrgConnectionAdmissionRejectedUserVCPU OrgConnectionAdmissionRejectionReason = "user_vcpu"
)
// OrgConnectionAdmissionRejectedError carries the stable reason and values
// needed to return an actionable PostgreSQL configuration-limit error.
type OrgConnectionAdmissionRejectedError struct {
- Reason OrgConnectionAdmissionRejectionReason
- RequestedVCPUs int
- MaximumVCPUs int
+ Reason OrgConnectionAdmissionRejectionReason
+ RequestedVCPUs int
+ MaximumVCPUs int
+ RequestedMemoryBytes int64
+ MaximumMemoryBytes int64
}
func (e *OrgConnectionAdmissionRejectedError) Error() string {
if e == nil {
return ErrOrgConnectionAdmissionRejected.Error()
}
+ if e.Reason == OrgConnectionAdmissionRejectedOrgMemory {
+ return fmt.Sprintf("%s: requested %d memory bytes exceeds %s maximum of %d memory bytes", ErrOrgConnectionAdmissionRejected, e.RequestedMemoryBytes, e.Reason, e.MaximumMemoryBytes)
+ }
return fmt.Sprintf("%s: requested %d vCPUs exceeds %s maximum of %d vCPUs", ErrOrgConnectionAdmissionRejected, e.RequestedVCPUs, e.Reason, e.MaximumVCPUs)
}
@@ -80,8 +88,8 @@ type OrgConnectionMonitoringStatus struct {
// EnqueueOrgConnectionRequest inserts a pending cluster-wide connection
// admission request. FIFO ordering is scoped to org_id and ordered by
-// enqueued_at, then request_id. RequestedVCPUs is charged against active
-// resource leases when the request is granted.
+// enqueued_at, then request_id. RequestedVCPUs and RequestedMemoryBytes are
+// charged against active resource leases when the request is granted.
func (cs *ConfigStore) EnqueueOrgConnectionRequest(entry *OrgConnectionQueueEntry) error {
return cs.EnqueueOrgConnectionRequestContext(context.Background(), entry)
}
@@ -136,10 +144,10 @@ func (cs *ConfigStore) EnqueueOrgConnectionRequestContext(ctx context.Context, e
}
// TryAcquireOrgConnectionLease attempts to grant one queued request under
-// cluster-wide per-org and per-user vCPU budgets. It is retained for callers
-// that do not yet pass their control-plane identity explicitly. New runtime
-// callers use ScheduleAndClaimOrgConnectionLeaseForRef so owner validation is
-// part of the claim transaction.
+// cluster-wide per-org resource budgets and per-user vCPU budgets. It is
+// retained for callers that do not yet pass their control-plane identity
+// explicitly. New runtime callers use ScheduleAndClaimOrgConnectionLeaseForRef
+// so owner validation is part of the claim transaction.
//
// Legacy compatibility: This ID-only adapter assumes randomly generated
// request IDs are globally unique and never reused. New callers use
@@ -408,6 +416,13 @@ func (cs *ConfigStore) scheduleAndClaimOrgConnectionLeaseOnce(ctx context.Contex
MaximumVCPUs: requestLimits.OrgMaxVCPUs,
}
outcome = orgConnectionAdmissionOutcomeRejectedOrgVCPU
+ case requestLimits.OrgMaxMemoryBytes > 0 && request.RequestedMemoryBytes > requestLimits.OrgMaxMemoryBytes:
+ rejection = &OrgConnectionAdmissionRejectedError{
+ Reason: OrgConnectionAdmissionRejectedOrgMemory,
+ RequestedMemoryBytes: request.RequestedMemoryBytes,
+ MaximumMemoryBytes: requestLimits.OrgMaxMemoryBytes,
+ }
+ outcome = orgConnectionAdmissionOutcomeRejectedOrgMemory
case requestLimits.UserMaxVCPUs > 0 && request.RequestedVCPUs > requestLimits.UserMaxVCPUs:
rejection = &OrgConnectionAdmissionRejectedError{
Reason: OrgConnectionAdmissionRejectedUserVCPU,
@@ -565,6 +580,7 @@ type authoritativeOrgConnectionUserLimit struct {
type authoritativeOrgConnectionLimitSet struct {
orgMaxVCPUs int64
+ orgMaxMemoryBytes int64
users map[string]authoritativeOrgConnectionUserLimit
serviceCredential map[string]struct{}
}
@@ -574,12 +590,16 @@ func (l authoritativeOrgConnectionLimitSet) lookup(username string) OrgResourceL
// Service credentials are root-shaped org identities. They share the
// org budget but never inherit an ordinary user's per-user cap, even if
// a legacy/user row happens to have the same name.
- return OrgResourceLimits{OrgMaxVCPUs: int(l.orgMaxVCPUs)}
+ return OrgResourceLimits{
+ OrgMaxVCPUs: int(l.orgMaxVCPUs),
+ OrgMaxMemoryBytes: l.orgMaxMemoryBytes,
+ }
}
user := l.users[username]
return OrgResourceLimits{
- OrgMaxVCPUs: int(l.orgMaxVCPUs),
- UserMaxVCPUs: int(user.maxVCPUs),
+ OrgMaxVCPUs: int(l.orgMaxVCPUs),
+ OrgMaxMemoryBytes: l.orgMaxMemoryBytes,
+ UserMaxVCPUs: int(user.maxVCPUs),
}
}
@@ -594,11 +614,12 @@ func (l authoritativeOrgConnectionLimitSet) userAllowed(username string) bool {
func (cs *ConfigStore) authoritativeOrgConnectionLimits(tx *gorm.DB, orgID string) (authoritativeOrgConnectionLimitSet, bool, error) {
type orgLimitRow struct {
- MaxVCPUs *int64 `gorm:"column:max_vcpus"`
+ MaxVCPUs *int64 `gorm:"column:max_vcpus"`
+ MaxMemory string `gorm:"column:max_memory"`
}
var orgRow orgLimitRow
if err := tx.Model(&Org{}).
- Select("max_vcpus").
+ Select("max_vcpus, max_memory").
Where("name = ?", orgID).
Take(&orgRow).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
@@ -617,6 +638,11 @@ func (cs *ConfigStore) authoritativeOrgConnectionLimits(tx *gorm.DB, orgID strin
if limits.orgMaxVCPUs < 0 {
return authoritativeOrgConnectionLimitSet{}, false, fmt.Errorf("org %q has invalid negative max_vcpus", orgID)
}
+ maxMemoryBytes, err := ParseOrgMaxMemoryBytes(orgRow.MaxMemory)
+ if err != nil {
+ return authoritativeOrgConnectionLimitSet{}, false, fmt.Errorf("org %q has invalid max_memory: %w", orgID, err)
+ }
+ limits.orgMaxMemoryBytes = maxMemoryBytes
type userLimitRow struct {
Username string `gorm:"column:username"`
@@ -724,6 +750,7 @@ func (cs *ConfigStore) nextEligibleOrgConnectionRequestLocked(tx *gorm.DB, table
limits = targetLimits
}
requested := int64(head.RequestedVCPUs)
+ requestedMemoryBytes := head.RequestedMemoryBytes
// A permanently impossible foreign head must not block unrelated
// requests. Only that request's owner deletes and receives the typed
// rejection; other evaluators simply skip it.
@@ -733,6 +760,9 @@ func (cs *ConfigStore) nextEligibleOrgConnectionRequestLocked(tx *gorm.DB, table
if limits.OrgMaxVCPUs > 0 && requested > int64(limits.OrgMaxVCPUs) {
continue
}
+ if limits.OrgMaxMemoryBytes > 0 && requestedMemoryBytes > limits.OrgMaxMemoryBytes {
+ continue
+ }
if limits.UserMaxVCPUs > 0 {
used := legacyUserUsed + userUsed[head.Username]
if used+requested > int64(limits.UserMaxVCPUs) {
@@ -742,6 +772,15 @@ func (cs *ConfigStore) nextEligibleOrgConnectionRequestLocked(tx *gorm.DB, table
if limits.OrgMaxVCPUs > 0 && orgUsed+requested > int64(limits.OrgMaxVCPUs) {
return nil, orgConnectionAdmissionOutcomeBlockedOrgVCPU, nil
}
+ if limits.OrgMaxMemoryBytes > 0 {
+ orgMemoryUsed, orgMemoryUsageKnown, err := cs.activeOrgConnectionLeaseMemoryUsage(tx, target.OrgID)
+ if err != nil {
+ return nil, orgConnectionAdmissionOutcomeError, err
+ }
+ if orgConnectionMemoryLimitExceeded(limits.OrgMaxMemoryBytes, orgMemoryUsed, orgMemoryUsageKnown, requestedMemoryBytes) {
+ return nil, orgConnectionAdmissionOutcomeBlockedOrgMemory, nil
+ }
+ }
if sameOrgConnectionAdmissionRequest(head, target) {
return head, orgConnectionAdmissionOutcomeGranted, nil
@@ -752,6 +791,16 @@ func (cs *ConfigStore) nextEligibleOrgConnectionRequestLocked(tx *gorm.DB, table
return nil, orgConnectionAdmissionOutcomeWaiting, nil
}
+func orgConnectionMemoryLimitExceeded(limit, used int64, usageKnown bool, requested int64) bool {
+ if limit <= 0 {
+ return false
+ }
+ if !usageKnown || requested <= 0 || used >= limit {
+ return true
+ }
+ return requested > limit-used
+}
+
func sameOrgConnectionAdmissionRequest(a, b *OrgConnectionQueueEntry) bool {
return a != nil && b != nil &&
a.RequestID == b.RequestID &&
@@ -809,18 +858,42 @@ func (cs *ConfigStore) activeOrgConnectionLeaseVCPUUsage(tx *gorm.DB, orgID stri
return orgUsed, userUsed, legacyUserUsed, nil
}
+func (cs *ConfigStore) activeOrgConnectionLeaseMemoryUsage(tx *gorm.DB, orgID string) (int64, bool, error) {
+ type memoryUsageRow struct {
+ MemoryBytes int64 `gorm:"column:memory_bytes"`
+ UnknownCount int64 `gorm:"column:unknown_count"`
+ }
+ var usage memoryUsageRow
+ leaseTable := cs.runtimeTable((&OrgConnectionLease{}).TableName())
+ cpTable := cs.runtimeTable((&ControlPlaneInstance{}).TableName())
+ // PostgreSQL promotes SUM(bigint) to numeric, so a corrupt or extreme set
+ // of rows can exceed int64 even though each requested value cannot. Clamp
+ // before scanning; the capped value is at least every representable limit,
+ // which keeps admission fail-closed without an overflowing Go addition.
+ if err := tx.Table(leaseTable+" AS l").
+ Select("LEAST(COALESCE(SUM(CASE WHEN l.requested_memory_bytes > 0 THEN l.requested_memory_bytes ELSE 0 END), 0), ?)::bigint AS memory_bytes, COUNT(*) FILTER (WHERE l.requested_memory_bytes <= 0) AS unknown_count", maximumOrgConnectionMemoryBytes).
+ Joins("LEFT JOIN "+cpTable+" AS cp ON cp.id = l.cp_instance_id").
+ Where("l.org_id = ?", orgID).
+ Where("cp.id IS NULL OR cp.state <> ?", ControlPlaneInstanceStateExpired).
+ Scan(&usage).Error; err != nil {
+ return 0, false, err
+ }
+ return usage.MemoryBytes, usage.UnknownCount == 0, nil
+}
+
func (cs *ConfigStore) createOrgConnectionLease(tx *gorm.DB, request *OrgConnectionQueueEntry, now time.Time) (*OrgConnectionLease, error) {
granted := now
created := &OrgConnectionLease{
- LeaseID: request.RequestID,
- RequestID: request.RequestID,
- OrgID: request.OrgID,
- Username: request.Username,
- CPInstanceID: request.CPInstanceID,
- PID: request.PID,
- Protocol: request.Protocol,
- RequestedVCPUs: request.RequestedVCPUs,
- AcquiredAt: now,
+ LeaseID: request.RequestID,
+ RequestID: request.RequestID,
+ OrgID: request.OrgID,
+ Username: request.Username,
+ CPInstanceID: request.CPInstanceID,
+ PID: request.PID,
+ Protocol: request.Protocol,
+ RequestedVCPUs: request.RequestedVCPUs,
+ RequestedMemoryBytes: request.RequestedMemoryBytes,
+ AcquiredAt: now,
}
if err := tx.Table(cs.runtimeTable(created.TableName())).Create(created).Error; err != nil {
return nil, err
diff --git a/controlplane/configstore/store.go b/controlplane/configstore/store.go
index 5f3d076b..2e9ba984 100644
--- a/controlplane/configstore/store.go
+++ b/controlplane/configstore/store.go
@@ -9,6 +9,7 @@ import (
"fmt"
"hash/fnv"
"log/slog"
+ "math"
"sort"
"strings"
"sync"
@@ -240,6 +241,51 @@ func (cs *ConfigStore) Start(ctx context.Context) {
}()
}
+// ParseOrgMaxMemoryBytes parses an org max_memory Kubernetes resource
+// quantity into the byte count used by admission and monitoring. Empty and
+// zero quantities mean unlimited. Negative quantities are never meaningful
+// for a ceiling and are rejected rather than silently treated as unlimited.
+func ParseOrgMaxMemoryBytes(raw string) (int64, error) {
+ raw = strings.TrimSpace(raw)
+ if raw == "" {
+ return 0, nil
+ }
+ quantity, err := resource.ParseQuantity(raw)
+ if err != nil {
+ return 0, fmt.Errorf("invalid memory quantity %q: %w", raw, err)
+ }
+ if quantity.Sign() < 0 {
+ return 0, fmt.Errorf("memory quantity %q must be >= 0", raw)
+ }
+ max := resource.NewQuantity(math.MaxInt64, resource.DecimalSI)
+ // DecimalSI values can remain larger than int64 and wrap in Value().
+ // ParseQuantity instead caps overflowing BinarySI values (for example 8Ei)
+ // at MaxInt64; no exact positive BinarySI quantity can equal that sentinel.
+ if quantity.Cmp(*max) > 0 || (quantity.Format == resource.BinarySI && quantity.Value() == math.MaxInt64) {
+ return 0, fmt.Errorf("memory quantity %q exceeds %d bytes", raw, int64(math.MaxInt64))
+ }
+ return quantity.Value(), nil
+}
+
+// NormalizeOrgMaxMemory validates max_memory and returns its canonical
+// Kubernetes quantity spelling. Empty and every valid zero quantity normalize
+// to the single persisted unlimited representation: an empty string.
+func NormalizeOrgMaxMemory(raw string) (string, error) {
+ raw = strings.TrimSpace(raw)
+ bytes, err := ParseOrgMaxMemoryBytes(raw)
+ if err != nil {
+ return "", err
+ }
+ if bytes == 0 {
+ return "", nil
+ }
+ quantity, err := resource.ParseQuantity(raw)
+ if err != nil {
+ return "", err // ParseOrgMaxMemoryBytes already validated this input.
+ }
+ return quantity.String(), nil
+}
+
// load fetches all config from the database and builds a Snapshot.
func (cs *ConfigStore) load() (*Snapshot, error) {
var orgs []Org
@@ -280,6 +326,10 @@ func (cs *ConfigStore) load() (*Snapshot, error) {
}
for _, o := range orgs {
+ maxMemoryBytes, err := ParseOrgMaxMemoryBytes(o.MaxMemory)
+ if err != nil {
+ return nil, fmt.Errorf("load org %q max_memory: %w", o.Name, err)
+ }
alias := ""
if o.HostnameAlias != nil {
alias = *o.HostnameAlias
@@ -302,6 +352,7 @@ func (cs *ConfigStore) load() (*Snapshot, error) {
HostnameAlias: alias,
MaxWorkers: o.MaxWorkers,
MaxVCPUs: o.MaxVCPUs,
+ MaxMemoryBytes: maxMemoryBytes,
DefaultWorkerCPU: o.DefaultWorkerCPU,
DefaultWorkerMemory: o.DefaultWorkerMemory,
DefaultWorkerTTL: o.DefaultWorkerTTL,
diff --git a/controlplane/configstore/store_test.go b/controlplane/configstore/store_test.go
index da126c71..67fe1611 100644
--- a/controlplane/configstore/store_test.go
+++ b/controlplane/configstore/store_test.go
@@ -16,6 +16,42 @@ func mustHash(t *testing.T, password string) string {
return hash
}
+func TestParseOrgMaxMemoryBytes(t *testing.T) {
+ tests := []struct {
+ name string
+ raw string
+ want int64
+ wantErr bool
+ }{
+ {name: "unset", raw: "", want: 0},
+ {name: "zero", raw: "0", want: 0},
+ {name: "zero quantity", raw: "0Gi", want: 0},
+ {name: "binary quantity", raw: "120Gi", want: 120 * 1024 * 1024 * 1024},
+ {name: "trimmed", raw: " 1500Mi ", want: 1500 * 1024 * 1024},
+ {name: "invalid", raw: "lots", wantErr: true},
+ {name: "negative", raw: "-1Gi", wantErr: true},
+ {name: "decimal overflow", raw: "10000000000000000000", wantErr: true},
+ {name: "binary overflow capped by quantity parser", raw: "8Ei", wantErr: true},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got, err := ParseOrgMaxMemoryBytes(tt.raw)
+ if tt.wantErr {
+ if err == nil {
+ t.Fatalf("ParseOrgMaxMemoryBytes(%q) = %d, nil; want error", tt.raw, got)
+ }
+ return
+ }
+ if err != nil {
+ t.Fatalf("ParseOrgMaxMemoryBytes(%q): %v", tt.raw, err)
+ }
+ if got != tt.want {
+ t.Fatalf("ParseOrgMaxMemoryBytes(%q) = %d, want %d", tt.raw, got, tt.want)
+ }
+ })
+ }
+}
+
func TestSnapshotBuild(t *testing.T) {
// Verify OrgConfig construction from models
hash1 := mustHash(t, "secret1")
diff --git a/controlplane/connection_limiter.go b/controlplane/connection_limiter.go
index e4cb79a8..01d90670 100644
--- a/controlplane/connection_limiter.go
+++ b/controlplane/connection_limiter.go
@@ -21,10 +21,11 @@ type connectionLease interface {
}
type connectionAdmissionRequest struct {
- PID int32
- Username string
- Protocol string
- RequestedVCPUs int
+ PID int32
+ Username string
+ Protocol string
+ RequestedVCPUs int
+ RequestedMemoryBytes int64
}
type connectionLimiter interface {
@@ -75,7 +76,7 @@ type runtimeOrgConnectionLimiter struct {
newID func() (string, error)
}
-// NewRuntimeOrgConnectionLimiter builds the runtime vCPU admission gate. The
+// NewRuntimeOrgConnectionLimiter builds the runtime resource-admission gate. The
// optional form preserves source compatibility for older callers, but Acquire
// deliberately fails closed unless the control-plane-wide reclaimer is
// supplied; durable admission must never start without reserved cleanup
@@ -119,15 +120,16 @@ func (l *runtimeOrgConnectionLimiter) Acquire(ctx context.Context, request conne
enqueuedAt := l.now()
expiresAt := enqueuedAt.Add(l.queueTTL)
entry := &configstore.OrgConnectionQueueEntry{
- RequestID: requestID,
- OrgID: l.orgID,
- Username: request.Username,
- CPInstanceID: l.cpInstanceID,
- PID: request.PID,
- Protocol: request.Protocol,
- RequestedVCPUs: request.RequestedVCPUs,
- EnqueuedAt: enqueuedAt,
- ExpiresAt: expiresAt,
+ RequestID: requestID,
+ OrgID: l.orgID,
+ Username: request.Username,
+ CPInstanceID: l.cpInstanceID,
+ PID: request.PID,
+ Protocol: request.Protocol,
+ RequestedVCPUs: request.RequestedVCPUs,
+ RequestedMemoryBytes: request.RequestedMemoryBytes,
+ EnqueuedAt: enqueuedAt,
+ ExpiresAt: expiresAt,
}
if err := ctx.Err(); err != nil {
return nil, runtimeAdmissionContextError(err)
@@ -228,11 +230,13 @@ func (l *runtimeOrgConnectionLimiter) Acquire(ctx context.Context, request conne
if lease != nil {
outcome = "granted"
sessionAdmissionActiveVCPUsGauge.WithLabelValues(l.orgID).Add(float64(request.RequestedVCPUs))
+ sessionAdmissionActiveMemoryBytesGauge.WithLabelValues(l.orgID).Add(float64(request.RequestedMemoryBytes))
cleanupArmed = false
return &runtimeOrgConnectionLease{
- reservation: reservation,
- orgID: l.orgID,
- requestedVCPUs: request.RequestedVCPUs,
+ reservation: reservation,
+ orgID: l.orgID,
+ requestedVCPUs: request.RequestedVCPUs,
+ requestedMemoryBytes: request.RequestedMemoryBytes,
}, nil
}
@@ -279,6 +283,8 @@ func sessionAdmissionRejectionReason(rejection *configstore.OrgConnectionAdmissi
switch rejection.Reason {
case configstore.OrgConnectionAdmissionRejectedOrgVCPU:
return "org_vcpu"
+ case configstore.OrgConnectionAdmissionRejectedOrgMemory:
+ return "org_memory"
case configstore.OrgConnectionAdmissionRejectedUserVCPU:
return "user_vcpu"
default:
@@ -294,10 +300,11 @@ func runtimeAdmissionContextError(err error) error {
}
type runtimeOrgConnectionLease struct {
- reservation AdmissionReclaimReservation
- orgID string
- requestedVCPUs int
- released sync.Once
+ reservation AdmissionReclaimReservation
+ orgID string
+ requestedVCPUs int
+ requestedMemoryBytes int64
+ released sync.Once
}
func (l *runtimeOrgConnectionLease) Release(context.Context) error {
@@ -307,6 +314,7 @@ func (l *runtimeOrgConnectionLease) Release(context.Context) error {
l.released.Do(func() {
l.reservation.Reclaim(admissionReclaimCauseLeaseRelease)
sessionAdmissionActiveVCPUsGauge.WithLabelValues(l.orgID).Sub(float64(l.requestedVCPUs))
+ sessionAdmissionActiveMemoryBytesGauge.WithLabelValues(l.orgID).Sub(float64(l.requestedMemoryBytes))
})
return nil
}
diff --git a/controlplane/connection_limiter_test.go b/controlplane/connection_limiter_test.go
index 4e42ac31..f922e445 100644
--- a/controlplane/connection_limiter_test.go
+++ b/controlplane/connection_limiter_test.go
@@ -581,10 +581,14 @@ func resetSessionAdmissionMetricGauges(t *testing.T, org string) {
sessionAdmissionQueueDepthGauge.DeleteLabelValues(org)
sessionAdmissionActiveVCPUsGauge.DeleteLabelValues(org)
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(org)
+ sessionAdmissionActiveMemoryBytesGauge.DeleteLabelValues(org)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(org)
t.Cleanup(func() {
sessionAdmissionQueueDepthGauge.DeleteLabelValues(org)
sessionAdmissionActiveVCPUsGauge.DeleteLabelValues(org)
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(org)
+ sessionAdmissionActiveMemoryBytesGauge.DeleteLabelValues(org)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(org)
})
}
@@ -592,6 +596,7 @@ func TestRuntimeOrgConnectionLimiterObservesGrantedRequestLifecycle(t *testing.T
const org = "org-admission-metrics-granted"
resetSessionAdmissionMetricGauges(t, org)
sessionAdmissionLimitVCPUsGauge.WithLabelValues(org).Set(7)
+ sessionAdmissionLimitMemoryBytesGauge.WithLabelValues(org).Set(240 << 30)
store := &admissionMetricsScheduleStore{results: []admissionMetricsScheduleResult{
{evaluation: configstore.OrgConnectionAdmissionEvaluation{Decision: "blocked", Reason: "org_vcpu"}},
{evaluation: configstore.OrgConnectionAdmissionEvaluation{Decision: "granted_current", Reason: "none"}},
@@ -602,7 +607,7 @@ func TestRuntimeOrgConnectionLimiterObservesGrantedRequestLifecycle(t *testing.T
requestsBefore := counterVecLabelValue(t, sessionAdmissionRequestsCounter, org, "granted", "org_vcpu")
waitsBefore := histogramVecLabelSampleCount(t, sessionAdmissionWaitHistogram, org, "granted", "org_vcpu")
lease, err := limiter.Acquire(context.Background(), connectionAdmissionRequest{
- PID: 1001, Username: "alice", Protocol: "postgres", RequestedVCPUs: 2,
+ PID: 1001, Username: "alice", Protocol: "postgres", RequestedVCPUs: 2, RequestedMemoryBytes: 120 << 30,
}, func(string) configstore.OrgResourceLimits { return configstore.OrgResourceLimits{} })
if err != nil || lease == nil {
t.Fatalf("Acquire = (%v, %v), want lease", lease, err)
@@ -619,9 +624,15 @@ func TestRuntimeOrgConnectionLimiterObservesGrantedRequestLifecycle(t *testing.T
if got := gaugeVecLabelValue(t, sessionAdmissionActiveVCPUsGauge, org); got != 2 {
t.Fatalf("active vCPUs = %v, want 2", got)
}
+ if got := gaugeVecLabelValue(t, sessionAdmissionActiveMemoryBytesGauge, org); got != 120<<30 {
+ t.Fatalf("active memory bytes = %v, want %d", got, int64(120<<30))
+ }
if got := gaugeVecLabelValue(t, sessionAdmissionLimitVCPUsGauge, org); got != 7 {
t.Fatalf("limit vCPUs = %v, want config-reconciled value 7", got)
}
+ if got := gaugeVecLabelValue(t, sessionAdmissionLimitMemoryBytesGauge, org); got != 240<<30 {
+ t.Fatalf("limit memory bytes = %v, want %d", got, int64(240<<30))
+ }
if err := lease.Release(context.Background()); err != nil {
t.Fatalf("Release: %v", err)
@@ -632,6 +643,9 @@ func TestRuntimeOrgConnectionLimiterObservesGrantedRequestLifecycle(t *testing.T
if got := gaugeVecLabelValue(t, sessionAdmissionActiveVCPUsGauge, org); got != 0 {
t.Fatalf("active vCPUs after repeated release = %v, want 0", got)
}
+ if got := gaugeVecLabelValue(t, sessionAdmissionActiveMemoryBytesGauge, org); got != 0 {
+ t.Fatalf("active memory bytes after repeated release = %v, want 0", got)
+ }
submissions, _ := reclaimer.snapshot()
if len(submissions) != 1 || submissions[0].cause != admissionReclaimCauseLeaseRelease {
t.Fatalf("lease cleanup submissions = %#v, want exactly one release", submissions)
diff --git a/controlplane/control.go b/controlplane/control.go
index b80695a8..e3925b65 100644
--- a/controlplane/control.go
+++ b/controlplane/control.go
@@ -30,6 +30,7 @@ import (
"github.com/posthog/duckgres/server/observe"
"github.com/posthog/duckgres/server/sessionmeta"
"github.com/prometheus/client_golang/prometheus/promhttp"
+ "k8s.io/apimachinery/pkg/api/resource"
)
// DefaultAdmissionReclaimerMaxReservations bounds cleanup ownership retained by
@@ -795,6 +796,11 @@ func sessionCreationErrorResponse(err error) (code string, message string) {
case errors.As(err, &capacityErr):
return "53300", capacityMissPolicyForReason(capacityErr.missReason()).sqlMessage(capacityErr.RetryAfter)
case errors.As(err, &rejection):
+ if rejection.Reason == configstore.OrgConnectionAdmissionRejectedOrgMemory {
+ requested := resource.NewQuantity(rejection.RequestedMemoryBytes, resource.BinarySI)
+ maximum := resource.NewQuantity(rejection.MaximumMemoryBytes, resource.BinarySI)
+ return "53400", fmt.Sprintf("requested worker requires %s memory, exceeding the organization limit of %s; request a smaller worker or raise the limit", requested.String(), maximum.String())
+ }
scope := "organization"
if rejection.Reason == configstore.OrgConnectionAdmissionRejectedUserVCPU {
scope = "user"
@@ -1622,7 +1628,7 @@ func (cp *ControlPlane) handleConnection(conn net.Conn) {
if err != nil {
// Classified with the same logic the eager connect path uses, so
// the client sees the real SQLSTATE + message (53300 with a retry
- // hint at capacity, 53400 for a vCPU-admission rejection, 57P03
+ // hint at capacity, 53400 for a resource-admission rejection, 57P03
// while draining) instead of a substring guess.
return nil, 0, "", fmt.Errorf("escalate to standard worker: %w", newSessionAcquireError(err))
}
diff --git a/controlplane/control_cancel_test.go b/controlplane/control_cancel_test.go
index 55309fcb..4faac8c5 100644
--- a/controlplane/control_cancel_test.go
+++ b/controlplane/control_cancel_test.go
@@ -148,6 +148,21 @@ func TestSessionCreationErrorResponse(t *testing.T) {
}
})
+ t.Run("request exceeds org memory limit", func(t *testing.T) {
+ code, message := sessionCreationErrorResponse(&configstore.OrgConnectionAdmissionRejectedError{
+ Reason: configstore.OrgConnectionAdmissionRejectedOrgMemory,
+ RequestedMemoryBytes: 120 * 1024 * 1024 * 1024,
+ MaximumMemoryBytes: 60 * 1024 * 1024 * 1024,
+ })
+ if code != "53400" {
+ t.Fatalf("code = %q, want 53400", code)
+ }
+ want := "requested worker requires 120Gi memory, exceeding the organization limit of 60Gi; request a smaller worker or raise the limit"
+ if message != want {
+ t.Fatalf("message = %q, want %q", message, want)
+ }
+ })
+
t.Run("worker capacity exhausted", func(t *testing.T) {
code, message := sessionCreationErrorResponse(NewWorkerCapacityExhaustedError(45 * time.Second))
if code != "53300" {
diff --git a/controlplane/k8s_pool_acquire.go b/controlplane/k8s_pool_acquire.go
index 64f2dc32..b7b97ba6 100644
--- a/controlplane/k8s_pool_acquire.go
+++ b/controlplane/k8s_pool_acquire.go
@@ -18,6 +18,7 @@ import (
"github.com/posthog/duckgres/server/flightclient"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
+ corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)
@@ -568,6 +569,14 @@ func (p *K8sWorkerPool) adoptClaimedWorker(ctx context.Context, claimed *configs
if err != nil {
return nil, fmt.Errorf("get claimed worker pod %s: %w", claimed.PodName, err)
}
+ profile := WorkerProfile{
+ CPU: claimed.ProfileCPU,
+ Memory: claimed.ProfileMemory,
+ TTL: time.Duration(claimed.TTLMinutes) * time.Minute,
+ }
+ if err := p.validateClaimedWorkerPodShape(pod, profile); err != nil {
+ return nil, err
+ }
// For hot-idle workers, skip the epoch-validated health check. The worker's
// epoch and CP instance ID are from the previous owner, and ClaimHotIdleWorker
@@ -595,18 +604,52 @@ func (p *K8sWorkerPool) adoptClaimedWorker(ctx context.Context, claimed *configs
// worker keeps its shape. Without this, re-adopting a sized worker reset its
// profile to the default and the next hot-idle persist dropped its
// cpu/mem/ttl, so a same-size request could no longer reuse it.
- profile: WorkerProfile{
- CPU: claimed.ProfileCPU,
- Memory: claimed.ProfileMemory,
- TTL: time.Duration(claimed.TTLMinutes) * time.Minute,
- },
- done: make(chan struct{}),
+ profile: profile,
+ done: make(chan struct{}),
}
worker.SetOwnerCPInstanceID(claimed.OwnerCPInstanceID)
worker.SetOwnerEpoch(claimed.OwnerEpoch)
return worker, nil
}
+// validateClaimedWorkerPodShape prevents a control plane from adopting a pod
+// whose concrete resource requests differ from the shape admission will charge.
+// This matters most for default-profile runtime rows: they intentionally persist
+// empty CPU/memory fields, so a pod created before a rolling default change can
+// otherwise be mistaken for the new (larger or smaller) default shape.
+func (p *K8sWorkerPool) validateClaimedWorkerPodShape(pod *corev1.Pod, profile WorkerProfile) error {
+ if pod == nil {
+ return fmt.Errorf("claimed worker pod is missing")
+ }
+
+ var workerContainer *corev1.Container
+ for i := range pod.Spec.Containers {
+ if pod.Spec.Containers[i].Name == "duckdb-worker" {
+ workerContainer = &pod.Spec.Containers[i]
+ break
+ }
+ }
+ if workerContainer == nil {
+ return fmt.Errorf("claimed worker pod %s has no duckdb-worker container", pod.Name)
+ }
+
+ expected := p.workerResourcesForProfile(profile).Requests
+ for _, resourceName := range []corev1.ResourceName{corev1.ResourceCPU, corev1.ResourceMemory} {
+ actual, ok := workerContainer.Resources.Requests[resourceName]
+ if !ok {
+ return fmt.Errorf("claimed worker pod %s duckdb-worker container is missing %s request", pod.Name, resourceName)
+ }
+ if actual.Sign() <= 0 {
+ return fmt.Errorf("claimed worker pod %s duckdb-worker container has invalid %s request %q", pod.Name, resourceName, actual.String())
+ }
+ want := expected[resourceName]
+ if actual.Cmp(want) != 0 {
+ return fmt.Errorf("claimed worker pod %s duckdb-worker container %s request %q does not match effective request %q", pod.Name, resourceName, actual.String(), want.String())
+ }
+ }
+ return nil
+}
+
func (p *K8sWorkerPool) connectWorker(ctx context.Context, podName, podIP, bearerToken string) (*flightsql.Client, error) {
return p.connectWorkerWithHealthCheck(ctx, podName, podIP, bearerToken, server.WorkerHealthCheckPayload{})
}
diff --git a/controlplane/k8s_pool_test.go b/controlplane/k8s_pool_test.go
index dedacd25..41721711 100644
--- a/controlplane/k8s_pool_test.go
+++ b/controlplane/k8s_pool_test.go
@@ -1401,6 +1401,95 @@ func TestK8sPoolReserveClaimedWorkerRejectsStaleInMemoryEpoch(t *testing.T) {
}
}
+// A default-profile runtime row does not record the concrete pod shape. During
+// a rolling config change, the pod may therefore have been spawned with the old
+// defaults while a new control plane would charge admission using the new
+// defaults. Adoption must compare the live pod requests with the effective
+// shape before handing the worker to a session.
+func TestK8sPoolAdoptClaimedWorkerValidatesEffectivePodShape(t *testing.T) {
+ tests := []struct {
+ name string
+ profileCPU string
+ profileMemory string
+ podCPU string
+ podMemory string
+ includeWorker bool
+ wantErr bool
+ }{
+ {name: "default cpu changed", podCPU: "8", podMemory: "120Gi", includeWorker: true, wantErr: true},
+ {name: "default memory changed", podCPU: "15", podMemory: "16Gi", includeWorker: true, wantErr: true},
+ {name: "cpu request missing", podMemory: "120Gi", includeWorker: true, wantErr: true},
+ {name: "memory request missing", podCPU: "15", includeWorker: true, wantErr: true},
+ {name: "zero request is invalid", podCPU: "0", podMemory: "120Gi", includeWorker: true, wantErr: true},
+ {name: "worker container missing", podCPU: "15", podMemory: "120Gi", wantErr: true},
+ {name: "default matches current config", podCPU: "15", podMemory: "120Gi", includeWorker: true},
+ {name: "explicit profile ignores changed defaults", profileCPU: "8", profileMemory: "16Gi", podCPU: "8", podMemory: "16Gi", includeWorker: true},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ pool, cs := newTestK8sPool(t, 5)
+ pool.workerCPURequest = "15"
+ pool.workerMemoryRequest = "120Gi"
+ podName := "duckgres-worker-test-cp-41"
+ if _, err := pool.ensureWorkerRPCSecret(context.Background(), podName); err != nil {
+ t.Fatalf("ensure worker RPC secret: %v", err)
+ }
+
+ pod := &corev1.Pod{
+ ObjectMeta: metav1.ObjectMeta{Name: podName, Namespace: pool.namespace},
+ Status: corev1.PodStatus{PodIP: "10.0.0.41"},
+ }
+ if tt.includeWorker {
+ requests := corev1.ResourceList{}
+ if tt.podCPU != "" {
+ requests[corev1.ResourceCPU] = resource.MustParse(tt.podCPU)
+ }
+ if tt.podMemory != "" {
+ requests[corev1.ResourceMemory] = resource.MustParse(tt.podMemory)
+ }
+ pod.Spec.Containers = []corev1.Container{{
+ Name: "duckdb-worker",
+ Resources: corev1.ResourceRequirements{Requests: requests},
+ }}
+ } else {
+ pod.Spec.Containers = []corev1.Container{{Name: "sidecar"}}
+ }
+ if _, err := cs.CoreV1().Pods(pool.namespace).Create(context.Background(), pod, metav1.CreateOptions{}); err != nil {
+ t.Fatalf("create claimed worker pod: %v", err)
+ }
+
+ connectCalled := false
+ pool.connectWorkerFunc = func(context.Context, string, string, string) (*flightsql.Client, error) {
+ connectCalled = true
+ return nil, nil
+ }
+ _, err := pool.adoptClaimedWorker(context.Background(), &configstore.WorkerRecord{
+ WorkerID: 41,
+ PodName: podName,
+ ProfileCPU: tt.profileCPU,
+ ProfileMemory: tt.profileMemory,
+ OwnerEpoch: 2,
+ })
+ if tt.wantErr {
+ if err == nil {
+ t.Fatal("adoptClaimedWorker succeeded for an incompatible pod shape")
+ }
+ if connectCalled {
+ t.Fatal("adoptClaimedWorker connected before validating the pod shape")
+ }
+ return
+ }
+ if err != nil {
+ t.Fatalf("adoptClaimedWorker: %v", err)
+ }
+ if !connectCalled {
+ t.Fatal("adoptClaimedWorker did not connect after validating the pod shape")
+ }
+ })
+ }
+}
+
func TestK8sPoolReserveSharedWorkerReturnsOrgCapFromHotIdleClaim(t *testing.T) {
pool, _ := newTestK8sPool(t, 5)
store := &captureRuntimeWorkerStore{
diff --git a/controlplane/org_router.go b/controlplane/org_router.go
index ef859eed..b9011f14 100644
--- a/controlplane/org_router.go
+++ b/controlplane/org_router.go
@@ -167,6 +167,9 @@ func (tr *OrgRouter) createOrgStackWhileMutationHeld(tc *configstore.OrgConfig)
sessions.SetRequestedVCPUsResolver(func(profile *WorkerProfile) (int, error) {
return requestedWorkerVCPUs(profile, tr.baseCfg.WorkerCPURequest)
})
+ sessions.SetRequestedMemoryBytesResolver(func(profile *WorkerProfile) (int64, error) {
+ return requestedWorkerMemoryBytes(profile, tr.baseCfg.WorkerMemoryRequest)
+ })
sessions.SetConnectionLimiter(NewRuntimeOrgConnectionLimiter(
tr.configStore,
tc.Name,
@@ -383,10 +386,12 @@ func (tr *OrgRouter) publishOrgStack(orgID string, stack *OrgStack) orgStackPubl
tr.orgs[orgID] = stack
if stack != nil && stack.Config != nil {
sessionAdmissionLimitVCPUsGauge.WithLabelValues(orgID).Set(float64(stack.Config.MaxVCPUs))
+ sessionAdmissionLimitMemoryBytesGauge.WithLabelValues(orgID).Set(float64(stack.Config.MaxMemoryBytes))
} else {
// Minimal stacks are valid in lifecycle and concurrency paths. Without a
// config there is no authoritative limit value to publish.
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(orgID)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(orgID)
}
tr.mu.Unlock()
return orgStackPublishAccepted
@@ -401,6 +406,7 @@ func (tr *OrgRouter) resourceLimitsForOrg(orgID string) func(username string) co
limits := configstore.OrgResourceLimits{}
if org, ok := snap.Orgs[orgID]; ok && org != nil {
limits.OrgMaxVCPUs = org.MaxVCPUs
+ limits.OrgMaxMemoryBytes = org.MaxMemoryBytes
}
limits.UserMaxVCPUs = snap.OrgUserMaxVCPUs[configstore.OrgUserKey{OrgID: orgID, Username: username}]
return limits
@@ -504,10 +510,11 @@ func (tr *OrgRouter) refreshOrgStackWhileMutationHeld(orgID string, stack *OrgSt
oldTC := stack.Config
stack.Config = tc
sessionAdmissionLimitVCPUsGauge.WithLabelValues(orgID).Set(float64(tc.MaxVCPUs))
+ sessionAdmissionLimitMemoryBytesGauge.WithLabelValues(orgID).Set(float64(tc.MaxMemoryBytes))
tr.mu.Unlock()
limitsChanged := oldTC == nil || oldTC.MaxWorkers != tc.MaxWorkers
- resourceLimitChanged := oldTC == nil || oldTC.MaxVCPUs != tc.MaxVCPUs
+ resourceLimitChanged := oldTC == nil || oldTC.MaxVCPUs != tc.MaxVCPUs || oldTC.MaxMemoryBytes != tc.MaxMemoryBytes
floorChanged := oldTC == nil || oldTC.DefaultWorkerMinHotIdle != tc.DefaultWorkerMinHotIdle
oldImage := tr.baseCfg.WorkerImage
if oldTC != nil {
@@ -518,11 +525,14 @@ func (tr *OrgRouter) refreshOrgStackWhileMutationHeld(orgID string, stack *OrgSt
if resourceLimitChanged {
oldMaxVCPUs := 0
+ var oldMaxMemoryBytes int64
if oldTC != nil {
oldMaxVCPUs = oldTC.MaxVCPUs
+ oldMaxMemoryBytes = oldTC.MaxMemoryBytes
}
slog.Info("Org resource limit changed.", "org", orgID,
- "old_max_vcpus", oldMaxVCPUs, "new_max_vcpus", tc.MaxVCPUs)
+ "old_max_vcpus", oldMaxVCPUs, "new_max_vcpus", tc.MaxVCPUs,
+ "old_max_memory_bytes", oldMaxMemoryBytes, "new_max_memory_bytes", tc.MaxMemoryBytes)
}
if limitsChanged && stack.Pool != nil {
oldMaxWorkers := 0
@@ -556,6 +566,7 @@ func (tr *OrgRouter) removeOrgStackWhileMutationHeld(orgID string, expected *Org
stack, ok := tr.orgs[orgID]
if !ok {
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(orgID)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(orgID)
tr.mu.Unlock()
return
}
@@ -565,6 +576,7 @@ func (tr *OrgRouter) removeOrgStackWhileMutationHeld(orgID string, expected *Org
}
delete(tr.orgs, orgID)
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(orgID)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(orgID)
tr.mu.Unlock()
slog.Info("Destroying org stack.", "org", orgID)
@@ -811,6 +823,7 @@ func (tr *OrgRouter) shutdownAll() {
for k, v := range tr.orgs {
orgs[k] = v
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(k)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(k)
}
tr.orgs = make(map[string]*OrgStack)
tr.mu.Unlock()
diff --git a/controlplane/org_router_test.go b/controlplane/org_router_test.go
index e2d8c030..17218a32 100644
--- a/controlplane/org_router_test.go
+++ b/controlplane/org_router_test.go
@@ -179,8 +179,13 @@ func TestOrgRouterShutdownAllDrainsAdmissionReclaimerAfterAllSessions(t *testing
for i, org := range orgs {
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(org)
sessionAdmissionLimitVCPUsGauge.WithLabelValues(org).Set(float64(i + 1))
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(org)
+ sessionAdmissionLimitMemoryBytesGauge.WithLabelValues(org).Set(float64((i + 1) << 30))
org := org
- t.Cleanup(func() { sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(org) })
+ t.Cleanup(func() {
+ sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(org)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(org)
+ })
}
events := []string{}
firstPool := &recordingOrgRouterPool{events: &events}
@@ -226,6 +231,9 @@ func TestOrgRouterShutdownAllDrainsAdmissionReclaimerAfterAllSessions(t *testing
if _, ok := metricGaugeFamilyLabelValue(t, "duckgres_session_admission_limit_vcpus", map[string]string{"org": org}); ok {
t.Fatalf("shutdown left admission limit series for org %q", org)
}
+ if _, ok := metricGaugeFamilyLabelValue(t, "duckgres_session_admission_limit_memory_bytes", map[string]string{"org": org}); ok {
+ t.Fatalf("shutdown left memory admission limit series for org %q", org)
+ }
}
}
@@ -1090,9 +1098,10 @@ func TestOrgRouterHandleConfigChangeRefreshesRuntimeOnlyUpdates(t *testing.T) {
func TestOrgRouterReconcilesAdmissionLimitGaugeWithOrgStackLifecycle(t *testing.T) {
const org = "admission-limit-lifecycle"
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(org)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(org)
sharedPool, _ := newTestK8sPool(t, 10)
- initial := &configstore.OrgConfig{Name: org, MaxVCPUs: 16}
+ initial := &configstore.OrgConfig{Name: org, MaxVCPUs: 16, MaxMemoryBytes: 240 << 30}
initialSnapshot := &configstore.Snapshot{Orgs: map[string]*configstore.OrgConfig{org: initial}}
store := newTestConfigStoreWithSnapshot(initialSnapshot)
tr := &OrgRouter{
@@ -1105,6 +1114,7 @@ func TestOrgRouterReconcilesAdmissionLimitGaugeWithOrgStackLifecycle(t *testing.
t.Cleanup(func() {
tr.DestroyOrgStack(org)
sessionAdmissionLimitVCPUsGauge.DeleteLabelValues(org)
+ sessionAdmissionLimitMemoryBytesGauge.DeleteLabelValues(org)
})
if _, err := tr.createOrgStack(initial); err != nil {
@@ -1113,22 +1123,31 @@ func TestOrgRouterReconcilesAdmissionLimitGaugeWithOrgStackLifecycle(t *testing.
if got := gaugeVecLabelValue(t, sessionAdmissionLimitVCPUsGauge, org); got != 16 {
t.Fatalf("initial admission limit = %v, want 16", got)
}
+ if got := gaugeVecLabelValue(t, sessionAdmissionLimitMemoryBytesGauge, org); got != 240<<30 {
+ t.Fatalf("initial memory admission limit = %v, want %d", got, int64(240<<30))
+ }
- lowered := &configstore.OrgConfig{Name: org, MaxVCPUs: 8}
+ lowered := &configstore.OrgConfig{Name: org, MaxVCPUs: 8, MaxMemoryBytes: 120 << 30}
loweredSnapshot := &configstore.Snapshot{Orgs: map[string]*configstore.OrgConfig{org: lowered}}
setTestConfigStoreSnapshot(store, loweredSnapshot)
tr.HandleConfigChange(initialSnapshot, loweredSnapshot)
if got := gaugeVecLabelValue(t, sessionAdmissionLimitVCPUsGauge, org); got != 8 {
t.Fatalf("lowered admission limit = %v, want 8", got)
}
+ if got := gaugeVecLabelValue(t, sessionAdmissionLimitMemoryBytesGauge, org); got != 120<<30 {
+ t.Fatalf("lowered memory admission limit = %v, want %d", got, int64(120<<30))
+ }
- unlimited := &configstore.OrgConfig{Name: org, MaxVCPUs: 0}
+ unlimited := &configstore.OrgConfig{Name: org, MaxVCPUs: 0, MaxMemoryBytes: 0}
unlimitedSnapshot := &configstore.Snapshot{Orgs: map[string]*configstore.OrgConfig{org: unlimited}}
setTestConfigStoreSnapshot(store, unlimitedSnapshot)
tr.HandleConfigChange(loweredSnapshot, unlimitedSnapshot)
if got := gaugeVecLabelValue(t, sessionAdmissionLimitVCPUsGauge, org); got != 0 {
t.Fatalf("unlimited admission limit = %v, want 0", got)
}
+ if got := gaugeVecLabelValue(t, sessionAdmissionLimitMemoryBytesGauge, org); got != 0 {
+ t.Fatalf("unlimited memory admission limit = %v, want 0", got)
+ }
removedSnapshot := &configstore.Snapshot{Orgs: map[string]*configstore.OrgConfig{}}
setTestConfigStoreSnapshot(store, removedSnapshot)
@@ -1136,6 +1155,31 @@ func TestOrgRouterReconcilesAdmissionLimitGaugeWithOrgStackLifecycle(t *testing.
if _, ok := metricGaugeFamilyLabelValue(t, "duckgres_session_admission_limit_vcpus", map[string]string{"org": org}); ok {
t.Fatal("expected removed org's admission limit series to be deleted")
}
+ if _, ok := metricGaugeFamilyLabelValue(t, "duckgres_session_admission_limit_memory_bytes", map[string]string{"org": org}); ok {
+ t.Fatal("expected removed org's memory admission limit series to be deleted")
+ }
+}
+
+func TestOrgRouterResourceLimitsForOrgIncludesDynamicMemoryLimit(t *testing.T) {
+ const org = "dynamic-memory-limit"
+ initial := &configstore.OrgConfig{Name: org, MaxVCPUs: 4, MaxMemoryBytes: 120 << 30}
+ store := newTestConfigStoreWithSnapshot(&configstore.Snapshot{
+ Orgs: map[string]*configstore.OrgConfig{org: initial},
+ })
+ router := &OrgRouter{configStore: store}
+ lookup := router.resourceLimitsForOrg(org)
+
+ if got := lookup("alice"); got.OrgMaxVCPUs != 4 || got.OrgMaxMemoryBytes != 120<<30 {
+ t.Fatalf("initial resource limits = %#v, want 4 vCPU/120Gi", got)
+ }
+
+ updated := &configstore.OrgConfig{Name: org, MaxVCPUs: 8, MaxMemoryBytes: 240 << 30}
+ setTestConfigStoreSnapshot(store, &configstore.Snapshot{
+ Orgs: map[string]*configstore.OrgConfig{org: updated},
+ })
+ if got := lookup("alice"); got.OrgMaxVCPUs != 8 || got.OrgMaxMemoryBytes != 240<<30 {
+ t.Fatalf("updated resource limits = %#v, want 8 vCPU/240Gi", got)
+ }
}
func TestOrgRouterHandleConfigChangeRefreshesOrgWorkerImage(t *testing.T) {
diff --git a/controlplane/session_activation.go b/controlplane/session_activation.go
index 20e31653..9039d4b4 100644
--- a/controlplane/session_activation.go
+++ b/controlplane/session_activation.go
@@ -77,7 +77,7 @@ func activationOutcomeForCode(code string) sessionActivationOutcome {
// logic the eager connect path uses (sessionCreationErrorResponse) and packages
// the result for the server package, which cannot see the control plane's
// sentinels. Without this the server had to guess the SQLSTATE from the error
-// text, which silently degraded a vCPU-admission rejection or a draining
+// text, which silently degraded a resource-admission rejection or a draining
// control plane into a generic 53400.
func newSessionAcquireError(err error) *server.SessionAcquireError {
code, message := sessionCreationErrorResponse(err)
diff --git a/controlplane/session_mgr.go b/controlplane/session_mgr.go
index 90bca626..7f33c046 100644
--- a/controlplane/session_mgr.go
+++ b/controlplane/session_mgr.go
@@ -87,12 +87,13 @@ type SessionManager struct {
// org stack) so every session/worker lifecycle line is org-filterable.
log *slog.Logger
- maxConnections int
- activeSlots int
- waiters []*connectionWaiter
- limiter connectionLimiter
- resourceLimits func(username string) configstore.OrgResourceLimits
- requestedVCPUs func(profile *WorkerProfile) (int, error)
+ maxConnections int
+ activeSlots int
+ waiters []*connectionWaiter
+ limiter connectionLimiter
+ resourceLimits func(username string) configstore.OrgResourceLimits
+ requestedVCPUs func(profile *WorkerProfile) (int, error)
+ requestedMemoryBytes func(profile *WorkerProfile) (int64, error)
// userSecretLoader returns the user's persistent CREATE SECRET statements
// (decrypted) to replay on a worker at session creation. nil outside the
@@ -169,6 +170,16 @@ func (sm *SessionManager) SetRequestedVCPUsResolver(fn func(profile *WorkerProfi
sm.requestedVCPUs = fn
}
+// SetRequestedMemoryBytesResolver installs the worker-profile-to-memory-byte
+// resolver used for resource admission. nil uses the effective profile memory
+// with the built-in default, which keeps every durable request positive even
+// for compatibility callers that do not install a deployment-aware resolver.
+func (sm *SessionManager) SetRequestedMemoryBytesResolver(fn func(profile *WorkerProfile) (int64, error)) {
+ sm.mu.Lock()
+ defer sm.mu.Unlock()
+ sm.requestedMemoryBytes = fn
+}
+
// ReservePID generates a new unique PID for a session.
func (sm *SessionManager) ReservePID() int32 {
return reservePID(globalNextPID)
@@ -188,6 +199,7 @@ func (sm *SessionManager) acquireConnectionSlot(ctx context.Context, pid int32,
limiter := sm.limiter
resourceLimits := sm.resourceLimits
requestedVCPUs := sm.requestedVCPUs
+ requestedMemoryBytes := sm.requestedMemoryBytes
sm.mu.Unlock()
if limiter != nil {
vcpus := 1
@@ -201,6 +213,16 @@ func (sm *SessionManager) acquireConnectionSlot(ctx context.Context, pid int32,
if vcpus <= 0 {
return nil, fmt.Errorf("requested vcpus must be positive, got %d", vcpus)
}
+ memoryBytes, err := requestedWorkerMemoryBytes(profile, "")
+ if requestedMemoryBytes != nil {
+ memoryBytes, err = requestedMemoryBytes(profile)
+ }
+ if err != nil {
+ return nil, err
+ }
+ if memoryBytes <= 0 {
+ return nil, fmt.Errorf("requested memory bytes must be positive, got %d", memoryBytes)
+ }
limits := func(user string) configstore.OrgResourceLimits {
if resourceLimits == nil {
return configstore.OrgResourceLimits{}
@@ -208,10 +230,11 @@ func (sm *SessionManager) acquireConnectionSlot(ctx context.Context, pid int32,
return resourceLimits(user)
}
lease, err := limiter.Acquire(ctx, connectionAdmissionRequest{
- PID: pid,
- Username: username,
- Protocol: protocol,
- RequestedVCPUs: vcpus,
+ PID: pid,
+ Username: username,
+ Protocol: protocol,
+ RequestedVCPUs: vcpus,
+ RequestedMemoryBytes: memoryBytes,
}, limits)
if err != nil {
return nil, err
diff --git a/controlplane/session_mgr_drain_test.go b/controlplane/session_mgr_drain_test.go
index 44e0fef2..5ed12b28 100644
--- a/controlplane/session_mgr_drain_test.go
+++ b/controlplane/session_mgr_drain_test.go
@@ -217,9 +217,13 @@ type observingConnectionLimiter struct {
firstRead chan configstore.OrgResourceLimits
readAgain chan struct{}
secondRead chan configstore.OrgResourceLimits
+ request chan connectionAdmissionRequest
}
func (l *observingConnectionLimiter) Acquire(ctx context.Context, request connectionAdmissionRequest, limits func(string) configstore.OrgResourceLimits) (connectionLease, error) {
+ if l.request != nil {
+ l.request <- request
+ }
first := limits(request.Username)
l.firstRead <- first
@@ -1203,21 +1207,28 @@ func countEvents(events []string, want string) int {
func TestSessionManager_RuntimeLimiterObservesDynamicResourceLimitWhileQueued(t *testing.T) {
sm := NewSessionManager(nil, nil)
var orgMaxVCPUs atomic.Int32
+ var orgMaxMemoryBytes atomic.Int64
orgMaxVCPUs.Store(4)
+ orgMaxMemoryBytes.Store(120 << 30)
sm.SetResourceLimitsProvider(func(username string) configstore.OrgResourceLimits {
return configstore.OrgResourceLimits{
- OrgMaxVCPUs: int(orgMaxVCPUs.Load()),
- UserMaxVCPUs: 2,
+ OrgMaxVCPUs: int(orgMaxVCPUs.Load()),
+ OrgMaxMemoryBytes: orgMaxMemoryBytes.Load(),
+ UserMaxVCPUs: 2,
}
})
sm.SetRequestedVCPUsResolver(func(profile *WorkerProfile) (int, error) {
return 2, nil
})
+ sm.SetRequestedMemoryBytesResolver(func(profile *WorkerProfile) (int64, error) {
+ return 60 << 30, nil
+ })
limiter := &observingConnectionLimiter{
firstRead: make(chan configstore.OrgResourceLimits, 1),
readAgain: make(chan struct{}),
secondRead: make(chan configstore.OrgResourceLimits, 1),
+ request: make(chan connectionAdmissionRequest, 1),
}
sm.SetConnectionLimiter(limiter)
@@ -1229,20 +1240,30 @@ func TestSessionManager_RuntimeLimiterObservesDynamicResourceLimitWhileQueued(t
select {
case got := <-limiter.firstRead:
- if got.OrgMaxVCPUs != 4 || got.UserMaxVCPUs != 2 {
- t.Fatalf("expected first limits org=4/user=2, got %#v", got)
+ if got.OrgMaxVCPUs != 4 || got.OrgMaxMemoryBytes != 120<<30 || got.UserMaxVCPUs != 2 {
+ t.Fatalf("expected first limits org_cpu=4/org_memory=120Gi/user_cpu=2, got %#v", got)
}
case <-time.After(time.Second):
t.Fatal("timed out waiting for first limit read")
}
+ select {
+ case got := <-limiter.request:
+ if got.RequestedVCPUs != 2 || got.RequestedMemoryBytes != 60<<30 {
+ t.Fatalf("admission request resources = %#v, want 2 vCPUs/60Gi", got)
+ }
+ case <-time.After(time.Second):
+ t.Fatal("timed out waiting for admission request")
+ }
+
orgMaxVCPUs.Store(8)
+ orgMaxMemoryBytes.Store(240 << 30)
close(limiter.readAgain)
select {
case got := <-limiter.secondRead:
- if got.OrgMaxVCPUs != 8 || got.UserMaxVCPUs != 2 {
- t.Fatalf("expected queued limiter read to observe updated org limit 8, got %#v", got)
+ if got.OrgMaxVCPUs != 8 || got.OrgMaxMemoryBytes != 240<<30 || got.UserMaxVCPUs != 2 {
+ t.Fatalf("expected queued limiter read to observe updated org limits 8 vCPU/240Gi, got %#v", got)
}
case <-time.After(time.Second):
t.Fatal("timed out waiting for second limit read")
@@ -1258,6 +1279,59 @@ func TestSessionManager_RuntimeLimiterObservesDynamicResourceLimitWhileQueued(t
}
}
+func TestSessionManager_PropagatesWorkerResourcesForEverySessionProtocol(t *testing.T) {
+ for _, protocol := range []string{"postgres", "flight"} {
+ t.Run(protocol, func(t *testing.T) {
+ profile := &WorkerProfile{CPU: "4", Memory: "60Gi"}
+ sm := NewSessionManager(nil, nil)
+ sm.SetRequestedVCPUsResolver(func(got *WorkerProfile) (int, error) {
+ if got != profile {
+ t.Fatalf("vCPU resolver profile = %p, want %p", got, profile)
+ }
+ return 4, nil
+ })
+ sm.SetRequestedMemoryBytesResolver(func(got *WorkerProfile) (int64, error) {
+ if got != profile {
+ t.Fatalf("memory resolver profile = %p, want %p", got, profile)
+ }
+ return 60 << 30, nil
+ })
+ limiter := &observingConnectionLimiter{
+ firstRead: make(chan configstore.OrgResourceLimits, 1),
+ readAgain: make(chan struct{}),
+ secondRead: make(chan configstore.OrgResourceLimits, 1),
+ request: make(chan connectionAdmissionRequest, 1),
+ }
+ close(limiter.readAgain)
+ sm.SetConnectionLimiter(limiter)
+
+ if _, err := sm.acquireConnectionSlot(context.Background(), 1001, "alice", protocol, profile); err != nil {
+ t.Fatalf("acquireConnectionSlot: %v", err)
+ }
+ got := <-limiter.request
+ if got.Protocol != protocol || got.RequestedVCPUs != 4 || got.RequestedMemoryBytes != 60<<30 {
+ t.Fatalf("admission request = %#v, want protocol %q with 4 vCPU/60Gi", got, protocol)
+ }
+ })
+ }
+}
+
+func TestSessionManager_RejectsMemoryOverflowBeforeAdmissionEnqueue(t *testing.T) {
+ sm := NewSessionManager(nil, nil)
+ requests := make(chan connectionAdmissionRequest, 1)
+ sm.SetConnectionLimiter(&observingConnectionLimiter{request: requests})
+ sm.SetRequestedMemoryBytesResolver(func(profile *WorkerProfile) (int64, error) {
+ return requestedWorkerMemoryBytes(profile, "8Ei")
+ })
+
+ if _, err := sm.acquireConnectionSlot(context.Background(), 1001, "alice", "postgres", nil); err == nil {
+ t.Fatal("expected overflowing worker memory to be rejected")
+ }
+ if len(requests) != 0 {
+ t.Fatalf("overflowing memory reached admission limiter: %#v", <-requests)
+ }
+}
+
func TestRuntimeOrgConnectionLimiterKeepsQueuedLeaseWhenResourceLimitBecomesUnlimited(t *testing.T) {
store := &runtimeLimiterTestStore{
firstTry: make(chan struct{}),
@@ -1357,12 +1431,13 @@ func TestRuntimeOrgConnectionLimiterUsesQueue(t *testing.T) {
}
lease, err := limiter.Acquire(context.Background(), connectionAdmissionRequest{
- PID: 2002,
- Username: "alice",
- Protocol: "postgres",
- RequestedVCPUs: 4,
+ PID: 2002,
+ Username: "alice",
+ Protocol: "postgres",
+ RequestedVCPUs: 4,
+ RequestedMemoryBytes: 120 << 30,
}, func(username string) configstore.OrgResourceLimits {
- return configstore.OrgResourceLimits{OrgMaxVCPUs: 4, UserMaxVCPUs: 4}
+ return configstore.OrgResourceLimits{OrgMaxVCPUs: 4, OrgMaxMemoryBytes: 120 << 30, UserMaxVCPUs: 4}
})
if err != nil {
t.Fatalf("Acquire: %v", err)
@@ -1380,7 +1455,7 @@ func TestRuntimeOrgConnectionLimiterUsesQueue(t *testing.T) {
if store.entry == nil {
t.Fatal("expected reconnect admission to enqueue a request")
}
- if store.entry.Username != "alice" || store.entry.Protocol != "postgres" || store.entry.RequestedVCPUs != 4 || store.entry.PID != 2002 {
+ if store.entry.Username != "alice" || store.entry.Protocol != "postgres" || store.entry.RequestedVCPUs != 4 || store.entry.RequestedMemoryBytes != 120<<30 || store.entry.PID != 2002 {
t.Fatalf("unexpected queued request: %#v", store.entry)
}
if err := lease.Release(context.Background()); err != nil {
diff --git a/controlplane/worker_profile.go b/controlplane/worker_profile.go
index a7356c06..ad8413f0 100644
--- a/controlplane/worker_profile.go
+++ b/controlplane/worker_profile.go
@@ -3,6 +3,7 @@ package controlplane
import (
"errors"
"fmt"
+ "math"
"strings"
"time"
@@ -243,13 +244,11 @@ func firstNonEmpty(a, b string) string {
}
func requestedWorkerVCPUs(profile *WorkerProfile, workerCPURequest string) (int, error) {
- cpu := strings.TrimSpace(workerCPURequest)
- if profile != nil && strings.TrimSpace(profile.CPU) != "" {
- cpu = strings.TrimSpace(profile.CPU)
- }
- if cpu == "" {
- cpu = defaultWorkerCPU
+ profileCPU := ""
+ if profile != nil {
+ profileCPU = profile.CPU
}
+ cpu := effectiveWorkerResourceRequest(profileCPU, workerCPURequest, defaultWorkerCPU)
q, err := resource.ParseQuantity(cpu)
if err != nil {
return 0, fmt.Errorf("invalid worker cpu quantity %q: %w", cpu, err)
@@ -261,6 +260,52 @@ func requestedWorkerVCPUs(profile *WorkerProfile, workerCPURequest string) (int,
return int((millis + 999) / 1000), nil
}
+// requestedWorkerMemoryBytes returns the exact pod memory request charged to
+// session admission. Its precedence mirrors workerResourcesForProfile: an
+// effective profile value, then the deployment-wide worker request, then the
+// built-in default. Kubernetes quantities round positive sub-byte values up to
+// one byte; values that cannot be represented by the durable int64 accounting
+// columns are rejected before a queue entry can be written.
+func requestedWorkerMemoryBytes(profile *WorkerProfile, workerMemoryRequest string) (int64, error) {
+ profileMemory := ""
+ if profile != nil {
+ profileMemory = profile.Memory
+ }
+ memory := effectiveWorkerResourceRequest(profileMemory, workerMemoryRequest, defaultWorkerMemory)
+ q, err := resource.ParseQuantity(memory)
+ if err != nil {
+ return 0, fmt.Errorf("invalid worker memory quantity %q: %w", memory, err)
+ }
+ if q.Sign() <= 0 {
+ return 0, fmt.Errorf("worker memory quantity %q must be positive", memory)
+ }
+ max := resource.NewQuantity(math.MaxInt64, resource.DecimalSI)
+ // ParseQuantity caps overflowing BinarySI values (for example 8Ei) at
+ // MaxInt64. No exact positive BinarySI quantity can equal MaxInt64 because
+ // binary suffixes are powers of two, so that sentinel also means overflow.
+ if q.Cmp(*max) > 0 || (q.Format == resource.BinarySI && q.Value() == math.MaxInt64) {
+ return 0, fmt.Errorf("worker memory quantity %q exceeds %d bytes", memory, int64(math.MaxInt64))
+ }
+ bytes := q.Value()
+ if bytes <= 0 {
+ return 0, fmt.Errorf("worker memory quantity %q must resolve to a positive int64 byte count", memory)
+ }
+ return bytes, nil
+}
+
+// effectiveWorkerResourceRequest is the shared profile/deployment/built-in
+// precedence used by resource admission. Keep this aligned with pod resource
+// construction in workerResourcesForProfile.
+func effectiveWorkerResourceRequest(profileValue, deploymentValue, builtIn string) string {
+ if value := strings.TrimSpace(profileValue); value != "" {
+ return value
+ }
+ if value := strings.TrimSpace(deploymentValue); value != "" {
+ return value
+ }
+ return builtIn
+}
+
// useExploratoryTier decides whether a connection starts on the small
// exploratory worker rather than directly on its target shape. Every exclusion
// here is a case where starting small is a guaranteed-wasted acquire +
diff --git a/controlplane/worker_profile_k8s_test.go b/controlplane/worker_profile_k8s_test.go
new file mode 100644
index 00000000..7c0400a6
--- /dev/null
+++ b/controlplane/worker_profile_k8s_test.go
@@ -0,0 +1,38 @@
+//go:build kubernetes
+
+package controlplane
+
+import "testing"
+
+func TestRequestedWorkerMemoryBytesMatchesWorkerPodRequest(t *testing.T) {
+ tests := []struct {
+ name string
+ profile *WorkerProfile
+ deploymentMemory string
+ }{
+ {name: "profile", profile: &WorkerProfile{Memory: "60Gi"}, deploymentMemory: "120Gi"},
+ {name: "deployment default", deploymentMemory: "120Gi"},
+ {name: "built-in default"},
+ {name: "Kubernetes sub-byte rounding", deploymentMemory: "500m"},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ pool := &K8sWorkerPool{workerMemoryRequest: tt.deploymentMemory}
+ podProfile := WorkerProfile{}
+ if tt.profile != nil {
+ podProfile = *tt.profile
+ }
+ podResources := pool.workerResourcesForProfile(podProfile)
+ podMemory := podResources.Requests.Memory().Value()
+
+ admissionMemory, err := requestedWorkerMemoryBytes(tt.profile, tt.deploymentMemory)
+ if err != nil {
+ t.Fatalf("requestedWorkerMemoryBytes: %v", err)
+ }
+ if admissionMemory != podMemory {
+ t.Fatalf("admission memory = %d, pod request = %d", admissionMemory, podMemory)
+ }
+ })
+ }
+}
diff --git a/controlplane/worker_profile_test.go b/controlplane/worker_profile_test.go
index 7db3a444..657c014f 100644
--- a/controlplane/worker_profile_test.go
+++ b/controlplane/worker_profile_test.go
@@ -160,6 +160,45 @@ func TestRequestedWorkerVCPUs(t *testing.T) {
}
}
+func TestRequestedWorkerMemoryBytes(t *testing.T) {
+ tests := []struct {
+ name string
+ profile *WorkerProfile
+ defaultMemory string
+ want int64
+ wantErr bool
+ }{
+ {name: "profile wins", profile: &WorkerProfile{Memory: "32Gi"}, defaultMemory: "120Gi", want: 32 << 30},
+ {name: "deployment default", defaultMemory: "120Gi", want: 120 << 30},
+ {name: "built in default", want: 16 << 30},
+ {name: "binary quantity", defaultMemory: "1536Mi", want: 1536 << 20},
+ {name: "sub-byte quantity rounds to Kubernetes byte request", defaultMemory: "500m", want: 1},
+ {name: "bad memory", defaultMemory: "a lot", wantErr: true},
+ {name: "zero memory", defaultMemory: "0", wantErr: true},
+ {name: "negative memory", defaultMemory: "-1Gi", wantErr: true},
+ {name: "decimal overflow", defaultMemory: "9223372036854775808", wantErr: true},
+ {name: "binary overflow capped by quantity parser", defaultMemory: "8Ei", wantErr: true},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ got, err := requestedWorkerMemoryBytes(tt.profile, tt.defaultMemory)
+ if tt.wantErr {
+ if err == nil {
+ t.Fatalf("expected error, got %d", got)
+ }
+ return
+ }
+ if err != nil {
+ t.Fatalf("unexpected error: %v", err)
+ }
+ if got != tt.want {
+ t.Fatalf("requestedWorkerMemoryBytes() = %d, want %d", got, tt.want)
+ }
+ })
+ }
+}
+
// When the client DOES set sizing and no default size is configured, the built-in
// 8/16Gi/20m applies for the omitted fields.
func TestResolveWorkerProfileSizing_BuiltinDefaults(t *testing.T) {
diff --git a/docs/metrics.md b/docs/metrics.md
index 30a4deae..e1f9c609 100644
--- a/docs/metrics.md
+++ b/docs/metrics.md
@@ -45,7 +45,7 @@ warehouse disappears from the metric on the next snapshot refresh.
|---|---|---|
| Admission evaluation | `duckgres_session_admission_evaluation_*` | One request-owned DB admission poll, including its serialized transaction and lock wait. |
| Admission queue | `duckgres_session_admission_wait_seconds`, `duckgres_session_admission_requests_total` | After a request is successfully enqueued until grant, hard rejection, timeout, cancellation, or evaluation error. Enqueue failures are excluded. |
-| Admission state | `duckgres_session_admission_queue_depth`, `duckgres_session_admission_active_vcpus`, `duckgres_session_admission_limit_vcpus` | Local waiting callers, live local lease handles, and the effective org cap reconciled for active org stacks from each control-plane process's current config snapshot. |
+| Admission state | `duckgres_session_admission_queue_depth`, `duckgres_session_admission_active_vcpus`, `duckgres_session_admission_limit_vcpus`, `duckgres_session_admission_active_memory_bytes`, `duckgres_session_admission_limit_memory_bytes` | Local waiting callers, resources held by live local lease handles, and effective org caps reconciled for active org stacks from each control-plane process's current config snapshot. |
| Worker acquisition | `duckgres_worker_acquire_*` | After admission grants until an existing, hot-idle, or newly spawned worker is allocated. |
| Session start | `duckgres_session_start_duration_seconds`, `duckgres_postgres_session_start_total` | After successful PostgreSQL authentication until `ReadyForQuery` is flushed or session bootstrap terminates. The counter records exactly one terminal result after server-side retries. |
| Query | `duckgres_query_total`, `duckgres_query_duration_seconds` | One non-empty query attempt and its execution duration. |
@@ -135,12 +135,14 @@ counted with `outcome="not_local"`; the owning session remains active.
| `duckgres_session_admission_queue_depth` | Gauge | `org` | In-process callers still waiting after successful durable enqueue. It is not a count of durable queue rows. |
| `duckgres_session_admission_active_vcpus` | Gauge | `org` | Requested vCPUs held by live local lease handles. It is admitted capacity, not measured CPU usage or the exact durable lease-row total. |
| `duckgres_session_admission_limit_vcpus` | Gauge | `org` | Effective org cap for an active org stack, reconciled from this process's current config snapshot. `0` means unlimited. |
+| `duckgres_session_admission_active_memory_bytes` | Gauge | `org` | Requested memory bytes held by live local lease handles. It is admitted capacity, not measured memory usage or the exact durable lease-row total. |
+| `duckgres_session_admission_limit_memory_bytes` | Gauge | `org` | Effective org memory cap in bytes for an active org stack, reconciled from this process's current config snapshot. `0` means unlimited. |
| `duckgres_session_start_duration_seconds` | Histogram | `org`, `protocol`, `outcome` | Authenticated PostgreSQL create-to-ready latency. |
| `duckgres_postgres_session_start_total` | Counter | `org`, `outcome`, `reason` | Exactly one terminal authenticated PostgreSQL session-start result after server-side retries. |
Admission request outcomes are `granted`, `rejected`, `timeout`, `canceled`,
and `error`. `rejected` means the requested worker shape can never fit its hard
-organization or user vCPU ceiling.
+organization memory/vCPU ceiling or user vCPU ceiling.
Session-start outcomes are `success`, `timeout`, `canceled`, `capacity`,
`draining`, and `error`.
@@ -152,26 +154,27 @@ operator can usually alleviate. The remaining reasons let alerts exclude bad
client input, planned lifecycle transitions, client disconnects, wire errors,
and newly added paths that have not yet been classified. `capacity` covers
runtime worker exhaustion and admission
-timeouts; requests that exceed a configured hard org or user vCPU limit are
+timeouts; requests that exceed a configured hard org resource or user vCPU limit are
reported with reason `client`.
Evaluation decisions are `granted_current`, `already_granted`, `rejected`,
`blocked`, `waiting`, `inactive`, `missing`, `canceled`, `timeout`, and `error`.
Each evaluation describes only the polling request. Evaluation reasons are
-`none`, `org_vcpu`, `user_vcpu`, `org_user_vcpu`, `user_ineligible`,
+`none`, `org_vcpu`, `org_memory`, `user_vcpu`, `org_user_vcpu`, `user_ineligible`,
`resharding`, `fifo`, and `store_error`.
-A terminal request retains vCPU-cap attribution across every admission poll.
-No blocking poll produces `reason="none"`. If the request encountered an org
-cap, a user cap, or both, its terminal reason is `org_vcpu`, `user_vcpu`, or
-`org_user_vcpu`, respectively, even when it also encountered another reason.
-Without a vCPU-cap reason, one distinct reason is kept as-is and multiple
+A terminal request retains resource-cap attribution across every admission
+poll. No blocking poll produces `reason="none"`. A single observed cap is
+reported as `org_vcpu`, `org_memory`, or `user_vcpu`; observing both vCPU caps
+produces `org_user_vcpu`, while observing memory and a vCPU cap produces
+`resource_mixed`. Resource-cap reasons take precedence over unrelated blocking reasons.
+Without a resource-cap reason, one distinct reason is kept as-is and multiple
distinct reasons become `mixed`. An admission store failure contributes
`store_error`; interruption by the caller is classified as cancellation or
timeout instead.
-Queue depth and active vCPUs are process-local logical contributions. Use
-`sum by (org)` across control-plane replicas. Active vCPUs drop when a live
+Queue depth and active resource gauges are process-local logical contributions.
+Use `sum by (org)` across control-plane replicas. Active resources drop when a live
lease handle transfers cleanup ownership to the reclaimer; durable rows still
awaiting cleanup are excluded, so database-enforced usage can temporarily be
higher. Each replica reconciles the limit when it creates, updates, or removes
@@ -311,19 +314,24 @@ histogram_quantile(
)
```
-Terminal requests affected by an org or user vCPU cap:
+Terminal requests affected by an org memory/vCPU or user vCPU cap:
```promql
sum by (org, reason) (
- rate(duckgres_session_admission_requests_total{reason=~"org_vcpu|user_vcpu|org_user_vcpu"}[5m])
+ rate(duckgres_session_admission_requests_total{reason=~"org_vcpu|org_memory|user_vcpu|org_user_vcpu|resource_mixed"}[5m])
)
```
-Current queue depth and admitted vCPUs:
+`resource_mixed` means the request encountered both a memory cap and a vCPU
+cap. The separate `mixed` reason remains reserved for multiple unrelated
+non-resource blockers, so this query has no ambiguity.
+
+Current queue depth and admitted resources:
```promql
sum by (org) (duckgres_session_admission_queue_depth)
sum by (org) (duckgres_session_admission_active_vcpus)
+sum by (org) (duckgres_session_admission_active_memory_bytes)
```
Live admitted-session utilization for capped orgs (`limit=0` is deliberately
@@ -335,6 +343,14 @@ sum by (org) (duckgres_session_admission_active_vcpus)
(max by (org) (duckgres_session_admission_limit_vcpus) > 0)
```
+The equivalent memory utilization ratio is:
+
+```promql
+sum by (org) (duckgres_session_admission_active_memory_bytes)
+ / on (org)
+(max by (org) (duckgres_session_admission_limit_memory_bytes) > 0)
+```
+
Authenticated session-start p95:
```promql
diff --git a/docs/runbooks/org-connection-admission.md b/docs/runbooks/org-connection-admission.md
index 4d56d6df..994155e9 100644
--- a/docs/runbooks/org-connection-admission.md
+++ b/docs/runbooks/org-connection-admission.md
@@ -1,8 +1,18 @@
# Org Connection Admission
-The multitenant control plane admits each connection by its requested worker
-vCPUs. `duckgres_orgs.max_vcpus` is the org ceiling and
-`duckgres_org_users.max_vcpus` is the user ceiling; `0` means unlimited.
+The multitenant control plane admits each connection by its resolved worker
+vCPUs and memory. `duckgres_orgs.max_vcpus` and `duckgres_orgs.max_memory` are
+the org ceilings, while `duckgres_org_users.max_vcpus` is the user vCPU
+ceiling. CPU and memory are independent: neither limit derives the other.
+`max_memory` is a Kubernetes quantity such as `240Gi`; empty or `0` means
+unlimited. Newly provisioned orgs default to unlimited memory.
+
+Admission charges the worker shape resolved from client options, org defaults,
+deployment defaults, or built-in defaults. Only live session leases count.
+Hot-idle workers do not count, just as they do not count toward `max_vcpus`.
+For example, two simultaneous `15` vCPU / `120Gi` sessions require ceilings of
+at least `30` vCPUs and `240Gi`; setting either ceiling lower blocks the second
+session.
## Admission invariants
@@ -16,9 +26,12 @@ vCPUs. `duckgres_orgs.max_vcpus` is the org ceiling and
- Admission selection can create only the caller's lease and reject only the
caller's request. Serialized housekeeping may prune expired or inactive
foreign rows, but never grants or reserves capacity for them.
-- A request larger than its hard org or user ceiling is rejected. Temporary
+- A request larger than its hard org memory/vCPU or user vCPU ceiling is rejected. Temporary
saturation remains queued until capacity becomes available or the request
times out.
+- Lowering a limit below current usage does not terminate existing sessions.
+ Existing leases remain active and new requests wait until usage falls within
+ both ceilings.
- Resharding takes the same org lock. No lease can be granted after the
ready-to-resharding transition commits.
@@ -28,12 +41,23 @@ that head is the polling request. Otherwise the caller remains queued and the
head's owner admits itself on its next poll.
During a rolling deployment from the previous admission implementation, old
-replicas may
-still grant a foreign queue head and evaluate limits from their local config
-snapshot. Capacity remains protected by the shared org lock, but the strict
-request-owned and authoritative-limit invariants begin only after every old
-replica has exited. Avoid changing vCPU limits during that overlap when an
-exact change boundary matters.
+replicas may still grant a foreign queue head, evaluate limits from their local
+config snapshot, and grant leases without recording requested memory. Capacity
+remains protected by the shared org lock for limits every replica understands,
+but the strict request-owned, authoritative-limit, and memory-limit invariants
+begin only after every old replica has exited. Keep `max_memory` unlimited until
+all control-plane replicas run the new version and all legacy active leases
+whose `requested_memory_bytes` is zero have drained. A new replica deliberately
+fails closed on such an unknown active lease when a non-zero memory limit is
+enabled. Avoid changing vCPU limits during the overlap when an exact change
+boundary matters.
+
+After any org has a non-zero `max_memory`, do not roll an old control-plane
+binary back into the fleet: it does not enforce the cap and writes leases with
+unknown memory. Prefer roll-forward. If rollback is unavoidable, first clear
+every `max_memory` cap, complete the rollback, and leave the caps unlimited.
+After rolling forward again, wait until every replica is upgraded and every
+zero-memory active lease has drained before re-enabling the caps.
The connection queue timeout is configured by
`DUCKGRES_WORKER_QUEUE_TIMEOUT` (default `60s`). Owners poll every `100ms` while
@@ -71,8 +95,8 @@ admission reclaimer.
- A committed lease is authoritative. Its queue row is only a lifecycle mirror
and is removed with the lease when the session ends.
- If admission is blocked, inspect active leases and unexpired queue rows for
- the org. Confirm the owning control-plane instance is active before removing
- any row manually.
+ the org, including `requested_vcpus` and `requested_memory_bytes`. Confirm the
+ owning control-plane instance is active before removing any row manually.
- Do not delete a lease for a live session. If an owner is gone, expire its
control-plane runtime record and let the serialized cleanup path reclaim its
admission rows.
@@ -90,12 +114,16 @@ non-zero pending count can be healthy during steady connection churn. Reclaim
logs include the request, org, retry count, and age; the metrics deliberately
omit request and org labels.
-The org-labeled admission queue and active-vCPU gauges are logical local
-contributions, not exact durable row counts. Active vCPUs drop when cleanup is
+The org-labeled admission queue and active-resource gauges are logical local
+contributions, not exact durable row counts. Active resources drop when cleanup is
transferred to the reclaimer, before the durable lease row is necessarily
deleted. Use the reclaim backlog and attempt metrics above when that distinction
matters; [the metrics reference](../metrics.md) documents aggregation rules.
+Monitor `duckgres_session_admission_active_memory_bytes` with `sum by (org)`
+across replicas and `duckgres_session_admission_limit_memory_bytes` with
+`max by (org)`. The corresponding vCPU gauges use the same aggregation rules.
+
For local verification, run `just test-configstore-integration`; it exercises
cross-replica ordering, cancellation races, eventual live-owner reclamation,
hard-limit rejection, resharding, and stale-owner cleanup against PostgreSQL.
diff --git a/docs/runbooks/replenish-capacity.md b/docs/runbooks/replenish-capacity.md
index efcbbed1..da35cbdd 100644
--- a/docs/runbooks/replenish-capacity.md
+++ b/docs/runbooks/replenish-capacity.md
@@ -24,8 +24,10 @@ workers)`) and the size following the largest worker shape spawned in the last
7 days — both derived from the worker spawn log, no config knobs; a real worker
spawn preempts a placeholder and schedules immediately. So "low capacity" now means one of: spawns
are failing, the cluster can't schedule pods, headroom is exhausted, or an org
-has hit its per-org worker cap (`Org.MaxWorkers`; 0 = unbounded). There is no
-global/cluster worker cap — the node pool / autoscaler is the only shared ceiling.
+has hit its per-org worker cap (`Org.MaxWorkers`; 0 = unbounded). Session
+admission also enforces independent per-org vCPU and memory ceilings before
+worker acquisition. There is no global/cluster worker cap — the node pool /
+autoscaler is the only shared ceiling.
## Metrics to watch
@@ -34,6 +36,8 @@ global/cluster worker cap — the node pool / autoscaler is the only shared ceil
| `sum(duckgres_worker_lifecycle_count{state="spawning"})` sustained > 0 | Spawns are slow/stuck (image pull, Pending pods, cold nodes) |
| `sum by (reason)(rate(duckgres_worker_spawn_failures_total[5m]))` | Pod spawn errors by stage |
| `duckgres_control_plane_worker_acquire_seconds` (if present) | Session acquire latency |
+| `sum by (org) (duckgres_session_admission_active_memory_bytes)` | Memory requested by local live admission leases; sum across replicas |
+| `max by (org) (duckgres_session_admission_limit_memory_bytes)` | Effective org memory ceiling; 0 means unlimited |
| Crash-rate query (procedure step 1) | Workers crashing → repeated respawns |
Note: `binding="neutral"` lifecycle series are legacy (no worker is spawned
@@ -88,12 +92,18 @@ unassigned anymore); production capacity lives in `binding="org_bound"`.
binds in steady state, revisit the constants in `controlplane/headroom.go`.
4. **Check the cap.** If you hit `worker capacity exhausted for organization`,
- the request is at a real cap, not a scheduling problem:
+ a `53400` resource-limit error, or an admission timeout, the request may be
+ at a real cap rather than a scheduling problem:
- Per-org: the org reached its configured max concurrent workers
(`Org.MaxWorkers`) and all are busy — expected backpressure; the client
should retry as queries finish. Raise the org's `max_workers` (or set it to
0 = unbounded) if it is legitimately undersized; otherwise scale the worker
node pool. There is no global worker cap to raise.
+ - Resource admission: compare the active and limit vCPU/memory gauges and
+ inspect durable leases. A worker shape larger than the hard ceiling is
+ rejected immediately; temporary aggregate saturation queues. Follow the
+ [org connection admission runbook](org-connection-admission.md) before
+ changing a limit or reclaiming a stale lease.
5. **Verify recovery.** New sessions stop getting capacity errors;
`duckgres_worker_lifecycle_count{state="spawning"}` settles back toward 0.
diff --git a/tests/configstore/migrations_postgres_test.go b/tests/configstore/migrations_postgres_test.go
index edfa56ab..b58a679c 100644
--- a/tests/configstore/migrations_postgres_test.go
+++ b/tests/configstore/migrations_postgres_test.go
@@ -55,7 +55,8 @@ func TestConfigStoreRunsVersionedSQLMigrations(t *testing.T) {
requireGooseMigrationRecorded(t, db, 34)
requireGooseMigrationRecorded(t, db, 35)
requireGooseMigrationRecorded(t, db, 36)
- requireGooseLatestVersion(t, db, 36)
+ requireGooseMigrationRecorded(t, db, 37)
+ requireGooseLatestVersion(t, db, 37)
requireTableAbsent(t, db, "duckgres_schema_migrations")
// Migration 000018 added the reshard operation + verbose log tables.
@@ -193,6 +194,9 @@ func TestConfigStoreRunsVersionedSQLMigrations(t *testing.T) {
t.Fatalf("default_worker_min_hot_idle column count = %d, want 1", columnCount)
}
requireColumnDefault(t, db, "duckgres_orgs", "max_vcpus", "0")
+ requireColumnType(t, db, "duckgres_orgs", "max_memory", "character varying")
+ requireColumnNotNull(t, db, "duckgres_orgs", "max_memory")
+ requireColumnDefault(t, db, "duckgres_orgs", "max_memory", "''::character varying")
requireColumnDefault(t, db, "duckgres_org_users", "max_vcpus", "0")
// Migration 000011 added the per-user kill-switch column.
requireColumnDefault(t, db, "duckgres_org_users", "disabled", "false")
@@ -230,6 +234,34 @@ func TestConfigStoreRunsVersionedSQLMigrations(t *testing.T) {
requireColumnAbsent(t, db, "duckgres_org_users", "service_grant_expires_at")
}
+func TestConfigStoreSnapshotParsesOrgMaxMemory(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ if err := store.DB().Create(&cpconfigstore.Org{
+ Name: "memory-limited", DatabaseName: "memory-limited", MaxMemory: "120Gi",
+ }).Error; err != nil {
+ t.Fatalf("create memory-limited org: %v", err)
+ }
+ if err := store.Reload(); err != nil {
+ t.Fatalf("reload memory-limited org: %v", err)
+ }
+ const wantBytes = int64(120 * 1024 * 1024 * 1024)
+ if got := store.Snapshot().Orgs["memory-limited"].MaxMemoryBytes; got != wantBytes {
+ t.Fatalf("snapshot max memory bytes = %d, want %d", got, wantBytes)
+ }
+
+ if err := store.DB().Model(&cpconfigstore.Org{}).
+ Where("name = ?", "memory-limited").
+ Update("max_memory", "invalid").Error; err != nil {
+ t.Fatalf("seed invalid max_memory: %v", err)
+ }
+ if err := store.Reload(); err == nil || !strings.Contains(err.Error(), "max_memory") {
+ t.Fatalf("reload invalid max_memory error = %v, want max_memory validation failure", err)
+ }
+ if got := store.Snapshot().Orgs["memory-limited"].MaxMemoryBytes; got != wantBytes {
+ t.Fatalf("failed reload replaced last good snapshot: max memory bytes = %d, want %d", got, wantBytes)
+ }
+}
+
func TestConfigStoreSQLMigrationsUpgradeVersion8Schema(t *testing.T) {
_, connStr := newIsolatedConfigStoreSchema(t)
store, err := cpconfigStoreNew(connStr)
@@ -243,6 +275,7 @@ func TestConfigStoreSQLMigrationsUpgradeVersion8Schema(t *testing.T) {
if err := store.DB().Exec(`
ALTER TABLE duckgres_orgs DROP COLUMN data_imports_table_naming_version;
ALTER TABLE duckgres_orgs DROP COLUMN max_vcpus;
+ ALTER TABLE duckgres_orgs DROP COLUMN max_memory;
ALTER TABLE duckgres_org_users DROP COLUMN max_vcpus;
ALTER TABLE duckgres_org_users DROP COLUMN disabled;
ALTER TABLE duckgres_orgs ADD COLUMN IF NOT EXISTS max_connections BIGINT DEFAULT 0;
@@ -272,7 +305,7 @@ func TestConfigStoreSQLMigrationsUpgradeVersion8Schema(t *testing.T) {
DROP TABLE IF EXISTS duckgres_reshard_operation_log;
DROP TABLE IF EXISTS duckgres_reshard_operations;
DROP TABLE IF EXISTS duckgres_service_grants;
- DELETE FROM goose_db_version WHERE version_id IN (9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36);
+ DELETE FROM goose_db_version WHERE version_id IN (9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37);
`).Error; err != nil {
t.Fatalf("downgrade baseline schema to pre-v9 shape: %v", err)
}
@@ -322,10 +355,12 @@ func TestConfigStoreSQLMigrationsUpgradeVersion8Schema(t *testing.T) {
requireGooseMigrationRecorded(t, upgradedDB, 34)
requireGooseMigrationRecorded(t, upgradedDB, 35)
requireGooseMigrationRecorded(t, upgradedDB, 36)
- requireGooseLatestVersion(t, upgradedDB, 36)
+ requireGooseMigrationRecorded(t, upgradedDB, 37)
+ requireGooseLatestVersion(t, upgradedDB, 37)
requireColumnPresent(t, upgradedDB, "duckgres_reshard_operations", "password_url")
requireTablePresent(t, upgradedDB, "duckgres_worker_spawn_log")
requireColumnDefault(t, upgradedDB, "duckgres_orgs", "max_vcpus", "0")
+ requireColumnDefault(t, upgradedDB, "duckgres_orgs", "max_memory", "''::character varying")
requireColumnDefault(t, upgradedDB, "duckgres_org_users", "max_vcpus", "0")
requireColumnDefault(t, upgradedDB, "duckgres_org_users", "disabled", "false")
requireColumnAbsent(t, upgradedDB, "duckgres_orgs", "default_team_id")
@@ -365,7 +400,8 @@ func TestConfigStoreSQLMigration34VersionsExistingAndNewOrgs(t *testing.T) {
ALTER TABLE duckgres_orgs DROP COLUMN data_imports_table_naming_version;
ALTER TABLE duckgres_org_users DROP COLUMN IF EXISTS service_grant_expires_at;
DROP TABLE IF EXISTS duckgres_service_grants;
- DELETE FROM goose_db_version WHERE version_id IN (34, 35, 36);
+ ALTER TABLE duckgres_orgs DROP COLUMN max_memory;
+ DELETE FROM goose_db_version WHERE version_id IN (34, 35, 36, 37);
`).Error; err != nil {
t.Fatalf("restore pre-migration-34 schema: %v", err)
}
diff --git a/tests/configstore/org_connection_limiter_test.go b/tests/configstore/org_connection_limiter_test.go
index 99834b43..4c010d27 100644
--- a/tests/configstore/org_connection_limiter_test.go
+++ b/tests/configstore/org_connection_limiter_test.go
@@ -165,6 +165,48 @@ func seedAuthoritativeOrgConnectionLimits(t *testing.T, store *cpconfigstore.Con
}
}
+func seedAuthoritativeOrgConnectionMemoryLimit(t *testing.T, store *cpconfigstore.ConfigStore, orgID, orgMaxMemory string, usernames ...string) {
+ t.Helper()
+
+ if err := store.DB().Create(&cpconfigstore.Org{
+ Name: orgID,
+ DatabaseName: orgID,
+ MaxMemory: orgMaxMemory,
+ }).Error; err != nil {
+ t.Fatalf("seed org %q memory limit: %v", orgID, err)
+ }
+ for _, username := range usernames {
+ if err := store.DB().Create(&cpconfigstore.OrgUser{
+ OrgID: orgID,
+ Username: username,
+ Password: "test-password-hash",
+ }).Error; err != nil {
+ t.Fatalf("seed user %s/%s: %v", orgID, username, err)
+ }
+ }
+}
+
+func enqueueOrgConnectionMemoryRequest(t *testing.T, store *cpconfigstore.ConfigStore, requestID, orgID, username, cpInstanceID string, pid int32, requestedMemoryBytes int64, enqueuedAt time.Time) *cpconfigstore.OrgConnectionQueueEntry {
+ t.Helper()
+
+ request := &cpconfigstore.OrgConnectionQueueEntry{
+ RequestID: requestID,
+ OrgID: orgID,
+ Username: username,
+ CPInstanceID: cpInstanceID,
+ PID: pid,
+ Protocol: "postgres",
+ RequestedVCPUs: 1,
+ RequestedMemoryBytes: requestedMemoryBytes,
+ EnqueuedAt: enqueuedAt,
+ ExpiresAt: enqueuedAt.Add(time.Minute),
+ }
+ if err := store.EnqueueOrgConnectionRequest(request); err != nil {
+ t.Fatalf("enqueue %s: %v", requestID, err)
+ }
+ return request
+}
+
func TestOrgConnectionLeasesEnforceOrgVCPUBudget(t *testing.T) {
store := newIsolatedConfigStore(t)
upsertActiveCP(t, store, "cp-a")
@@ -229,6 +271,326 @@ func TestOrgConnectionLeasesEnforceOrgVCPUBudget(t *testing.T) {
}
}
+func TestOrgConnectionLeasesEnforceOrgMemoryBudget(t *testing.T) {
+ t.Run("exact boundary succeeds", func(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+ upsertActiveCP(t, store, "cp-b")
+
+ now := time.Now()
+ first := enqueueOrgConnectionMemoryRequest(t, store, "request-a", "org-memory", "alice", "cp-a", 1001, 80, now)
+ second := enqueueOrgConnectionMemoryRequest(t, store, "request-b", "org-memory", "bob", "cp-b", 1002, 40, now.Add(time.Millisecond))
+ limits := cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 120}
+
+ firstLease, err := store.TryAcquireOrgConnectionLease(first.RequestID, limits, now)
+ if err != nil || firstLease == nil {
+ t.Fatalf("acquire first lease = %#v, err %v", firstLease, err)
+ }
+ if firstLease.RequestedMemoryBytes != first.RequestedMemoryBytes {
+ t.Fatalf("first lease requested memory = %d, want %d", firstLease.RequestedMemoryBytes, first.RequestedMemoryBytes)
+ }
+ secondLease, err := store.TryAcquireOrgConnectionLease(second.RequestID, limits, now)
+ if err != nil || secondLease == nil {
+ t.Fatalf("acquire exact-boundary second lease = %#v, err %v", secondLease, err)
+ }
+ })
+
+ t.Run("one byte over blocks until release", func(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+ upsertActiveCP(t, store, "cp-b")
+
+ now := time.Now()
+ first := enqueueOrgConnectionMemoryRequest(t, store, "request-a", "org-memory", "alice", "cp-a", 1001, 80, now)
+ second := enqueueOrgConnectionMemoryRequest(t, store, "request-b", "org-memory", "bob", "cp-b", 1002, 41, now.Add(time.Millisecond))
+ limits := cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 120}
+
+ firstLease, err := store.TryAcquireOrgConnectionLease(first.RequestID, limits, now)
+ if err != nil || firstLease == nil {
+ t.Fatalf("acquire first lease = %#v, err %v", firstLease, err)
+ }
+ blocked, err := store.TryAcquireOrgConnectionLease(second.RequestID, limits, now)
+ if err != nil {
+ t.Fatalf("acquire one-byte-over request: %v", err)
+ }
+ if blocked != nil {
+ t.Fatalf("one-byte-over request acquired lease %#v", blocked)
+ }
+ assertOrgConnectionRequestPending(t, store, second.RequestID)
+
+ if err := store.ReleaseOrgConnectionLease(firstLease.LeaseID); err != nil {
+ t.Fatalf("release first lease: %v", err)
+ }
+ secondLease, err := store.TryAcquireOrgConnectionLease(second.RequestID, limits, now)
+ if err != nil || secondLease == nil {
+ t.Fatalf("acquire second lease after release = %#v, err %v", secondLease, err)
+ }
+ })
+
+ t.Run("zero is unlimited", func(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+
+ now := time.Now()
+ request := enqueueOrgConnectionMemoryRequest(t, store, "request-a", "org-memory", "alice", "cp-a", 1001, 1<<50, now)
+ lease, err := store.TryAcquireOrgConnectionLease(request.RequestID, cpconfigstore.OrgResourceLimits{}, now)
+ if err != nil || lease == nil {
+ t.Fatalf("acquire with unlimited memory = %#v, err %v", lease, err)
+ }
+ })
+}
+
+func TestOrgConnectionLeasesMemoryBudgetIgnoresExpiredOwner(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+ upsertActiveCP(t, store, "cp-b")
+
+ now := time.Now()
+ first := enqueueOrgConnectionMemoryRequest(t, store, "request-a", "org-memory-expired", "alice", "cp-a", 1001, 120, now)
+ limits := cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 120}
+ if lease, err := store.TryAcquireOrgConnectionLease(first.RequestID, limits, now); err != nil || lease == nil {
+ t.Fatalf("acquire first lease = %#v, err %v", lease, err)
+ }
+
+ expiredAt := now.Add(time.Second)
+ if err := store.UpsertControlPlaneInstance(&cpconfigstore.ControlPlaneInstance{
+ ID: "cp-a",
+ PodName: "cp-a",
+ State: cpconfigstore.ControlPlaneInstanceStateExpired,
+ StartedAt: now,
+ LastHeartbeatAt: now,
+ ExpiredAt: &expiredAt,
+ }); err != nil {
+ t.Fatalf("expire cp-a: %v", err)
+ }
+
+ second := enqueueOrgConnectionMemoryRequest(t, store, "request-b", "org-memory-expired", "bob", "cp-b", 1002, 120, now.Add(time.Millisecond))
+ lease, err := store.TryAcquireOrgConnectionLease(second.RequestID, limits, now)
+ if err != nil || lease == nil {
+ t.Fatalf("acquire after owner expiry = %#v, err %v", lease, err)
+ }
+}
+
+func TestOrgConnectionLeasesMemoryBudgetFailsClosedForLegacyUnknownMemory(t *testing.T) {
+ t.Run("active lease", func(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+ upsertActiveCP(t, store, "cp-b")
+
+ now := time.Now()
+ if err := store.DB().Exec(
+ "INSERT INTO "+store.RuntimeSchema()+".org_connection_leases (lease_id, request_id, org_id, username, cp_instance_id, p_id, protocol, requested_vcpus, requested_memory_bytes, acquired_at, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+ "legacy-lease", "legacy-request", "org-legacy-memory", "alice", "cp-a", 1001, "postgres", 1, 0, now, now, now,
+ ).Error; err != nil {
+ t.Fatalf("insert legacy lease: %v", err)
+ }
+
+ request := enqueueOrgConnectionMemoryRequest(t, store, "request-b", "org-legacy-memory", "bob", "cp-b", 1002, 1, now.Add(time.Millisecond))
+ blocked, err := store.TryAcquireOrgConnectionLease(request.RequestID, cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 120}, now)
+ if err != nil {
+ t.Fatalf("acquire with unknown active usage: %v", err)
+ }
+ if blocked != nil {
+ t.Fatalf("request acquired despite unknown active memory: %#v", blocked)
+ }
+ assertOrgConnectionRequestPending(t, store, request.RequestID)
+
+ lease, err := store.TryAcquireOrgConnectionLease(request.RequestID, cpconfigstore.OrgResourceLimits{}, now)
+ if err != nil || lease == nil {
+ t.Fatalf("acquire with unlimited memory = %#v, err %v", lease, err)
+ }
+ })
+
+ for _, requestedMemoryBytes := range []int64{0, -1} {
+ t.Run(fmt.Sprintf("queued request %d", requestedMemoryBytes), func(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+
+ now := time.Now()
+ request := enqueueOrgConnectionMemoryRequest(t, store, "legacy-request", "org-legacy-memory", "alice", "cp-a", 1001, requestedMemoryBytes, now)
+ blocked, err := store.TryAcquireOrgConnectionLease(request.RequestID, cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 120}, now)
+ if err != nil {
+ t.Fatalf("acquire request with unknown memory: %v", err)
+ }
+ if blocked != nil {
+ t.Fatalf("request with unknown memory acquired lease %#v", blocked)
+ }
+ assertOrgConnectionRequestPending(t, store, request.RequestID)
+ })
+ }
+}
+
+func TestOrgConnectionLeasesMemoryBudgetPreservesFIFO(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ for _, cpID := range []string{"cp-active", "cp-oldest", "cp-later"} {
+ upsertActiveCP(t, store, cpID)
+ }
+
+ now := time.Now()
+ active := enqueueOrgConnectionMemoryRequest(t, store, "request-active", "org-memory-fifo", "carol", "cp-active", 1001, 60, now)
+ limits := cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 100}
+ activeLease, err := store.TryAcquireOrgConnectionLease(active.RequestID, limits, now)
+ if err != nil || activeLease == nil {
+ t.Fatalf("acquire active lease = %#v, err %v", activeLease, err)
+ }
+
+ oldest := enqueueOrgConnectionMemoryRequest(t, store, "request-oldest", active.OrgID, "alice", "cp-oldest", 1002, 50, now.Add(time.Millisecond))
+ later := enqueueOrgConnectionMemoryRequest(t, store, "request-later", active.OrgID, "bob", "cp-later", 1003, 40, now.Add(2*time.Millisecond))
+ lease, err := store.TryAcquireOrgConnectionLease(later.RequestID, limits, now)
+ if err != nil {
+ t.Fatalf("poll later request: %v", err)
+ }
+ if lease != nil {
+ t.Fatalf("later smaller request bypassed memory-blocked FIFO head: %#v", lease)
+ }
+ assertOrgConnectionRequestPending(t, store, oldest.RequestID)
+ assertOrgConnectionRequestPending(t, store, later.RequestID)
+
+ if err := store.ReleaseOrgConnectionLease(activeLease.LeaseID); err != nil {
+ t.Fatalf("release active lease: %v", err)
+ }
+ lease, err = store.TryAcquireOrgConnectionLease(later.RequestID, limits, now)
+ if err != nil {
+ t.Fatalf("poll later request after release: %v", err)
+ }
+ if lease != nil {
+ t.Fatalf("later request bypassed now-runnable FIFO head: %#v", lease)
+ }
+ oldestLease, err := store.TryAcquireOrgConnectionLease(oldest.RequestID, limits, now)
+ if err != nil || oldestLease == nil {
+ t.Fatalf("acquire oldest request = %#v, err %v", oldestLease, err)
+ }
+ laterLease, err := store.TryAcquireOrgConnectionLease(later.RequestID, limits, now)
+ if err != nil || laterLease == nil {
+ t.Fatalf("acquire later request = %#v, err %v", laterLease, err)
+ }
+}
+
+func TestOrgConnectionLeasesMemoryBudgetHandlesForeignUnserviceableHead(t *testing.T) {
+ t.Run("permanently oversized head is skipped", func(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-oldest")
+ upsertActiveCP(t, store, "cp-later")
+
+ now := time.Now()
+ oldest := enqueueOrgConnectionMemoryRequest(t, store, "request-oldest", "org-memory-head", "alice", "cp-oldest", 1001, 121, now)
+ later := enqueueOrgConnectionMemoryRequest(t, store, "request-later", oldest.OrgID, "bob", "cp-later", 1002, 1, now.Add(time.Millisecond))
+ limits := cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 120}
+
+ lease, err := store.TryAcquireOrgConnectionLease(later.RequestID, limits, now)
+ if err != nil || lease == nil || lease.RequestID != later.RequestID {
+ t.Fatalf("acquire behind permanently oversized head = %#v, err %v", lease, err)
+ }
+ assertOrgConnectionRequestPending(t, store, oldest.RequestID)
+
+ lease, err = store.TryAcquireOrgConnectionLease(oldest.RequestID, limits, now)
+ if lease != nil || !errors.Is(err, cpconfigstore.ErrOrgConnectionAdmissionRejected) {
+ t.Fatalf("oversized head rejection = lease %#v, err %v", lease, err)
+ }
+ assertOrgConnectionRequestAbsent(t, store, oldest.RequestID)
+ })
+
+ t.Run("unknown head fails closed", func(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-oldest")
+ upsertActiveCP(t, store, "cp-later")
+
+ now := time.Now()
+ oldest := enqueueOrgConnectionMemoryRequest(t, store, "request-oldest", "org-memory-head", "alice", "cp-oldest", 1001, 0, now)
+ later := enqueueOrgConnectionMemoryRequest(t, store, "request-later", oldest.OrgID, "bob", "cp-later", 1002, 1, now.Add(time.Millisecond))
+
+ lease, err := store.TryAcquireOrgConnectionLease(later.RequestID, cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: 120}, now)
+ if err != nil {
+ t.Fatalf("poll behind unknown-memory head: %v", err)
+ }
+ if lease != nil {
+ t.Fatalf("later request bypassed unknown-memory head: %#v", lease)
+ }
+ assertOrgConnectionRequestPending(t, store, oldest.RequestID)
+ assertOrgConnectionRequestPending(t, store, later.RequestID)
+ })
+}
+
+func TestOrgConnectionLeasesMemoryBudgetFailsClosedWhenActiveSumExceedsInt64(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ for _, cpID := range []string{"cp-a", "cp-b", "cp-c"} {
+ upsertActiveCP(t, store, cpID)
+ }
+
+ now := time.Now()
+ const maxMemoryBytes int64 = 1<<63 - 1
+ leases := []*cpconfigstore.OrgConnectionLease{
+ {LeaseID: "lease-a", RequestID: "active-a", OrgID: "org-memory-overflow", Username: "alice", CPInstanceID: "cp-a", PID: 1001, Protocol: "postgres", RequestedVCPUs: 1, RequestedMemoryBytes: maxMemoryBytes, AcquiredAt: now},
+ {LeaseID: "lease-b", RequestID: "active-b", OrgID: "org-memory-overflow", Username: "bob", CPInstanceID: "cp-b", PID: 1002, Protocol: "postgres", RequestedVCPUs: 1, RequestedMemoryBytes: 1, AcquiredAt: now},
+ }
+ if err := store.DB().Table(store.RuntimeSchema() + ".org_connection_leases").Create(leases).Error; err != nil {
+ t.Fatalf("seed overflowing active memory sum: %v", err)
+ }
+ request := enqueueOrgConnectionMemoryRequest(t, store, "request-c", "org-memory-overflow", "carol", "cp-c", 1003, 1, now)
+
+ lease, err := store.TryAcquireOrgConnectionLease(request.RequestID, cpconfigstore.OrgResourceLimits{OrgMaxMemoryBytes: maxMemoryBytes}, now)
+ if err != nil {
+ t.Fatalf("evaluate overflowing active memory sum: %v", err)
+ }
+ if lease != nil {
+ t.Fatalf("request acquired over overflowing memory usage: %#v", lease)
+ }
+ assertOrgConnectionRequestPending(t, store, request.RequestID)
+}
+
+func TestOrgConnectionAdmissionUsesAuthoritativeOrgMemoryLimit(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+ upsertActiveCP(t, store, "cp-b")
+ const orgID = "org-authoritative-memory"
+ seedAuthoritativeOrgConnectionMemoryLimit(t, store, orgID, "120Gi", "alice", "bob")
+
+ now := time.Now()
+ const gibibyte = int64(1024 * 1024 * 1024)
+ active := enqueueOrgConnectionMemoryRequest(t, store, "request-active", orgID, "alice", "cp-a", 1001, 80*gibibyte, now)
+ activeRef := cpconfigstore.OrgConnectionAdmissionRef{RequestID: active.RequestID, OrgID: orgID, CPInstanceID: active.CPInstanceID}
+ if lease, _, err := store.ScheduleAndClaimOrgConnectionLeaseForRefWithEvaluationContext(t.Context(), activeRef); err != nil || lease == nil {
+ t.Fatalf("grant active request = %#v, err %v", lease, err)
+ }
+
+ blocked := enqueueOrgConnectionMemoryRequest(t, store, "request-blocked", orgID, "bob", "cp-b", 1002, 40*gibibyte+1, now.Add(time.Millisecond))
+ blockedRef := cpconfigstore.OrgConnectionAdmissionRef{RequestID: blocked.RequestID, OrgID: orgID, CPInstanceID: blocked.CPInstanceID}
+ lease, evaluation, err := store.ScheduleAndClaimOrgConnectionLeaseForRefWithEvaluationContext(t.Context(), blockedRef)
+ if err != nil || lease != nil {
+ t.Fatalf("evaluate blocked request = lease %#v, err %v", lease, err)
+ }
+ if evaluation.Decision != "blocked" || evaluation.Reason != "org_memory" {
+ t.Fatalf("evaluation = %#v, want blocked/org_memory", evaluation)
+ }
+}
+
+func TestOrgConnectionAdmissionRejectsRequestLargerThanOrgMemoryLimit(t *testing.T) {
+ store := newIsolatedConfigStore(t)
+ upsertActiveCP(t, store, "cp-a")
+ const orgID = "org-rejected-memory"
+ seedAuthoritativeOrgConnectionMemoryLimit(t, store, orgID, "120Gi", "alice")
+
+ now := time.Now()
+ const gibibyte = int64(1024 * 1024 * 1024)
+ request := enqueueOrgConnectionMemoryRequest(t, store, "request-rejected", orgID, "alice", "cp-a", 1001, 120*gibibyte+1, now)
+ ref := cpconfigstore.OrgConnectionAdmissionRef{RequestID: request.RequestID, OrgID: orgID, CPInstanceID: request.CPInstanceID}
+ lease, evaluation, err := store.ScheduleAndClaimOrgConnectionLeaseForRefWithEvaluationContext(t.Context(), ref)
+ if lease != nil || !errors.Is(err, cpconfigstore.ErrOrgConnectionAdmissionRejected) {
+ t.Fatalf("hard rejection = lease %#v, err %v", lease, err)
+ }
+ var rejection *cpconfigstore.OrgConnectionAdmissionRejectedError
+ if !errors.As(err, &rejection) {
+ t.Fatalf("hard rejection error %T, want *OrgConnectionAdmissionRejectedError", err)
+ }
+ if rejection.Reason != cpconfigstore.OrgConnectionAdmissionRejectedOrgMemory || rejection.RequestedMemoryBytes != request.RequestedMemoryBytes || rejection.MaximumMemoryBytes != 120*gibibyte {
+ t.Fatalf("hard rejection = %#v", rejection)
+ }
+ if evaluation.Decision != "rejected" || evaluation.Reason != "org_memory" {
+ t.Fatalf("hard-rejection evaluation = %#v, want rejected/org_memory", evaluation)
+ }
+ assertOrgConnectionRequestAbsent(t, store, request.RequestID)
+}
+
func TestOrgConnectionLeasesEnforceUserVCPUBudget(t *testing.T) {
store := newIsolatedConfigStore(t)
upsertActiveCP(t, store, "cp-a")
diff --git a/tests/configstore/runtime_store_postgres_test.go b/tests/configstore/runtime_store_postgres_test.go
index efc8496b..44a378bb 100644
--- a/tests/configstore/runtime_store_postgres_test.go
+++ b/tests/configstore/runtime_store_postgres_test.go
@@ -68,6 +68,19 @@ func TestRuntimeStorePostgres(t *testing.T) {
}
}
}
+ for _, table := range []string{"org_connection_queue", "org_connection_leases"} {
+ var nullable, defaultValue string
+ if err := store.DB().Raw(
+ "SELECT is_nullable, column_default FROM information_schema.columns WHERE table_schema = ? AND table_name = ? AND column_name = 'requested_memory_bytes'",
+ runtimeSchema,
+ table,
+ ).Row().Scan(&nullable, &defaultValue); err != nil {
+ t.Fatalf("lookup runtime memory column %s.requested_memory_bytes: %v", table, err)
+ }
+ if nullable != "NO" || defaultValue != "0" {
+ t.Fatalf("%s.requested_memory_bytes metadata = nullable %q default %q, want NO/0", table, nullable, defaultValue)
+ }
+ }
requireRuntimeIndexDefinition(t, store, "org_connection_queue", "idx_org_connection_queue_user_heads",
"(org_id, username, enqueued_at, request_id)", "WHERE (granted_at IS NULL)")
requireRuntimeIndexDefinition(t, store, "org_connection_leases", "idx_org_connection_leases_org_user",