diff --git a/docs/concepts/tokens-and-billing.md b/docs/concepts/tokens-and-billing.md index f64e09472..bf5e18525 100644 --- a/docs/concepts/tokens-and-billing.md +++ b/docs/concepts/tokens-and-billing.md @@ -60,7 +60,10 @@ health/quota signal. an estimated start. The sweep's ends and backfilled starts carry the `estimated` flag. Other events are stamped when the server records the binding change; a Runner-reconnect reap records the reap time. The log is - auditable and reconstructable. + auditable and reconstructable. Closed intervals feed derived hourly and daily + duration rollups, while open intervals are excluded. The raw compute log + follows the configured event retention window; rollups remain available after + pruning. - **LLM token usage and spend — recorded for display, not billed day-1.** Every model call's tokens-in/out and cost, captured at the gateway. It powers the in-product usage/spend charts the user sees, and it is *recorded* even though diff --git a/go/cmd/compass-server/main.go b/go/cmd/compass-server/main.go index 2f14db8f5..75b0bfa0f 100644 --- a/go/cmd/compass-server/main.go +++ b/go/cmd/compass-server/main.go @@ -368,7 +368,7 @@ func registerServeFlags(fs *flag.FlagSet) serveFlags { transcriptSafetyValveCapBytes: fs.Int("transcript-safety-valve-cap-bytes", 0, "Hot-tail safety-valve cap in bytes. Defaults to $COMPASS_TRANSCRIPT_SAFETY_VALVE_CAP_BYTES."), usageRetention: fs.String("usage-event-retention", "", - "How long raw token-usage events are kept before the daily prune "+ + "How long raw token- and compute-usage events are kept before the daily prune "+ "deletes them, as a Go duration (e.g. 720h). The hourly and daily "+ "usage totals are kept. Falls back to $COMPASS_USAGE_EVENT_RETENTION, "+ "then 2160h (90 days). 0 disables the prune."), diff --git a/go/internal/store/compute_usage_pgtest_test.go b/go/internal/store/compute_usage_pgtest_test.go index 4ef91ab89..72a24cff1 100644 --- a/go/internal/store/compute_usage_pgtest_test.go +++ b/go/internal/store/compute_usage_pgtest_test.go @@ -444,3 +444,118 @@ func execAsSystem(t *testing.T, s *Store, sql string, args ...any) { t.Fatalf("commit: %v", err) } } + +// Two system sweeps insert ends for the same tenants in opposite natural scan +// orders. The held first-tenant lock makes the inversion deterministic. +func TestComputeUsageCrossTenantSweepsLockTenantsInOrder(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + tenantA := seedTenant(t, s, "compute-lock-a") + tenantB := seedTenant(t, s, "compute-lock-b") + first, second := tenantA, tenantB + if string(first) > string(second) { + first, second = second, first + } + + // The runner index retains heap insertion order for its duplicate key, so + // seed bindings in reverse tenant order; orphan starts follow tenant order. + for i, tenant := range []TenantID{second, first} { + suffix := []string{"b", "a"}[i] + tenantCtx := WithTenant(ctx, tenant) + owner, err := s.CreateUser(tenantCtx, NewUser{ + Handle: "compute-lock-owner-" + suffix, DisplayName: "compute-lock-owner-" + suffix, + }) + if err != nil { + t.Fatalf("CreateUser(%s): %v", tenant, err) + } + agent, err := s.CreateAgent(tenantCtx, owner.ID, NewAgent{ + Handle: "compute-lock-agent-" + suffix, DisplayName: "compute-lock-agent-" + suffix, + }) + if err != nil { + t.Fatalf("CreateAgent(%s): %v", tenant, err) + } + if _, err := s.RecordSessionBinding(tenantCtx, "compute-lock-session-"+suffix, agent.ID, "compute-lock-runner"); err != nil { + t.Fatalf("RecordSessionBinding(%s): %v", tenant, err) + } + } + for i, tenant := range []TenantID{first, second} { + suffix := []string{"a", "b"}[i] + if _, err := s.pool.Exec(ctx, ` + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id, estimated + ) VALUES ($1, $2, $3, 'start', clock_timestamp(), $4, $5, $6, $7, TRUE)`, + string(tenant), "compute-lock-orphan-start-"+suffix, "compute-lock-orphan-"+suffix, + "compute-lock-orphan-agent-"+suffix, "compute-lock-orphan-owner-"+suffix, + "compute-lock-orphan-session-"+suffix, "compute-lock-orphan-runner"); err != nil { + t.Fatalf("seed orphan start for %s: %v", tenant, err) + } + } + + gate, err := s.pool.Begin(ctx) + if err != nil { + t.Fatalf("begin advisory-lock gate: %v", err) + } + defer func() { + if err := gate.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) { + t.Errorf("rollback advisory-lock gate: %v", err) + } + }() + if _, err := gate.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtext('compute_usage:' || $1))`, string(first)); err != nil { + t.Fatalf("hold first-tenant compute lock: %v", err) + } + + type result struct { + name string + err error + } + done := make(chan result, 2) + go func() { + _, err := s.CloseOrphanedComputeIntervals(WithSystemRole(ctx)) + done <- result{name: "orphan sweep", err: err} + }() + waitForWaiters := func(wantWaiters int) { + t.Helper() + deadline := time.After(5 * time.Second) + tick := time.NewTicker(10 * time.Millisecond) + defer tick.Stop() + for { + var waiters int + if err := gate.QueryRow(ctx, ` + WITH key AS (SELECT hashtext('compute_usage:' || $1)::bigint AS value) + SELECT count(*) + FROM pg_locks, key + WHERE locktype = 'advisory' AND NOT granted + AND classid = ((key.value >> 32) & 4294967295)::oid + AND objid = (key.value & 4294967295)::oid`, string(first)).Scan(&waiters); err != nil { + t.Fatalf("read compute advisory lock waiters: %v", err) + } + if waiters >= wantWaiters { + return + } + select { + case result := <-done: + t.Fatalf("%s completed before lock gate release: %v", result.name, result.err) + case <-deadline: + t.Fatalf("%d sweeps wait on first-tenant lock; want at least %d", waiters, wantWaiters) + case <-tick.C: + } + } + } + waitForWaiters(1) + go func() { + _, err := s.DeleteSessionBindingsForRunner(WithSystemRole(ctx), "compute-lock-runner") + done <- result{name: "runner sweep", err: err} + }() + waitForWaiters(2) + + if err := gate.Rollback(ctx); err != nil { + t.Fatalf("release first-tenant compute lock: %v", err) + } + for range 2 { + result := <-done + if result.err != nil { + t.Errorf("%s: %v", result.name, result.err) + } + } +} diff --git a/go/internal/store/compute_usage_rollup_migration_pgtest_test.go b/go/internal/store/compute_usage_rollup_migration_pgtest_test.go new file mode 100644 index 000000000..9b0de13a0 --- /dev/null +++ b/go/internal/store/compute_usage_rollup_migration_pgtest_test.go @@ -0,0 +1,161 @@ +//go:build pgtest + +package store + +import ( + "context" + "slices" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/RigelBuild/compass/go/internal/pgtest" +) + +func TestComputeUsageRollupMigrationBackfillsClosedIntervals(t *testing.T) { + ctx := t.Context() + dsn := pgtest.RequireDSN(t) + // Look the migration up by file name, so its version is never hard-coded. + migs, err := loadMigrations() + if err != nil { + t.Fatalf("load migrations: %v", err) + } + i := slices.IndexFunc(migs, func(m migration) bool { return m.name == "0003_compute_usage_rollups.sql" }) + if i < 0 { + t.Fatal("migration 0003_compute_usage_rollups.sql is not embedded") + } + version := migs[i].version + applyMigrationsThrough(t, dsn, version-1) + + seed, err := pgxpool.New(ctx, dsn) + if err != nil { + t.Fatalf("connect to seed pre-rollup data: %v", err) + } + start := time.Date(2026, time.March, 10, 23, 30, 0, 0, time.UTC) + end := start.Add(time.Hour) + _, err = seed.Exec(ctx, + "INSERT INTO tenants (id, slug, display_name, created_at_unix_ms) VALUES ('rollup-tenant-a', 'rollup-tenant-a', 'A', 0), ('rollup-tenant-b', 'rollup-tenant-b', 'B', 0)") + if err == nil { + _, err = seed.Exec(ctx, ` + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id + ) VALUES + ('rollup-tenant-a', 'start', 'interval', 'start', $1, 'agent', 'owner', 'session', 'runner'), + ('rollup-tenant-a', 'end', 'interval', 'end', $2, 'agent', 'owner', 'session', 'runner'), + ('rollup-tenant-b', 'start', 'interval', 'start', $1, 'agent', 'owner', 'session', 'runner'), + ('rollup-tenant-b', 'end', 'interval', 'end', $2, 'agent', 'owner', 'session', 'runner')`, start, end) + } + seed.Close() + if err != nil { + t.Fatalf("seed pre-rollup closed intervals: %v", err) + } + + applyMigrationsThrough(t, dsn, version) + s := reopenStore(t, dsn) + var horizon pgtype.Timestamptz + if err := s.pool.QueryRow(ctx, "SELECT horizon FROM compute_usage_prune_horizon").Scan(&horizon); err != nil { + t.Fatalf("read compute prune horizon: %v", err) + } + if horizon.InfinityModifier != pgtype.NegativeInfinity { + t.Fatalf("compute prune horizon = %+v, want -infinity", horizon) + } + for _, tc := range []struct { + name string + table string + buckets []time.Time + activeMS []int64 + intervals []int64 + }{ + { + name: "hourly", + table: "compute_usage_rollups_hourly", + buckets: []time.Time{time.Date(2026, time.March, 10, 23, 0, 0, 0, time.UTC), time.Date(2026, time.March, 11, 0, 0, 0, 0, time.UTC)}, + activeMS: []int64{30 * 60 * 1000, 30 * 60 * 1000}, + intervals: []int64{1, 0}, + }, + { + name: "daily", + table: "compute_usage_rollups_daily", + buckets: []time.Time{time.Date(2026, time.March, 10, 0, 0, 0, 0, time.UTC), time.Date(2026, time.March, 11, 0, 0, 0, 0, time.UTC)}, + activeMS: []int64{30 * 60 * 1000, 30 * 60 * 1000}, + intervals: []int64{1, 0}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + rows, err := s.pool.Query(ctx, "SELECT tenant_id, bucket_start, active_ms, intervals FROM "+tc.table+" ORDER BY tenant_id, bucket_start") + if err != nil { + t.Fatalf("read %s backfill: %v", tc.name, err) + } + defer rows.Close() + for _, tenant := range []string{"rollup-tenant-a", "rollup-tenant-b"} { + for bucketIndex, bucket := range tc.buckets { + if !rows.Next() { + t.Fatalf("%s backfill has no row for %s at %s", tc.name, tenant, bucket) + } + var gotTenant string + var gotBucket time.Time + var activeMS, intervals int64 + if err := rows.Scan(&gotTenant, &gotBucket, &activeMS, &intervals); err != nil { + t.Fatalf("scan %s backfill: %v", tc.name, err) + } + if gotTenant != tenant || !gotBucket.Equal(bucket) || activeMS != tc.activeMS[bucketIndex] || intervals != tc.intervals[bucketIndex] { + t.Errorf("%s backfill row = %s, %s, %d ms, %d intervals; want %s, %s, %d ms, %d intervals", + tc.name, gotTenant, gotBucket, activeMS, intervals, + tenant, bucket, tc.activeMS[bucketIndex], tc.intervals[bucketIndex]) + } + } + } + if rows.Next() { + t.Errorf("%s backfill has extra rows", tc.name) + } + if err := rows.Err(); err != nil { + t.Fatalf("iterate %s backfill: %v", tc.name, err) + } + }) + } +} + +func applyMigrationsThrough(t *testing.T, dsn string, version int) { + t.Helper() + ctx := t.Context() + migs, err := loadMigrations() + if err != nil { + t.Fatalf("load migrations: %v", err) + } + pool, err := pgxpool.New(ctx, dsn) + if err != nil { + t.Fatalf("connect to apply migrations: %v", err) + } + defer pool.Close() + conn, err := pool.Acquire(ctx) + if err != nil { + t.Fatalf("acquire to apply migrations: %v", err) + } + defer conn.Release() + if _, err := conn.Exec(ctx, "SELECT pg_advisory_lock($1)", migrationLockKey); err != nil { + t.Fatalf("acquire migration lock: %v", err) + } + defer func() { + if _, err := conn.Exec(context.WithoutCancel(ctx), "SELECT pg_advisory_unlock($1)", migrationLockKey); err != nil { + t.Errorf("release migration lock: %v", err) + } + }() + if err := ensureMigrationsTable(ctx, conn); err != nil { + t.Fatal(err) + } + applied, err := appliedChecksums(ctx, conn) + if err != nil { + t.Fatal(err) + } + for _, migration := range migs { + if _, done := applied[migration.version]; migration.version > version || done { + continue + } + if err := applyMigration(ctx, conn, migration); err != nil { + t.Fatal(err) + } + } +} diff --git a/go/internal/store/compute_usage_rollup_pgtest_test.go b/go/internal/store/compute_usage_rollup_pgtest_test.go new file mode 100644 index 000000000..9a162333e --- /dev/null +++ b/go/internal/store/compute_usage_rollup_pgtest_test.go @@ -0,0 +1,163 @@ +//go:build pgtest + +package store + +import ( + "testing" + + "github.com/jackc/pgx/v5" +) + +// computeRollupTotals is one agent's rollup sums at one granularity. +type computeRollupTotals struct { + ActiveMs int64 + Intervals int64 +} + +// readComputeRollups sums one agent's hourly and daily rollup rows as owner. +func readComputeRollups(t *testing.T, s *Store, tenant TenantID, accountID AccountID) (hourly, daily computeRollupTotals) { + t.Helper() + for _, row := range []struct { + table string + dst *computeRollupTotals + }{ + {"compute_usage_rollups_hourly", &hourly}, + {"compute_usage_rollups_daily", &daily}, + } { + if err := s.pool.QueryRow(t.Context(), + "SELECT coalesce(sum(active_ms), 0)::bigint, coalesce(sum(intervals), 0)::bigint FROM "+row.table+ + " WHERE tenant_id = $1 AND agent_account_id = $2", + string(tenant), string(accountID)).Scan(&row.dst.ActiveMs, &row.dst.Intervals); err != nil { + t.Fatalf("read %s: %v", row.table, err) + } + } + return hourly, daily +} + +// wantRollupsMatchEvents checks the rollups hold every closed interval's full +// duration exactly once, at both granularities. +func wantRollupsMatchEvents(t *testing.T, s *Store, tenant TenantID, accountID AccountID, wantClosed int64) { + t.Helper() + starts := map[string]computeUsageEvent{} + var want computeRollupTotals + events := computeEvents(t, s, tenant, accountID) + for _, event := range events { + if event.Kind == "start" { + starts[event.IntervalID] = event + } + } + for _, event := range events { + start, ok := starts[event.IntervalID] + if event.Kind != "end" || !ok { + continue + } + want.ActiveMs += event.OccurredAt.UnixMilli() - start.OccurredAt.UnixMilli() + want.Intervals++ + } + if want.Intervals != wantClosed { + t.Fatalf("closed intervals = %d in %+v, want %d", want.Intervals, events, wantClosed) + } + hourly, daily := readComputeRollups(t, s, tenant, accountID) + if hourly != want || daily != want { + t.Fatalf("rollups hourly %+v daily %+v, want both %+v from events %+v", hourly, daily, want, events) + } +} + +// advanceComputeHorizonPastOpenIntervals moves the prune horizon past every +// open start, as a prune that ran while they were open would. +func advanceComputeHorizonPastOpenIntervals(t *testing.T, s *Store) { + t.Helper() + execAsSystem(t, s, "UPDATE compute_usage_prune_horizon SET horizon = date_trunc('day', now(), 'UTC') + interval '1 day'") +} + +func TestComputeUsageRollupsFromRepoint(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "rollup-repoint-owner") + agent := mustAgent(t, s, owner.ID, "rollup-repoint-agent") + tenant := s.EffectiveTenant(ctx) + mustBind(t, ctx, s, "rollup-before", agent.ID, "runner-before") + advanceComputeHorizonPastOpenIntervals(t, s) + + mustBind(t, ctx, s, "rollup-after", agent.ID, "runner-after") + wantRollupsMatchEvents(t, s, tenant, agent.ID, 1) +} + +func TestComputeUsageRollupsFromDeleteSessionBinding(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "rollup-delete-owner") + agent := mustAgent(t, s, owner.ID, "rollup-delete-agent") + tenant := s.EffectiveTenant(ctx) + mustBind(t, ctx, s, "rollup-delete", agent.ID, "rollup-runner") + advanceComputeHorizonPastOpenIntervals(t, s) + + if err := s.DeleteSessionBinding(ctx, "rollup-delete"); err != nil { + t.Fatalf("DeleteSessionBinding: %v", err) + } + wantRollupsMatchEvents(t, s, tenant, agent.ID, 1) + if err := s.DeleteSessionBinding(ctx, "rollup-delete"); err != nil { + t.Fatalf("second DeleteSessionBinding: %v", err) + } + wantRollupsMatchEvents(t, s, tenant, agent.ID, 1) +} + +func TestComputeUsageRollupsFromRunnerSweep(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "rollup-sweep-owner") + agentA := mustAgent(t, s, owner.ID, "rollup-sweep-a") + agentB := mustAgent(t, s, owner.ID, "rollup-sweep-b") + tenant := s.EffectiveTenant(ctx) + mustBind(t, ctx, s, "rollup-sweep-a", agentA.ID, "rollup-sweep-runner") + mustBind(t, ctx, s, "rollup-sweep-b", agentB.ID, "rollup-sweep-runner") + advanceComputeHorizonPastOpenIntervals(t, s) + + if _, err := s.DeleteSessionBindingsForRunner(ctx, "rollup-sweep-runner"); err != nil { + t.Fatalf("DeleteSessionBindingsForRunner: %v", err) + } + wantRollupsMatchEvents(t, s, tenant, agentA.ID, 1) + wantRollupsMatchEvents(t, s, tenant, agentB.ID, 1) +} + +func TestComputeUsageRollupsFromOrphanClose(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "rollup-orphan-owner") + orphan := mustAgent(t, s, owner.ID, "rollup-orphan-agent") + live := mustAgent(t, s, owner.ID, "rollup-live-agent") + tenant := s.EffectiveTenant(ctx) + mustBind(t, ctx, s, "rollup-orphan", orphan.ID, "rollup-orphan-runner") + mustBind(t, ctx, s, "rollup-live", live.ID, "rollup-live-runner") + execAsSystem(t, s, "DELETE FROM session_bindings WHERE tenant_id = $1 AND agent_account_id = $2", + string(tenant), string(orphan.ID)) + advanceComputeHorizonPastOpenIntervals(t, s) + + if _, err := s.CloseOrphanedComputeIntervals(WithSystemRole(ctx)); err != nil { + t.Fatalf("CloseOrphanedComputeIntervals: %v", err) + } + wantRollupsMatchEvents(t, s, tenant, orphan.ID, 1) + wantRollupsMatchEvents(t, s, tenant, live.ID, 0) +} + +// The request role appends to the raw log but must never delete from it; only +// the system-role prune may. +func TestComputeUsageAppRoleCannotDeleteEvents(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "rollup-grant-owner") + agent := mustAgent(t, s, owner.ID, "rollup-grant-agent") + tenant := s.EffectiveTenant(ctx) + mustBind(t, ctx, s, "rollup-grant", agent.ID, "rollup-grant-runner") + + err := s.WithTx(ctx, func(tx pgx.Tx) error { + _, err := tx.Exec(ctx, "DELETE FROM compute_usage_events") + return err + }) + if !pgErrIs(err, "42501") { + t.Fatalf("DELETE as %s error = %v, want insufficient_privilege", appRole, err) + } + if got := computeEvents(t, s, tenant, agent.ID); len(got) != 1 { + t.Fatalf("events after refused DELETE = %+v, want the start kept", got) + } +} diff --git a/go/internal/store/db/compute_usage.sql.go b/go/internal/store/db/compute_usage.sql.go index f50d460b9..df023cd42 100644 --- a/go/internal/store/db/compute_usage.sql.go +++ b/go/internal/store/db/compute_usage.sql.go @@ -7,8 +7,24 @@ package db import ( "context" + + "github.com/jackc/pgx/v5/pgtype" ) +const advanceComputeUsagePruneHorizon = `-- name: AdvanceComputeUsagePruneHorizon :execrows +UPDATE compute_usage_prune_horizon AS state + SET horizon = GREATEST(state.horizon, $1::timestamptz) +` + +// AdvanceComputeUsagePruneHorizon moves the single global horizon forward. +func (q *Queries) AdvanceComputeUsagePruneHorizon(ctx context.Context, cutoff pgtype.Timestamptz) (int64, error) { + result, err := q.db.Exec(ctx, advanceComputeUsagePruneHorizon, cutoff) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + const closeOrphanedComputeIntervals = `-- name: CloseOrphanedComputeIntervals :execrows WITH orphaned AS ( SELECT starts.tenant_id, starts.interval_id, starts.agent_account_id, @@ -33,15 +49,17 @@ INSERT INTO compute_usage_events ( tenant_id, id, interval_id, kind, occurred_at, agent_account_id, owner_user_id, session_id, runner_id, estimated ) -SELECT orphaned.tenant_id, gen_random_uuid()::text AS id, orphaned.interval_id, 'end' AS kind, now() AS occurred_at, - orphaned.agent_account_id, orphaned.owner_user_id, orphaned.session_id, - orphaned.runner_id, true AS estimated +SELECT orphaned.tenant_id, gen_random_uuid()::text AS id, orphaned.interval_id, + 'end' AS kind, now() AS occurred_at, orphaned.agent_account_id, + orphaned.owner_user_id, orphaned.session_id, orphaned.runner_id, + true AS estimated FROM orphaned + ORDER BY orphaned.tenant_id, orphaned.interval_id ON CONFLICT DO NOTHING ` -// Close intervals only after both binding deletion and its end event are absent. -// The tenant id comes from each start because the system role has no tenant GUC. +// Close intervals only after the binding and matching end event are absent. +// The tenant id comes from each start because this sweep has no tenant GUC. func (q *Queries) CloseOrphanedComputeIntervals(ctx context.Context) (int64, error) { result, err := q.db.Exec(ctx, closeOrphanedComputeIntervals) if err != nil { @@ -49,3 +67,230 @@ func (q *Queries) CloseOrphanedComputeIntervals(ctx context.Context) (int64, err } return result.RowsAffected(), nil } + +const computeUsagePruneHorizon = `-- name: ComputeUsagePruneHorizon :one +SELECT horizon.horizon FROM compute_usage_prune_horizon AS horizon FOR SHARE +` + +// ComputeUsagePruneHorizon pins the horizon while the rebuild replaces newer rows. +func (q *Queries) ComputeUsagePruneHorizon(ctx context.Context) (pgtype.Timestamptz, error) { + row := q.db.QueryRow(ctx, computeUsagePruneHorizon) + var horizon pgtype.Timestamptz + err := row.Scan(&horizon) + return horizon, err +} + +const computeUsageSeries = `-- name: ComputeUsageSeries :many +SELECT r.bucket_start::timestamptz AS bucket_start, + sum(r.active_ms)::bigint AS active_ms, + sum(r.intervals)::bigint AS intervals + FROM (SELECT rollups.bucket_start, rollups.agent_account_id, rollups.active_ms, rollups.intervals + FROM compute_usage_rollups_hourly AS rollups + WHERE $1::integer = 1 + UNION ALL + SELECT rollups.bucket_start, rollups.agent_account_id, rollups.active_ms, rollups.intervals + FROM compute_usage_rollups_daily AS rollups + WHERE $1::integer = 2) AS r + WHERE r.bucket_start >= $2::timestamptz + AND r.bucket_start < $3::timestamptz + AND (coalesce(cardinality($4::text[]), 0) = 0 + OR r.agent_account_id = ANY ($4::text[])) + GROUP BY r.bucket_start + ORDER BY r.bucket_start +` + +type ComputeUsageSeriesParams struct { + Granularity int32 + StartAt pgtype.Timestamptz + EndAt pgtype.Timestamptz + AgentAccountIds []string +} + +type ComputeUsageSeriesRow struct { + BucketStart pgtype.Timestamptz + ActiveMs int64 + Intervals int64 +} + +// ComputeUsageSeries sums matching rows at one pre-aggregated granularity. +func (q *Queries) ComputeUsageSeries(ctx context.Context, arg ComputeUsageSeriesParams) ([]ComputeUsageSeriesRow, error) { + rows, err := q.db.Query(ctx, computeUsageSeries, + arg.Granularity, + arg.StartAt, + arg.EndAt, + arg.AgentAccountIds, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ComputeUsageSeriesRow + for rows.Next() { + var i ComputeUsageSeriesRow + if err := rows.Scan(&i.BucketStart, &i.ActiveMs, &i.Intervals); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const deleteComputeUsageIntervalsBefore = `-- name: DeleteComputeUsageIntervalsBefore :execrows +WITH closed AS ( + SELECT ends.tenant_id, ends.interval_id + FROM compute_usage_events AS ends + WHERE ends.tenant_id = $1::text + AND ends.kind = 'end' + AND ends.occurred_at < $2::timestamptz + AND EXISTS ( + SELECT 1 + FROM compute_usage_events AS starts + WHERE starts.tenant_id = ends.tenant_id + AND starts.interval_id = ends.interval_id + AND starts.kind = 'start' + ) +) +DELETE FROM compute_usage_events AS events + USING closed + WHERE events.tenant_id = $1::text + AND events.tenant_id = closed.tenant_id + AND events.interval_id = closed.interval_id +` + +type DeleteComputeUsageIntervalsBeforeParams struct { + TenantID string + Cutoff pgtype.Timestamptz +} + +// DeleteComputeUsageIntervalsBefore removes both events only when the end is old. +// It runs as the system role, so the tenant predicate is the only tenant scope. +func (q *Queries) DeleteComputeUsageIntervalsBefore(ctx context.Context, arg DeleteComputeUsageIntervalsBeforeParams) (int64, error) { + result, err := q.db.Exec(ctx, deleteComputeUsageIntervalsBefore, arg.TenantID, arg.Cutoff) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const deleteComputeUsageRollupsFrom = `-- name: DeleteComputeUsageRollupsFrom :exec +WITH hourly AS ( + DELETE FROM compute_usage_rollups_hourly AS rollups + WHERE rollups.bucket_start >= $1::timestamptz +) +DELETE FROM compute_usage_rollups_daily AS rollups + WHERE rollups.bucket_start >= $1::timestamptz +` + +// DeleteComputeUsageRollupsFrom clears rows the rebuild can reconstruct. +func (q *Queries) DeleteComputeUsageRollupsFrom(ctx context.Context, horizon pgtype.Timestamptz) error { + _, err := q.db.Exec(ctx, deleteComputeUsageRollupsFrom, horizon) + return err +} + +const lockComputeUsage = `-- name: LockComputeUsage :exec + +SELECT pg_advisory_xact_lock(hashtext('compute_usage:' || current_setting('compass.tenant_id', TRUE))) +` + +// Compute-usage rollups are derived from closed start/end event pairs. +// Each query runs under the tenant role and transaction scope supplied by Store. +// LockComputeUsage serializes a tenant rebuild with end-event rollup triggers. +func (q *Queries) LockComputeUsage(ctx context.Context) error { + _, err := q.db.Exec(ctx, lockComputeUsage) + return err +} + +const lockComputeUsageTenant = `-- name: LockComputeUsageTenant :exec +SELECT pg_advisory_xact_lock(hashtext('compute_usage:' || $1::text)) +` + +// LockComputeUsageTenant takes the same lock for a system-role caller. +// The key must stay byte-identical to LockComputeUsage and the end trigger. +func (q *Queries) LockComputeUsageTenant(ctx context.Context, tenantID string) error { + _, err := q.db.Exec(ctx, lockComputeUsageTenant, tenantID) + return err +} + +const rollUpComputeUsageFrom = `-- name: RollUpComputeUsageFrom :exec +WITH intervals AS ( + SELECT starts.tenant_id, starts.interval_id, starts.occurred_at AS start_at, + ends.occurred_at AS end_at, starts.owner_user_id, starts.agent_account_id, + floor(extract(epoch FROM starts.occurred_at) * 1000)::bigint AS start_ms, + floor(extract(epoch FROM ends.occurred_at) * 1000)::bigint AS end_ms + FROM compute_usage_events AS starts + JOIN compute_usage_events AS ends + ON ends.tenant_id = starts.tenant_id + AND ends.interval_id = starts.interval_id + AND ends.kind = 'end' + WHERE starts.kind = 'start' + AND ends.occurred_at >= $1::timestamptz +), bounds AS ( + SELECT intervals.tenant_id, intervals.interval_id, intervals.start_at, intervals.end_at, intervals.owner_user_id, intervals.agent_account_id, intervals.start_ms, intervals.end_ms, + CASE WHEN $1::timestamptz = '-infinity'::timestamptz + THEN date_trunc('hour', intervals.start_at, 'UTC') + ELSE greatest(date_trunc('hour', intervals.start_at, 'UTC'), + date_trunc('hour', $1::timestamptz, 'UTC')) + END AS first_hour, + greatest(date_trunc('hour', intervals.start_at, 'UTC'), + date_trunc('hour', intervals.end_at - interval '1 millisecond', 'UTC')) AS last_hour, + CASE WHEN $1::timestamptz = '-infinity'::timestamptz + THEN date_trunc('day', intervals.start_at, 'UTC') + ELSE greatest(date_trunc('day', intervals.start_at, 'UTC'), + date_trunc('day', $1::timestamptz, 'UTC')) + END AS first_day, + greatest(date_trunc('day', intervals.start_at, 'UTC'), + date_trunc('day', intervals.end_at - interval '1 millisecond', 'UTC')) AS last_day + FROM intervals +), hourly AS ( + INSERT INTO compute_usage_rollups_hourly ( + bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) + SELECT buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id, + sum(greatest(0, least(bounds.end_ms, epoch.bucket_ms + 3600000) - + greatest(bounds.start_ms, epoch.bucket_ms)))::bigint, + sum(CASE WHEN buckets.bucket_start = date_trunc('hour', bounds.start_at, 'UTC') + AND bounds.start_at >= $1::timestamptz + THEN 1 ELSE 0 END)::bigint + FROM bounds + CROSS JOIN LATERAL generate_series( + bounds.first_hour, + bounds.last_hour, + interval '1 hour' + ) AS buckets(bucket_start) + CROSS JOIN LATERAL ( + SELECT floor(extract(epoch FROM buckets.bucket_start) * 1000)::bigint AS bucket_ms + ) AS epoch + GROUP BY buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id +), daily AS ( + INSERT INTO compute_usage_rollups_daily ( + bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) + SELECT buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id, + sum(greatest(0, least(bounds.end_ms, epoch.bucket_ms + 86400000) - + greatest(bounds.start_ms, epoch.bucket_ms)))::bigint, + sum(CASE WHEN buckets.bucket_start = date_trunc('day', bounds.start_at, 'UTC') + AND bounds.start_at >= $1::timestamptz + THEN 1 ELSE 0 END)::bigint + FROM bounds + CROSS JOIN LATERAL generate_series( + bounds.first_day, + bounds.last_day, + interval '24 hours' + ) AS buckets(bucket_start) + CROSS JOIN LATERAL ( + SELECT floor(extract(epoch FROM buckets.bucket_start) * 1000)::bigint AS bucket_ms + ) AS epoch + GROUP BY buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id +) +SELECT 1 +` + +// RollUpComputeUsageFrom rebuilds only buckets at or after the horizon, which are +// the rows it deleted; older buckets are frozen. The sentinel skips clipping. +func (q *Queries) RollUpComputeUsageFrom(ctx context.Context, horizon pgtype.Timestamptz) error { + _, err := q.db.Exec(ctx, rollUpComputeUsageFrom, horizon) + return err +} diff --git a/go/internal/store/db/models.go b/go/internal/store/db/models.go index 45cf884d5..740bf7f53 100644 --- a/go/internal/store/db/models.go +++ b/go/internal/store/db/models.go @@ -182,6 +182,35 @@ type ComputeUsageEvent struct { Estimated bool } +type ComputeUsagePruneHorizon struct { + Singleton bool + Horizon pgtype.Timestamptz + CreatedAt pgtype.Timestamptz + UpdatedAt pgtype.Timestamptz +} + +type ComputeUsageRollupsDaily struct { + TenantID string + BucketStart pgtype.Timestamptz + OwnerUserID string + AgentAccountID string + ActiveMs int64 + Intervals int64 + CreatedAt pgtype.Timestamptz + UpdatedAt pgtype.Timestamptz +} + +type ComputeUsageRollupsHourly struct { + TenantID string + BucketStart pgtype.Timestamptz + OwnerUserID string + AgentAccountID string + ActiveMs int64 + Intervals int64 + CreatedAt pgtype.Timestamptz + UpdatedAt pgtype.Timestamptz +} + type ForgeArtifactCursor struct { ForgeProvider int16 ForgeHost string diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 2c921b45f..65721af86 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -17,6 +17,8 @@ type Querier interface { AccountVisibleTo(ctx context.Context, arg AccountVisibleToParams) (bool, error) AcquireOwnerTreeLock(ctx context.Context, hashtext string) error ActivityFor(ctx context.Context, dollar_1 []string) ([]ActivityForRow, error) + // AdvanceComputeUsagePruneHorizon moves the single global horizon forward. + AdvanceComputeUsagePruneHorizon(ctx context.Context, cutoff pgtype.Timestamptz) (int64, error) AdvanceDeliveryCursor(ctx context.Context, arg AdvanceDeliveryCursorParams) error AdvanceForgeDeliveredRevision(ctx context.Context, arg AdvanceForgeDeliveredRevisionParams) (int64, error) // Compare-and-set advance for the notify-router suppress path: the write lands @@ -83,10 +85,14 @@ type Querier interface { // The unique key makes concurrent first-run attempts a single-winner claim. ClaimTourStart(ctx context.Context, arg ClaimTourStartParams) (string, error) ClearOwedMention(ctx context.Context, arg ClearOwedMentionParams) (int64, error) - // Close intervals only after both binding deletion and its end event are absent. - // The tenant id comes from each start because the system role has no tenant GUC. + // Close intervals only after the binding and matching end event are absent. + // The tenant id comes from each start because this sweep has no tenant GUC. CloseOrphanedComputeIntervals(ctx context.Context) (int64, error) CollectSegment(ctx context.Context, arg CollectSegmentParams) ([]CollectSegmentRow, error) + // ComputeUsagePruneHorizon pins the horizon while the rebuild replaces newer rows. + ComputeUsagePruneHorizon(ctx context.Context) (pgtype.Timestamptz, error) + // ComputeUsageSeries sums matching rows at one pre-aggregated granularity. + ComputeUsageSeries(ctx context.Context, arg ComputeUsageSeriesParams) ([]ComputeUsageSeriesRow, error) // Single-statement clear-and-return: the row is claimed and its actor returned // in ONE UPDATE, so a memo attributes at most one event and a concurrent second // reader matches nothing (consumed_at is no longer NULL). A state mismatch or a @@ -120,6 +126,11 @@ type Querier interface { DeleteChannelMember(ctx context.Context, arg DeleteChannelMemberParams) (int64, error) DeleteChannelPin(ctx context.Context, arg DeleteChannelPinParams) error DeleteChannelPinReturningPosition(ctx context.Context, arg DeleteChannelPinReturningPositionParams) (int32, error) + // DeleteComputeUsageIntervalsBefore removes both events only when the end is old. + // It runs as the system role, so the tenant predicate is the only tenant scope. + DeleteComputeUsageIntervalsBefore(ctx context.Context, arg DeleteComputeUsageIntervalsBeforeParams) (int64, error) + // DeleteComputeUsageRollupsFrom clears rows the rebuild can reconstruct. + DeleteComputeUsageRollupsFrom(ctx context.Context, horizon pgtype.Timestamptz) error DeleteModelRegistry(ctx context.Context) error // DeleteSecret addresses one scope coordinate — a name alone no longer // identifies a row (composite PK). @@ -383,6 +394,13 @@ type Querier interface { LockChannelForPins(ctx context.Context, id string) (LockChannelForPinsRow, error) LockChannelMandatoryKind(ctx context.Context, id string) (LockChannelMandatoryKindRow, error) LockChannelPolicy(ctx context.Context, id string) (LockChannelPolicyRow, error) + // Compute-usage rollups are derived from closed start/end event pairs. + // Each query runs under the tenant role and transaction scope supplied by Store. + // LockComputeUsage serializes a tenant rebuild with end-event rollup triggers. + LockComputeUsage(ctx context.Context) error + // LockComputeUsageTenant takes the same lock for a system-role caller. + // The key must stay byte-identical to LockComputeUsage and the end trigger. + LockComputeUsageTenant(ctx context.Context, tenantID string) error LockLinearRouting(ctx context.Context, dollar_1 pgtype.Text) error LockOwnerCoordination(ctx context.Context, dollar_1 pgtype.Text) error LockOwnerDM(ctx context.Context, dollar_1 pgtype.Text) error @@ -518,6 +536,9 @@ type Querier interface { ReviveTopic(ctx context.Context, id string) error RevokeForgeScope(ctx context.Context, arg RevokeForgeScopeParams) error RevokeToken(ctx context.Context, hash []byte) (int64, error) + // RollUpComputeUsageFrom rebuilds only buckets at or after the horizon, which are + // the rows it deleted; older buckets are frozen. The sentinel skips clipping. + RollUpComputeUsageFrom(ctx context.Context, horizon pgtype.Timestamptz) error // RollUpTokenUsageFrom rebuilds both rollups from the events at or after the // horizon. RollUpTokenUsageFrom(ctx context.Context, horizon pgtype.Timestamptz) error diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go index 5608f5704..b320e0a24 100644 --- a/go/internal/store/db/session_bindings.sql.go +++ b/go/internal/store/db/session_bindings.sql.go @@ -37,6 +37,7 @@ SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_t d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id FROM d JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ORDER BY d.tenant_id, d.usage_interval_id ON CONFLICT DO NOTHING ` @@ -73,6 +74,7 @@ WITH d AS ( d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id FROM d JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ORDER BY d.tenant_id, d.usage_interval_id ON CONFLICT DO NOTHING RETURNING 1 ) diff --git a/go/internal/store/migrations/0003_compute_usage_rollups.sql b/go/internal/store/migrations/0003_compute_usage_rollups.sql new file mode 100644 index 000000000..69f8e5b4a --- /dev/null +++ b/go/internal/store/migrations/0003_compute_usage_rollups.sql @@ -0,0 +1,229 @@ +-- 0009_compute_usage_rollups adds derived compute duration buckets and a raw-log horizon. + +CREATE TABLE compute_usage_rollups_hourly ( + tenant_id TEXT NOT NULL DEFAULT current_setting('compass.tenant_id', TRUE) REFERENCES tenants (id) ON DELETE RESTRICT, + bucket_start TIMESTAMPTZ NOT NULL, + owner_user_id TEXT NOT NULL, + agent_account_id TEXT NOT NULL, + active_ms BIGINT NOT NULL, + intervals BIGINT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (tenant_id, bucket_start, owner_user_id, agent_account_id) +); + +CREATE TABLE compute_usage_rollups_daily ( + tenant_id TEXT NOT NULL DEFAULT current_setting('compass.tenant_id', TRUE) REFERENCES tenants (id) ON DELETE RESTRICT, + bucket_start TIMESTAMPTZ NOT NULL, + owner_user_id TEXT NOT NULL, + agent_account_id TEXT NOT NULL, + active_ms BIGINT NOT NULL, + intervals BIGINT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (tenant_id, bucket_start, owner_user_id, agent_account_id) +); + +CREATE TABLE compute_usage_prune_horizon ( + singleton BOOLEAN PRIMARY KEY DEFAULT TRUE CHECK (singleton), + horizon TIMESTAMPTZ NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +INSERT INTO compute_usage_prune_horizon (horizon) VALUES ('-infinity'); + +GRANT DELETE ON compute_usage_events TO compass_system; +GRANT SELECT, INSERT, UPDATE, DELETE + ON compute_usage_rollups_hourly, compute_usage_rollups_daily, + compute_usage_prune_horizon + TO compass_app, compass_system; +REVOKE INSERT, DELETE ON compute_usage_prune_horizon FROM compass_app, compass_system; + +DO $$ +DECLARE + t text; + tenant_tables text[] := ARRAY[ + 'compute_usage_rollups_hourly', 'compute_usage_rollups_daily' + ]; +BEGIN + FOREACH t IN ARRAY tenant_tables LOOP + EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', t); + EXECUTE format('ALTER TABLE %I FORCE ROW LEVEL SECURITY', t); + EXECUTE format($f$ + CREATE POLICY tenant_isolation ON %I + USING ((SELECT current_setting('compass.tenant_id', TRUE)) <> '' + AND tenant_id = (SELECT current_setting('compass.tenant_id', TRUE))) + WITH CHECK ((SELECT current_setting('compass.tenant_id', TRUE)) <> '' + AND tenant_id = (SELECT current_setting('compass.tenant_id', TRUE))) + $f$, t); + END LOOP; +END $$; + +DO $$ +DECLARE + t text; + updated_at_tables text[] := ARRAY[ + 'compute_usage_rollups_hourly', + 'compute_usage_rollups_daily', + 'compute_usage_prune_horizon' + ]; +BEGIN + FOREACH t IN ARRAY updated_at_tables LOOP + EXECUTE format( + 'CREATE TRIGGER set_updated_at BEFORE UPDATE ON %I + FOR EACH ROW EXECUTE FUNCTION set_updated_at()', t); + END LOOP; +END $$; + +-- Keep this trigger unpinned so isolated schemas resolve their own tables. +CREATE FUNCTION roll_up_compute_usage_end() RETURNS TRIGGER +LANGUAGE plpgsql +AS $$ +DECLARE + started compute_usage_events%ROWTYPE; + start_ms BIGINT; + end_ms BIGINT; + width_ms BIGINT; + rollup_bucket TIMESTAMPTZ; + first_bucket TIMESTAMPTZ; + last_bucket TIMESTAMPTZ; + bucket_ms BIGINT; + active_ms BIGINT; +BEGIN + IF NEW.kind <> 'end' THEN + RETURN NEW; + END IF; + + PERFORM pg_advisory_xact_lock(hashtext('compute_usage:' || NEW.tenant_id)); + -- No horizon clipping: an interval open at prune time is in no frozen bucket yet. + SELECT events.* INTO started + FROM compute_usage_events AS events + WHERE events.tenant_id = NEW.tenant_id + AND events.interval_id = NEW.interval_id + AND events.kind = 'start'; + IF NOT FOUND THEN + RETURN NEW; + END IF; + + start_ms := floor(extract(epoch FROM started.occurred_at) * 1000)::bigint; + end_ms := floor(extract(epoch FROM NEW.occurred_at) * 1000)::bigint; + + FOREACH width_ms IN ARRAY ARRAY[3600000::bigint, 86400000::bigint] LOOP + IF width_ms = 3600000 THEN + first_bucket := date_trunc('hour', started.occurred_at, 'UTC'); + last_bucket := greatest(first_bucket, date_trunc('hour', NEW.occurred_at - interval '1 millisecond', 'UTC')); + ELSE + first_bucket := date_trunc('day', started.occurred_at, 'UTC'); + last_bucket := greatest(first_bucket, date_trunc('day', NEW.occurred_at - interval '1 millisecond', 'UTC')); + END IF; + + rollup_bucket := first_bucket; + WHILE rollup_bucket <= last_bucket LOOP + bucket_ms := floor(extract(epoch FROM rollup_bucket) * 1000)::bigint; + active_ms := greatest(0, least(end_ms, bucket_ms + width_ms) - greatest(start_ms, bucket_ms)); + IF width_ms = 3600000 THEN + INSERT INTO compute_usage_rollups_hourly ( + tenant_id, bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) VALUES ( + NEW.tenant_id, rollup_bucket, started.owner_user_id, started.agent_account_id, + active_ms, CASE WHEN rollup_bucket = first_bucket THEN 1 ELSE 0 END + ) + ON CONFLICT (tenant_id, bucket_start, owner_user_id, agent_account_id) + DO UPDATE SET + active_ms = compute_usage_rollups_hourly.active_ms + EXCLUDED.active_ms, + intervals = compute_usage_rollups_hourly.intervals + EXCLUDED.intervals; + ELSE + INSERT INTO compute_usage_rollups_daily ( + tenant_id, bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) VALUES ( + NEW.tenant_id, rollup_bucket, started.owner_user_id, started.agent_account_id, + active_ms, CASE WHEN rollup_bucket = first_bucket THEN 1 ELSE 0 END + ) + ON CONFLICT (tenant_id, bucket_start, owner_user_id, agent_account_id) + DO UPDATE SET + active_ms = compute_usage_rollups_daily.active_ms + EXCLUDED.active_ms, + intervals = compute_usage_rollups_daily.intervals + EXCLUDED.intervals; + END IF; + rollup_bucket := rollup_bucket + CASE WHEN width_ms = 3600000 THEN interval '1 hour' ELSE interval '24 hours' END; + END LOOP; + END LOOP; + RETURN NEW; +END $$; + +CREATE TRIGGER compute_usage_rollup_after_end +AFTER INSERT ON compute_usage_events +FOR EACH ROW EXECUTE FUNCTION roll_up_compute_usage_end(); + +-- Seed rollups for closed intervals written before the trigger existed. +SET LOCAL ROLE compass_system; + +WITH intervals AS ( + SELECT starts.tenant_id, starts.interval_id, starts.occurred_at AS start_at, + ends.occurred_at AS end_at, starts.owner_user_id, starts.agent_account_id, + floor(extract(EPOCH FROM starts.occurred_at) * 1000)::BIGINT AS start_ms, + floor(extract(EPOCH FROM ends.occurred_at) * 1000)::BIGINT AS end_ms + FROM compute_usage_events AS starts + JOIN compute_usage_events AS ends + ON ends.tenant_id = starts.tenant_id + AND ends.interval_id = starts.interval_id + AND ends.kind = 'end' + WHERE starts.kind = 'start' +), + +bounds AS ( + SELECT intervals.*, + date_trunc('hour', intervals.start_at, 'UTC') AS first_hour, + greatest(date_trunc('hour', intervals.start_at, 'UTC'), + date_trunc('hour', intervals.end_at - INTERVAL '1 millisecond', 'UTC')) AS last_hour, + date_trunc('day', intervals.start_at, 'UTC') AS first_day, + greatest(date_trunc('day', intervals.start_at, 'UTC'), + date_trunc('day', intervals.end_at - INTERVAL '1 millisecond', 'UTC')) AS last_day + FROM intervals +), + +hourly AS ( + INSERT INTO compute_usage_rollups_hourly ( + tenant_id, bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) + SELECT bounds.tenant_id, buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id, + sum(greatest(0, least(bounds.end_ms, epoch.bucket_ms + 3600000) - + greatest(bounds.start_ms, epoch.bucket_ms)))::BIGINT AS active_ms, + sum(CASE WHEN buckets.bucket_start = bounds.first_hour THEN 1 ELSE 0 END)::BIGINT AS intervals + FROM bounds + CROSS JOIN LATERAL generate_series( + bounds.first_hour, + bounds.last_hour, + INTERVAL '1 hour' + ) AS buckets(bucket_start) + CROSS JOIN LATERAL ( + SELECT floor(extract(EPOCH FROM buckets.bucket_start) * 1000)::BIGINT AS bucket_ms + ) AS epoch + GROUP BY bounds.tenant_id, buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id + RETURNING 1 +), + +daily AS ( + INSERT INTO compute_usage_rollups_daily ( + tenant_id, bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) + SELECT bounds.tenant_id, buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id, + sum(greatest(0, least(bounds.end_ms, epoch.bucket_ms + 86400000) - + greatest(bounds.start_ms, epoch.bucket_ms)))::BIGINT AS active_ms, + sum(CASE WHEN buckets.bucket_start = bounds.first_day THEN 1 ELSE 0 END)::BIGINT AS intervals + FROM bounds + CROSS JOIN LATERAL generate_series( + bounds.first_day, + bounds.last_day, + INTERVAL '24 hours' + ) AS buckets(bucket_start) + CROSS JOIN LATERAL ( + SELECT floor(extract(EPOCH FROM buckets.bucket_start) * 1000)::BIGINT AS bucket_ms + ) AS epoch + GROUP BY bounds.tenant_id, buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id + RETURNING 1 +) + +SELECT (SELECT count(*) FROM hourly) + (SELECT count(*) FROM daily); + +RESET ROLE; diff --git a/go/internal/store/queries/compute_usage.sql b/go/internal/store/queries/compute_usage.sql index 1d9096a06..130258957 100644 --- a/go/internal/store/queries/compute_usage.sql +++ b/go/internal/store/queries/compute_usage.sql @@ -1,5 +1,5 @@ --- Close intervals only after both binding deletion and its end event are absent. --- The tenant id comes from each start because the system role has no tenant GUC. +-- Close intervals only after the binding and matching end event are absent. +-- The tenant id comes from each start because this sweep has no tenant GUC. -- name: CloseOrphanedComputeIntervals :execrows WITH orphaned AS ( SELECT starts.tenant_id, starts.interval_id, starts.agent_account_id, @@ -24,8 +24,157 @@ INSERT INTO compute_usage_events ( tenant_id, id, interval_id, kind, occurred_at, agent_account_id, owner_user_id, session_id, runner_id, estimated ) -SELECT orphaned.tenant_id, gen_random_uuid()::text AS id, orphaned.interval_id, 'end' AS kind, now() AS occurred_at, - orphaned.agent_account_id, orphaned.owner_user_id, orphaned.session_id, - orphaned.runner_id, true AS estimated +SELECT orphaned.tenant_id, gen_random_uuid()::text AS id, orphaned.interval_id, + 'end' AS kind, now() AS occurred_at, orphaned.agent_account_id, + orphaned.owner_user_id, orphaned.session_id, orphaned.runner_id, + true AS estimated FROM orphaned + ORDER BY orphaned.tenant_id, orphaned.interval_id ON CONFLICT DO NOTHING; + +-- Compute-usage rollups are derived from closed start/end event pairs. +-- Each query runs under the tenant role and transaction scope supplied by Store. + +-- LockComputeUsage serializes a tenant rebuild with end-event rollup triggers. +-- name: LockComputeUsage :exec +SELECT pg_advisory_xact_lock(hashtext('compute_usage:' || current_setting('compass.tenant_id', TRUE))); + +-- LockComputeUsageTenant takes the same lock for a system-role caller. +-- The key must stay byte-identical to LockComputeUsage and the end trigger. +-- name: LockComputeUsageTenant :exec +SELECT pg_advisory_xact_lock(hashtext('compute_usage:' || @tenant_id::text)); + +-- ComputeUsageSeries sums matching rows at one pre-aggregated granularity. +-- name: ComputeUsageSeries :many +SELECT r.bucket_start::timestamptz AS bucket_start, + sum(r.active_ms)::bigint AS active_ms, + sum(r.intervals)::bigint AS intervals + FROM (SELECT rollups.bucket_start, rollups.agent_account_id, rollups.active_ms, rollups.intervals + FROM compute_usage_rollups_hourly AS rollups + WHERE @granularity::integer = 1 + UNION ALL + SELECT rollups.bucket_start, rollups.agent_account_id, rollups.active_ms, rollups.intervals + FROM compute_usage_rollups_daily AS rollups + WHERE @granularity::integer = 2) AS r + WHERE r.bucket_start >= @start_at::timestamptz + AND r.bucket_start < @end_at::timestamptz + AND (coalesce(cardinality(@agent_account_ids::text[]), 0) = 0 + OR r.agent_account_id = ANY (@agent_account_ids::text[])) + GROUP BY r.bucket_start + ORDER BY r.bucket_start; + +-- ComputeUsagePruneHorizon pins the horizon while the rebuild replaces newer rows. +-- name: ComputeUsagePruneHorizon :one +SELECT horizon.horizon FROM compute_usage_prune_horizon AS horizon FOR SHARE; + +-- DeleteComputeUsageRollupsFrom clears rows the rebuild can reconstruct. +-- name: DeleteComputeUsageRollupsFrom :exec +WITH hourly AS ( + DELETE FROM compute_usage_rollups_hourly AS rollups + WHERE rollups.bucket_start >= @horizon::timestamptz +) +DELETE FROM compute_usage_rollups_daily AS rollups + WHERE rollups.bucket_start >= @horizon::timestamptz; + +-- RollUpComputeUsageFrom rebuilds only buckets at or after the horizon, which are +-- the rows it deleted; older buckets are frozen. The sentinel skips clipping. +-- name: RollUpComputeUsageFrom :exec +WITH intervals AS ( + SELECT starts.tenant_id, starts.interval_id, starts.occurred_at AS start_at, + ends.occurred_at AS end_at, starts.owner_user_id, starts.agent_account_id, + floor(extract(epoch FROM starts.occurred_at) * 1000)::bigint AS start_ms, + floor(extract(epoch FROM ends.occurred_at) * 1000)::bigint AS end_ms + FROM compute_usage_events AS starts + JOIN compute_usage_events AS ends + ON ends.tenant_id = starts.tenant_id + AND ends.interval_id = starts.interval_id + AND ends.kind = 'end' + WHERE starts.kind = 'start' + AND ends.occurred_at >= @horizon::timestamptz +), bounds AS ( + SELECT intervals.*, + CASE WHEN @horizon::timestamptz = '-infinity'::timestamptz + THEN date_trunc('hour', intervals.start_at, 'UTC') + ELSE greatest(date_trunc('hour', intervals.start_at, 'UTC'), + date_trunc('hour', @horizon::timestamptz, 'UTC')) + END AS first_hour, + greatest(date_trunc('hour', intervals.start_at, 'UTC'), + date_trunc('hour', intervals.end_at - interval '1 millisecond', 'UTC')) AS last_hour, + CASE WHEN @horizon::timestamptz = '-infinity'::timestamptz + THEN date_trunc('day', intervals.start_at, 'UTC') + ELSE greatest(date_trunc('day', intervals.start_at, 'UTC'), + date_trunc('day', @horizon::timestamptz, 'UTC')) + END AS first_day, + greatest(date_trunc('day', intervals.start_at, 'UTC'), + date_trunc('day', intervals.end_at - interval '1 millisecond', 'UTC')) AS last_day + FROM intervals +), hourly AS ( + INSERT INTO compute_usage_rollups_hourly ( + bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) + SELECT buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id, + sum(greatest(0, least(bounds.end_ms, epoch.bucket_ms + 3600000) - + greatest(bounds.start_ms, epoch.bucket_ms)))::bigint, + sum(CASE WHEN buckets.bucket_start = date_trunc('hour', bounds.start_at, 'UTC') + AND bounds.start_at >= @horizon::timestamptz + THEN 1 ELSE 0 END)::bigint + FROM bounds + CROSS JOIN LATERAL generate_series( + bounds.first_hour, + bounds.last_hour, + interval '1 hour' + ) AS buckets(bucket_start) + CROSS JOIN LATERAL ( + SELECT floor(extract(epoch FROM buckets.bucket_start) * 1000)::bigint AS bucket_ms + ) AS epoch + GROUP BY buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id +), daily AS ( + INSERT INTO compute_usage_rollups_daily ( + bucket_start, owner_user_id, agent_account_id, active_ms, intervals + ) + SELECT buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id, + sum(greatest(0, least(bounds.end_ms, epoch.bucket_ms + 86400000) - + greatest(bounds.start_ms, epoch.bucket_ms)))::bigint, + sum(CASE WHEN buckets.bucket_start = date_trunc('day', bounds.start_at, 'UTC') + AND bounds.start_at >= @horizon::timestamptz + THEN 1 ELSE 0 END)::bigint + FROM bounds + CROSS JOIN LATERAL generate_series( + bounds.first_day, + bounds.last_day, + interval '24 hours' + ) AS buckets(bucket_start) + CROSS JOIN LATERAL ( + SELECT floor(extract(epoch FROM buckets.bucket_start) * 1000)::bigint AS bucket_ms + ) AS epoch + GROUP BY buckets.bucket_start, bounds.owner_user_id, bounds.agent_account_id +) +SELECT 1; + +-- AdvanceComputeUsagePruneHorizon moves the single global horizon forward. +-- name: AdvanceComputeUsagePruneHorizon :execrows +UPDATE compute_usage_prune_horizon AS state + SET horizon = GREATEST(state.horizon, @cutoff::timestamptz); + +-- DeleteComputeUsageIntervalsBefore removes both events only when the end is old. +-- It runs as the system role, so the tenant predicate is the only tenant scope. +-- name: DeleteComputeUsageIntervalsBefore :execrows +WITH closed AS ( + SELECT ends.tenant_id, ends.interval_id + FROM compute_usage_events AS ends + WHERE ends.tenant_id = @tenant_id::text + AND ends.kind = 'end' + AND ends.occurred_at < @cutoff::timestamptz + AND EXISTS ( + SELECT 1 + FROM compute_usage_events AS starts + WHERE starts.tenant_id = ends.tenant_id + AND starts.interval_id = ends.interval_id + AND starts.kind = 'start' + ) +) +DELETE FROM compute_usage_events AS events + USING closed + WHERE events.tenant_id = @tenant_id::text + AND events.tenant_id = closed.tenant_id + AND events.interval_id = closed.interval_id; diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql index 8b850a916..45d6cf932 100644 --- a/go/internal/store/queries/session_bindings.sql +++ b/go/internal/store/queries/session_bindings.sql @@ -156,6 +156,7 @@ SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_t d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id FROM d JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ORDER BY d.tenant_id, d.usage_interval_id ON CONFLICT DO NOTHING; -- The reconnect sweep, run by Hub.enroll under the system role because a Runner @@ -188,6 +189,7 @@ WITH d AS ( d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id FROM d JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ORDER BY d.tenant_id, d.usage_interval_id ON CONFLICT DO NOTHING RETURNING 1 ) diff --git a/go/internal/store/rls_pgtest_test.go b/go/internal/store/rls_pgtest_test.go index 154795750..07e285257 100644 --- a/go/internal/store/rls_pgtest_test.go +++ b/go/internal/store/rls_pgtest_test.go @@ -606,7 +606,7 @@ func TestRLSCatalogEnabledAndForced(t *testing.T) { "issues", "forge_repo_subscriptions", "forge_artifact_cursors", "forge_state_transitions", "token_usage_events", "token_usage_rollups_hourly", "token_usage_rollups_daily", - "compute_usage_events", + "compute_usage_events", "compute_usage_rollups_hourly", "compute_usage_rollups_daily", } for _, tbl := range tenantOwned { if !enumerated[tbl] { diff --git a/go/internal/store/server_secrets_pgtest_test.go b/go/internal/store/server_secrets_pgtest_test.go index 2f956c24c..37b72ceaf 100644 --- a/go/internal/store/server_secrets_pgtest_test.go +++ b/go/internal/store/server_secrets_pgtest_test.go @@ -59,18 +59,20 @@ func TestT0ServerSecretsShape(t *testing.T) { // server_key_state's withheld DELETE is asserted ABSENT — that omission is a // deliberate least-privilege choice (the tripwire digest must not be // droppable), so it is pinned, not left to chance. The usage prune horizon and - // the compute log withhold the writes 0001_init.sql revokes for them. + // the compute log withhold writes; only the system role may prune the log. for _, tc := range []struct { tbl string + roles []string granted []string denied []string }{ - {"server_secrets", []string{"SELECT", "INSERT", "UPDATE", "DELETE"}, nil}, - {"server_key_state", []string{"SELECT", "INSERT", "UPDATE"}, []string{"DELETE"}}, - {"token_usage_prune_horizon", []string{"SELECT", "UPDATE"}, []string{"INSERT", "DELETE"}}, - {"compute_usage_events", []string{"SELECT", "INSERT"}, []string{"UPDATE", "DELETE"}}, + {"server_secrets", []string{"compass_app", "compass_system"}, []string{"SELECT", "INSERT", "UPDATE", "DELETE"}, nil}, + {"server_key_state", []string{"compass_app", "compass_system"}, []string{"SELECT", "INSERT", "UPDATE"}, []string{"DELETE"}}, + {"token_usage_prune_horizon", []string{"compass_app", "compass_system"}, []string{"SELECT", "UPDATE"}, []string{"INSERT", "DELETE"}}, + {"compute_usage_events", []string{"compass_app"}, []string{"SELECT", "INSERT"}, []string{"UPDATE", "DELETE"}}, + {"compute_usage_events", []string{"compass_system"}, []string{"SELECT", "INSERT", "DELETE"}, []string{"UPDATE"}}, } { - for _, role := range []string{"compass_app", "compass_system"} { + for _, role := range tc.roles { for _, priv := range tc.granted { if !hasTablePrivilege(t, s, role, tc.tbl, priv) { t.Fatalf("%s: %s lacks %s", tc.tbl, role, priv) diff --git a/go/internal/usage/export_test.go b/go/internal/usage/export_test.go index 0b2858563..e6df6d0af 100644 --- a/go/internal/usage/export_test.go +++ b/go/internal/usage/export_test.go @@ -2,6 +2,7 @@ package usage import ( "context" + "time" "github.com/jackc/pgx/v5" @@ -17,10 +18,54 @@ func ClearMemoryRollups(s Store) { s.(*Memory).clearRollups() } func ClearPostgresRollups(ctx context.Context, s Store) error { ctx = store.WithSystemRole(ctx) return s.(*Postgres).st.WithTx(ctx, func(tx pgx.Tx) error { - if _, err := tx.Exec(ctx, "DELETE FROM token_usage_rollups_hourly"); err != nil { - return err + for _, table := range []string{ + "token_usage_rollups_hourly", "token_usage_rollups_daily", + "compute_usage_rollups_hourly", "compute_usage_rollups_daily", + } { + if _, err := tx.Exec(ctx, "DELETE FROM "+table); err != nil { + return err + } } - _, err := tx.Exec(ctx, "DELETE FROM token_usage_rollups_daily") - return err + return nil + }) +} + +// SeedPostgresComputeIntervals writes test intervals through the production trigger. +func SeedPostgresComputeIntervals(ctx context.Context, s Store, intervals ...ComputeInterval) error { + p := s.(*Postgres) + return p.st.WithTx(ctx, func(tx pgx.Tx) error { + for _, interval := range intervals { + if err := validateComputeInterval(interval); err != nil { + return err + } + args := []any{ + interval.IntervalID + "-start", interval.IntervalID, "start", + time.UnixMilli(interval.StartUnixMs), interval.AgentAccountID, + interval.OwnerUserID, "session-" + interval.IntervalID, "runner-1", + } + if _, err := tx.Exec(ctx, ` + INSERT INTO compute_usage_events ( + id, interval_id, kind, occurred_at, agent_account_id, owner_user_id, + session_id, runner_id + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + ON CONFLICT (tenant_id, interval_id, kind) DO NOTHING`, args...); err != nil { + return err + } + if interval.EndUnixMs == 0 { + continue + } + args[0] = interval.IntervalID + "-end" + args[2] = "end" + args[3] = time.UnixMilli(interval.EndUnixMs) + if _, err := tx.Exec(ctx, ` + INSERT INTO compute_usage_events ( + id, interval_id, kind, occurred_at, agent_account_id, owner_user_id, + session_id, runner_id + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + ON CONFLICT (tenant_id, interval_id, kind) DO NOTHING`, args...); err != nil { + return err + } + } + return nil }) } diff --git a/go/internal/usage/memory.go b/go/internal/usage/memory.go index 021af15a1..260332548 100644 --- a/go/internal/usage/memory.go +++ b/go/internal/usage/memory.go @@ -3,6 +3,7 @@ package usage import ( "cmp" "context" + "fmt" "math" "slices" "sync" @@ -13,7 +14,7 @@ import ( // rollupGranularities are the widths every event is rolled up at. var rollupGranularities = [...]Granularity{GranularityHour, GranularityDay} -// rollupKey is one rollup row's identity, the same key the Postgres tables use. +// rollupKey is one token rollup row's identity. type rollupKey struct { granularity Granularity startUnixMs int64 @@ -23,10 +24,20 @@ type rollupKey struct { model string } -// memoryTenant is one tenant's raw events and rollups. +// computeRollupKey is one compute rollup row identity. +type computeRollupKey struct { + granularity Granularity + startUnixMs int64 + ownerUserID string + agentAccountID string +} + +// memoryTenant holds one tenant's raw usage and derived rollups. type memoryTenant struct { - events map[string]TokenUsageEvent - rollups map[rollupKey]Bucket + events map[string]TokenUsageEvent + rollups map[rollupKey]Bucket + computeEvents map[string]ComputeInterval + computeRollups map[computeRollupKey]ComputeBucket } // Memory is the in-memory reference Store. It keeps the tenant from ctx, or "" @@ -34,16 +45,167 @@ type memoryTenant struct { type Memory struct { mu sync.Mutex tenants map[store.TenantID]*memoryTenant - // horizonUnixMs is the UTC day the latest prune cut at. Rollups before it - // may count pruned events, so a rebuild must keep them. - horizonUnixMs int64 + // The horizons are the UTC day each raw log was last pruned through. + horizonUnixMs int64 + computeHorizonUnixMs int64 } var _ Store = (*Memory)(nil) // NewMemory returns an empty in-memory Store. func NewMemory() *Memory { - return &Memory{tenants: map[store.TenantID]*memoryTenant{}, horizonUnixMs: math.MinInt64} + return &Memory{ + tenants: map[store.TenantID]*memoryTenant{}, + horizonUnixMs: math.MinInt64, + computeHorizonUnixMs: math.MinInt64, + } +} + +// AppendComputeInterval adds a test/reference interval and incrementally rolls +// up its duration when it is closed. +func (m *Memory) AppendComputeInterval(ctx context.Context, interval ComputeInterval) error { + if err := validateComputeInterval(interval); err != nil { + return err + } + m.mu.Lock() + defer m.mu.Unlock() + t := m.tenant(ctx) + existing, exists := t.computeEvents[interval.IntervalID] + if exists { + if existing.EndUnixMs != 0 || interval.EndUnixMs == 0 { + return nil + } + existing.EndUnixMs = interval.EndUnixMs + if err := validateComputeInterval(existing); err != nil { + return err + } + t.computeEvents[interval.IntervalID] = existing + interval = existing + } else { + t.computeEvents[interval.IntervalID] = interval + } + if interval.EndUnixMs != 0 { + // Unclipped: an interval open at prune time is in no frozen bucket yet. + t.rollUpCompute(interval, math.MinInt64) + } + return nil +} + +func (t *memoryTenant) rollUpCompute(interval ComputeInterval, horizon int64) { + if interval.EndUnixMs == 0 || (horizon != math.MinInt64 && interval.EndUnixMs < horizon) { + return + } + for _, g := range rollupGranularities { + width := g.widthMs() + startBucket := g.BucketStart(interval.StartUnixMs) + horizonStart := int64(math.MinInt64) + if horizon != math.MinInt64 { + horizonStart = g.BucketStart(horizon) + } + first := max(startBucket, horizonStart) + last := max(startBucket, g.BucketStart(interval.EndUnixMs-1)) + for start := first; start <= last; start += width { + activeStart := max(interval.StartUnixMs, start) + activeEnd := min(interval.EndUnixMs, start+width) + count := int64(0) + if interval.StartUnixMs >= horizon && startBucket == start { + count = 1 + } + key := computeRollupKey{ + granularity: g, + startUnixMs: start, + ownerUserID: interval.OwnerUserID, + agentAccountID: interval.AgentAccountID, + } + bucket := t.computeRollups[key] + bucket.StartUnixMs = start + bucket.ActiveMs += max(0, activeEnd-activeStart) + bucket.Intervals += count + t.computeRollups[key] = bucket + } + } +} + +// validateComputeInterval rejects data that cannot be represented in UTC buckets. +func validateComputeInterval(interval ComputeInterval) error { + switch { + case interval.IntervalID == "": + return fmt.Errorf("%w: empty compute interval id", ErrInvalidArgument) + case interval.StartUnixMs <= 0 || interval.StartUnixMs > maxTimestamptzMs: + return fmt.Errorf("%w: compute interval start is outside the supported range", ErrInvalidArgument) + case interval.EndUnixMs < 0 || interval.EndUnixMs > maxTimestamptzMs || + (interval.EndUnixMs != 0 && interval.EndUnixMs < interval.StartUnixMs): + return fmt.Errorf("%w: compute interval end is outside the supported range", ErrInvalidArgument) + case interval.AgentAccountID == "" || interval.OwnerUserID == "": + return fmt.Errorf("%w: empty compute interval rollup key", ErrInvalidArgument) + } + return nil +} + +// ComputeUsageSeries implements Store. +func (m *Memory) ComputeUsageSeries(ctx context.Context, q SeriesQuery) ([]ComputeBucket, error) { + if err := q.Validate(); err != nil { + return nil, err + } + m.mu.Lock() + defer m.mu.Unlock() + sums := map[int64]ComputeBucket{} + for key, bucket := range m.tenant(ctx).computeRollups { + if key.granularity != q.Granularity || key.startUnixMs < q.StartUnixMs || key.startUnixMs >= q.EndUnixMs { + continue + } + if len(q.AgentAccountIDs) > 0 && !slices.Contains(q.AgentAccountIDs, key.agentAccountID) { + continue + } + sum := sums[key.startUnixMs] + sum.StartUnixMs = key.startUnixMs + sum.ActiveMs += bucket.ActiveMs + sum.Intervals += bucket.Intervals + sums[key.startUnixMs] = sum + } + series := make([]ComputeBucket, 0, len(sums)) + for _, bucket := range sums { + series = append(series, bucket) + } + slices.SortFunc(series, func(a, b ComputeBucket) int { return cmp.Compare(a.StartUnixMs, b.StartUnixMs) }) + return series, nil +} + +// RebuildComputeUsageRollups implements Store. +func (m *Memory) RebuildComputeUsageRollups(ctx context.Context) error { + m.mu.Lock() + defer m.mu.Unlock() + t := m.tenant(ctx) + for key := range t.computeRollups { + if key.startUnixMs >= GranularityDay.BucketStart(m.computeHorizonUnixMs) { + delete(t.computeRollups, key) + } + } + for _, interval := range t.computeEvents { + if interval.EndUnixMs != 0 && GranularityDay.BucketStart(interval.EndUnixMs) >= m.computeHorizonUnixMs { + // Clip to the horizon: older buckets were kept above, not cleared. + t.rollUpCompute(interval, m.computeHorizonUnixMs) + } + } + return nil +} + +// PruneComputeUsageBefore deletes both events of closed intervals before cutoff. +func (m *Memory) PruneComputeUsageBefore(_ context.Context, beforeUnixMs int64) (int64, error) { + cutoff := GranularityDay.BucketStart(beforeUnixMs) + m.mu.Lock() + defer m.mu.Unlock() + var deleted int64 + for _, t := range m.tenants { + for id, interval := range t.computeEvents { + if interval.EndUnixMs != 0 && GranularityDay.BucketStart(interval.EndUnixMs) < cutoff { + delete(t.computeEvents, id) + deleted += 2 + } + } + } + m.computeHorizonUnixMs = max(m.computeHorizonUnixMs, cutoff) + return deleted, nil } // AppendTokenUsage implements Store. @@ -153,6 +315,7 @@ func (m *Memory) clearRollups() { defer m.mu.Unlock() for _, t := range m.tenants { clear(t.rollups) + clear(t.computeRollups) } } @@ -161,7 +324,12 @@ func (m *Memory) tenant(ctx context.Context) *memoryTenant { id, _ := store.TenantFromContext(ctx) t, ok := m.tenants[id] if !ok { - t = &memoryTenant{events: map[string]TokenUsageEvent{}, rollups: map[rollupKey]Bucket{}} + t = &memoryTenant{ + events: map[string]TokenUsageEvent{}, + rollups: map[rollupKey]Bucket{}, + computeEvents: map[string]ComputeInterval{}, + computeRollups: map[computeRollupKey]ComputeBucket{}, + } m.tenants[id] = t } return t diff --git a/go/internal/usage/memory_test.go b/go/internal/usage/memory_test.go index 7a2856c22..794e37cc2 100644 --- a/go/internal/usage/memory_test.go +++ b/go/internal/usage/memory_test.go @@ -18,5 +18,14 @@ func TestMemory(t *testing.T) { return store.WithTenant(t.Context(), store.TenantID(fmt.Sprint("t", tenant))) }, CorruptRollups: func(_ *testing.T, s usage.Store) { usage.ClearMemoryRollups(s) }, + AppendComputeInterval: func(t *testing.T, s usage.Store, ctx context.Context, intervals ...usage.ComputeInterval) { + t.Helper() + memory := s.(*usage.Memory) + for _, interval := range intervals { + if err := memory.AppendComputeInterval(ctx, interval); err != nil { + t.Fatalf("AppendComputeInterval: %v", err) + } + } + }, }) } diff --git a/go/internal/usage/postgres.go b/go/internal/usage/postgres.go index 620a5c2a2..990eeea0a 100644 --- a/go/internal/usage/postgres.go +++ b/go/internal/usage/postgres.go @@ -131,7 +131,7 @@ func (p *Postgres) PruneTokenUsageBefore(ctx context.Context, beforeUnixMs int64 } var deleted int64 for _, id := range tenants { - n, err := p.pruneTenant(store.WithTenant(ctx, store.TenantID(id)), cutoff) + n, err := p.pruneTenant(store.WithTenant(store.WithoutSystemRole(ctx), store.TenantID(id)), cutoff) if err != nil { // The earlier tenants' deletes have committed, so the count keeps them. return deleted, fmt.Errorf("usage: prune token usage of tenant %s: %w", id, err) diff --git a/go/internal/usage/postgres_compute.go b/go/internal/usage/postgres_compute.go new file mode 100644 index 000000000..aeca455f5 --- /dev/null +++ b/go/internal/usage/postgres_compute.go @@ -0,0 +1,117 @@ +package usage + +import ( + "context" + "errors" + "fmt" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/RigelBuild/compass/go/internal/store" + "github.com/RigelBuild/compass/go/internal/store/db" +) + +// ComputeUsageSeries implements Store. +func (p *Postgres) ComputeUsageSeries(ctx context.Context, q SeriesQuery) ([]ComputeBucket, error) { + if err := q.Validate(); err != nil { + return nil, err + } + var rows []db.ComputeUsageSeriesRow + err := p.st.WithTx(ctx, func(tx pgx.Tx) error { + var err error + rows, err = db.New(tx).ComputeUsageSeries(ctx, db.ComputeUsageSeriesParams{ + Granularity: int32(q.Granularity), + StartAt: bound(q.StartUnixMs), + EndAt: bound(q.EndUnixMs), + AgentAccountIds: q.AgentAccountIDs, + }) + return err + }) + if err != nil { + return nil, fmt.Errorf("usage: read compute usage series: %w", err) + } + series := make([]ComputeBucket, len(rows)) + for i := range rows { + series[i] = ComputeBucket{ + StartUnixMs: rows[i].BucketStart.Time.UnixMilli(), + ActiveMs: rows[i].ActiveMs, + Intervals: rows[i].Intervals, + } + } + return series, nil +} + +// RebuildComputeUsageRollups implements Store. +func (p *Postgres) RebuildComputeUsageRollups(ctx context.Context) error { + err := p.st.WithTx(ctx, func(tx pgx.Tx) error { + q := db.New(tx) + if err := q.LockComputeUsage(ctx); err != nil { + return err + } + horizon, err := q.ComputeUsagePruneHorizon(ctx) + if err != nil { + return err + } + if err := q.DeleteComputeUsageRollupsFrom(ctx, horizon); err != nil { + return err + } + return q.RollUpComputeUsageFrom(ctx, horizon) + }) + if err != nil { + return fmt.Errorf("usage: rebuild compute usage rollups: %w", err) + } + return nil +} + +// PruneComputeUsageBefore deletes closed intervals before the UTC day cutoff. +func (p *Postgres) PruneComputeUsageBefore(ctx context.Context, beforeUnixMs int64) (int64, error) { + cutoff := bound(GranularityDay.BucketStart(beforeUnixMs)) + var tenants []string + err := p.st.WithTx(store.WithSystemRole(ctx), func(tx pgx.Tx) error { + q := db.New(tx) + n, err := q.AdvanceComputeUsagePruneHorizon(ctx, cutoff) + if err != nil { + return err + } + if n == 0 { + return errors.New("prune horizon row is missing") + } + tenants, err = q.ListTenantIDs(ctx) + return err + }) + if err != nil { + return 0, fmt.Errorf("usage: advance compute usage prune horizon: %w", err) + } + var deleted int64 + for _, id := range tenants { + n, err := p.pruneComputeTenant(ctx, id, cutoff) + if err != nil { + return deleted, fmt.Errorf("usage: prune compute usage of tenant %s: %w", id, err) + } + deleted += n + } + return deleted, nil +} + +// pruneComputeTenant deletes only this tenant's closed intervals. It runs as +// the system role so the request role never needs DELETE on the raw log. +func (p *Postgres) pruneComputeTenant(ctx context.Context, tenantID string, cutoff pgtype.Timestamptz) (int64, error) { + var deleted int64 + err := p.st.WithTx(store.WithSystemRole(ctx), func(tx pgx.Tx) error { + q := db.New(tx) + if err := q.LockComputeUsageTenant(ctx, tenantID); err != nil { + return err + } + var err error + deleted, err = q.DeleteComputeUsageIntervalsBefore(ctx, db.DeleteComputeUsageIntervalsBeforeParams{ + TenantID: tenantID, + Cutoff: cutoff, + }) + return err + }) + if err != nil { + return 0, err + } + return deleted, nil +} diff --git a/go/internal/usage/postgres_pgtest_test.go b/go/internal/usage/postgres_pgtest_test.go index e23ef8a2f..6e1d91a93 100644 --- a/go/internal/usage/postgres_pgtest_test.go +++ b/go/internal/usage/postgres_pgtest_test.go @@ -10,6 +10,7 @@ import ( "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" "github.com/RigelBuild/compass/go/internal/pgtest" "github.com/RigelBuild/compass/go/internal/store" @@ -34,6 +35,7 @@ func TestPostgres(t *testing.T) { t.Fatalf("clear rollups: %v", err) } }, + AppendComputeInterval: appendPostgresComputeIntervals, }) } @@ -72,6 +74,13 @@ func openPostgres(t *testing.T) (*store.Store, usage.Store) { return st, usage.NewPostgres(st) } +func appendPostgresComputeIntervals(t *testing.T, s usage.Store, ctx context.Context, intervals ...usage.ComputeInterval) { + t.Helper() + if err := usage.SeedPostgresComputeIntervals(ctx, s, intervals...); err != nil { + t.Fatalf("seed compute intervals: %v", err) + } +} + // A rebuild reads the horizon FOR SHARE, so a prune must not advance it, and // so must not delete events the rebuild would miss, until the rebuild commits. func TestPostgresPruneWaitsForHeldHorizon(t *testing.T) { @@ -114,6 +123,131 @@ func TestPostgresPruneWaitsForHeldHorizon(t *testing.T) { } } +func TestPostgresComputeRebuildUsesPruneHorizon(t *testing.T) { + _, s := openPostgres(t) + ctx := store.WithTenant(t.Context(), pgTenants[0]) + day0 := time.Date(2026, time.March, 10, 0, 0, 0, 0, time.UTC) + appendPostgresComputeIntervals(t, s, ctx, + usage.ComputeInterval{IntervalID: "horizon", StartUnixMs: day0.Add(23 * time.Hour).UnixMilli(), EndUnixMs: day0.Add(25 * time.Hour).UnixMilli(), AgentAccountID: "a1", OwnerUserID: "u1"}, + ) + if _, err := s.PruneComputeUsageBefore(ctx, day0.Add(24*time.Hour).UnixMilli()); err != nil { + t.Fatalf("PruneComputeUsageBefore: %v", err) + } + query := usage.SeriesQuery{ + Granularity: usage.GranularityHour, + StartUnixMs: day0.UnixMilli(), + EndUnixMs: day0.Add(48 * time.Hour).UnixMilli(), + } + want := []usage.ComputeBucket{ + {StartUnixMs: day0.Add(23 * time.Hour).UnixMilli(), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + {StartUnixMs: day0.Add(24 * time.Hour).UnixMilli(), ActiveMs: int64(time.Hour / time.Millisecond)}, + } + for range 2 { + if err := s.RebuildComputeUsageRollups(ctx); err != nil { + t.Fatalf("RebuildComputeUsageRollups: %v", err) + } + got, err := s.ComputeUsageSeries(ctx, query) + if err != nil { + t.Fatalf("ComputeUsageSeries: %v", err) + } + if !slices.Equal(got, want) { + t.Fatalf("compute usage series after rebuild = %+v, want %+v", got, want) + } + } +} + +func TestPostgresComputePruneRemovesOnlyClosedIntervalsBeforeCutoffDay(t *testing.T) { + st, s := openPostgres(t) + ctx := store.WithTenant(t.Context(), pgTenants[0]) + day0 := time.Date(2026, time.March, 10, 0, 0, 0, 0, time.UTC) + appendPostgresComputeIntervals(t, s, ctx, + usage.ComputeInterval{IntervalID: "old-compute", StartUnixMs: day0.Add(time.Hour).UnixMilli(), EndUnixMs: day0.Add(2 * time.Hour).UnixMilli(), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "new-compute", StartUnixMs: day0.Add(24 * time.Hour).UnixMilli(), EndUnixMs: day0.Add(25 * time.Hour).UnixMilli(), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "open-compute", StartUnixMs: day0.Add(3 * time.Hour).UnixMilli(), AgentAccountID: "a1", OwnerUserID: "u1"}, + ) + if deleted, err := s.PruneComputeUsageBefore(ctx, day0.Add(24*time.Hour).UnixMilli()); err != nil || deleted != 2 { + t.Fatalf("PruneComputeUsageBefore = %d, %v; want two events, nil", deleted, err) + } + var old, newer, open int + if err := st.WithTx(ctx, func(tx pgx.Tx) error { + return tx.QueryRow(ctx, ` + SELECT count(*) FILTER (WHERE interval_id = 'old-compute'), + count(*) FILTER (WHERE interval_id = 'new-compute'), + count(*) FILTER (WHERE interval_id = 'open-compute') + FROM compute_usage_events`).Scan(&old, &newer, &open) + }); err != nil { + t.Fatalf("count compute events: %v", err) + } + if old != 0 || newer != 2 || open != 1 { + t.Fatalf("compute event counts = old %d, new %d, open %d; want 0, 2, 1", old, newer, open) + } +} + +// Daily buckets are fixed UTC days, so a non-UTC session crossing a DST change +// must not shift later buckets off UTC midnight on either rollup path. +func TestPostgresComputeDailyBucketsIgnoreSessionTimeZone(t *testing.T) { + st, _ := openPostgres(t) + ctx := store.WithTenant(t.Context(), pgTenants[0]) + // America/New_York springs forward on 2026-03-08. + start := time.Date(2026, time.March, 7, 12, 0, 0, 0, time.UTC) + end := start.Add(72 * time.Hour) + var misaligned []string + err := st.WithTx(ctx, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, "SET LOCAL timezone = 'America/New_York'"); err != nil { + return err + } + for _, ev := range []struct { + kind string + at time.Time + }{{"start", start}, {"end", end}} { + if _, err := tx.Exec(ctx, ` + INSERT INTO compute_usage_events ( + id, interval_id, kind, occurred_at, agent_account_id, owner_user_id, + session_id, runner_id + ) VALUES ($1, 'dst', $2, $3, 'a1', 'u1', 'session-dst', 'runner-1')`, + "dst-"+ev.kind, ev.kind, ev.at); err != nil { + return err + } + } + check := func(path string) error { + rows, err := tx.Query(ctx, ` + SELECT rollups.bucket_start AT TIME ZONE 'UTC' + FROM compute_usage_rollups_daily AS rollups + WHERE rollups.bucket_start <> date_trunc('day', rollups.bucket_start, 'UTC')`) + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var b time.Time + if err := rows.Scan(&b); err != nil { + return err + } + misaligned = append(misaligned, path+" "+b.Format(time.DateTime)) + } + return rows.Err() + } + if err := check("trigger"); err != nil { + return err + } + q := db.New(tx) + horizon := pgtype.Timestamptz{InfinityModifier: pgtype.NegativeInfinity, Valid: true} + if err := q.DeleteComputeUsageRollupsFrom(ctx, horizon); err != nil { + return err + } + if err := q.RollUpComputeUsageFrom(ctx, horizon); err != nil { + return err + } + return check("rebuild") + }) + if err != nil { + t.Fatalf("roll up across DST: %v", err) + } + if len(misaligned) != 0 { + t.Fatalf("daily buckets off UTC midnight: %v", misaligned) + } +} + // A rebuild beside an append would count the append twice or lose it, so both // wait for whoever holds the tenant's usage lock. func TestPostgresWritersWaitForUsageLock(t *testing.T) { diff --git a/go/internal/usage/retention.go b/go/internal/usage/retention.go index 16d498aa8..8b2f4d0c4 100644 --- a/go/internal/usage/retention.go +++ b/go/internal/usage/retention.go @@ -14,9 +14,10 @@ const DefaultRetention = 90 * 24 * time.Hour // time, so a faster sweep would only rescan the same events. const sweepInterval = 24 * time.Hour -// pruner is the part of Store the sweeper drives. +// pruner is the raw-event retention surface the sweeper drives. type pruner interface { PruneTokenUsageBefore(ctx context.Context, beforeUnixMs int64) (int64, error) + PruneComputeUsageBefore(ctx context.Context, beforeUnixMs int64) (int64, error) } // RetentionConfig configures a RetentionSweeper. @@ -28,7 +29,7 @@ type RetentionConfig struct { Log *slog.Logger } -// RetentionSweeper deletes the raw events older than the retention window. +// RetentionSweeper deletes raw usage events older than the configured window. type RetentionSweeper struct { store pruner retention time.Duration @@ -48,8 +49,7 @@ func NewRetentionSweeper(s pruner, cfg RetentionConfig) *RetentionSweeper { // never returns a prune error: in the serve group that would stop the server. func (w *RetentionSweeper) Run(ctx context.Context) error { if w.retention <= 0 { - // The serve group treats a nil return as a clean exit, not a failure. - w.log.InfoContext(ctx, "usage retention: raw token-usage prune disabled") + w.log.InfoContext(ctx, "usage retention: raw-event prune disabled") return nil } w.sweep(ctx) @@ -65,16 +65,23 @@ func (w *RetentionSweeper) Run(ctx context.Context) error { } } -// sweep runs one prune. A failed prune is logged, and the next tick retries it. +// sweep runs one prune for each raw usage log. func (w *RetentionSweeper) sweep(ctx context.Context) { - deleted, err := w.store.PruneTokenUsageBefore(ctx, time.Now().Add(-w.retention).UnixMilli()) - switch { - case ctx.Err() != nil: - // Shutdown interrupted the prune; the next start sweeps again. - case err != nil: - w.log.ErrorContext(ctx, "usage retention: prune raw token-usage events", "error", err) - case deleted > 0: + cutoff := time.Now().Add(-w.retention).UnixMilli() + if deleted, err := w.store.PruneTokenUsageBefore(ctx, cutoff); err != nil { + if ctx.Err() == nil { + w.log.ErrorContext(ctx, "usage retention: prune raw token-usage events", "error", err) + } + } else if deleted > 0 { w.log.InfoContext(ctx, "usage retention: pruned raw token-usage events", "deleted", deleted, "retention", w.retention) } + if deleted, err := w.store.PruneComputeUsageBefore(ctx, cutoff); err != nil { + if ctx.Err() == nil { + w.log.ErrorContext(ctx, "usage retention: prune raw compute-usage events", "error", err) + } + } else if deleted > 0 { + w.log.InfoContext(ctx, "usage retention: pruned raw compute-usage events", + "deleted", deleted, "retention", w.retention) + } } diff --git a/go/internal/usage/retention_test.go b/go/internal/usage/retention_test.go index 34a12dd48..d3bdbe854 100644 --- a/go/internal/usage/retention_test.go +++ b/go/internal/usage/retention_test.go @@ -14,8 +14,9 @@ import ( // chanPruner hands each prune cutoff to the test. Blocking on the receive is // what lets the synctest clock run forward to the next tick. type chanPruner struct { - cutoffs chan int64 - err error + cutoffs chan int64 + computeCutoffs chan int64 + err error } func (p *chanPruner) PruneTokenUsageBefore(ctx context.Context, beforeUnixMs int64) (int64, error) { @@ -27,6 +28,15 @@ func (p *chanPruner) PruneTokenUsageBefore(ctx context.Context, beforeUnixMs int } } +func (p *chanPruner) PruneComputeUsageBefore(ctx context.Context, beforeUnixMs int64) (int64, error) { + select { + case p.computeCutoffs <- beforeUnixMs: + return 0, p.err + case <-ctx.Done(): + return 0, ctx.Err() + } +} + func TestRetentionSweeperPrunesAtStartAndDaily(t *testing.T) { const day = 24 * time.Hour tests := []struct { @@ -41,7 +51,7 @@ func TestRetentionSweeperPrunesAtStartAndDaily(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { synctest.Test(t, func(t *testing.T) { - p := &chanPruner{cutoffs: make(chan int64), err: tt.err} + p := &chanPruner{cutoffs: make(chan int64), computeCutoffs: make(chan int64), err: tt.err} w := usage.NewRetentionSweeper(p, usage.RetentionConfig{ Retention: tt.retention, Log: slog.New(slog.DiscardHandler), @@ -52,13 +62,16 @@ func TestRetentionSweeperPrunesAtStartAndDaily(t *testing.T) { go func() { errc <- w.Run(ctx) }() for _, at := range []time.Time{start, start.Add(day)} { - got := <-p.cutoffs + want := at.Add(-tt.window).UnixMilli() + if got := <-p.cutoffs; got != want { + t.Fatalf("token prune cutoff = %d, want %d", got, want) + } + if got := <-p.computeCutoffs; got != want { + t.Fatalf("compute prune cutoff = %d, want %d", got, want) + } if now := time.Now(); !now.Equal(at) { t.Fatalf("prune ran at %v, want %v", now, at) } - if want := at.Add(-tt.window).UnixMilli(); got != want { - t.Fatalf("prune cutoff = %d, want %d (%v before %v)", got, want, tt.window, at) - } } cancel() if err := <-errc; err != nil { diff --git a/go/internal/usage/usage.go b/go/internal/usage/usage.go index d425b7ec7..b31ee2aba 100644 --- a/go/internal/usage/usage.go +++ b/go/internal/usage/usage.go @@ -92,22 +92,35 @@ type Bucket struct { CostMicroUSD int64 } +// ComputeInterval is one completed, or currently open, agent-active interval. +// EndUnixMs is zero while open; only closed intervals are rolled up. +type ComputeInterval struct { + IntervalID string + StartUnixMs int64 + EndUnixMs int64 + AgentAccountID string + OwnerUserID string +} + +// ComputeBucket holds active duration and the number of intervals that started +// in one UTC bucket. +type ComputeBucket struct { + StartUnixMs int64 + ActiveMs int64 + Intervals int64 +} + // Store is the Plane-A usage store seam: an append-only event write plus rollup -// reads. Each call but PruneTokenUsageBefore is scoped to ctx's tenant, which -// the backend resolves. +// reads. Each call is scoped to ctx's tenant, which the backend resolves, except +// the two Prune calls: they advance a global horizon and prune every tenant. type Store interface { - // AppendTokenUsage records a batch atomically. It is idempotent on event ID - // for events still inside the retention window, so a retried batch never - // double-counts. A pruned event's ID is forgotten and counts again. AppendTokenUsage(ctx context.Context, events []TokenUsageEvent) error - // TokenUsageSeries returns the non-empty buckets in the window, oldest first. TokenUsageSeries(ctx context.Context, query SeriesQuery) ([]Bucket, error) - // RebuildTokenUsageRollups recomputes every rollup from the prune horizon on. - // Older rollups outlive their pruned events, so it keeps them. RebuildTokenUsageRollups(ctx context.Context) error - // PruneTokenUsageBefore deletes every tenant's raw events before the UTC day - // that holds beforeUnixMs and keeps the rollups. It returns the delete count. PruneTokenUsageBefore(ctx context.Context, beforeUnixMs int64) (int64, error) + ComputeUsageSeries(ctx context.Context, query SeriesQuery) ([]ComputeBucket, error) + RebuildComputeUsageRollups(ctx context.Context) error + PruneComputeUsageBefore(ctx context.Context, beforeUnixMs int64) (int64, error) } // BucketStart returns the start of the UTC-aligned bucket that holds unixMs, diff --git a/go/internal/usage/usagetest/usagetest.go b/go/internal/usage/usagetest/usagetest.go index 0a1be1a41..26b44a3b9 100644 --- a/go/internal/usage/usagetest/usagetest.go +++ b/go/internal/usage/usagetest/usagetest.go @@ -20,6 +20,8 @@ type Harness struct { Ctx func(t *testing.T, tenant int) context.Context // CorruptRollups deletes every rollup of s and keeps the raw events. CorruptRollups func(t *testing.T, s usage.Store) + // AppendComputeInterval seeds the backend's compute fixture events. + AppendComputeInterval func(t *testing.T, s usage.Store, ctx context.Context, intervals ...usage.ComputeInterval) } // Run runs the contract suite against the backend h adapts. @@ -38,6 +40,16 @@ func Run(t *testing.T, h Harness) { t.Run("tenants_are_isolated", tenantsAreIsolated(h)) t.Run("invalid_events_are_rejected_and_write_nothing", invalidEventsWriteNothing(h)) t.Run("invalid_queries_are_rejected", invalidQueriesAreRejected(h)) + t.Run("compute_intervals_split_across_hour_and_day_boundaries", computeIntervalsSplitAcrossBoundaries(h)) + t.Run("open_compute_interval_is_excluded", openComputeIntervalIsExcluded(h)) + t.Run("compute_rebuild_matches_incremental_rollups", computeRebuildMatchesIncremental(h)) + t.Run("compute_open_interval_ending_at_prune_horizon", computeOpenIntervalEndingAtPruneHorizon(h)) + t.Run("compute_open_interval_crossing_prune_horizon_rebuilds_consistently", computeOpenIntervalCrossingPruneHorizon(h)) + t.Run("compute_interval_open_across_two_prunes_counts_once", computeIntervalOpenAcrossTwoPrunesCountsOnce(h)) + t.Run("compute_multiple_intervals_crossing_bucket_edges_rebuild_consistently", computeMultipleIntervalsCrossingBucketEdges(h)) + t.Run("compute_zero_duration_at_prune_horizon_matches_sql", computeZeroDurationAtPruneHorizon(h)) + t.Run("compute_prune_keeps_rollups_and_open_intervals", computePruneKeepsRollups(h)) + t.Run("compute_series_is_tenant_scoped", computeSeriesIsTenantScoped(h)) } // day0 is a UTC midnight, so the fixtures sit on known bucket boundaries. @@ -395,3 +407,274 @@ func invalidQueriesAreRejected(h Harness) func(*testing.T) { } } } + +func computeIntervalsSplitAcrossBoundaries(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + h.AppendComputeInterval(t, s, ctx, + usage.ComputeInterval{IntervalID: "hour", StartUnixMs: at(hour - time.Minute), EndUnixMs: at(hour + 2*time.Minute), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "day", StartUnixMs: at(day - time.Minute), EndUnixMs: at(day + 2*time.Minute), AgentAccountID: "a1", OwnerUserID: "u1"}, + ) + wantComputeSeries(t, ctx, s, query(usage.GranularityHour), + usage.ComputeBucket{StartUnixMs: at(0), ActiveMs: int64(time.Minute / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(hour), ActiveMs: int64(2 * time.Minute / time.Millisecond)}, + usage.ComputeBucket{StartUnixMs: at(23 * hour), ActiveMs: int64(time.Minute / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(day), ActiveMs: int64(2 * time.Minute / time.Millisecond)}, + ) + wantComputeSeries(t, ctx, s, query(usage.GranularityDay), + usage.ComputeBucket{StartUnixMs: at(0), ActiveMs: int64(4 * time.Minute / time.Millisecond), Intervals: 2}, + usage.ComputeBucket{StartUnixMs: at(day), ActiveMs: int64(2 * time.Minute / time.Millisecond)}, + ) + } +} + +func openComputeIntervalIsExcluded(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + h.AppendComputeInterval(t, s, ctx, usage.ComputeInterval{ + IntervalID: "open", StartUnixMs: at(hour), AgentAccountID: "a1", OwnerUserID: "u1", + }) + wantComputeSeries(t, ctx, s, query(usage.GranularityHour)) + mustComputeRebuild(t, ctx, s) + wantComputeSeries(t, ctx, s, query(usage.GranularityDay)) + } +} + +func computeRebuildMatchesIncremental(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + h.AppendComputeInterval(t, s, ctx, + usage.ComputeInterval{IntervalID: "cross", StartUnixMs: at(hour - time.Minute), EndUnixMs: at(day + time.Minute), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "same", StartUnixMs: at(hour), EndUnixMs: at(hour + 30*time.Minute), AgentAccountID: "a2", OwnerUserID: "u1"}, + ) + beforeHour := computeSeries(t, ctx, s, query(usage.GranularityHour)) + beforeDay := computeSeries(t, ctx, s, query(usage.GranularityDay)) + h.CorruptRollups(t, s) + mustComputeRebuild(t, ctx, s) + if got := computeSeries(t, ctx, s, query(usage.GranularityHour)); !slices.Equal(got, beforeHour) { + t.Fatalf("hourly rebuild = %+v, want incremental %+v", got, beforeHour) + } + if got := computeSeries(t, ctx, s, query(usage.GranularityDay)); !slices.Equal(got, beforeDay) { + t.Fatalf("daily rebuild = %+v, want incremental %+v", got, beforeDay) + } + } +} + +func computeOpenIntervalEndingAtPruneHorizon(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + interval := usage.ComputeInterval{ + IntervalID: "ends-at-prune-horizon", StartUnixMs: at(day - hour), + AgentAccountID: "a1", OwnerUserID: "u1", + } + h.AppendComputeInterval(t, s, ctx, interval) + mustComputePrune(t, ctx, s, at(day), 0) + + interval.EndUnixMs = at(day) + h.AppendComputeInterval(t, s, ctx, interval) + want := []usage.ComputeBucket{{StartUnixMs: at(day - hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}} + wantComputeSeries(t, ctx, s, query(usage.GranularityHour), want...) + wantComputeSeries(t, ctx, s, query(usage.GranularityDay), + usage.ComputeBucket{StartUnixMs: at(0), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + ) + assertComputeRebuildMatchesIncremental(t, s, ctx) + } +} + +func computeOpenIntervalCrossingPruneHorizon(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + interval := usage.ComputeInterval{ + IntervalID: "cross-pruned-horizon", StartUnixMs: at(day - hour), + AgentAccountID: "a1", OwnerUserID: "u1", + } + h.AppendComputeInterval(t, s, ctx, interval) + mustComputePrune(t, ctx, s, at(day), 0) + + interval.EndUnixMs = at(day + 2*hour) + h.AppendComputeInterval(t, s, ctx, interval) + wantComputeSeries(t, ctx, s, query(usage.GranularityHour), + usage.ComputeBucket{StartUnixMs: at(day - hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(day), ActiveMs: int64(time.Hour / time.Millisecond)}, + usage.ComputeBucket{StartUnixMs: at(day + hour), ActiveMs: int64(time.Hour / time.Millisecond)}, + ) + wantComputeSeries(t, ctx, s, query(usage.GranularityDay), + usage.ComputeBucket{StartUnixMs: at(0), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(day), ActiveMs: int64(2 * time.Hour / time.Millisecond)}, + ) + assertComputeRebuildMatchesIncremental(t, s, ctx) + } +} + +// An interval open at one prune and closed before the next must bill its full +// time once, whatever rebuilds and prunes run around the close. +func computeIntervalOpenAcrossTwoPrunesCountsOnce(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + h.AppendComputeInterval(t, s, ctx, + usage.ComputeInterval{IntervalID: "frozen", StartUnixMs: at(hour), EndUnixMs: at(2 * hour), AgentAccountID: "a1", OwnerUserID: "u1"}, + ) + interval := usage.ComputeInterval{ + IntervalID: "open-across-prunes", StartUnixMs: at(day - hour), + AgentAccountID: "a1", OwnerUserID: "u1", + } + h.AppendComputeInterval(t, s, ctx, interval) + mustComputePrune(t, ctx, s, at(day), 2) + mustComputeRebuild(t, ctx, s) + + interval.EndUnixMs = at(day + 2*hour) + h.AppendComputeInterval(t, s, ctx, interval) + hourly := []usage.ComputeBucket{ + {StartUnixMs: at(hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + {StartUnixMs: at(day - hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + {StartUnixMs: at(day), ActiveMs: int64(time.Hour / time.Millisecond)}, + {StartUnixMs: at(day + hour), ActiveMs: int64(time.Hour / time.Millisecond)}, + } + daily := []usage.ComputeBucket{ + {StartUnixMs: at(0), ActiveMs: int64(2 * time.Hour / time.Millisecond), Intervals: 2}, + {StartUnixMs: at(day), ActiveMs: int64(2 * time.Hour / time.Millisecond)}, + } + check := func(step string) { + t.Helper() + if got := computeSeries(t, ctx, s, query(usage.GranularityHour)); !slices.Equal(got, hourly) { + t.Fatalf("%s: hourly = %+v, want %+v", step, got, hourly) + } + if got := computeSeries(t, ctx, s, query(usage.GranularityDay)); !slices.Equal(got, daily) { + t.Fatalf("%s: daily = %+v, want %+v", step, got, daily) + } + } + check("after close") + mustComputeRebuild(t, ctx, s) + check("rebuild under the first horizon") + mustComputePrune(t, ctx, s, at(2*day), 2) + check("second prune") + mustComputeRebuild(t, ctx, s) + check("rebuild under the second horizon") + } +} + +func computeZeroDurationAtPruneHorizon(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + interval := usage.ComputeInterval{ + IntervalID: "zero-at-prune-horizon", StartUnixMs: at(day), + AgentAccountID: "a1", OwnerUserID: "u1", + } + h.AppendComputeInterval(t, s, ctx, interval) + mustComputePrune(t, ctx, s, at(day), 0) + + interval.EndUnixMs = at(day) + h.AppendComputeInterval(t, s, ctx, interval) + want := []usage.ComputeBucket{{StartUnixMs: at(day), ActiveMs: 0, Intervals: 1}} + wantComputeSeries(t, ctx, s, query(usage.GranularityHour), want...) + wantComputeSeries(t, ctx, s, query(usage.GranularityDay), want...) + mustComputeRebuild(t, ctx, s) + wantComputeSeries(t, ctx, s, query(usage.GranularityHour), want...) + wantComputeSeries(t, ctx, s, query(usage.GranularityDay), want...) + } +} + +func computeMultipleIntervalsCrossingBucketEdges(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + h.AppendComputeInterval(t, s, ctx, + usage.ComputeInterval{IntervalID: "cross-midnight", StartUnixMs: at(day - 30*time.Minute), EndUnixMs: at(day + 30*time.Minute), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "cross-hours", StartUnixMs: at(day - 15*time.Minute), EndUnixMs: at(day + 2*time.Hour + 15*time.Minute), AgentAccountID: "a2", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "cross-next-midnight", StartUnixMs: at(2*day - 30*time.Minute), EndUnixMs: at(2*day + 30*time.Minute), AgentAccountID: "a1", OwnerUserID: "u1"}, + ) + assertComputeRebuildMatchesIncremental(t, s, ctx) + } +} + +func assertComputeRebuildMatchesIncremental(t *testing.T, s usage.Store, ctx context.Context) { + t.Helper() + beforeHour := computeSeries(t, ctx, s, query(usage.GranularityHour)) + beforeDay := computeSeries(t, ctx, s, query(usage.GranularityDay)) + mustComputeRebuild(t, ctx, s) + if got := computeSeries(t, ctx, s, query(usage.GranularityHour)); !slices.Equal(got, beforeHour) { + t.Fatalf("hourly rebuild = %+v, want incremental %+v", got, beforeHour) + } + if got := computeSeries(t, ctx, s, query(usage.GranularityDay)); !slices.Equal(got, beforeDay) { + t.Fatalf("daily rebuild = %+v, want incremental %+v", got, beforeDay) + } +} + +func computePruneKeepsRollups(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx := h.New(t), h.Ctx(t, 0) + h.AppendComputeInterval(t, s, ctx, + usage.ComputeInterval{IntervalID: "old", StartUnixMs: at(hour), EndUnixMs: at(2 * hour), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "cross-horizon", StartUnixMs: at(day - hour), EndUnixMs: at(day + hour), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "closed-new", StartUnixMs: at(day + hour), EndUnixMs: at(day + 2*hour), AgentAccountID: "a1", OwnerUserID: "u1"}, + usage.ComputeInterval{IntervalID: "open", StartUnixMs: at(hour), AgentAccountID: "a1", OwnerUserID: "u1"}, + ) + mustComputePrune(t, ctx, s, at(day+12*hour), 2) + wantComputeSeries(t, ctx, s, query(usage.GranularityHour), + usage.ComputeBucket{StartUnixMs: at(hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(day - hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(day), ActiveMs: int64(time.Hour / time.Millisecond)}, + usage.ComputeBucket{StartUnixMs: at(day + hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + ) + mustComputeRebuild(t, ctx, s) + wantComputeSeries(t, ctx, s, query(usage.GranularityHour), + usage.ComputeBucket{StartUnixMs: at(hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(day - hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(day), ActiveMs: int64(time.Hour / time.Millisecond)}, + usage.ComputeBucket{StartUnixMs: at(day + hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + ) + } +} + +func computeSeriesIsTenantScoped(h Harness) func(*testing.T) { + return func(t *testing.T) { + s, ctx0, ctx1 := h.New(t), h.Ctx(t, 0), h.Ctx(t, 1) + h.AppendComputeInterval(t, s, ctx0, usage.ComputeInterval{ + IntervalID: "t0", StartUnixMs: at(hour), EndUnixMs: at(2 * hour), AgentAccountID: "a1", OwnerUserID: "u1", + }) + wantComputeSeries(t, ctx1, s, query(usage.GranularityHour)) + h.AppendComputeInterval(t, s, ctx1, usage.ComputeInterval{ + IntervalID: "t1", StartUnixMs: at(hour), EndUnixMs: at(3 * hour), AgentAccountID: "a1", OwnerUserID: "u1", + }) + wantComputeSeries(t, ctx0, s, query(usage.GranularityHour), + usage.ComputeBucket{StartUnixMs: at(hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + ) + wantComputeSeries(t, ctx1, s, query(usage.GranularityHour), + usage.ComputeBucket{StartUnixMs: at(hour), ActiveMs: int64(time.Hour / time.Millisecond), Intervals: 1}, + usage.ComputeBucket{StartUnixMs: at(2 * hour), ActiveMs: int64(time.Hour / time.Millisecond)}, + ) + } +} + +func computeSeries(t *testing.T, ctx context.Context, s usage.Store, q usage.SeriesQuery) []usage.ComputeBucket { + t.Helper() + got, err := s.ComputeUsageSeries(ctx, q) + if err != nil { + t.Fatalf("ComputeUsageSeries(%+v): %v", q, err) + } + return got +} + +func wantComputeSeries(t *testing.T, ctx context.Context, s usage.Store, q usage.SeriesQuery, want ...usage.ComputeBucket) { + t.Helper() + if got := computeSeries(t, ctx, s, q); !slices.Equal(got, want) { + t.Fatalf("ComputeUsageSeries(%+v)\n got %+v\n want %+v", q, got, want) + } +} + +func mustComputeRebuild(t *testing.T, ctx context.Context, s usage.Store) { + t.Helper() + if err := s.RebuildComputeUsageRollups(ctx); err != nil { + t.Fatalf("RebuildComputeUsageRollups: %v", err) + } +} + +func mustComputePrune(t *testing.T, ctx context.Context, s usage.Store, beforeUnixMs, wantDeleted int64) { + t.Helper() + got, err := s.PruneComputeUsageBefore(ctx, beforeUnixMs) + if err != nil { + t.Fatalf("PruneComputeUsageBefore: %v", err) + } + if got != wantDeleted { + t.Fatalf("PruneComputeUsageBefore deleted %d events, want %d", got, wantDeleted) + } +} diff --git a/go/server/serve.go b/go/server/serve.go index 29654f617..e6fb4be4a 100644 --- a/go/server/serve.go +++ b/go/server/serve.go @@ -141,8 +141,8 @@ type ServeConfig struct { // points the whole deployment's custody — including the master key — at a // managed store, per the record's A2 KMS-by-provider-URI custody note. SecretProvider string - // UsageEventRetention is how long raw token-usage events are kept; 0 turns - // the daily prune off. The CLI defaults it to 90 days. + // UsageEventRetention is how long raw token- and compute-usage events are kept. + // 0 turns the daily prune off. The CLI defaults it to 90 days. UsageEventRetention time.Duration // OnBound, when set, is called once the dev and network-door listeners are // bound, with each bound address (nil for a door that is off). It lets a diff --git a/go/server/sinks.go b/go/server/sinks.go index 7e8536eff..399b37122 100644 --- a/go/server/sinks.go +++ b/go/server/sinks.go @@ -186,8 +186,8 @@ func startForgeIngestLanes(gctx context.Context, g *errgroup.Group, board *board } } -// startUsageSweepers starts the usage retention prune and the orphaned compute -// interval close on the serve group; binding changes close intervals inline. +// startUsageSweepers starts the token- and compute-usage retention prune and the +// orphaned compute interval close; binding changes close intervals inline. func startUsageSweepers(gctx context.Context, g *errgroup.Group, st *store.Store, retention time.Duration, log *slog.Logger) { retain := usage.NewRetentionSweeper(usage.NewPostgres(st), usage.RetentionConfig{Retention: retention, Log: log}) closer := usage.NewComputeUsageSweeper(computeUsageCloser{st: st}, log)