diff --git a/coderd/database/check_constraint.go b/coderd/database/check_constraint.go index 0402e4b8ee7..2f791242f5e 100644 --- a/coderd/database/check_constraint.go +++ b/coderd/database/check_constraint.go @@ -57,6 +57,7 @@ const ( CheckTelemetryLockEventTypeConstraint CheckConstraint = "telemetry_lock_event_type_constraint" // telemetry_locks CheckValidationMonotonicOrder CheckConstraint = "validation_monotonic_order" // template_version_parameters CheckUsageEventTypeCheck CheckConstraint = "usage_event_type_check" // usage_events + CheckUsageEventsAgentRuntimeHourAligned CheckConstraint = "usage_events_agent_runtime_hour_aligned" // usage_events CheckUserAIBudgetOverridesSpendLimitMicrosCheck CheckConstraint = "user_ai_budget_overrides_spend_limit_micros_check" // user_ai_budget_overrides CheckUserAIProviderKeysAPIKeyCheck CheckConstraint = "user_ai_provider_keys_api_key_check" // user_ai_provider_keys CheckUserSecretsEnabledRequiresTarget CheckConstraint = "user_secrets_enabled_requires_target" // user_secrets diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index d7595115c90..d70220b91ee 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -3552,7 +3552,8 @@ CREATE TABLE usage_events ( publish_started_at timestamp with time zone, published_at timestamp with time zone, failure_message text, - CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_v1'::text, 'hb_agent_runtime_v1'::text]))) + CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_v1'::text, 'hb_agent_runtime_v1'::text]))), + CONSTRAINT usage_events_agent_runtime_hour_aligned CHECK (((event_type <> 'hb_agent_runtime_v1'::text) OR (date_trunc('hour'::text, timezone('UTC'::text, created_at)) = timezone('UTC'::text, created_at)))) ); COMMENT ON TABLE usage_events IS 'usage_events contains usage data that is collected from the product and potentially shipped to the usage collector service.'; @@ -4899,7 +4900,7 @@ CREATE INDEX idx_template_versions_has_ai_task ON template_versions USING btree CREATE UNIQUE INDEX idx_unique_preset_name ON template_version_presets USING btree (name, template_version_id); -CREATE INDEX idx_usage_events_agent_runtime ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_agent_runtime_v1'::text); +CREATE UNIQUE INDEX idx_usage_events_agent_runtime ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_agent_runtime_v1'::text); CREATE INDEX idx_usage_events_ai_seats ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_ai_seats_v1'::text); diff --git a/coderd/database/migrations/000570_usage_events_agent_runtime_invariants.down.sql b/coderd/database/migrations/000570_usage_events_agent_runtime_invariants.down.sql new file mode 100644 index 00000000000..55d3b173185 --- /dev/null +++ b/coderd/database/migrations/000570_usage_events_agent_runtime_invariants.down.sql @@ -0,0 +1,9 @@ +-- IF EXISTS tolerates the index already being gone (e.g. rolling back out +-- of order during an incident) instead of failing. +DROP INDEX IF EXISTS idx_usage_events_agent_runtime; +CREATE INDEX idx_usage_events_agent_runtime + ON usage_events (event_type, created_at) + WHERE event_type = 'hb_agent_runtime_v1'; + +ALTER TABLE usage_events + DROP CONSTRAINT IF EXISTS usage_events_agent_runtime_hour_aligned; diff --git a/coderd/database/migrations/000570_usage_events_agent_runtime_invariants.up.sql b/coderd/database/migrations/000570_usage_events_agent_runtime_invariants.up.sql new file mode 100644 index 00000000000..7577b386f02 --- /dev/null +++ b/coderd/database/migrations/000570_usage_events_agent_runtime_invariants.up.sql @@ -0,0 +1,24 @@ +-- The usage generator writes hb_agent_runtime_v1 rows with created_at at +-- the UTC hourly bucket start and exactly one row per bucket. Uniqueness +-- keeps any consumer that sums runtime_ms from counting a bucket twice; +-- the alignment CHECK protects the attribution model, which charges a +-- bucket to the usage period containing its start. +-- +-- Both statements validate existing rows. Every supported writer has always +-- produced conforming data, so a pre-existing violator is anomalous and +-- failing the migration loudly beats silently rewriting usage rows. +ALTER TABLE usage_events + ADD CONSTRAINT usage_events_agent_runtime_hour_aligned + CHECK ( + event_type <> 'hb_agent_runtime_v1' + OR date_trunc('hour', (created_at AT TIME ZONE 'UTC')) = (created_at AT TIME ZONE 'UTC') + ); + +-- Inserts keep their (id) arbiter: re-inserting a bucket under its +-- deterministic id stays a silent no-op, while a duplicate bucket row under +-- a different id raises a unique violation (generateBucket in +-- enterprise/coderd/usage/generator.go handles it). +DROP INDEX idx_usage_events_agent_runtime; +CREATE UNIQUE INDEX idx_usage_events_agent_runtime + ON usage_events (event_type, created_at) + WHERE event_type = 'hb_agent_runtime_v1'; diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 44fc384b619..39d27c4475f 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -10911,11 +10911,10 @@ func TestUsageEventsTrigger(t *testing.T) { insert("hb_agent_runtime_v1:2025-01-02_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 250}`, day2) requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`) - // Re-inserting a bucket must not double-count it. The daily rollup - // sums runtime_ms, so idempotency rests on the aggregate trigger - // being AFTER INSERT: Postgres does not fire it for rows suppressed - // by ON CONFLICT (id) DO NOTHING. Concurrent replicas and backfill - // re-runs both take this path. + // Re-inserting a bucket under its deterministic id must not + // double-count it: the daily rollup's AFTER INSERT trigger does not + // fire for rows suppressed by the insert's ON CONFLICT (id) + // arbiter. insert("hb_agent_runtime_v1:2025-01-01_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 1000}`, day1) requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`) @@ -10923,6 +10922,38 @@ func TestUsageEventsTrigger(t *testing.T) { insert("hb-seats-1", "hb_ai_seats_v1", `{"count": 3}`, day2) rows := getDailyRows(ctx, sqlDB) require.Len(t, rows, 3) + + // The same bucket under a different id is not an idempotent + // re-insert but a duplicate that would double any aggregate summing + // runtime_ms; the unique partial index + // idx_usage_events_agent_runtime rejects it loudly instead of the + // (id) arbiter silently dropping it. + err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: "different-id-same-bucket", + EventType: "hb_agent_runtime_v1", + EventData: []byte(`{"runtime_ms": 9999}`), + CreatedAt: day1, + }) + require.True(t, database.IsUniqueViolation(err, database.UniqueIndexUsageEventsAgentRuntime), + "expected unique violation on idx_usage_events_agent_runtime, got %v", err) + // The rejected row must not have reached the daily rollup either. + rows = getDailyRows(ctx, sqlDB) + require.Len(t, rows, 3) + require.JSONEq(t, `{"runtime_ms": 1500}`, string(rows[0].UsageData)) + + // created_at must be the exact UTC hourly bucket start; + // usage_events_agent_runtime_hour_aligned rejects a misaligned row + // so it cannot skew the period a bucket is attributed to. + err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: "hb_agent_runtime_v1:misaligned", + EventType: "hb_agent_runtime_v1", + EventData: []byte(`{"runtime_ms": 100}`), + CreatedAt: day1.Add(30 * time.Minute), + }) + require.ErrorContains(t, err, string(database.CheckUsageEventsAgentRuntimeHourAligned)) + rows = getDailyRows(ctx, sqlDB) + require.Len(t, rows, 3) + require.JSONEq(t, `{"runtime_ms": 1500}`, string(rows[0].UsageData)) }) t.Run("UnknownEventType", func(t *testing.T) { diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 4b1a4376f2d..23256ed3b86 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -160,6 +160,7 @@ const ( UniqueIndexProvisionerDaemonsOrgNameOwnerKey UniqueConstraint = "idx_provisioner_daemons_org_name_owner_key" // CREATE UNIQUE INDEX idx_provisioner_daemons_org_name_owner_key ON provisioner_daemons USING btree (organization_id, name, lower(COALESCE((tags ->> 'owner'::text), ''::text))); UniqueIndexTemplateVersionPresetsDefault UniqueConstraint = "idx_template_version_presets_default" // CREATE UNIQUE INDEX idx_template_version_presets_default ON template_version_presets USING btree (template_version_id) WHERE (is_default = true); UniqueIndexUniquePresetName UniqueConstraint = "idx_unique_preset_name" // CREATE UNIQUE INDEX idx_unique_preset_name ON template_version_presets USING btree (name, template_version_id); + UniqueIndexUsageEventsAgentRuntime UniqueConstraint = "idx_usage_events_agent_runtime" // CREATE UNIQUE INDEX idx_usage_events_agent_runtime ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_agent_runtime_v1'::text); UniqueIndexUsersEmail UniqueConstraint = "idx_users_email" // CREATE UNIQUE INDEX idx_users_email ON users USING btree (email) WHERE ((deleted = false) AND (email <> ''::text)); UniqueIndexUsersUsername UniqueConstraint = "idx_users_username" // CREATE UNIQUE INDEX idx_users_username ON users USING btree (username) WHERE (deleted = false); UniqueNotificationMessagesDedupeHashIndex UniqueConstraint = "notification_messages_dedupe_hash_idx" // CREATE UNIQUE INDEX notification_messages_dedupe_hash_idx ON notification_messages USING btree (dedupe_hash); diff --git a/enterprise/coderd/usage/generator.go b/enterprise/coderd/usage/generator.go index 91d911b57a7..74e8ae9ac19 100644 --- a/enterprise/coderd/usage/generator.go +++ b/enterprise/coderd/usage/generator.go @@ -48,8 +48,14 @@ const ( // Cron jobs, which sample live state when they fire, the Generator derives // events from data already persisted in the database, so it can // deterministically backfill hours missed while the deployment was down, -// zero-filling idle hours. Deterministic event IDs plus the database's -// ON CONFLICT (id) DO NOTHING make concurrent replicas safe without locking. +// zero-filling idle hours. Deterministic event IDs make concurrent replicas +// safe without locking: the insert's ON CONFLICT (id) arbiter turns a +// re-insert of a bucket into a no-op, even when the competing insert is +// still in flight (once its arbiter index entry is visible, PostgreSQL +// waits on that transaction and takes the DO NOTHING path if it commits). +// Only the narrow speculative-insertion race, before the competing row's +// arbiter entry exists, surfaces a bucket unique violation instead, which +// generateBucket recognizes as the other replica winning. // // Events are generated unconditionally in enterprise builds; the // publish_usage_data license flag only gates publishing to Tallyman. @@ -157,7 +163,7 @@ func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { // A row marks its bucket complete regardless of publish outcome, so a // bucket whose event Tallyman permanently rejected is never // regenerated (re-inserting under the deterministic ID is a no-op via - // ON CONFLICT (id) DO NOTHING). + // the insert's ON CONFLICT (id) arbiter). // // The runtime is not lost locally: the row still holds it, and the // event can be re-queued for publishing with @@ -235,6 +241,12 @@ func (g *Generator) generateBucket(ctx context.Context, bucket time.Time) error // time) so daily rollups attribute backfilled hours to the correct day. stableID := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(usageEventIDTimeFormat) err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, stableID, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}) + if database.IsUniqueViolation(err, database.UniqueIndexUsageEventsAgentRuntime) { + // Another replica already created this bucket's row. The Generator + // doc comment explains why this race reaches the bucket unique + // index instead of the insert's ON CONFLICT (id) arbiter. + return nil + } if err != nil { return xerrors.Errorf("insert usage event: %w", err) } diff --git a/enterprise/coderd/usage/generator_internal_test.go b/enterprise/coderd/usage/generator_internal_test.go new file mode 100644 index 00000000000..fa6aeb63440 --- /dev/null +++ b/enterprise/coderd/usage/generator_internal_test.go @@ -0,0 +1,43 @@ +package usage + +import ( + "testing" + "time" + + "github.com/lib/pq" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +// TestGenerateBucketUniqueViolation pins that a unique violation on the +// bucket index resolves the bucket as complete: another writer already +// recorded it. TestGeneratorConcurrentReplicas also reaches this path, but +// only when its goroutines actually interleave; this case cannot pass by +// scheduling accident. +func TestGenerateBucketUniqueViolation(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + mDB := dbmock.NewMockStore(ctrl) + gen := NewGenerator(quartz.NewMock(t), slogtest.Make(t, nil), mDB, NewDBInserter()) + + mDB.EXPECT(). + GetTotalChatMessageRuntimeMsInRange(gomock.Any(), gomock.Any()). + Return(int64(1000), nil) + mDB.EXPECT(). + InsertUsageEvent(gomock.Any(), gomock.Any()). + Return(&pq.Error{ + Code: "23505", // unique_violation + Constraint: string(database.UniqueIndexUsageEventsAgentRuntime), + }) + + bucket := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + require.NoError(t, gen.generateBucket(ctx, bucket)) +}