From 71faf7326e075346099d82c8a60c9310e4ed2b7b Mon Sep 17 00:00:00 2001 From: Benjamin Knofe-Vider Date: Mon, 6 Jul 2026 09:29:23 +0200 Subject: [PATCH 1/2] Replace compute-billing push with the pull API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Billing now pulls compute usage from duckgres instead of duckgres pushing capture events to PostHog ingestion (docs/design/billing-pull-api.md). The per-connection metering is unchanged; what changes is the bucket key and how data leaves: - Migration 000015 widens the buffer key to (org_id, team_id, query_source, cpu, mem_gib, bucket_start) — team_id backfilled from the org's default_team_id, worker size as exact NUMERIC decimals — adds the single global ack cursor and drops the per-org push-drain state table. - The meter records the org's default team (config-snapshot lookup) and the session's duckgres.query_source GUC (via server.ConnectionBilling) alongside the provisioned worker size; distinct keys accumulate and bill separately. Metering is now always on for the remote backend. - GET /api/v1/billing/usage aggregates all closed buckets since the last ack into one row per key per UTC day, with watermark_low (cursor) and watermark_high (newest closed minute: now − 60s − 30s grace). POST /api/v1/billing/ack advances the cursor monotonically and deletes buckets ≤ watermark_high in one transaction; idempotent, and an ack beyond the latest closed bucket is rejected so unserved buckets can never be deleted. Both routes are internal-secret/admin-gated inside the audited /api/v1 group. - A leader-only safety GC hard-deletes buckets older than 30 days and WARNs with the dropped count (billing not keeping up = alertable). - The push path is gone: leader drain + capture client, and the --billing-ingest-url/--billing-ingest-token flags with their DUCKGRES_BILLING_INGEST_URL/_TOKEN envs. e2e: compute_usage_metering_wired becomes compute_usage_pull_api — two real connections (standard + endpoints GUC) must surface as separate usage rows carrying the provisioned default_team_id and worker size, then ack must advance the cursor, delete the acked buckets, stay idempotent on re-ack, and reject a future watermark. --- CLAUDE.md | 81 +++--- README.md | 2 - configresolve/cliflags.go | 4 - configresolve/cliflags_test.go | 2 - configresolve/resolve.go | 22 -- controlplane/compute_billing_api.go | 117 ++++++++ controlplane/compute_billing_api_test.go | 270 ++++++++++++++++++ controlplane/compute_capture.go | 120 -------- controlplane/compute_capture_test.go | 81 ------ controlplane/compute_drain.go | 111 ------- controlplane/compute_drain_test.go | 178 ------------ controlplane/compute_gc.go | 61 ++++ controlplane/compute_meter.go | 53 +++- controlplane/compute_meter_test.go | 91 ++++-- controlplane/configstore/compute_usage.go | 192 ++++++++----- .../000015_compute_usage_pull_key.sql | 63 ++++ controlplane/control.go | 29 +- controlplane/multitenant.go | 35 ++- docs/design/billing-compute-seconds-plan.md | 6 +- docs/design/billing-pull-api.md | 4 +- main.go | 2 - server/exports.go | 17 +- tests/configstore/migrations_postgres_test.go | 19 +- tests/e2e-mw-dev/harness.sh | 87 ++++-- 24 files changed, 910 insertions(+), 737 deletions(-) create mode 100644 controlplane/compute_billing_api.go create mode 100644 controlplane/compute_billing_api_test.go delete mode 100644 controlplane/compute_capture.go delete mode 100644 controlplane/compute_capture_test.go delete mode 100644 controlplane/compute_drain.go delete mode 100644 controlplane/compute_drain_test.go create mode 100644 controlplane/compute_gc.go create mode 100644 controlplane/configstore/migrations/000015_compute_usage_pull_key.sql diff --git a/CLAUDE.md b/CLAUDE.md index b5637c51..e5426fe2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -78,7 +78,7 @@ Key CLI flags for control-plane mode: - Config store: `--config-store`, `--config-poll-interval`, `--internal-secret` - K8s pool: `--k8s-worker-image`, `--k8s-worker-namespace`, `--k8s-control-plane-id`, `--k8s-worker-port`, `--k8s-worker-secret`, `--k8s-worker-configmap`, `--k8s-worker-image-pull-policy`, `--k8s-worker-service-account` (no global worker cap — per-org `Org.MaxWorkers`, 0=unbounded, is the only cap) - AWS / STS: `--aws-region` - - Compute-usage billing (emit side): `--billing-ingest-url` / `--billing-ingest-token` (env `DUCKGRES_BILLING_INGEST_URL` / `DUCKGRES_BILLING_INGEST_TOKEN`). Both must be set to enable per-org compute-usage metering+emit; either unset disables it (usage ships nowhere, queries never fail). See `docs/design/billing-compute-seconds-plan.md` and "Compute-Usage Billing" below. + - Compute-usage billing needs no config: metering is always on for the remote backend and billing PULLS usage over the internal-secret-authed HTTP API (`GET /api/v1/billing/usage` + `POST /api/v1/billing/ack`). See `docs/design/billing-pull-api.md` and "Compute-Usage Billing" below. - Pod scheduling knobs (CPU/memory requests, node selector, tolerations) are env-only — see `config_resolution.go`. Key CLI flags for duckdb-service mode: @@ -438,58 +438,73 @@ impersonation, audit log; sliceable by org + user). Design + decisions: ## Compute-Usage Billing (managed-warehouse, remote backend only) -duckgres meters per-org compute usage of worker pods and ships it to PostHog's -public ingestion as `"managed warehouse compute usage"` capture events (billing -sums the two raw metrics externally). Full design + decisions: -`docs/design/billing-compute-seconds-plan.md`. Scope is the **emit side**, and -**only** the remote/k8s backend (per-org worker pod with a known -`WorkerProfile` size). Pipeline: +duckgres meters per-org compute usage of worker pods into 60s buckets in the +config store; the billing service **pulls** the accumulated usage over an HTTP +API and acks a watermark, at which point duckgres deletes the acked buckets. +Full design + decisions: `docs/design/billing-pull-api.md` (supersedes the +push/capture reporting hop of `billing-compute-seconds-plan.md`; the metering +side of that doc still applies). Scope is **only** the remote/k8s backend +(per-org worker pod with a known `WorkerProfile` size). Pipeline: ``` -conn end → in-proc per-org counter (best-effort; never fails teardown) +conn end → in-proc counter keyed (org, default team, query_source, worker size) │ flusher (~15s) UPSERT-increment → config-store buffer (cross-CP sum) - ▼ duckgres_org_compute_usage / duckgres_org_compute_drain_state - leader drain (~60s) → POST {DUCKGRES_BILLING_INGEST_URL}/capture (ship-then-delete) + ▼ duckgres_org_compute_usage (+ duckgres_compute_billing_cursor) +billing: GET /api/v1/billing/usage (aggregated per key per UTC day, watermarks) + → POST /api/v1/billing/ack {watermark_high} → cursor advance + delete ≤ it +safety: leader-only GC hard-deletes buckets older than 30 days (WARN, alertable) ``` Two raw metrics per connection over its full lifetime, using the **provisioned** worker size: `cpu_seconds = vCPU × ceil(conn_secs)`, `memory_seconds = GiB × ceil(conn_secs)`. Counted internally in integer **millicore-seconds** / **MiB-seconds** (`compute_meter.go`) to avoid truncating a fractional-core / -sub-GiB worker. Invariants for anyone touching this path: +sub-GiB worker; worker size is stored in the bucket key as exact NUMERIC +decimals (vCPU / GiB). `team_id` is the org\'s `default_team_id` (resolved from +the config snapshot at connection end); `query_source` is the +`duckgres.query_source` session GUC (`standard` unless set; a mid-connection +change bills the whole connection under the final value). Invariants for anyone +touching this path: - **Metering is strictly best-effort and off the hot path.** A metering error - (counter, flush, drain) must NEVER block or fail a query or connection - teardown. The connection-end record is added to an in-process per-org counter - (map+mutex, microseconds, no I/O); everything downstream is async/retried. - `cp.computeMeter` is nil unless the remote backend is configured with both - ingest URL+token — every call site is nil-safe. -- **Enable gate = both `DUCKGRES_BILLING_INGEST_URL` and - `DUCKGRES_BILLING_INGEST_TOKEN` set** (`ControlPlaneConfig.BillingMeteringEnabled`). - Either unset → metering disabled (logged once at startup, ships nowhere). The - env names are fixed (infra wires them); keep them exact. + (counter, flush) must NEVER block or fail a query or connection teardown. The + connection-end record is added to an in-process counter (map+mutex, + microseconds, no I/O); the flush is async. `cp.computeMeter` is nil outside + the remote backend — every call site is nil-safe. There is no enable knob: + the remote backend always meters. - **Worker size is plumbed onto the connection** (`server.SetConnectionWorkerSize` → `clientConn.workerMillicores/workerMiB`, set in `control.go::handleConnection` from `workerBillingSize(workerProfile)`, remote-only). `workerMillicores==0` (non-remote / unknown) → metering skipped. The metric is computed once at the - SAME teardown point as `CloseConnectionMetrics` (the `#841` lifetime defer). + SAME teardown point as `CloseConnectionMetrics` (the `#841` lifetime defer), + via `server.ConnectionBilling` (which also carries the query source). - **Bucket = connection-end time floored to 60s.** Flush carries the sub-unit remainder forward so rounding never loses counts across flushes. Buffer flush - is UPSERT-increment so all CP pods sum into one row. -- **Drain is leader-only** (co-located under the janitor lease via - `JanitorLeaderManager.AttachLeaderLoop`), **ship-then-delete / at-least-once**: - only buckets `bucket_start ≤ now-60s-30s grace AND > last_drained_bucket` ship; - on HTTP 2xx → TXN advance high-water + DELETE row; on failure → keep + retry. - `event_uuid = hash(org_id, bucket_start)` + `timestamp = bucket_start` are - deterministic so a re-ship collapses onto the same ClickHouse row. + is UPSERT-increment so all CP pods sum into one row per key. +- **Serve only closed buckets.** `watermark_high` = the newest bucket with + `bucket_start ≤ now − 60s − 30s grace` (grace > flush interval, so every + CP\'s contribution has landed before a minute is served). The GET aggregates + the window `(cursor, watermark_high]` into one row per + `(org, team, query_source, cpu, mem_gib)` per **UTC day** — response size is + bounded by active keys × days, so billing downtime can\'t make it explode. +- **Ack is the only deletion path (plus the 30d GC).** `POST /billing/ack` + advances the single global cursor monotonically and deletes buckets + `≤ watermark_high` in one TXN (`AckComputeUsage`). Idempotent — re-acks and + stale acks are no-ops. An ack beyond the latest closed bucket is rejected + (400) so it can never delete buckets that were never served. Auth is the + admin internal secret (`RequireAdmin` on both routes, registered inside the + audited `/api/v1` group in `multitenant.go`). +- **Safety GC is leader-only** (`runComputeUsageGC`, attached under the janitor + lease): hard-deletes buckets older than 30 days regardless of ack and logs a + WARN with the dropped count — nonzero means billing stopped pulling (alert). - **Graceful shutdown does a final flush** after connections drain to their natural end (`shutdown`/`drainAndShutdown`), so a departing CP pod lands its last interval before exit. -- Touching the meter/flush/drain/capture, the worker-size plumbing, or the - config knobs → update `controlplane/compute_meter_test.go`, - `compute_drain_test.go`, `compute_size_test.go`, the migration assertion in - `tests/configstore/migrations_postgres_test.go`, and the - `compute_usage_metering_wired` assertion in `tests/e2e-mw-dev/harness.sh`. +- Touching the meter/flush/API/GC, the worker-size or query-source plumbing, or + the bucket key → update `controlplane/compute_meter_test.go`, + `compute_billing_api_test.go`, `compute_size_test.go`, the migration assertion + in `tests/configstore/migrations_postgres_test.go`, and the + `compute_usage_pull_api` assertion in `tests/e2e-mw-dev/harness.sh`. ## TODO Reference diff --git a/README.md b/README.md index d273a8d1..a5bb0dd8 100644 --- a/README.md +++ b/README.md @@ -280,8 +280,6 @@ Run with config file: | `DUCKGRES_QUERY_LOG_KAFKA_TOPIC` | Kafka topic for query-log events; required for Kafka mode | - | | `DUCKGRES_QUERY_LOG_KAFKA_CLIENT_ID` | Kafka client ID used by the query-log producer | `duckgres-query-log` | | `DUCKGRES_QUERY_LOG_KAFKA_GROUP_ID` | Kafka consumer group used by `duckgres --mode query-log-writer` | `duckgres-query-log-writer` | -| `DUCKGRES_BILLING_INGEST_URL` | PostHog public ingestion base URL for managed-warehouse compute-usage events (e.g. `https://us.i.posthog.com`). Remote backend only. Unset (or token unset) disables metering — usage ships nowhere and a query is never failed on its account. | - | -| `DUCKGRES_BILLING_INGEST_TOKEN` | PostHog project API token stamped on the compute-usage capture event. Remote backend only; both URL and token must be set to enable metering. | - | | `POSTHOG_API_KEY` | PostHog project API key (`phc_...`); enables log export **and product-analytics events** | - | | `POSTHOG_HOST` | PostHog ingest host | `us.i.posthog.com` | | `ADDITIONAL_POSTHOG_API_KEYS` | **(Experimental)** Comma-separated list of additional PostHog API keys to publish logs to. Requires `POSTHOG_API_KEY` to be set. | - | diff --git a/configresolve/cliflags.go b/configresolve/cliflags.go index ff5d7bfa..4b38a567 100644 --- a/configresolve/cliflags.go +++ b/configresolve/cliflags.go @@ -72,8 +72,6 @@ func RegisterCLIInputsFlags(fs *flag.FlagSet) func() CLIInputs { k8sWorkerServiceAccount := fs.String("k8s-worker-service-account", "", "Neutral ServiceAccount name for K8s worker pods (default: duckgres-worker) (env: DUCKGRES_K8S_WORKER_SERVICE_ACCOUNT)") awsRegion := fs.String("aws-region", "", "AWS region for STS client (env: DUCKGRES_AWS_REGION)") queryLog := fs.Bool("query-log", true, "Enable/disable DuckLake query log (use --query-log=false to disable; env: DUCKGRES_QUERY_LOG_ENABLED)") - billingIngestURL := fs.String("billing-ingest-url", "", "PostHog public ingestion base URL for managed-warehouse compute-usage events, e.g. https://us.i.posthog.com (remote backend only; unset disables metering) (env: DUCKGRES_BILLING_INGEST_URL)") - billingIngestToken := fs.String("billing-ingest-token", "", "PostHog project API token for managed-warehouse compute-usage events (remote backend only; unset disables metering) (env: DUCKGRES_BILLING_INGEST_TOKEN)") return func() CLIInputs { cli := CLIInputs{Set: map[string]bool{}} @@ -130,8 +128,6 @@ func RegisterCLIInputsFlags(fs *flag.FlagSet) func() CLIInputs { cli.K8sWorkerServiceAccount = *k8sWorkerServiceAccount cli.AWSRegion = *awsRegion cli.QueryLog = *queryLog - cli.BillingIngestURL = *billingIngestURL - cli.BillingIngestToken = *billingIngestToken return cli } } diff --git a/configresolve/cliflags_test.go b/configresolve/cliflags_test.go index e4fa2781..f8ce8569 100644 --- a/configresolve/cliflags_test.go +++ b/configresolve/cliflags_test.go @@ -51,8 +51,6 @@ func fieldNameToFlagName(name string) string { {"AWSRegion", "aws-region"}, {"SNIRoutingMode", "sni-routing-mode"}, {"ManagedHostnameSuffixes", "managed-hostname-suffixes"}, - {"BillingIngestURL", "billing-ingest-url"}, - {"BillingIngestToken", "billing-ingest-token"}, {"ConfigStoreConn", "config-store"}, {"ConfigPollInterval", "config-poll-interval"}, {"InternalSecretFallbacks", "internal-secret-fallbacks"}, diff --git a/configresolve/resolve.go b/configresolve/resolve.go index b68092c8..bf2afc79 100644 --- a/configresolve/resolve.go +++ b/configresolve/resolve.go @@ -82,8 +82,6 @@ type CLIInputs struct { K8sWorkerTolerationValue string AWSRegion string QueryLog bool - BillingIngestURL string - BillingIngestToken string } type Resolved struct { @@ -131,8 +129,6 @@ type Resolved struct { ManagedHostnameSuffixes []string DucklingBucketSuffix string DuckLakeDefaultSpecVersion string - BillingIngestURL string - BillingIngestToken string } func intPtr(n int) *int { return &n } @@ -222,7 +218,6 @@ func ResolveEffective(fileCfg *configloader.FileConfig, cli CLIInputs, getenv fu var sniRoutingMode string var managedHostnameSuffixes []string var ducklingBucketSuffix string - var billingIngestURL, billingIngestToken string if fileCfg != nil { if fileCfg.Host != "" { @@ -792,15 +787,6 @@ func ResolveEffective(fileCfg *configloader.FileConfig, cli CLIInputs, getenv fu if v := getenv("DUCKGRES_DUCKLING_BUCKET_SUFFIX"); v != "" { ducklingBucketSuffix = v } - // Managed-warehouse compute-usage billing ingest (remote backend only). Both - // must be set for metering to ship anything; either unset disables metering. - // Env names are fixed — infra wires these exact keys. - if v := getenv("DUCKGRES_BILLING_INGEST_URL"); v != "" { - billingIngestURL = v - } - if v := getenv("DUCKGRES_BILLING_INGEST_TOKEN"); v != "" { - billingIngestToken = v - } if v := getenv("DUCKGRES_WORKER_BACKEND"); v != "" { workerBackend = v } @@ -1112,12 +1098,6 @@ func ResolveEffective(fileCfg *configloader.FileConfig, cli CLIInputs, getenv fu // JWT_SIGNING_KEY_FALLBACKS. internalSecretFallbacks = splitAndTrim(cli.InternalSecretFallbacks, ",") } - if cli.Set["billing-ingest-url"] { - billingIngestURL = cli.BillingIngestURL - } - if cli.Set["billing-ingest-token"] { - billingIngestToken = cli.BillingIngestToken - } if cli.Set["sni-routing-mode"] { sniRoutingMode = cli.SNIRoutingMode } @@ -1273,8 +1253,6 @@ func ResolveEffective(fileCfg *configloader.FileConfig, cli CLIInputs, getenv fu ManagedHostnameSuffixes: managedHostnameSuffixes, DucklingBucketSuffix: ducklingBucketSuffix, DuckLakeDefaultSpecVersion: cfg.DuckLake.SpecVersion, - BillingIngestURL: billingIngestURL, - BillingIngestToken: billingIngestToken, } } diff --git a/controlplane/compute_billing_api.go b/controlplane/compute_billing_api.go new file mode 100644 index 00000000..429b3427 --- /dev/null +++ b/controlplane/compute_billing_api.go @@ -0,0 +1,117 @@ +package controlplane + +import ( + "net/http" + "time" + + "github.com/gin-gonic/gin" + + "github.com/posthog/duckgres/controlplane/configstore" +) + +// billingUsageStore is the config-store surface the billing pull API needs. +type billingUsageStore interface { + AggregateComputeUsage(low, high time.Time) ([]configstore.ComputeUsageRow, error) + ComputeBillingCursor() (time.Time, bool, error) + AckComputeUsage(watermarkHigh time.Time) (int64, error) +} + +// billingAPIHandler serves the pull-based compute-billing API +// (docs/design/billing-pull-api.md): billing GETs the usage accumulated since +// its last ack, processes it, and POSTs the watermark back; duckgres advances +// the cursor and deletes everything at or below it. Reads/deletes on the +// config store only — never on the query hot path. +type billingAPIHandler struct { + store billingUsageStore + now func() time.Time +} + +// registerBillingAPI mounts the billing pull API on an authenticated route +// group. The caller passes the admin-gating middleware (the billing service +// authenticates with the internal secret, which resolves to admin). +func registerBillingAPI(r gin.IRouter, store billingUsageStore, requireAdmin gin.HandlerFunc) { + h := &billingAPIHandler{store: store, now: time.Now} + r.GET("/billing/usage", requireAdmin, h.getUsage) + r.POST("/billing/ack", requireAdmin, h.postAck) +} + +// latestClosedBucket returns the newest bucket_start that is fully closed: +// every contribution for it has landed (grace exceeds the in-process flush +// interval), so it is safe to serve — and later delete on ack. +func (h *billingAPIHandler) latestClosedBucket() time.Time { + return h.now().UTC().Add(-computeBucketWidth - computeBucketGrace).Truncate(computeBucketWidth) +} + +// getUsage returns usage aggregated since the last ack — one row per key +// (org, team, query_source, worker size) per UTC day — plus the watermark +// window. watermark_low is the server cursor (what billing last acked; +// billing should cross-check it against its own record), watermark_high is +// what billing acks after processing. +func (h *billingAPIHandler) getUsage(c *gin.Context) { + low, _, err := h.store.ComputeBillingCursor() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "read billing cursor: " + err.Error()}) + return + } + low = low.UTC() // zero time (never acked) serves everything buffered + + high := h.latestClosedBucket() + if !high.After(low) { + // Nothing closed beyond the cursor yet (fresh deploy or a pull racing + // right behind an ack). An empty window with high == low is a valid + // response: billing acks it as a no-op. + c.JSON(http.StatusOK, gin.H{ + "watermark_low": low.Format(time.RFC3339), + "watermark_high": low.Format(time.RFC3339), + "usage": []configstore.ComputeUsageRow{}, + }) + return + } + + rows, err := h.store.AggregateComputeUsage(low, high) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "aggregate usage: " + err.Error()}) + return + } + if rows == nil { + rows = []configstore.ComputeUsageRow{} + } + c.JSON(http.StatusOK, gin.H{ + "watermark_low": low.Format(time.RFC3339), + "watermark_high": high.Format(time.RFC3339), + "usage": rows, + }) +} + +type billingAckRequest struct { + WatermarkHigh time.Time `json:"watermark_high" binding:"required"` +} + +// postAck advances the cursor to watermark_high and deletes every buffered +// bucket at or below it. Idempotent: re-acking an already-acked (or older) +// watermark is a no-op. The watermark must be a closed bucket boundary the +// server could have served — acking into the still-accumulating present is +// rejected so an ack can never delete buckets that were never returned. +func (h *billingAPIHandler) postAck(c *gin.Context) { + var req billingAckRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid ack body (want {\"watermark_high\": RFC3339}): " + err.Error()}) + return + } + high := req.WatermarkHigh.UTC() + if latest := h.latestClosedBucket(); high.After(latest) { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "watermark_high is beyond the latest closed bucket; ack exactly the watermark_high returned by GET /billing/usage", + }) + return + } + deleted, err := h.store.AckComputeUsage(high) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "ack usage: " + err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "acked": high.Format(time.RFC3339), + "deleted": deleted, + }) +} diff --git a/controlplane/compute_billing_api_test.go b/controlplane/compute_billing_api_test.go new file mode 100644 index 00000000..0df999b0 --- /dev/null +++ b/controlplane/compute_billing_api_test.go @@ -0,0 +1,270 @@ +package controlplane + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + + "github.com/posthog/duckgres/controlplane/configstore" +) + +// fakeBillingStore implements billingUsageStore + computeGCStore in memory. +type fakeBillingStore struct { + mu sync.Mutex + cursor time.Time + hasAck bool + rows []configstore.ComputeUsageRow + aggLow time.Time + aggHigh time.Time + ackedTo time.Time + deleted int64 + gcCutoff time.Time + gcDrop int64 + failWith error +} + +func (f *fakeBillingStore) AggregateComputeUsage(low, high time.Time) ([]configstore.ComputeUsageRow, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.failWith != nil { + return nil, f.failWith + } + f.aggLow, f.aggHigh = low, high + return f.rows, nil +} + +func (f *fakeBillingStore) ComputeBillingCursor() (time.Time, bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.failWith != nil { + return time.Time{}, false, f.failWith + } + return f.cursor, f.hasAck, nil +} + +func (f *fakeBillingStore) AckComputeUsage(high time.Time) (int64, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.failWith != nil { + return 0, f.failWith + } + if high.After(f.cursor) { + f.cursor, f.hasAck = high, true + } + f.ackedTo = high + return f.deleted, nil +} + +func (f *fakeBillingStore) GCComputeUsage(olderThan time.Time) (int64, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.failWith != nil { + return 0, f.failWith + } + f.gcCutoff = olderThan + return f.gcDrop, nil +} + +func newBillingTestRouter(store *fakeBillingStore, now time.Time) *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + h := &billingAPIHandler{store: store, now: func() time.Time { return now }} + grp := r.Group("/api/v1") + grp.GET("/billing/usage", h.getUsage) + grp.POST("/billing/ack", h.postAck) + return r +} + +type usageResponse struct { + WatermarkLow time.Time `json:"watermark_low"` + WatermarkHigh time.Time `json:"watermark_high"` + Usage []configstore.ComputeUsageRow `json:"usage"` +} + +func TestBillingUsageWindowAndRows(t *testing.T) { + now := time.Date(2026, 7, 1, 12, 40, 47, 0, time.UTC) + cursor := time.Date(2026, 7, 1, 12, 30, 0, 0, time.UTC) + store := &fakeBillingStore{ + cursor: cursor, hasAck: true, + rows: []configstore.ComputeUsageRow{{ + Date: "2026-07-01", OrgID: "org_abc", TeamID: "12345", QuerySource: "standard", + CPU: "8", MemGiB: "16", CPUSeconds: 4800, MemorySeconds: 9600, + }}, + } + router := newBillingTestRouter(store, now) + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/billing/usage", nil)) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d: %s", rec.Code, rec.Body.String()) + } + var resp usageResponse + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode: %v", err) + } + if !resp.WatermarkLow.Equal(cursor) { + t.Fatalf("watermark_low = %v, want cursor %v", resp.WatermarkLow, cursor) + } + // Latest closed bucket at 12:40:47 with 60s width + 30s grace: + // 12:40:47 − 90s = 12:39:17 → truncated to 12:39:00. + wantHigh := time.Date(2026, 7, 1, 12, 39, 0, 0, time.UTC) + if !resp.WatermarkHigh.Equal(wantHigh) { + t.Fatalf("watermark_high = %v, want %v", resp.WatermarkHigh, wantHigh) + } + // The aggregate window must be exactly (cursor, high]. + if !store.aggLow.Equal(cursor) || !store.aggHigh.Equal(wantHigh) { + t.Fatalf("aggregate window = (%v, %v], want (%v, %v]", store.aggLow, store.aggHigh, cursor, wantHigh) + } + if len(resp.Usage) != 1 || resp.Usage[0].OrgID != "org_abc" || resp.Usage[0].CPUSeconds != 4800 { + t.Fatalf("usage = %+v", resp.Usage) + } + // The exact-decimal sizes must serialize as unquoted JSON numbers. + if !bytes.Contains(rec.Body.Bytes(), []byte(`"cpu":8`)) || !bytes.Contains(rec.Body.Bytes(), []byte(`"mem_gib":16`)) { + t.Fatalf("cpu/mem_gib not serialized as JSON numbers: %s", rec.Body.String()) + } +} + +func TestBillingUsageNeverAckedServesEverything(t *testing.T) { + // No cursor row yet → the window starts at the zero time (serves all + // buffered history). + now := time.Date(2026, 7, 1, 12, 40, 47, 0, time.UTC) + store := &fakeBillingStore{} + router := newBillingTestRouter(store, now) + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/billing/usage", nil)) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d: %s", rec.Code, rec.Body.String()) + } + if !store.aggLow.IsZero() { + t.Fatalf("window low = %v, want zero time when never acked", store.aggLow) + } +} + +func TestBillingUsageEmptyWindowWhenCursorCurrent(t *testing.T) { + // Cursor already at the latest closed bucket → nothing to serve; the + // response must be a valid empty window with high == low, not an error. + now := time.Date(2026, 7, 1, 12, 40, 47, 0, time.UTC) + cursor := time.Date(2026, 7, 1, 12, 39, 0, 0, time.UTC) + store := &fakeBillingStore{cursor: cursor, hasAck: true} + router := newBillingTestRouter(store, now) + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/billing/usage", nil)) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d: %s", rec.Code, rec.Body.String()) + } + var resp usageResponse + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode: %v", err) + } + if !resp.WatermarkHigh.Equal(resp.WatermarkLow) || !resp.WatermarkLow.Equal(cursor) { + t.Fatalf("want empty window at cursor, got (%v, %v]", resp.WatermarkLow, resp.WatermarkHigh) + } + if len(resp.Usage) != 0 { + t.Fatalf("usage = %+v, want empty", resp.Usage) + } + if !store.aggLow.IsZero() || !store.aggHigh.IsZero() { + t.Fatal("aggregate must not run for an empty window") + } +} + +func TestBillingAckAdvancesAndIsBounded(t *testing.T) { + now := time.Date(2026, 7, 1, 12, 40, 47, 0, time.UTC) + store := &fakeBillingStore{deleted: 7} + router := newBillingTestRouter(store, now) + + // A valid ack (≤ latest closed bucket) advances and reports deletions. + body := []byte(`{"watermark_high":"2026-07-01T12:39:00Z"}`) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/api/v1/billing/ack", bytes.NewReader(body))) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d: %s", rec.Code, rec.Body.String()) + } + if want := time.Date(2026, 7, 1, 12, 39, 0, 0, time.UTC); !store.ackedTo.Equal(want) { + t.Fatalf("acked = %v, want %v", store.ackedTo, want) + } + var resp struct { + Deleted int64 `json:"deleted"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil || resp.Deleted != 7 { + t.Fatalf("deleted = %d (err=%v), want 7", resp.Deleted, err) + } + + // An ack beyond the latest closed bucket must be rejected — it could + // delete buckets that were never served. + store.ackedTo = time.Time{} + body = []byte(`{"watermark_high":"2026-07-01T12:40:00Z"}`) + rec = httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/api/v1/billing/ack", bytes.NewReader(body))) + if rec.Code != http.StatusBadRequest { + t.Fatalf("future ack status = %d, want 400: %s", rec.Code, rec.Body.String()) + } + if !store.ackedTo.IsZero() { + t.Fatal("future ack must not reach the store") + } + + // Garbage body → 400. + rec = httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/api/v1/billing/ack", bytes.NewReader([]byte(`{}`)))) + if rec.Code != http.StatusBadRequest { + t.Fatalf("empty ack status = %d, want 400: %s", rec.Code, rec.Body.String()) + } +} + +func TestBillingStoreErrorsSurfaceAs500(t *testing.T) { + now := time.Date(2026, 7, 1, 12, 40, 47, 0, time.UTC) + store := &fakeBillingStore{failWith: errors.New("db down")} + router := newBillingTestRouter(store, now) + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/billing/usage", nil)) + if rec.Code != http.StatusInternalServerError { + t.Fatalf("usage status = %d, want 500", rec.Code) + } + rec = httptest.NewRecorder() + router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/api/v1/billing/ack", bytes.NewReader([]byte(`{"watermark_high":"2026-07-01T12:39:00Z"}`)))) + if rec.Code != http.StatusInternalServerError { + t.Fatalf("ack status = %d, want 500", rec.Code) + } +} + +func TestComputeUsageGCRunsAndStops(t *testing.T) { + store := &fakeBillingStore{gcDrop: 3} + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { runComputeUsageGC(ctx, store); close(done) }() + // The first tick runs immediately; wait for it, then cancel. + deadline := time.After(2 * time.Second) + for { + store.mu.Lock() + got := store.gcCutoff + store.mu.Unlock() + if !got.IsZero() { + if d := time.Since(got); d < computeUsageRetention-time.Minute || d > computeUsageRetention+time.Minute { + t.Fatalf("gc cutoff %v not ~retention (30d) ago", got) + } + break + } + select { + case <-deadline: + t.Fatal("gc first tick never ran") + default: + time.Sleep(10 * time.Millisecond) + } + } + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("gc loop did not stop on cancel") + } +} diff --git a/controlplane/compute_capture.go b/controlplane/compute_capture.go deleted file mode 100644 index a8a2e783..00000000 --- a/controlplane/compute_capture.go +++ /dev/null @@ -1,120 +0,0 @@ -package controlplane - -import ( - "bytes" - "context" - "crypto/sha256" - "encoding/hex" - "encoding/json" - "fmt" - "net/http" - "strings" - "time" -) - -// computeUsageEventName is the PostHog capture event carrying the two raw -// compute-usage metrics. Billing derives its two usage_keys by summing each -// property. Must match the posthog-side gather query. -const computeUsageEventName = "managed warehouse compute usage" - -// computeCaptureClient ships closed compute-usage buckets to PostHog's public -// ingestion endpoint via the capture API, exactly like any SDK. Authed by a -// project API token stamped on the event `token` property. -type computeCaptureClient struct { - baseURL string - token string - http *http.Client -} - -func newComputeCaptureClient(baseURL, token string) *computeCaptureClient { - return &computeCaptureClient{ - baseURL: strings.TrimRight(baseURL, "/"), - token: token, - http: &http.Client{Timeout: 15 * time.Second}, - } -} - -// computeEventUUID derives a deterministic, stable event UUID from -// (org_id, bucket_start) so a re-ship after an ingestion ack but before the -// delete commits collapses onto the same ClickHouse row (ReplacingMergeTree -// dedups on cityHash64(uuid)). Formatted as a UUID string for ingestion. -func computeEventUUID(orgID string, bucketStart time.Time) string { - h := sha256.Sum256([]byte(orgID + "|" + fmt.Sprintf("%d", bucketStart.UTC().Unix()))) - // Stamp RFC-4122 version (4) and variant bits so it is a well-formed UUID. - var b [16]byte - copy(b[:], h[:16]) - b[6] = (b[6] & 0x0f) | 0x40 - b[8] = (b[8] & 0x3f) | 0x80 - return fmt.Sprintf("%s-%s-%s-%s-%s", - hex.EncodeToString(b[0:4]), - hex.EncodeToString(b[4:6]), - hex.EncodeToString(b[6:8]), - hex.EncodeToString(b[8:10]), - hex.EncodeToString(b[10:16]), - ) -} - -// capturePayload is the PostHog capture API request body. -// -// APIKey is the project write key — it routes the event to its PostHog project -// (team); the ingestion host (base URL) only selects the region. The same token -// is ALSO stamped on properties.token as the billing validity marker the -// usage-report gather query filters on. Both uses carry the same value. -type capturePayload struct { - APIKey string `json:"api_key"` - Event string `json:"event"` - DistinctID string `json:"distinct_id"` - Properties map[string]interface{} `json:"properties"` - Timestamp string `json:"timestamp"` - UUID string `json:"uuid"` -} - -// Ship sends one bucket as a capture event. Any 2xx is success. On a non-2xx or -// transport error it returns an error so the drainer keeps the row and retries. -func (c *computeCaptureClient) Ship(ctx context.Context, b computeUsageBucket) error { - uuid := computeEventUUID(b.OrgID, b.BucketStart) - payload := capturePayload{ - APIKey: c.token, - Event: computeUsageEventName, - DistinctID: b.OrgID, - Properties: map[string]interface{}{ - "cpu_seconds": b.CPUSeconds, - "memory_seconds": b.MemorySeconds, - "bucket_start": b.BucketStart.UTC().Format(time.RFC3339), - "token": c.token, - }, - Timestamp: b.BucketStart.UTC().Format(time.RFC3339), - UUID: uuid, - } - body, err := json.Marshal(payload) - if err != nil { - return fmt.Errorf("marshal capture payload: %w", err) - } - - url := c.baseURL + "/capture/" - req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) - if err != nil { - return fmt.Errorf("build capture request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - - resp, err := c.http.Do(req) - if err != nil { - return fmt.Errorf("post capture event: %w", err) - } - defer func() { _ = resp.Body.Close() }() - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return fmt.Errorf("capture event rejected: status %d", resp.StatusCode) - } - return nil -} - -// computeUsageBucket is the drainer's local view of one buffer row. It mirrors -// configstore.ComputeUsageBucket so the capture client and drainer do not -// depend on the store package's concrete type in tests. -type computeUsageBucket struct { - OrgID string - BucketStart time.Time - CPUSeconds int64 - MemorySeconds int64 -} diff --git a/controlplane/compute_capture_test.go b/controlplane/compute_capture_test.go deleted file mode 100644 index f2d80cb5..00000000 --- a/controlplane/compute_capture_test.go +++ /dev/null @@ -1,81 +0,0 @@ -package controlplane - -import ( - "context" - "encoding/json" - "io" - "net/http" - "net/http/httptest" - "testing" - "time" -) - -// TestComputeCapturePayload asserts the capture request routes to a project via -// the top-level api_key (the team is set by the token, not the URL), carries the -// two raw metrics + the validity token property, and uses the deterministic uuid -// and bucket_start timestamp. -func TestComputeCapturePayload(t *testing.T) { - var gotPath string - var body map[string]any - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotPath = r.URL.Path - b, _ := io.ReadAll(r.Body) - _ = json.Unmarshal(b, &body) - w.WriteHeader(http.StatusOK) - })) - defer srv.Close() - - const token = "phc_test_token" - c := newComputeCaptureClient(srv.URL, token) - bucket := computeUsageBucket{ - OrgID: "org-123", - BucketStart: time.Unix(1_700_000_040, 0).UTC(), // aligned 60s bucket - CPUSeconds: 80, - MemorySeconds: 160, - } - if err := c.Ship(context.Background(), bucket); err != nil { - t.Fatalf("Ship: %v", err) - } - - if gotPath != "/capture/" { - t.Errorf("path = %q, want /capture/", gotPath) - } - // Team routing: api_key MUST be the token (not just a property). - if body["api_key"] != token { - t.Errorf("api_key = %v, want %q (routes event to the project/team)", body["api_key"], token) - } - if body["event"] != "managed warehouse compute usage" { - t.Errorf("event = %v", body["event"]) - } - if body["distinct_id"] != bucket.OrgID { - t.Errorf("distinct_id = %v, want %q", body["distinct_id"], bucket.OrgID) - } - if body["uuid"] != computeEventUUID(bucket.OrgID, bucket.BucketStart) { - t.Errorf("uuid = %v, want deterministic hash", body["uuid"]) - } - props, ok := body["properties"].(map[string]any) - if !ok { - t.Fatalf("properties missing/not an object: %v", body["properties"]) - } - if props["cpu_seconds"] != float64(80) || props["memory_seconds"] != float64(160) { - t.Errorf("metrics = cpu %v / mem %v, want 80 / 160", props["cpu_seconds"], props["memory_seconds"]) - } - // Validity marker for the billing gather query. - if props["token"] != token { - t.Errorf("properties.token = %v, want %q", props["token"], token) - } -} - -// TestComputeCaptureNon2xxIsError ensures a non-2xx keeps the bucket (drainer retries). -func TestComputeCaptureNon2xxIsError(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusInternalServerError) - })) - defer srv.Close() - - c := newComputeCaptureClient(srv.URL, "tok") - err := c.Ship(context.Background(), computeUsageBucket{OrgID: "o", BucketStart: time.Unix(1_700_000_040, 0).UTC()}) - if err == nil { - t.Fatal("expected error on non-2xx so the drainer keeps the row and retries") - } -} diff --git a/controlplane/compute_drain.go b/controlplane/compute_drain.go deleted file mode 100644 index b2f16cc1..00000000 --- a/controlplane/compute_drain.go +++ /dev/null @@ -1,111 +0,0 @@ -package controlplane - -import ( - "context" - "log/slog" - "time" - - "github.com/posthog/duckgres/controlplane/configstore" -) - -const ( - // computeDrainTick is how often the leader scans for closed buckets to ship. - computeDrainTick = 60 * time.Second - // computeBucketGrace waits out CP-pod flush lag + clock skew so every - // contribution to a bucket has landed before it is shipped. Must be >= - // in-proc flush interval + clock-skew margin (15s + 15s → 30s). - computeBucketGrace = 30 * time.Second -) - -// computeDrainStore is the durable-buffer surface the drainer needs. -type computeDrainStore interface { - ListDrainableComputeBuckets(closedBefore time.Time) ([]configstore.ComputeUsageBucket, error) - MarkComputeBucketDrained(orgID string, bucketStart time.Time) error - SweepDrainedComputeBuckets() (int64, error) -} - -// computeShipper ships one bucket to ingestion. Any error means "keep the row, -// retry next tick". -type computeShipper interface { - Ship(ctx context.Context, b computeUsageBucket) error -} - -// computeDrainer is the leader-only loop that ships closed compute-usage -// buckets to PostHog (ship-then-delete, at-least-once). Runs alongside the -// janitor under the same leader lease. -type computeDrainer struct { - store computeDrainStore - shipper computeShipper - now func() time.Time -} - -func newComputeDrainer(store computeDrainStore, shipper computeShipper) *computeDrainer { - return &computeDrainer{store: store, shipper: shipper, now: time.Now} -} - -// Run scans for closed buckets every tick until ctx is cancelled. Leader-only: -// the caller wires it through a leaderOnlyLoop so exactly one CP pod drains. -func (d *computeDrainer) Run(ctx context.Context) { - if d == nil || d.store == nil || d.shipper == nil { - return - } - d.runOnce(ctx) - ticker := time.NewTicker(computeDrainTick) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - d.runOnce(ctx) - } - } -} - -func (d *computeDrainer) runOnce(ctx context.Context) { - closedBefore := d.now().Add(-computeBucketWidth - computeBucketGrace) - buckets, err := d.store.ListDrainableComputeBuckets(closedBefore) - if err != nil { - slog.Warn("Compute-usage drain: failed to list drainable buckets.", "error", err) - return - } - - var shipped, failed int - for _, b := range buckets { - if ctx.Err() != nil { - return - } - evt := computeUsageBucket{ - OrgID: b.OrgID, - BucketStart: b.BucketStart, - CPUSeconds: b.CPUSeconds, - MemorySeconds: b.MemorySeconds, - } - if err := d.shipper.Ship(ctx, evt); err != nil { - // Ship-then-delete: a failure keeps the row for the next tick. - slog.Warn("Compute-usage drain: ship failed; will retry next tick.", "org", b.OrgID, "bucket", b.BucketStart, "error", err) - failed++ - continue - } - // 2xx confirmed: advance high-water + delete the row, atomically. - if err := d.store.MarkComputeBucketDrained(b.OrgID, b.BucketStart); err != nil { - // The event was accepted but we couldn't record it. The deterministic - // uuid makes a re-ship next tick the same ClickHouse row, so this is - // safe (at-least-once collapses to effectively-once). - slog.Warn("Compute-usage drain: shipped but failed to mark drained; may re-ship (idempotent).", "org", b.OrgID, "bucket", b.BucketStart, "error", err) - failed++ - continue - } - shipped++ - } - - if swept, err := d.store.SweepDrainedComputeBuckets(); err != nil { - slog.Warn("Compute-usage drain: sweep of already-drained buckets failed.", "error", err) - } else if swept > 0 { - slog.Info("Compute-usage drain: swept already-drained buffer rows.", "rows", swept) - } - - if shipped > 0 || failed > 0 { - slog.Info("Compute-usage drain tick complete.", "shipped", shipped, "failed", failed) - } -} diff --git a/controlplane/compute_drain_test.go b/controlplane/compute_drain_test.go deleted file mode 100644 index 87d26f8b..00000000 --- a/controlplane/compute_drain_test.go +++ /dev/null @@ -1,178 +0,0 @@ -package controlplane - -import ( - "context" - "errors" - "testing" - "time" - - "github.com/posthog/duckgres/controlplane/configstore" -) - -// fakeDrainStore is an in-memory stand-in for the durable buffer. -type fakeDrainStore struct { - buckets []configstore.ComputeUsageBucket - highWater map[string]time.Time - marked []configstore.ComputeUsageBucket - swept int64 -} - -func newFakeDrainStore(b ...configstore.ComputeUsageBucket) *fakeDrainStore { - return &fakeDrainStore{buckets: b, highWater: map[string]time.Time{}} -} - -func (f *fakeDrainStore) ListDrainableComputeBuckets(closedBefore time.Time) ([]configstore.ComputeUsageBucket, error) { - var out []configstore.ComputeUsageBucket - for _, b := range f.buckets { - if b.BucketStart.After(closedBefore) { - continue // not closed yet - } - if hw, ok := f.highWater[b.OrgID]; ok && !b.BucketStart.After(hw) { - continue // already drained - } - out = append(out, b) - } - return out, nil -} - -func (f *fakeDrainStore) MarkComputeBucketDrained(orgID string, bucketStart time.Time) error { - if hw, ok := f.highWater[orgID]; !ok || bucketStart.After(hw) { - f.highWater[orgID] = bucketStart - } - f.marked = append(f.marked, configstore.ComputeUsageBucket{OrgID: orgID, BucketStart: bucketStart}) - // Remove the row from the buffer (DELETE half of the txn). - var remaining []configstore.ComputeUsageBucket - for _, b := range f.buckets { - if b.OrgID == orgID && b.BucketStart.Equal(bucketStart) { - continue - } - remaining = append(remaining, b) - } - f.buckets = remaining - return nil -} - -func (f *fakeDrainStore) SweepDrainedComputeBuckets() (int64, error) { - return f.swept, nil -} - -// fakeShipper records what it ships and can be made to fail per-org. -type fakeShipper struct { - shipped []computeUsageBucket - failOrg string -} - -func (s *fakeShipper) Ship(_ context.Context, b computeUsageBucket) error { - if s.failOrg != "" && b.OrgID == s.failOrg { - return errors.New("ingestion down") - } - s.shipped = append(s.shipped, b) - return nil -} - -func TestComputeDrainerShipThenDelete(t *testing.T) { - base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) - store := newFakeDrainStore( - configstore.ComputeUsageBucket{OrgID: "orgA", BucketStart: base, CPUSeconds: 80, MemorySeconds: 160}, - ) - shipper := &fakeShipper{} - d := newComputeDrainer(store, shipper) - // Pin "now" well past the bucket close window. - d.now = func() time.Time { return base.Add(10 * time.Minute) } - - d.runOnce(context.Background()) - - if len(shipper.shipped) != 1 { - t.Fatalf("shipped = %d, want 1", len(shipper.shipped)) - } - if shipper.shipped[0].CPUSeconds != 80 || shipper.shipped[0].MemorySeconds != 160 { - t.Fatalf("shipped payload = %+v", shipper.shipped[0]) - } - // Ship succeeded → high-water advanced + row deleted. - if len(store.marked) != 1 { - t.Fatalf("marked drained = %d, want 1", len(store.marked)) - } - if len(store.buckets) != 0 { - t.Fatalf("buffer rows after drain = %d, want 0", len(store.buckets)) - } - if !store.highWater["orgA"].Equal(base) { - t.Fatalf("high-water = %v, want %v", store.highWater["orgA"], base) - } -} - -func TestComputeDrainerKeepsRowOnShipFailure(t *testing.T) { - base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) - store := newFakeDrainStore( - configstore.ComputeUsageBucket{OrgID: "orgA", BucketStart: base, CPUSeconds: 80, MemorySeconds: 160}, - ) - shipper := &fakeShipper{failOrg: "orgA"} - d := newComputeDrainer(store, shipper) - d.now = func() time.Time { return base.Add(10 * time.Minute) } - - d.runOnce(context.Background()) - - // Ship failed → no mark, row retained for retry, no high-water advance. - if len(store.marked) != 0 { - t.Fatalf("marked drained = %d, want 0 on ship failure", len(store.marked)) - } - if len(store.buckets) != 1 { - t.Fatalf("buffer rows = %d, want 1 retained for retry", len(store.buckets)) - } - if _, ok := store.highWater["orgA"]; ok { - t.Fatalf("high-water advanced despite ship failure") - } -} - -func TestComputeDrainerSkipsOpenBuckets(t *testing.T) { - base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) - store := newFakeDrainStore( - configstore.ComputeUsageBucket{OrgID: "orgA", BucketStart: base, CPUSeconds: 1, MemorySeconds: 1}, - ) - shipper := &fakeShipper{} - d := newComputeDrainer(store, shipper) - // "now" is only 30s past bucket start — within width(60s)+grace(30s), still open. - d.now = func() time.Time { return base.Add(30 * time.Second) } - - d.runOnce(context.Background()) - if len(shipper.shipped) != 0 { - t.Fatalf("shipped an open bucket: %d", len(shipper.shipped)) - } -} - -func TestComputeDrainerHighWaterSkipsDrained(t *testing.T) { - base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) - store := newFakeDrainStore( - configstore.ComputeUsageBucket{OrgID: "orgA", BucketStart: base, CPUSeconds: 1, MemorySeconds: 1}, - ) - store.highWater["orgA"] = base // already drained at or past this bucket - shipper := &fakeShipper{} - d := newComputeDrainer(store, shipper) - d.now = func() time.Time { return base.Add(10 * time.Minute) } - - d.runOnce(context.Background()) - if len(shipper.shipped) != 0 { - t.Fatalf("re-shipped an already-drained bucket: %d", len(shipper.shipped)) - } -} - -func TestComputeEventUUIDDeterministicAndDistinct(t *testing.T) { - base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) - u1 := computeEventUUID("orgA", base) - u2 := computeEventUUID("orgA", base) - if u1 != u2 { - t.Fatalf("uuid not stable across calls: %q vs %q", u1, u2) - } - if len(u1) != 36 { - t.Fatalf("uuid not UUID-shaped: %q", u1) - } - // Version nibble must be 4. - if u1[14] != '4' { - t.Fatalf("uuid version nibble = %c, want 4 (%q)", u1[14], u1) - } - if computeEventUUID("orgB", base) == u1 { - t.Fatal("different org produced same uuid") - } - if computeEventUUID("orgA", base.Add(time.Minute)) == u1 { - t.Fatal("different bucket produced same uuid") - } -} diff --git a/controlplane/compute_gc.go b/controlplane/compute_gc.go new file mode 100644 index 00000000..88de8cdc --- /dev/null +++ b/controlplane/compute_gc.go @@ -0,0 +1,61 @@ +package controlplane + +import ( + "context" + "log/slog" + "time" +) + +const ( + // computeBucketGrace waits out CP-pod flush lag + clock skew so every + // contribution to a bucket has landed before the bucket is served by the + // pull API (and can be deleted by an ack). Must be >= in-proc flush + // interval + clock-skew margin (15s + 15s → 30s). + computeBucketGrace = 30 * time.Second + + // computeUsageRetention is the safety-GC cap: buckets older than this are + // hard-deleted even if billing never acked them, bounding table growth if + // billing stops pulling entirely. + computeUsageRetention = 30 * 24 * time.Hour + + // computeGCTick is how often the leader runs the safety GC. + computeGCTick = time.Hour +) + +// computeGCStore is the config-store surface the safety GC needs. +type computeGCStore interface { + GCComputeUsage(olderThan time.Time) (int64, error) +} + +// runComputeUsageGC hard-deletes compute-usage buckets past the retention cap +// on a slow tick until ctx is cancelled. Leader-only: the caller wires it +// through the janitor leader lease so exactly one CP pod sweeps. A nonzero +// drop count means billing is not pulling/acking — that data is gone for +// billing, so it is logged at WARN (alertable). +func runComputeUsageGC(ctx context.Context, store computeGCStore) { + if store == nil { + return + } + tick := func() { + dropped, err := store.GCComputeUsage(time.Now().Add(-computeUsageRetention)) + if err != nil { + slog.Warn("Compute-usage safety GC failed.", "error", err) + return + } + if dropped > 0 { + slog.Warn("Compute-usage safety GC dropped unacked buckets older than retention — billing is not keeping up.", + "rows", dropped, "retention", computeUsageRetention.String()) + } + } + tick() + ticker := time.NewTicker(computeGCTick) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + tick() + } + } +} diff --git a/controlplane/compute_meter.go b/controlplane/compute_meter.go index 4b49eb48..2b3c0715 100644 --- a/controlplane/compute_meter.go +++ b/controlplane/compute_meter.go @@ -19,11 +19,18 @@ const ( computeFlushInterval = 15 * time.Second ) -// computeUsageKey identifies one accumulator bucket: an org and an aligned -// time-bucket (connection-end time floored to computeBucketWidth). +// computeUsageKey identifies one accumulator bucket: the full billing key — +// org, the org's default team, the session's query source and the provisioned +// worker size (milli-units) — plus an aligned time-bucket (connection-end time +// floored to computeBucketWidth). Distinct sizes/teams/sources accumulate (and +// bill) separately. type computeUsageKey struct { - orgID string - bucket time.Time + orgID string + teamID string + querySource string + millicores int64 + mib int64 + bucket time.Time } // computeUsageCounter is the in-process, best-effort per-org compute-usage @@ -74,7 +81,7 @@ func computeConnectionUsage(millicores, mib int64, dur time.Duration) (milliCPUS // Record adds one connection's usage to the in-process counter, bucketed by // connection-end time. Best-effort: a zero/unknown size is a no-op. Never // errors. -func (c *computeUsageCounter) Record(orgID string, millicores, mib int64, endTime time.Time, dur time.Duration) { +func (c *computeUsageCounter) Record(orgID, teamID, querySource string, millicores, mib int64, endTime time.Time, dur time.Duration) { if c == nil || orgID == "" || millicores <= 0 { return } @@ -82,7 +89,14 @@ func (c *computeUsageCounter) Record(orgID string, millicores, mib int64, endTim if milliCPU == 0 && mibSec == 0 { return } - key := computeUsageKey{orgID: orgID, bucket: endTime.UTC().Truncate(computeBucketWidth)} + key := computeUsageKey{ + orgID: orgID, + teamID: teamID, + querySource: querySource, + millicores: millicores, + mib: mib, + bucket: endTime.UTC().Truncate(computeBucketWidth), + } c.mu.Lock() a := c.buckets[key] @@ -112,6 +126,10 @@ func (c *computeUsageCounter) drainWholeUnits() []configstore.ComputeUsageDelta } out = append(out, configstore.ComputeUsageDelta{ OrgID: key.orgID, + TeamID: key.teamID, + QuerySource: key.querySource, + Millicores: key.millicores, + MiB: key.mib, BucketStart: key.bucket, CPUSeconds: cpuWhole, MemorySeconds: memWhole, @@ -132,23 +150,30 @@ type computeUsageStore interface { } // computeMeter wires the per-org counter to a flusher. Nil-safe: a disabled -// meter (no ingest config / non-remote backend) leaves this nil and every call -// site no-ops. +// meter (non-remote backend) leaves this nil and every call site no-ops. +// resolveTeam maps an org to its default PostHog team id at record time (a +// config-snapshot read — no I/O); empty is tolerated per the OrgDefaultTeamID +// contract and the bucket then carries an empty team_id. type computeMeter struct { - counter *computeUsageCounter - store computeUsageStore + counter *computeUsageCounter + store computeUsageStore + resolveTeam func(orgID string) string } -func newComputeMeter(store computeUsageStore) *computeMeter { - return &computeMeter{counter: newComputeUsageCounter(), store: store} +func newComputeMeter(store computeUsageStore, resolveTeam func(orgID string) string) *computeMeter { + return &computeMeter{counter: newComputeUsageCounter(), store: store, resolveTeam: resolveTeam} } // Record forwards a connection-end record to the in-process counter. Best-effort. -func (m *computeMeter) Record(orgID string, millicores, mib int64, endTime time.Time, dur time.Duration) { +func (m *computeMeter) Record(orgID, querySource string, millicores, mib int64, endTime time.Time, dur time.Duration) { if m == nil { return } - m.counter.Record(orgID, millicores, mib, endTime, dur) + teamID := "" + if m.resolveTeam != nil { + teamID = m.resolveTeam(orgID) + } + m.counter.Record(orgID, teamID, querySource, millicores, mib, endTime, dur) } // Flush drains the in-process counter into the durable buffer once. Best-effort: diff --git a/controlplane/compute_meter_test.go b/controlplane/compute_meter_test.go index 5c412803..5601b8b7 100644 --- a/controlplane/compute_meter_test.go +++ b/controlplane/compute_meter_test.go @@ -55,35 +55,56 @@ func TestComputeConnectionUsage(t *testing.T) { } } +// keyA is the billing key most tests share: orgA's default team, standard +// source, an 8-vCPU/16-GiB worker. +func keyA(bucket time.Time) computeUsageKey { + return computeUsageKey{orgID: "orgA", teamID: "42", querySource: "standard", millicores: 8000, mib: 16 * 1024, bucket: bucket} +} + func TestComputeUsageCounterBucketing(t *testing.T) { c := newComputeUsageCounter() base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) - // Two connections in the same minute-bucket for org A accumulate. - c.Record("orgA", 8000, 16*1024, base.Add(5*time.Second), 10*time.Second) // 80 cpu-s, 160 gib-s - c.Record("orgA", 8000, 16*1024, base.Add(40*time.Second), 10*time.Second) // +80, +160 + // Two connections in the same minute-bucket for the same key accumulate. + c.Record("orgA", "42", "standard", 8000, 16*1024, base.Add(5*time.Second), 10*time.Second) // 80 cpu-s, 160 gib-s + c.Record("orgA", "42", "standard", 8000, 16*1024, base.Add(40*time.Second), 10*time.Second) // +80, +160 // A connection in the next bucket is separate. - c.Record("orgA", 8000, 16*1024, base.Add(65*time.Second), 5*time.Second) // 40 cpu-s, 80 gib-s + c.Record("orgA", "42", "standard", 8000, 16*1024, base.Add(65*time.Second), 5*time.Second) // 40 cpu-s, 80 gib-s // Different org, same bucket. - c.Record("orgB", 1000, 1024, base.Add(5*time.Second), 2*time.Second) // 2 cpu-s, 2 gib-s + c.Record("orgB", "7", "standard", 1000, 1024, base.Add(5*time.Second), 2*time.Second) // 2 cpu-s, 2 gib-s + // Same org+bucket but a different query source is a separate key. + c.Record("orgA", "42", "endpoints", 8000, 16*1024, base.Add(10*time.Second), 10*time.Second) // 80, 160 + // Same org+bucket but a different worker size is a separate key. + c.Record("orgA", "42", "standard", 2000, 4*1024, base.Add(10*time.Second), 10*time.Second) // 20, 40 deltas := c.drainWholeUnits() got := map[computeUsageKey]configstore.ComputeUsageDelta{} for _, d := range deltas { - got[computeUsageKey{orgID: d.OrgID, bucket: d.BucketStart}] = d + k := computeUsageKey{orgID: d.OrgID, teamID: d.TeamID, querySource: d.QuerySource, millicores: d.Millicores, mib: d.MiB, bucket: d.BucketStart} + got[k] = d } bucket0 := base.Truncate(computeBucketWidth) bucket1 := base.Add(60 * time.Second).Truncate(computeBucketWidth) - if d := got[computeUsageKey{"orgA", bucket0}]; d.CPUSeconds != 160 || d.MemorySeconds != 320 { + if d := got[keyA(bucket0)]; d.CPUSeconds != 160 || d.MemorySeconds != 320 { t.Errorf("orgA bucket0 = (%d,%d), want (160,320)", d.CPUSeconds, d.MemorySeconds) } - if d := got[computeUsageKey{"orgA", bucket1}]; d.CPUSeconds != 40 || d.MemorySeconds != 80 { + if d := got[keyA(bucket1)]; d.CPUSeconds != 40 || d.MemorySeconds != 80 { t.Errorf("orgA bucket1 = (%d,%d), want (40,80)", d.CPUSeconds, d.MemorySeconds) } - if d := got[computeUsageKey{"orgB", bucket0}]; d.CPUSeconds != 2 || d.MemorySeconds != 2 { + if d := got[computeUsageKey{orgID: "orgB", teamID: "7", querySource: "standard", millicores: 1000, mib: 1024, bucket: bucket0}]; d.CPUSeconds != 2 || d.MemorySeconds != 2 { t.Errorf("orgB bucket0 = (%d,%d), want (2,2)", d.CPUSeconds, d.MemorySeconds) } + endpointsKey := keyA(bucket0) + endpointsKey.querySource = "endpoints" + if d := got[endpointsKey]; d.CPUSeconds != 80 || d.MemorySeconds != 160 { + t.Errorf("orgA endpoints bucket0 = (%d,%d), want (80,160) — query_source must be its own key", d.CPUSeconds, d.MemorySeconds) + } + smallKey := keyA(bucket0) + smallKey.millicores, smallKey.mib = 2000, 4*1024 + if d := got[smallKey]; d.CPUSeconds != 20 || d.MemorySeconds != 40 { + t.Errorf("orgA small-worker bucket0 = (%d,%d), want (20,40) — worker size must be its own key", d.CPUSeconds, d.MemorySeconds) + } } func TestComputeUsageCounterRemainderCarry(t *testing.T) { @@ -91,16 +112,20 @@ func TestComputeUsageCounterRemainderCarry(t *testing.T) { base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) // 500 millicores × 1s = 500 milli-cpu-seconds = 0 whole vCPU-seconds yet, // and 512 MiB × 1s = 512 MiB-seconds = 0 whole GiB-seconds yet. - c.Record("orgA", 500, 512, base, 1*time.Second) + c.Record("orgA", "42", "standard", 500, 512, base, 1*time.Second) if deltas := c.drainWholeUnits(); len(deltas) != 0 { t.Fatalf("expected no whole units yet, got %v", deltas) } // Add another 500 milli-cpu-s (→ 1000 = 1 vCPU-s) and 512 MiB-s (→ 1024 = 1 GiB-s). - c.Record("orgA", 500, 512, base, 1*time.Second) + c.Record("orgA", "42", "standard", 500, 512, base, 1*time.Second) deltas := c.drainWholeUnits() if len(deltas) != 1 || deltas[0].CPUSeconds != 1 || deltas[0].MemorySeconds != 1 { t.Fatalf("expected exactly 1 vCPU-s + 1 GiB-s after carry, got %v", deltas) } + // The delta must carry the worker size for the NUMERIC key columns. + if deltas[0].Millicores != 500 || deltas[0].MiB != 512 { + t.Fatalf("delta size = (%d,%d), want (500,512)", deltas[0].Millicores, deltas[0].MiB) + } // Remainder is now zero — the bucket should be gone. if deltas := c.drainWholeUnits(); len(deltas) != 0 { t.Fatalf("expected bucket drained to empty, got %v", deltas) @@ -121,11 +146,18 @@ func (f *fakeFlushStore) FlushComputeUsage(d []configstore.ComputeUsageDelta) er return nil } +func teamResolverA(orgID string) string { + if orgID == "orgA" { + return "42" + } + return "" +} + func TestComputeMeterFlush(t *testing.T) { store := &fakeFlushStore{} - m := newComputeMeter(store) + m := newComputeMeter(store, teamResolverA) base := time.Date(2026, 6, 30, 12, 0, 0, 0, time.UTC) - m.Record("orgA", 8000, 16*1024, base, 10*time.Second) + m.Record("orgA", "standard", 8000, 16*1024, base, 10*time.Second) if n := m.Flush(); n != 1 { t.Fatalf("Flush rows = %d, want 1", n) @@ -133,16 +165,34 @@ func TestComputeMeterFlush(t *testing.T) { if len(store.flushed) != 1 || store.flushed[0].CPUSeconds != 80 || store.flushed[0].MemorySeconds != 160 { t.Fatalf("flushed = %v, want one (80,160)", store.flushed) } + // The meter resolves the org's default team at record time. + if store.flushed[0].TeamID != "42" || store.flushed[0].QuerySource != "standard" { + t.Fatalf("flushed key = (team=%q, source=%q), want (42, standard)", store.flushed[0].TeamID, store.flushed[0].QuerySource) + } // Nothing left to flush. if n := m.Flush(); n != 0 { t.Fatalf("second Flush rows = %d, want 0", n) } } +func TestComputeMeterUnknownTeamTolerated(t *testing.T) { + // An org with no default_team_id (or a nil resolver) still meters — the + // bucket carries an empty team_id rather than dropping usage. + store := &fakeFlushStore{} + m := newComputeMeter(store, teamResolverA) + m.Record("orgB", "standard", 1000, 1024, time.Now(), 2*time.Second) + if n := m.Flush(); n != 1 { + t.Fatalf("Flush rows = %d, want 1", n) + } + if store.flushed[0].TeamID != "" { + t.Fatalf("team = %q, want empty for unknown org", store.flushed[0].TeamID) + } +} + func TestComputeMeterFlushErrorIsBestEffort(t *testing.T) { store := &fakeFlushStore{err: errors.New("db down")} - m := newComputeMeter(store) - m.Record("orgA", 8000, 16*1024, time.Now(), 10*time.Second) + m := newComputeMeter(store, teamResolverA) + m.Record("orgA", "standard", 8000, 16*1024, time.Now(), 10*time.Second) // A flush error must not panic; the count for this interval is dropped. _ = m.Flush() } @@ -150,14 +200,15 @@ func TestComputeMeterFlushErrorIsBestEffort(t *testing.T) { func TestComputeMeterRecordNeverPanicsOnNilOrDisabled(t *testing.T) { // A nil meter (metering disabled) is a no-op. var m *computeMeter - m.Record("orgA", 8000, 16*1024, time.Now(), time.Second) + m.Record("orgA", "standard", 8000, 16*1024, time.Now(), time.Second) if got := m.Flush(); got != 0 { t.Fatalf("nil meter Flush = %d, want 0", got) } - // A meter with a nil store flushes nothing. + // A meter with a nil store (and nil resolver) flushes nothing and never + // panics. m2 := &computeMeter{counter: newComputeUsageCounter()} - m2.Record("orgA", 8000, 16*1024, time.Now(), time.Second) + m2.Record("orgA", "standard", 8000, 16*1024, time.Now(), time.Second) if got := m2.Flush(); got != 0 { t.Fatalf("nil-store meter Flush = %d, want 0", got) } @@ -165,8 +216,8 @@ func TestComputeMeterRecordNeverPanicsOnNilOrDisabled(t *testing.T) { func TestComputeMeterRunFinalFlushOnCancel(t *testing.T) { store := &fakeFlushStore{} - m := newComputeMeter(store) - m.Record("orgA", 8000, 16*1024, time.Now(), 10*time.Second) + m := newComputeMeter(store, teamResolverA) + m.Record("orgA", "standard", 8000, 16*1024, time.Now(), 10*time.Second) ctx, cancel := context.WithCancel(context.Background()) done := make(chan struct{}) diff --git a/controlplane/configstore/compute_usage.go b/controlplane/configstore/compute_usage.go index a60376a1..ae1946e1 100644 --- a/controlplane/configstore/compute_usage.go +++ b/controlplane/configstore/compute_usage.go @@ -1,51 +1,93 @@ package configstore import ( + "database/sql" + "encoding/json" + "errors" "fmt" + "strings" "time" "gorm.io/gorm" ) -// ComputeUsageDelta is one org's accumulated compute-usage for a single -// time-bucket, ready to be flushed into the durable buffer. Counts are whole -// vCPU-seconds / GiB-seconds (the in-process counter rounds millicore-/MiB- -// seconds down to whole units at flush time). +// ComputeUsageDelta is one accumulator's compute-usage for a single billing +// key + time-bucket, ready to be flushed into the durable buffer. Counts are +// whole vCPU-seconds / GiB-seconds (the in-process counter rounds millicore-/ +// MiB-seconds down to whole units at flush time, carrying the remainder). The +// worker size travels in milli-units (millicores, MiB) and is stored as exact +// NUMERIC decimals (vCPU / GiB), so fractional sizes group exactly with no +// float-equality issues. type ComputeUsageDelta struct { OrgID string + TeamID string + QuerySource string + Millicores int64 + MiB int64 BucketStart time.Time CPUSeconds int64 MemorySeconds int64 } -// ComputeUsageBucket is one durable buffer row read back for draining. -type ComputeUsageBucket struct { - OrgID string - BucketStart time.Time - CPUSeconds int64 - MemorySeconds int64 +// ComputeUsageRow is one aggregated row served by the billing pull API: the +// sum of every closed bucket for one key on one UTC day. CPU and MemGiB are +// the exact NUMERIC decimal texts (e.g. "8", "1.5", "0.5") carried as +// json.Number so they serialize as unquoted JSON numbers with full precision. +type ComputeUsageRow struct { + Date string `json:"date"` + OrgID string `json:"org_id"` + TeamID string `json:"team_id"` + QuerySource string `json:"query_source"` + CPU json.Number `json:"cpu"` + MemGiB json.Number `json:"mem_gib"` + CPUSeconds int64 `json:"cpu_seconds"` + MemorySeconds int64 `json:"memory_seconds"` } -// FlushComputeUsage applies a batch of per-(org, bucket) deltas to the durable -// buffer via UPSERT-increment, so contributions from every CP pod sum into the -// same row. Best-effort by contract — the caller logs and drops on error; a -// flush failure must never reach the request/teardown path. +// milliUnitDecimal renders value/denom as an exact decimal string with no +// trailing zeros, using integer math only. denom must divide a power of 10 +// (1000 and 1024 both do), so the result is always finite: 1500/1000 → "1.5", +// 512/1024 → "0.5", 8000/1000 → "8". Used to feed worker sizes into the +// NUMERIC key columns in a canonical text form. +func milliUnitDecimal(value, denom int64) string { + if value <= 0 || denom <= 0 { + return "0" + } + whole := value / denom + rem := value % denom + if rem == 0 { + return fmt.Sprintf("%d", whole) + } + // 10^10 is divisible by both 1000 and 1024·5^10 — scale the remainder to a + // fixed 10 decimal digits, then trim trailing zeros. + const scale = int64(10_000_000_000) // 10^10 + frac := rem * (scale / denom) + digits := strings.TrimRight(fmt.Sprintf("%010d", frac), "0") + return fmt.Sprintf("%d.%s", whole, digits) +} + +// FlushComputeUsage applies a batch of per-key deltas to the durable buffer +// via UPSERT-increment, so contributions from every CP pod sum into the same +// row. Best-effort by contract — the caller logs and drops on error; a flush +// failure must never reach the request/teardown path. func (cs *ConfigStore) FlushComputeUsage(deltas []ComputeUsageDelta) error { if len(deltas) == 0 { return nil } return cs.db.Transaction(func(tx *gorm.DB) error { const stmt = ` -INSERT INTO duckgres_org_compute_usage (org_id, bucket_start, cpu_seconds, memory_seconds) -VALUES (?, ?, ?, ?) -ON CONFLICT (org_id, bucket_start) +INSERT INTO duckgres_org_compute_usage (org_id, team_id, query_source, cpu, mem_gib, bucket_start, cpu_seconds, memory_seconds) +VALUES (?, ?, ?, ?::numeric, ?::numeric, ?, ?, ?) +ON CONFLICT (org_id, team_id, query_source, cpu, mem_gib, bucket_start) DO UPDATE SET cpu_seconds = duckgres_org_compute_usage.cpu_seconds + EXCLUDED.cpu_seconds, memory_seconds = duckgres_org_compute_usage.memory_seconds + EXCLUDED.memory_seconds` for _, d := range deltas { if d.CPUSeconds == 0 && d.MemorySeconds == 0 { continue } - if err := tx.Exec(stmt, d.OrgID, d.BucketStart.UTC(), d.CPUSeconds, d.MemorySeconds).Error; err != nil { + cpu := milliUnitDecimal(d.Millicores, 1000) + mem := milliUnitDecimal(d.MiB, 1024) + if err := tx.Exec(stmt, d.OrgID, d.TeamID, d.QuerySource, cpu, mem, d.BucketStart.UTC(), d.CPUSeconds, d.MemorySeconds).Error; err != nil { return fmt.Errorf("upsert compute usage (org=%s bucket=%s): %w", d.OrgID, d.BucketStart, err) } } @@ -53,74 +95,90 @@ DO UPDATE SET cpu_seconds = duckgres_org_compute_usage.cpu_seconds + EXCLU }) } -// ListDrainableComputeBuckets returns every closed, not-yet-drained buffer row, -// oldest first. A bucket is drainable when it is closed (bucket_start <= -// closedBefore, i.e. now - width - grace) and strictly newer than the org's -// high-water mark (so a late re-INSERT of an already-shipped bucket is skipped). -// Ordering by bucket_start lets the drainer advance the per-org high-water mark -// monotonically. -func (cs *ConfigStore) ListDrainableComputeBuckets(closedBefore time.Time) ([]ComputeUsageBucket, error) { +// AggregateComputeUsage sums every buffered bucket in the half-open window +// (low, high] into one row per key per UTC day, for the billing pull API. +// Callers pass low = the ack cursor and high = the newest closed bucket, so a +// minute that is still accumulating is never served. The NUMERIC size columns +// are canonical decimal texts on insert, so ::text round-trips them exactly. +func (cs *ConfigStore) AggregateComputeUsage(low, high time.Time) ([]ComputeUsageRow, error) { const q = ` -SELECT u.org_id, u.bucket_start, u.cpu_seconds, u.memory_seconds -FROM duckgres_org_compute_usage u -LEFT JOIN duckgres_org_compute_drain_state s ON s.org_id = u.org_id -WHERE u.bucket_start <= ? - AND (s.last_drained_bucket IS NULL OR u.bucket_start > s.last_drained_bucket) -ORDER BY u.org_id, u.bucket_start` +SELECT to_char((bucket_start AT TIME ZONE 'UTC')::date, 'YYYY-MM-DD') AS date, + org_id, team_id, query_source, cpu::text, mem_gib::text, + SUM(cpu_seconds), SUM(memory_seconds) +FROM duckgres_org_compute_usage +WHERE bucket_start > ? AND bucket_start <= ? +GROUP BY 1, org_id, team_id, query_source, cpu, mem_gib +ORDER BY 1, org_id, team_id, query_source, cpu, mem_gib` - rows, err := cs.db.Raw(q, closedBefore.UTC()).Rows() + rows, err := cs.db.Raw(q, low.UTC(), high.UTC()).Rows() if err != nil { - return nil, fmt.Errorf("list drainable compute buckets: %w", err) + return nil, fmt.Errorf("aggregate compute usage: %w", err) } defer func() { _ = rows.Close() }() - var out []ComputeUsageBucket + var out []ComputeUsageRow for rows.Next() { - var b ComputeUsageBucket - if err := rows.Scan(&b.OrgID, &b.BucketStart, &b.CPUSeconds, &b.MemorySeconds); err != nil { - return nil, fmt.Errorf("scan drainable compute bucket: %w", err) + var r ComputeUsageRow + var cpu, mem string + if err := rows.Scan(&r.Date, &r.OrgID, &r.TeamID, &r.QuerySource, &cpu, &mem, &r.CPUSeconds, &r.MemorySeconds); err != nil { + return nil, fmt.Errorf("scan compute usage row: %w", err) } - out = append(out, b) + r.CPU, r.MemGiB = json.Number(cpu), json.Number(mem) + out = append(out, r) } return out, rows.Err() } -// MarkComputeBucketDrained commits the ship-then-delete step for one bucket -// after ingestion confirms success: it advances the org's high-water mark to -// this bucket and deletes the buffer row, atomically. Called ONLY after a 2xx -// from ingestion. -func (cs *ConfigStore) MarkComputeBucketDrained(orgID string, bucketStart time.Time) error { - bucket := bucketStart.UTC() - return cs.db.Transaction(func(tx *gorm.DB) error { - const upsertHW = ` -INSERT INTO duckgres_org_compute_drain_state (org_id, last_drained_bucket) -VALUES (?, ?) -ON CONFLICT (org_id) -DO UPDATE SET last_drained_bucket = GREATEST(duckgres_org_compute_drain_state.last_drained_bucket, EXCLUDED.last_drained_bucket)` - if err := tx.Exec(upsertHW, orgID, bucket).Error; err != nil { - return fmt.Errorf("advance compute drain high-water (org=%s): %w", orgID, err) +// ComputeBillingCursor returns the ack cursor (the watermark billing last +// acked). ok=false when no ack has ever happened — the pull window then +// starts at the epoch (i.e. "everything buffered"). +func (cs *ConfigStore) ComputeBillingCursor() (cursor time.Time, ok bool, err error) { + row := cs.db.Raw(`SELECT last_acked FROM duckgres_compute_billing_cursor WHERE id = 1`).Row() + var t time.Time + if scanErr := row.Scan(&t); scanErr != nil { + if errors.Is(scanErr, sql.ErrNoRows) { + return time.Time{}, false, nil + } + return time.Time{}, false, fmt.Errorf("read compute billing cursor: %w", scanErr) + } + return t.UTC(), true, nil +} + +// AckComputeUsage commits billing's watermark: it advances the cursor to +// watermarkHigh (monotonically — an older value never moves it backwards) and +// deletes every buffered bucket at or below it, atomically. Idempotent: a +// retried ack at or below the current cursor deletes nothing and leaves the +// cursor unchanged. Returns the number of buffer rows deleted. +func (cs *ConfigStore) AckComputeUsage(watermarkHigh time.Time) (int64, error) { + high := watermarkHigh.UTC() + var deleted int64 + err := cs.db.Transaction(func(tx *gorm.DB) error { + const upsert = ` +INSERT INTO duckgres_compute_billing_cursor (id, last_acked) +VALUES (1, ?) +ON CONFLICT (id) +DO UPDATE SET last_acked = GREATEST(duckgres_compute_billing_cursor.last_acked, EXCLUDED.last_acked)` + if err := tx.Exec(upsert, high).Error; err != nil { + return fmt.Errorf("advance compute billing cursor: %w", err) } - const del = `DELETE FROM duckgres_org_compute_usage WHERE org_id = ? AND bucket_start = ?` - if err := tx.Exec(del, orgID, bucket).Error; err != nil { - return fmt.Errorf("delete drained compute bucket (org=%s): %w", orgID, err) + res := tx.Exec(`DELETE FROM duckgres_org_compute_usage WHERE bucket_start <= ?`, high) + if res.Error != nil { + return fmt.Errorf("delete acked compute buckets: %w", res.Error) } + deleted = res.RowsAffected return nil }) + return deleted, err } -// SweepDrainedComputeBuckets hard-deletes any lingering buffer rows at or below -// each org's high-water mark — the safety net for a late re-INSERT of an -// already-shipped bucket (a CP pod that stalled past the grace window). Returns -// the number of rows removed. -func (cs *ConfigStore) SweepDrainedComputeBuckets() (int64, error) { - const del = ` -DELETE FROM duckgres_org_compute_usage u -USING duckgres_org_compute_drain_state s -WHERE s.org_id = u.org_id - AND u.bucket_start <= s.last_drained_bucket` - res := cs.db.Exec(del) +// GCComputeUsage hard-deletes buffered buckets older than the cutoff +// regardless of ack — the safety net that bounds table growth if billing +// stops pulling entirely. Returns the number of rows dropped; a nonzero +// count means billing is not keeping up (the caller alerts on it). +func (cs *ConfigStore) GCComputeUsage(olderThan time.Time) (int64, error) { + res := cs.db.Exec(`DELETE FROM duckgres_org_compute_usage WHERE bucket_start < ?`, olderThan.UTC()) if res.Error != nil { - return 0, fmt.Errorf("sweep drained compute buckets: %w", res.Error) + return 0, fmt.Errorf("gc compute usage: %w", res.Error) } return res.RowsAffected, nil } diff --git a/controlplane/configstore/migrations/000015_compute_usage_pull_key.sql b/controlplane/configstore/migrations/000015_compute_usage_pull_key.sql new file mode 100644 index 00000000..6d4401fb --- /dev/null +++ b/controlplane/configstore/migrations/000015_compute_usage_pull_key.sql @@ -0,0 +1,63 @@ +-- +goose Up + +-- Pull-based compute billing (docs/design/billing-pull-api.md): billing now +-- PULLS usage buckets from duckgres over HTTP instead of duckgres pushing +-- capture events to ingestion. The buffer key gains the billing dimensions — +-- team_id (the org's default team), query_source (standard|endpoints, from the +-- duckgres.query_source session GUC) and the provisioned worker size (cpu vCPU +-- / mem_gib GiB as exact NUMERIC decimals, so fractional sizes group exactly) — +-- and the per-org push-drain high-water mark is replaced by a single global +-- ack cursor. +ALTER TABLE duckgres_org_compute_usage + ADD COLUMN IF NOT EXISTS team_id TEXT NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS query_source TEXT NOT NULL DEFAULT 'standard', + ADD COLUMN IF NOT EXISTS cpu NUMERIC NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS mem_gib NUMERIC NOT NULL DEFAULT 0; + +-- Backfill team_id for buckets accumulated before this deploy from the org's +-- default team (all orgs are backfilled with default_team_id). Their cpu / +-- mem_gib stay 0 ("size unknown" legacy rows) — the values were already summed +-- without the size dimension. +UPDATE duckgres_org_compute_usage u +SET team_id = COALESCE(o.default_team_id, '') +FROM duckgres_orgs o +WHERE o.name = u.org_id AND u.team_id = ''; + +ALTER TABLE duckgres_org_compute_usage + DROP CONSTRAINT duckgres_org_compute_usage_pkey; +ALTER TABLE duckgres_org_compute_usage + ADD PRIMARY KEY (org_id, team_id, query_source, cpu, mem_gib, bucket_start); + +-- Single global ack cursor (one billing consumer): everything at or below +-- last_acked has been pulled AND acked by billing, and is deleted. Replaces +-- the per-org push-drain high-water mark. +CREATE TABLE IF NOT EXISTS duckgres_compute_billing_cursor ( + id SMALLINT PRIMARY KEY DEFAULT 1 CHECK (id = 1), + last_acked TIMESTAMPTZ NOT NULL +); + +DROP TABLE IF EXISTS duckgres_org_compute_drain_state; + +-- +goose Down + +-- The buffer is transient (unshipped counts only); collapsing the widened key +-- back into (org_id, bucket_start) would have to merge rows, so the down +-- migration clears it instead — acceptable for a buffer, matching the +-- best-effort metering contract. +DELETE FROM duckgres_org_compute_usage; +ALTER TABLE duckgres_org_compute_usage + DROP CONSTRAINT duckgres_org_compute_usage_pkey; +ALTER TABLE duckgres_org_compute_usage + DROP COLUMN team_id, + DROP COLUMN query_source, + DROP COLUMN cpu, + DROP COLUMN mem_gib; +ALTER TABLE duckgres_org_compute_usage + ADD PRIMARY KEY (org_id, bucket_start); + +DROP TABLE IF EXISTS duckgres_compute_billing_cursor; + +CREATE TABLE IF NOT EXISTS duckgres_org_compute_drain_state ( + org_id TEXT PRIMARY KEY, + last_drained_bucket TIMESTAMPTZ NOT NULL +); diff --git a/controlplane/control.go b/controlplane/control.go index 007fce60..439cc992 100644 --- a/controlplane/control.go +++ b/controlplane/control.go @@ -117,23 +117,6 @@ type ControlPlaneConfig struct { // DuckLakeDefaultSpecVersion is the global default DuckLake spec version // used for migration checks when an org doesn't specify an override. DuckLakeDefaultSpecVersion string - - // BillingIngestURL and BillingIngestToken configure managed-warehouse - // compute-usage metering (remote/k8s backend only). The URL is PostHog's - // public ingestion base (e.g. https://us.i.posthog.com); the token is the - // project API write key, stamped on the capture event's `token` property. - // If EITHER is empty, metering is disabled: usage is never shipped and a - // query is NEVER failed on its account. See - // docs/design/billing-compute-seconds-plan.md. - BillingIngestURL string - BillingIngestToken string -} - -// BillingMeteringEnabled reports whether compute-usage metering is configured. -// Metering also requires the remote backend; this only checks the ingest -// config (both URL and token present). -func (c ControlPlaneConfig) BillingMeteringEnabled() bool { - return c.BillingIngestURL != "" && c.BillingIngestToken != "" } type ProcessConfig struct { @@ -1357,15 +1340,15 @@ func (cp *ControlPlane) handleConnection(conn net.Conn) { defer func() { dur := server.CloseConnectionMetrics(cc) clog.Info("Client disconnected.", "duration_ms", dur.Milliseconds()) - // Best-effort compute-usage metering. cp.computeMeter is nil unless the - // remote backend is configured with billing ingest; Record is nil-safe - // and a zero worker size (non-remote/unknown) is a no-op. A panic here - // must never escape teardown. + // Best-effort compute-usage metering. cp.computeMeter is nil outside the + // remote backend; Record is nil-safe and a zero worker size + // (non-remote/unknown) is a no-op. A panic here must never escape + // teardown. if cp.computeMeter != nil { func() { defer func() { _ = recover() }() - billOrg, millicores, mib, billDur := server.ConnectionBilling(cc) - cp.computeMeter.Record(billOrg, millicores, mib, time.Now(), billDur) + billOrg, billSource, millicores, mib, billDur := server.ConnectionBilling(cc) + cp.computeMeter.Record(billOrg, billSource, millicores, mib, time.Now(), billDur) }() } }() diff --git a/controlplane/multitenant.go b/controlplane/multitenant.go index 2692c9d8..81ca2dd7 100644 --- a/controlplane/multitenant.go +++ b/controlplane/multitenant.go @@ -525,6 +525,10 @@ func SetupMultiTenant( ) admin.RegisterAPI(api, store, adpt, liveFetcher) provisioning.RegisterAPI(api, provisioning.NewGormStore(store), cfg.DucklingBucketSuffix) + // Pull-based compute-billing API (GET /billing/usage + POST /billing/ack). + // The billing service authenticates with the internal secret (→ admin); + // RequireAdmin keeps SSO viewers away from raw usage + the ack mutation. + registerBillingAPI(api, store, admin.RequireAdmin()) // Node-overview topology reads reuse the shared K8s pool's in-cluster // clientset (nil when there's no shared pool — leaves those routes off). var clusterClient kubernetes.Interface @@ -580,24 +584,19 @@ func SetupMultiTenant( } }() - // Compute-usage metering (managed-warehouse billing). Enabled only when both - // ingest URL and token are configured; otherwise the meter is nil and every - // call site no-ops (queries are NEVER failed on its account). The flusher - // runs on every CP pod (cross-pod UPSERT-increment); the drain loop is - // leader-only, co-located under the janitor lease so exactly one CP ships. - var meter *computeMeter - if cfg.BillingMeteringEnabled() { - meter = newComputeMeter(store) - go meter.Run(context.Background()) - - drainer := newComputeDrainer(store, newComputeCaptureClient(cfg.BillingIngestURL, cfg.BillingIngestToken)) - if janitorLeader != nil { - janitorLeader.AttachLeaderLoop(drainer.Run) - } - slog.Info("Managed-warehouse compute-usage metering enabled.", "ingest_url", cfg.BillingIngestURL) - } else { - slog.Info("Managed-warehouse compute-usage metering disabled (DUCKGRES_BILLING_INGEST_URL/TOKEN not both set).") - } + // Compute-usage metering (managed-warehouse billing, pull model — see + // docs/design/billing-pull-api.md). Always on for the remote backend: every + // CP pod meters connection-end usage and flushes it into the config-store + // buffer (cross-pod UPSERT-increment); billing pulls it via the + // GET/ack API registered above. Best-effort throughout — queries are NEVER + // failed on its account. The 30-day safety GC is leader-only, co-located + // under the janitor lease so exactly one CP pod sweeps. + meter := newComputeMeter(store, store.OrgDefaultTeamID) + go meter.Run(context.Background()) + if janitorLeader != nil { + janitorLeader.AttachLeaderLoop(func(ctx context.Context) { runComputeUsageGC(ctx, store) }) + } + slog.Info("Managed-warehouse compute-usage metering enabled (pull API).") return store, adpt, apiServer, runtimeTracker, janitorLeader, meter, nil } diff --git a/docs/design/billing-compute-seconds-plan.md b/docs/design/billing-compute-seconds-plan.md index 5abd9f6f..19b808f9 100644 --- a/docs/design/billing-compute-seconds-plan.md +++ b/docs/design/billing-compute-seconds-plan.md @@ -1,6 +1,10 @@ # Compute-Seconds Billing — Design / Implementation Plan -Status: **DRAFT — iterating**. Last updated 2026-06-16. +Status: **PARTIALLY SUPERSEDED**. The metering side (connection wall-clock × +provisioned worker size → 60s config-store buckets) shipped and still applies; +the reporting hop (leader drain → PostHog ingestion capture events, +`DUCKGRES_BILLING_INGEST_URL/_TOKEN`) was replaced by the pull API — see +[`billing-pull-api.md`](./billing-pull-api.md). Last updated 2026-07-06. This is a living design doc. Decisions marked ✅ are locked; ❓ are open forks. diff --git a/docs/design/billing-pull-api.md b/docs/design/billing-pull-api.md index ed454d0e..b5aff9a1 100644 --- a/docs/design/billing-pull-api.md +++ b/docs/design/billing-pull-api.md @@ -1,6 +1,8 @@ # Managed-warehouse compute billing — pull API -Status: **DRAFT**. Supersedes the reporting hop of +Status: **IMPLEMENTED** (duckgres side; endpoints live at +`GET /api/v1/billing/usage` + `POST /api/v1/billing/ack`, internal-secret +bearer auth). Supersedes the reporting hop of [`billing-compute-seconds-plan.md`](./billing-compute-seconds-plan.md): billing now **pulls** usage from duckgres instead of duckgres pushing capture events to PostHog ingestion. The per-connection metering is unchanged — only how the data diff --git a/main.go b/main.go index 3d51de36..1fffd615 100644 --- a/main.go +++ b/main.go @@ -367,8 +367,6 @@ func main() { ManagedHostnameSuffixes: resolved.ManagedHostnameSuffixes, DucklingBucketSuffix: resolved.DucklingBucketSuffix, DuckLakeDefaultSpecVersion: resolved.DuckLakeDefaultSpecVersion, - BillingIngestURL: resolved.BillingIngestURL, - BillingIngestToken: resolved.BillingIngestToken, K8s: controlplane.K8sConfig{ WorkerImage: resolved.K8sWorkerImage, WorkerNamespace: resolved.K8sWorkerNamespace, diff --git a/server/exports.go b/server/exports.go index 287fe18f..3add4750 100644 --- a/server/exports.go +++ b/server/exports.go @@ -206,15 +206,18 @@ func SetConnectionWorkerSize(cc *clientConn, millicores, mib int64) { } // ConnectionBilling returns the data needed to meter one connection's -// compute-usage at teardown: the org, the provisioned worker size in -// milli-units, and the connection's elapsed lifetime. millicores == 0 means -// metering should be skipped (unknown worker size). Call at the same teardown -// point as CloseConnectionMetrics. -func ConnectionBilling(cc *clientConn) (orgID string, millicores, mib int64, dur time.Duration) { +// compute-usage at teardown: the org, the session's query source (the +// `duckgres.query_source` GUC — "standard" unless the client set it), the +// provisioned worker size in milli-units, and the connection's elapsed +// lifetime. millicores == 0 means metering should be skipped (unknown worker +// size). Call at the same teardown point as CloseConnectionMetrics. A +// mid-connection GUC change is not split: the whole connection is metered +// under the final value (documented in docs/design/billing-pull-api.md). +func ConnectionBilling(cc *clientConn) (orgID, querySource string, millicores, mib int64, dur time.Duration) { if cc == nil { - return "", 0, 0, 0 + return "", "", 0, 0, 0 } - return cc.orgID, cc.workerMillicores, cc.workerMiB, time.Since(cc.backendStart) + return cc.orgID, cc.QuerySource(), cc.workerMillicores, cc.workerMiB, time.Since(cc.backendStart) } // CancelClientConn cancels the context of a clientConn. diff --git a/tests/configstore/migrations_postgres_test.go b/tests/configstore/migrations_postgres_test.go index 7a9ab384..3f4db156 100644 --- a/tests/configstore/migrations_postgres_test.go +++ b/tests/configstore/migrations_postgres_test.go @@ -32,15 +32,23 @@ func TestConfigStoreRunsVersionedSQLMigrations(t *testing.T) { requireGooseMigrationRecorded(t, db, 12) requireGooseMigrationRecorded(t, db, 13) requireGooseMigrationRecorded(t, db, 14) - requireGooseLatestVersion(t, db, 14) + requireGooseMigrationRecorded(t, db, 15) + requireGooseLatestVersion(t, db, 15) requireTableAbsent(t, db, "duckgres_schema_migrations") // Migration 000013 added the nullable per-org default_team_id column. requireColumnPresent(t, db, "duckgres_orgs", "default_team_id") - // Migration 000007 added the compute-usage billing buffer + drain state. + // Migration 000007 added the compute-usage billing buffer; 000015 widened + // its key for pull-based billing (team_id, query_source, worker size), + // added the single ack cursor and dropped the push-drain state table. requireTablePresent(t, db, "duckgres_org_compute_usage") - requireTablePresent(t, db, "duckgres_org_compute_drain_state") + requireColumnPresent(t, db, "duckgres_org_compute_usage", "team_id") + requireColumnPresent(t, db, "duckgres_org_compute_usage", "query_source") + requireColumnPresent(t, db, "duckgres_org_compute_usage", "cpu") + requireColumnPresent(t, db, "duckgres_org_compute_usage", "mem_gib") + requireTablePresent(t, db, "duckgres_compute_billing_cursor") + requireTableAbsent(t, db, "duckgres_org_compute_drain_state") // Migration 000008 added the explicit Duckling CR name column on // managed warehouses, backfilled from lower(org_id). @@ -99,7 +107,7 @@ func TestConfigStoreSQLMigrationsUpgradeVersion8Schema(t *testing.T) { ALTER TABLE duckgres_managed_warehouses ADD COLUMN IF NOT EXISTS iceberg_enabled BOOLEAN DEFAULT false; ALTER TABLE duckgres_managed_warehouses ADD COLUMN IF NOT EXISTS iceberg_state VARCHAR(32); ALTER TABLE duckgres_org_users ADD COLUMN IF NOT EXISTS default_catalog VARCHAR(255); - DELETE FROM goose_db_version WHERE version_id IN (9, 10, 11, 12, 13, 14); + DELETE FROM goose_db_version WHERE version_id IN (9, 10, 11, 12, 13, 14, 15); `).Error; err != nil { t.Fatalf("downgrade baseline schema to pre-v9 shape: %v", err) } @@ -127,7 +135,8 @@ func TestConfigStoreSQLMigrationsUpgradeVersion8Schema(t *testing.T) { requireGooseMigrationRecorded(t, upgradedDB, 12) requireGooseMigrationRecorded(t, upgradedDB, 13) requireGooseMigrationRecorded(t, upgradedDB, 14) - requireGooseLatestVersion(t, upgradedDB, 14) + requireGooseMigrationRecorded(t, upgradedDB, 15) + requireGooseLatestVersion(t, upgradedDB, 15) requireColumnDefault(t, upgradedDB, "duckgres_orgs", "max_vcpus", "0") requireColumnDefault(t, upgradedDB, "duckgres_org_users", "max_vcpus", "0") requireColumnDefault(t, upgradedDB, "duckgres_org_users", "disabled", "false") diff --git a/tests/e2e-mw-dev/harness.sh b/tests/e2e-mw-dev/harness.sh index 0f249c4d..41dc2b60 100755 --- a/tests/e2e-mw-dev/harness.sh +++ b/tests/e2e-mw-dev/harness.sh @@ -451,30 +451,65 @@ connection_duration_logged() { # org password [ "$ms" -gt 0 ] || fail "disconnect duration_ms is $ms (want > 0 — backendStart not honoured?)" } -# Managed-warehouse compute-usage metering (billing emit side). At connection -# teardown the CP meters cpu_seconds/memory_seconds from the provisioned worker -# size over the connection lifetime into an in-proc per-org counter (best-effort), -# flushes it to the durable config-store buffer (~15s), and the leader drains -# closed buckets to PostHog ingestion (~60s, ship-then-delete). The :9090 metrics -# port + the ingestion HTTP path are not observable from this in-cluster Job, so -# we assert the CP's startup log line that proves the metering config knob is -# wired and reports its enabled/disabled state. When the deploy sets -# DUCKGRES_BILLING_INGEST_URL/TOKEN this line reads "enabled"; otherwise -# "disabled" — either way the knob is present and the emit path is compiled in. -# (NOTE: a full end-to-end "the event landed in PostHog ClickHouse" assertion -# needs a billing-analytics token + CH read access from the Job, which this lane -# does not have — see README. The buffer UPSERT/drain SQL is exercised by the -# tests/configstore Postgres migration test, and the meter/drain logic by the -# controlplane/ unit tests.) -compute_usage_metering_wired() { - log "compute-usage metering config knob wired in CP" - found="" - for p in $(k get pods -l app=duckgres-control-plane -o jsonpath='{.items[*].metadata.name}'); do - m="$(k logs "$p" 2>/dev/null | grep -E 'Managed-warehouse compute-usage metering (enabled|disabled)' | tail -1)" - if [ -n "$m" ]; then found="$m"; fi +# Managed-warehouse compute-usage billing (pull API, docs/design/billing-pull-api.md). +# At connection teardown the CP meters cpu_seconds/memory_seconds from the +# provisioned worker size over the connection lifetime into an in-proc counter +# keyed (org, default team, query_source, worker size), flushes it to the +# durable config-store buffer (~15s), and serves it aggregated over +# GET /api/v1/billing/usage; POST /api/v1/billing/ack advances the cursor and +# deletes acked buckets. This asserts the FULL round-trip against the real +# stack: two real connections (one standard, one with the duckgres.query_source +# GUC set to endpoints) must surface as separate usage rows carrying the org's +# provisioned default_team_id and a positive worker size, then an ack of the +# served watermark_high must 200 and the next GET's watermark_low must equal it +# (cursor advanced, acked buckets deleted). A bucket closes ~90s after the +# connection ends (60s width + 30s grace) plus ≤15s flush, so the poll allows +# ~4 minutes. This e2e stack has its own config store, so acking here cannot +# eat production usage. +compute_usage_pull_api() { # org password + org="$1"; pw="$2" + log "compute-usage pull API round-trip on $org" + + # Generate usage under both query sources. Each pg call is one connection; + # the batched SET applies to the SELECT's session (split-path, see #868). + pg "$org" "$pw" ducklake 'SELECT 1' >/dev/null + pg "$org" "$pw" ducklake "SET duckgres.query_source = 'endpoints'; SELECT 1" >/dev/null + + # Poll until both rows are served (bucket close + flush lag). + a=0 body="" + while [ "$a" -lt 30 ]; do + body="$(curl -fsS -H "$H" "$API/api/v1/billing/usage")" || body="" + if [ -n "$body" ] && echo "$body" | jq -e --arg o "$org" --arg t "$CNPG_DEFAULT_TEAM_ID" ' + (.usage | map(select(.org_id==$o and .team_id==$t and .query_source=="standard" and .cpu_seconds>0 and .cpu>0 and .mem_gib>0)) | length >= 1) + and + (.usage | map(select(.org_id==$o and .team_id==$t and .query_source=="endpoints" and .cpu_seconds>0)) | length >= 1)' >/dev/null 2>&1; then + break + fi + sleep 10; a=$((a + 1)) done - [ -n "$found" ] || fail "no compute-usage metering startup log line found (config knob not wired into SetupMultiTenant?)" - log "compute-usage metering state: $found" + [ "$a" -lt 30 ] || fail "compute-usage: rows for $org (standard+endpoints, team=$CNPG_DEFAULT_TEAM_ID) never appeared in GET /billing/usage: $(echo "$body" | head -c 600)" + wl="$(echo "$body" | jq -r '.watermark_low')" + wh="$(echo "$body" | jq -r '.watermark_high')" + log "compute-usage OK: usage served (low=$wl high=$wh)" + + # Ack the served watermark; the cursor must advance and acked buckets die. + ack="$(curl -fsS -X POST -H "$H" -H 'Content-Type: application/json' \ + -d "{\"watermark_high\":\"$wh\"}" "$API/api/v1/billing/ack")" \ + || fail "compute-usage: ack POST failed" + deleted="$(echo "$ack" | jq -r '.deleted')" + [ "${deleted:-0}" -ge 1 ] || fail "compute-usage: ack deleted=$deleted (want >=1): $ack" + low2="$(curl -fsS -H "$H" "$API/api/v1/billing/usage" | jq -r '.watermark_low')" + [ "$low2" = "$wh" ] || fail "compute-usage: after ack, watermark_low='$low2' want '$wh' (cursor did not advance)" + # Idempotency: re-acking the same watermark is a safe no-op (200). + curl -fsS -X POST -H "$H" -H 'Content-Type: application/json' \ + -d "{\"watermark_high\":\"$wh\"}" "$API/api/v1/billing/ack" >/dev/null \ + || fail "compute-usage: re-ack of the same watermark failed (must be idempotent)" + # Acking into the still-open present must be rejected (400), never delete. + future="$(jq -rn 'now + 3600 | todate')" + code="$(curl -s -o /tmp/ack_future -w '%{http_code}' -X POST -H "$H" -H 'Content-Type: application/json' \ + -d "{\"watermark_high\":\"$future\"}" "$API/api/v1/billing/ack")" + [ "$code" = "400" ] || fail "compute-usage: future ack -> HTTP $code want 400: $(cat /tmp/ack_future)" + log "compute-usage OK: ack advanced cursor (deleted=$deleted), idempotent re-ack, future ack rejected" } # jsonb || must keep Postgres concatenation semantics through the full CP @@ -2486,8 +2521,8 @@ main() { # many connect/disconnects, so the disconnect log is warm) ---- connection_duration_logged "$CNPG" "$cnpg_pw" - # ---- compute-usage metering wired (billing emit side) ---- - compute_usage_metering_wired + # ---- compute-usage billing pull API (meter → buffer → GET → ack) ---- + compute_usage_pull_api "$CNPG" "$cnpg_pw" # ---- cross-tenant isolation (cnpg vs ext) — needs both lanes done ---- tenant_isolation "$CNPG" "$cnpg_pw" "$EXT" "$ext_pw" @@ -2501,7 +2536,7 @@ main() { # mid-run image bump); it stays covered by the controlplane/ unit tests. log "SKIP version-reaper (needs an in-run image bump; see README)" - log "PASS: admin-no-query-token + models-explorer-api(redaction) + admin-console-api(me/live/metrics/auth-gate) + admin-rbac-viewer(403 mutate/audit) + admin-impersonation(round-trip+audit) + wire + malformed-startup-resilience + jsonb-concat + cold-burst-absorption + pipeline-error-recovery + cancel-reuse + activation(DuckLake) + ext-forks + worker-pod + concurrency + durability + crash-recovery + busy-only-do-not-disrupt + graceful-drain + one-session-per-worker + parallel-cold-burst-ramp + worker-sizing(cnpg DuckLake) + org-default-profile(ext) + persistent-user-secrets(cnpg+ext, cross-user isolation) + user-kill-switch(cnpg) + user-disable-block(cnpg+ext) + connection-duration-logged + isolation + lifecycle-teardown, on cnpg & ext (4 parallel lanes)" + log "PASS: admin-no-query-token + models-explorer-api(redaction) + admin-console-api(me/live/metrics/auth-gate) + admin-rbac-viewer(403 mutate/audit) + admin-impersonation(round-trip+audit) + wire + malformed-startup-resilience + jsonb-concat + cold-burst-absorption + pipeline-error-recovery + cancel-reuse + activation(DuckLake) + ext-forks + worker-pod + concurrency + durability + crash-recovery + busy-only-do-not-disrupt + graceful-drain + one-session-per-worker + parallel-cold-burst-ramp + worker-sizing(cnpg DuckLake) + org-default-profile(ext) + persistent-user-secrets(cnpg+ext, cross-user isolation) + user-kill-switch(cnpg) + user-disable-block(cnpg+ext) + connection-duration-logged + compute-usage-pull-api(cnpg) + isolation + lifecycle-teardown, on cnpg & ext (4 parallel lanes)" } main "$@" From 1f34ed1389e7109cda75bf122ec8843f421c758b Mon Sep 17 00:00:00 2001 From: Benjamin Knofe-Vider Date: Mon, 6 Jul 2026 09:32:46 +0200 Subject: [PATCH 2/2] Cover registerBillingAPI in the untagged build Its only production caller is kubernetes-tagged multitenant.go, so the plain build's linter flagged it unused. The new test exercises the real registration and locks in that both billing routes sit behind the admin-gating middleware. --- controlplane/compute_billing_api_test.go | 29 ++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/controlplane/compute_billing_api_test.go b/controlplane/compute_billing_api_test.go index 0df999b0..43202381 100644 --- a/controlplane/compute_billing_api_test.go +++ b/controlplane/compute_billing_api_test.go @@ -83,6 +83,35 @@ func newBillingTestRouter(store *fakeBillingStore, now time.Time) *gin.Engine { return r } +// TestRegisterBillingAPIGatesBothRoutes exercises the real registration used +// by multitenant.go: both routes are mounted and BOTH pass through the +// admin-gating middleware (a viewer/billing-secret mixup must never reach the +// ack mutation or raw usage). +func TestRegisterBillingAPIGatesBothRoutes(t *testing.T) { + gin.SetMode(gin.TestMode) + r := gin.New() + gateHits := 0 + deny := func(c *gin.Context) { + gateHits++ + c.AbortWithStatus(http.StatusForbidden) + } + registerBillingAPI(r.Group("/api/v1"), &fakeBillingStore{}, deny) + + for _, req := range []*http.Request{ + httptest.NewRequest(http.MethodGet, "/api/v1/billing/usage", nil), + httptest.NewRequest(http.MethodPost, "/api/v1/billing/ack", bytes.NewReader([]byte(`{"watermark_high":"2026-07-01T12:39:00Z"}`))), + } { + rec := httptest.NewRecorder() + r.ServeHTTP(rec, req) + if rec.Code != http.StatusForbidden { + t.Fatalf("%s %s = %d, want 403 from the admin gate", req.Method, req.URL.Path, rec.Code) + } + } + if gateHits != 2 { + t.Fatalf("admin gate ran %d times, want 2 (every billing route must be gated)", gateHits) + } +} + type usageResponse struct { WatermarkLow time.Time `json:"watermark_low"` WatermarkHigh time.Time `json:"watermark_high"`