diff --git a/coderd/coderdtest/usage.go b/coderd/coderdtest/usage.go index c7139128670..7af644fe015 100644 --- a/coderd/coderdtest/usage.go +++ b/coderd/coderdtest/usage.go @@ -3,6 +3,7 @@ package coderdtest import ( "context" "sync" + "time" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/usage" @@ -11,10 +12,16 @@ import ( var _ usage.Inserter = (*UsageInserter)(nil) +type HeartbeatEvent struct { + ID string + CreatedAt time.Time + Event usagetypes.HeartbeatEvent +} + type UsageInserter struct { sync.Mutex discreteEvents []usagetypes.DiscreteEvent - heartbeatEvents []usagetypes.HeartbeatEvent + heartbeatEvents []HeartbeatEvent seenHeartbeats map[string]struct{} } @@ -22,7 +29,7 @@ func NewUsageInserter() *UsageInserter { return &UsageInserter{ discreteEvents: []usagetypes.DiscreteEvent{}, seenHeartbeats: map[string]struct{}{}, - heartbeatEvents: []usagetypes.HeartbeatEvent{}, + heartbeatEvents: []HeartbeatEvent{}, } } @@ -33,7 +40,7 @@ func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.S return nil } -func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, event usagetypes.HeartbeatEvent) error { +func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error { u.Lock() defer u.Unlock() if _, seen := u.seenHeartbeats[id]; seen { @@ -41,14 +48,18 @@ func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database. } u.seenHeartbeats[id] = struct{}{} - u.heartbeatEvents = append(u.heartbeatEvents, event) + u.heartbeatEvents = append(u.heartbeatEvents, HeartbeatEvent{ + ID: id, + CreatedAt: createdAt, + Event: event, + }) return nil } -func (u *UsageInserter) GetHeartbeatEvents() []usagetypes.HeartbeatEvent { +func (u *UsageInserter) GetHeartbeatEvents() []HeartbeatEvent { u.Lock() defer u.Unlock() - eventsCopy := make([]usagetypes.HeartbeatEvent, len(u.heartbeatEvents)) + eventsCopy := make([]HeartbeatEvent, len(u.heartbeatEvents)) copy(eventsCopy, u.heartbeatEvents) return eventsCopy } @@ -72,5 +83,5 @@ func (u *UsageInserter) Reset() { defer u.Unlock() u.seenHeartbeats = map[string]struct{}{} u.discreteEvents = []usagetypes.DiscreteEvent{} - u.heartbeatEvents = []usagetypes.HeartbeatEvent{} + u.heartbeatEvents = []HeartbeatEvent{} } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index de153b24c2e..3b51ccceb65 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -671,9 +671,9 @@ var ( Site: rbac.Permissions(map[string][]policy.Action{ rbac.ResourceLicense.Type: {policy.ActionRead}, rbac.ResourceAiSeat.Type: {policy.ActionRead}, // Required for GetActiveAISeatCount. - // The usage publisher doesn't create events, just - // reads/processes them. - rbac.ResourceUsageEvent.Type: {policy.ActionRead, policy.ActionUpdate}, + // Create is required to insert heartbeat usage events + // under this subject. + rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate}, }), User: []rbac.Permission{}, ByOrgID: map[string]rbac.OrgPermissions{}, @@ -4992,6 +4992,16 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe return q.db.GetAuthorizedTemplates(ctx, arg, prep) } +func (q *querier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + // This query exists solely to compute hb_agent_runtime_v1 usage event + // payloads and returns a bare sum with no chat content, so it is gated + // on usage event creation rather than on reading chats. + if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceUsageEvent); err != nil { + return 0, err + } + return q.db.GetTotalChatMessageRuntimeMsInRange(ctx, arg) +} + func (q *querier) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil { return 0, err @@ -6901,6 +6911,13 @@ func (q *querier) ListTasks(ctx context.Context, arg database.ListTasksParams) ( return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ListTasks)(ctx, arg) } +func (q *querier) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil { + return nil, err + } + return q.db.ListUsageEventCreatedAtsByTypeSince(ctx, arg) +} + func (q *querier) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) { u, err := q.db.GetUserByID(ctx, userID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 07278858db4..40bf232df43 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6717,6 +6717,27 @@ func (s *MethodTestSuite) TestUsageEvents() { EndDate: time.Time{}, }).Asserts(rbac.ResourceUsageEvent, policy.ActionRead) })) + + s.Run("ListUsageEventCreatedAtsByTypeSince", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.ListUsageEventCreatedAtsByTypeSinceParams{ + EventType: "hb_agent_runtime_v1", + Since: dbtime.Now(), + } + db.EXPECT().ListUsageEventCreatedAtsByTypeSince(gomock.Any(), params).Return([]time.Time{}, nil) + check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionRead) + })) + + // GetTotalChatMessageRuntimeMsInRange exists solely to compute usage + // event payloads, so it asserts usage event creation rather than chat + // read permissions. + s.Run("GetTotalChatMessageRuntimeMsInRange", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: time.Time{}, + EndTime: time.Time{}, + } + db.EXPECT().GetTotalChatMessageRuntimeMsInRange(gomock.Any(), params).Return(int64(0), nil) + check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionCreate) + })) } // Ensures that the prebuilds actor may never insert an api key. @@ -6735,6 +6756,41 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) { require.True(t, dbauthz.IsNotAuthorizedError(err)) } +// TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied mechanically +// checks the invariant the query's authz gate relies on: it exposes a +// deployment-wide aggregate behind usage_event create at site scope, which no +// human-assignable role holds. Owner is excluded from usage_event via +// allPermsExcept in roles.go; org roles such as org-admin do carry +// usage_event permissions, but only at org scope, which cannot satisfy a +// site-scoped check. If either of those ever changes, this test fails. +func TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied(t *testing.T) { + t.Parallel() + + orgID := uuid.New() + var roles []rbac.RoleIdentifier + for _, role := range rbac.SiteBuiltInRoles() { + roles = append(roles, role.Identifier) + } + for _, role := range rbac.OrganizationRoles(orgID) { + roles = append(roles, role.Identifier) + } + require.NotEmpty(t, roles) + + for _, role := range roles { + subj := rbac.Subject{ + ID: uuid.NewString(), + Roles: rbac.RoleIdentifiers{role}, + Scope: rbac.ScopeAll, + } + ctx := dbauthz.As(testutil.Context(t, testutil.WaitShort), subj) + mDB := dbmock.NewMockStore(gomock.NewController(t)) + mDB.EXPECT().Wrappers().Times(1).Return([]string{}) + dbz := dbauthz.New(mDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) + _, err := dbz.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{}) + require.True(t, dbauthz.IsNotAuthorizedError(err), "role %s must be denied", role) + } +} + func (s *MethodTestSuite) TestAIBridge() { s.Run("InsertAIBridgeInterception", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { initID := uuid.UUID{3} diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 26f47ff51dd..a47adc4aae7 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3177,6 +3177,14 @@ func (m queryMetricsStore) GetTemplatesWithFilter(ctx context.Context, arg datab return r0, r1 } +func (m queryMetricsStore) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + start := time.Now() + r0, r1 := m.s.GetTotalChatMessageRuntimeMsInRange(ctx, arg) + m.queryLatencies.WithLabelValues("GetTotalChatMessageRuntimeMsInRange").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTotalChatMessageRuntimeMsInRange").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) { start := time.Now() r0, r1 := m.s.GetTotalUsageDCManagedAgentsV1(ctx, arg) @@ -4857,6 +4865,14 @@ func (m queryMetricsStore) ListTasks(ctx context.Context, arg database.ListTasks return r0, r1 } +func (m queryMetricsStore) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + start := time.Now() + r0, r1 := m.s.ListUsageEventCreatedAtsByTypeSince(ctx, arg) + m.queryLatencies.WithLabelValues("ListUsageEventCreatedAtsByTypeSince").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListUsageEventCreatedAtsByTypeSince").Inc() + return r0, r1 +} + func (m queryMetricsStore) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) { start := time.Now() r0, r1 := m.s.ListUserChatCompactionThresholds(ctx, userID) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index ac848e0f993..c63d297527b 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -5938,6 +5938,21 @@ func (mr *MockStoreMockRecorder) GetTemplatesWithFilter(ctx, arg any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplatesWithFilter", reflect.TypeOf((*MockStore)(nil).GetTemplatesWithFilter), ctx, arg) } +// GetTotalChatMessageRuntimeMsInRange mocks base method. +func (m *MockStore) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetTotalChatMessageRuntimeMsInRange", ctx, arg) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetTotalChatMessageRuntimeMsInRange indicates an expected call of GetTotalChatMessageRuntimeMsInRange. +func (mr *MockStoreMockRecorder) GetTotalChatMessageRuntimeMsInRange(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTotalChatMessageRuntimeMsInRange", reflect.TypeOf((*MockStore)(nil).GetTotalChatMessageRuntimeMsInRange), ctx, arg) +} + // GetTotalUsageDCManagedAgentsV1 mocks base method. func (m *MockStore) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) { m.ctrl.T.Helper() @@ -9146,6 +9161,21 @@ func (mr *MockStoreMockRecorder) ListTasks(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListTasks", reflect.TypeOf((*MockStore)(nil).ListTasks), ctx, arg) } +// ListUsageEventCreatedAtsByTypeSince mocks base method. +func (m *MockStore) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListUsageEventCreatedAtsByTypeSince", ctx, arg) + ret0, _ := ret[0].([]time.Time) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListUsageEventCreatedAtsByTypeSince indicates an expected call of ListUsageEventCreatedAtsByTypeSince. +func (mr *MockStoreMockRecorder) ListUsageEventCreatedAtsByTypeSince(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListUsageEventCreatedAtsByTypeSince", reflect.TypeOf((*MockStore)(nil).ListUsageEventCreatedAtsByTypeSince), ctx, arg) +} + // ListUserChatCompactionThresholds mocks base method. func (m *MockStore) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index fb7ab0b8e40..b63f1193958 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -728,7 +728,7 @@ CREATE FUNCTION aggregate_usage_event() RETURNS trigger AS $$ BEGIN -- Check for supported event types and throw error for unknown types. - IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1') THEN + IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1') THEN RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type; END IF; @@ -756,6 +756,13 @@ BEGIN COALESCE((NEW.event_data->>'count')::bigint, 0) ) ) + -- Hourly runtime heartbeats: sum the runtime per day. + WHEN NEW.event_type IN ('hb_agent_runtime_v1') THEN + jsonb_build_object( + 'runtime_ms', + COALESCE((usage_events_daily.usage_data->>'runtime_ms')::bigint, 0) + + COALESCE((NEW.event_data->>'runtime_ms')::bigint, 0) + ) END; RETURN NEW; @@ -3529,7 +3536,7 @@ 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]))) + 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]))) ); 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.'; @@ -3540,6 +3547,8 @@ COMMENT ON COLUMN usage_events.event_type IS 'The usage event type with version. COMMENT ON COLUMN usage_events.event_data IS 'Event payload. Determined by the matching usage struct for this event type.'; +COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; + COMMENT ON COLUMN usage_events.publish_started_at IS 'Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired.'; COMMENT ON COLUMN usage_events.published_at IS 'Set to a timestamp when the event is successfully (or permanently unsuccessfully) published to the usage collector service. If set, the event should never be attempted to be published again.'; @@ -4870,6 +4879,8 @@ 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 INDEX idx_usage_events_ai_seats ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_ai_seats_v1'::text); CREATE INDEX idx_usage_events_select_for_publishing ON usage_events USING btree (published_at, publish_started_at, created_at); diff --git a/coderd/database/migrations/000561_usage_events_agent_runtime.down.sql b/coderd/database/migrations/000561_usage_events_agent_runtime.down.sql new file mode 100644 index 00000000000..8c39020c1b7 --- /dev/null +++ b/coderd/database/migrations/000561_usage_events_agent_runtime.down.sql @@ -0,0 +1,50 @@ +COMMENT ON COLUMN usage_events.created_at IS NULL; + +DROP INDEX IF EXISTS idx_usage_events_agent_runtime; + +-- Remove hb_agent_runtime_v1 rows so the previous constraint can be restored. +DELETE FROM usage_events WHERE event_type = 'hb_agent_runtime_v1'; +DELETE FROM usage_events_daily WHERE event_type = 'hb_agent_runtime_v1'; + +ALTER TABLE usage_events + DROP CONSTRAINT usage_event_type_check, + ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1')); + +-- Restores the 000444 version of the function. +CREATE OR REPLACE FUNCTION aggregate_usage_event() +RETURNS TRIGGER AS $$ +BEGIN + -- Check for supported event types and throw error for unknown types. + IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1') THEN + RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type; + END IF; + + INSERT INTO usage_events_daily (day, event_type, usage_data) + VALUES ( + date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date, + NEW.event_type, + NEW.event_data + ) + ON CONFLICT (day, event_type) DO UPDATE SET + usage_data = CASE + -- Handle simple counter events by summing the count. + WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN + jsonb_build_object( + 'count', + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) + + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + -- Heartbeat events: keep the max value seen that day + WHEN NEW.event_type IN ('hb_ai_seats_v1') THEN + jsonb_build_object( + 'count', + GREATEST( + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0), + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + ) + END; + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; diff --git a/coderd/database/migrations/000561_usage_events_agent_runtime.up.sql b/coderd/database/migrations/000561_usage_events_agent_runtime.up.sql new file mode 100644 index 00000000000..1adb5d821a3 --- /dev/null +++ b/coderd/database/migrations/000561_usage_events_agent_runtime.up.sql @@ -0,0 +1,55 @@ +ALTER TABLE usage_events + DROP CONSTRAINT usage_event_type_check, + ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1')); + +COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.'; + +-- Serves the usage generator's scan for missing hourly buckets. +CREATE INDEX idx_usage_events_agent_runtime + ON usage_events (event_type, created_at) + WHERE event_type = 'hb_agent_runtime_v1'; + +CREATE OR REPLACE FUNCTION aggregate_usage_event() +RETURNS TRIGGER AS $$ +BEGIN + -- Check for supported event types and throw error for unknown types. + IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1') THEN + RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type; + END IF; + + INSERT INTO usage_events_daily (day, event_type, usage_data) + VALUES ( + date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date, + NEW.event_type, + NEW.event_data + ) + ON CONFLICT (day, event_type) DO UPDATE SET + usage_data = CASE + -- Handle simple counter events by summing the count. + WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN + jsonb_build_object( + 'count', + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) + + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + -- Heartbeat events: keep the max value seen that day + WHEN NEW.event_type IN ('hb_ai_seats_v1') THEN + jsonb_build_object( + 'count', + GREATEST( + COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0), + COALESCE((NEW.event_data->>'count')::bigint, 0) + ) + ) + -- Hourly runtime heartbeats: sum the runtime per day. + WHEN NEW.event_type IN ('hb_agent_runtime_v1') THEN + jsonb_build_object( + 'runtime_ms', + COALESCE((usage_events_daily.usage_data->>'runtime_ms')::bigint, 0) + + COALESCE((NEW.event_data->>'runtime_ms')::bigint, 0) + ) + END; + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; diff --git a/coderd/database/migrations/testdata/fixtures/000561_usage_events_agent_runtime.up.sql b/coderd/database/migrations/testdata/fixtures/000561_usage_events_agent_runtime.up.sql new file mode 100644 index 00000000000..71a42bd1a0a --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000561_usage_events_agent_runtime.up.sql @@ -0,0 +1,20 @@ +INSERT INTO usage_events ( + id, + event_type, + event_data, + created_at, + publish_started_at, + published_at, + failure_message +) +VALUES +-- Unpublished hb_agent_runtime_v1 event. +( + 'hb_agent_runtime_v1:2023-06-01_00:00:00', + 'hb_agent_runtime_v1', + '{"runtime_ms":3600000}', + '2023-06-01 00:00:00+00', + NULL, + NULL, + NULL +); diff --git a/coderd/database/models.go b/coderd/database/models.go index b3c598a2eda..4a979e94c8b 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -6142,7 +6142,8 @@ type UsageEvent struct { EventType string `db:"event_type" json:"event_type"` // Event payload. Determined by the matching usage struct for this event type. EventData json.RawMessage `db:"event_data" json:"event_data"` - CreatedAt time.Time `db:"created_at" json:"created_at"` + // The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp. + CreatedAt time.Time `db:"created_at" json:"created_at"` // Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired. PublishStartedAt sql.NullTime `db:"publish_started_at" json:"publish_started_at"` // Set to a timestamp when the event is successfully (or permanently unsuccessfully) published to the usage collector service. If set, the event should never be attempted to be published again. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index eb46ff6a8f8..7ed0f7bc02b 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -867,6 +867,9 @@ type sqlcQuerier interface { GetTemplateVersionsCreatedAfter(ctx context.Context, createdAt time.Time) ([]TemplateVersion, error) GetTemplates(ctx context.Context) ([]Template, error) GetTemplatesWithFilter(ctx context.Context, arg GetTemplatesWithFilterParams) ([]Template, error) + // Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes + // soft-deleted messages and messages from all chats. + GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) // Gets the total number of managed agents created between two dates. Uses the // aggregate table to avoid large scans or a complex index on the usage_events // table. @@ -1257,6 +1260,8 @@ type sqlcQuerier interface { ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListProvisionerKeysByOrganizationExcludeReserved(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListTasks(ctx context.Context, arg ListTasksParams) ([]Task, error) + // Used by the usage generator to find missing heartbeat buckets. + ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]UserConfig, error) ListUserChatPersonalModelOverrides(ctx context.Context, userID uuid.UUID) ([]ListUserChatPersonalModelOverridesRow, error) // Returns metadata only (no value or value_key_id) for the diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index a4e08e460b8..457512f1bde 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -10717,6 +10717,62 @@ func TestUsageEventsTrigger(t *testing.T) { require.Len(t, rows, 3) }) + t.Run("HeartbeatAgentRuntime", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + + insert := func(id, eventType, eventData string, createdAt time.Time) { + t.Helper() + err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: id, + EventType: eventType, + EventData: []byte(eventData), + CreatedAt: createdAt, + }) + require.NoError(t, err) + } + requireDaily := func(wantUsageData ...string) { + t.Helper() + rows := getDailyRows(ctx, sqlDB) + require.Len(t, rows, len(wantUsageData)) + for i, want := range wantUsageData { + require.JSONEq(t, want, string(rows[i].UsageData)) + } + } + + day1 := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + day2 := time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC) + + insert("hb_agent_runtime_v1:2025-01-01_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 1000}`, day1) + requireDaily(`{"runtime_ms": 1000}`) + + // Unlike hb_ai_seats_v1, hourly runtime events are summed per day. + insert("hb_agent_runtime_v1:2025-01-01_12:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 500}`, day1.Add(12*time.Hour)) + requireDaily(`{"runtime_ms": 1500}`) + + // Zero-valued events (idle hours) do not change the sum. + insert("hb_agent_runtime_v1:2025-01-01_18:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 0}`, day1.Add(18*time.Hour)) + requireDaily(`{"runtime_ms": 1500}`) + + 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. + 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}`) + + // A different event type on the same day gets its own daily row. + insert("hb-seats-1", "hb_ai_seats_v1", `{"count": 3}`, day2) + rows := getDailyRows(ctx, sqlDB) + require.Len(t, rows, 3) + }) + t.Run("UnknownEventType", func(t *testing.T) { t.Parallel() @@ -10750,6 +10806,117 @@ func TestUsageEventsTrigger(t *testing.T) { }) } +func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + + rangeStart := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + rangeEnd := rangeStart.Add(time.Hour) + + total, err := db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: rangeStart, + EndTime: rangeEnd, + }) + require.NoError(t, err) + require.EqualValues(t, 0, total) + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", + }) + mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "test-model", + ContextLimit: 8192, + }) + chat1 := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + chat2 := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + + insertMessage := func(chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) { + t.Helper() + msg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: mc.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + RuntimeMs: sql.NullInt64{Int64: runtimeMs, Valid: true}, + }) + _, err := sqlDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1, deleted = $2 WHERE id = $3", createdAt, deleted, msg.ID) + require.NoError(t, err) + } + + // Counted: on the inclusive start boundary, in the middle (across two + // chats), soft-deleted, and just before the exclusive end boundary. + insertMessage(chat1.ID, 1, rangeStart, false) + insertMessage(chat2.ID, 2, rangeStart.Add(30*time.Minute), false) + insertMessage(chat1.ID, 4, rangeStart.Add(45*time.Minute), true) + insertMessage(chat1.ID, 8, rangeEnd.Add(-time.Second), false) + // Not counted: before the range, on the exclusive end boundary, and a + // NULL runtime (runtime 0 is stored as NULL). + insertMessage(chat1.ID, 16, rangeStart.Add(-time.Second), false) + insertMessage(chat1.ID, 32, rangeEnd, false) + insertMessage(chat1.ID, 0, rangeStart.Add(10*time.Minute), false) + + total, err = db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: rangeStart, + EndTime: rangeEnd, + }) + require.NoError(t, err) + require.EqualValues(t, 15, total) +} + +func TestListUsageEventCreatedAtsByTypeSince(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, _ := dbtestutil.NewDB(t) + + since := time.Date(2025, 3, 10, 0, 0, 0, 0, time.UTC) + + insertEvent := func(id, eventType string, eventData string, createdAt time.Time) { + t.Helper() + err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: id, + EventType: eventType, + EventData: []byte(eventData), + CreatedAt: createdAt, + }) + require.NoError(t, err) + } + + // Matching type: one before since (excluded), one exactly at since + // (included), one after (included). + insertEvent("rt-old", "hb_agent_runtime_v1", `{"runtime_ms": 1}`, since.Add(-time.Hour)) + insertEvent("rt-at", "hb_agent_runtime_v1", `{"runtime_ms": 2}`, since) + insertEvent("rt-new", "hb_agent_runtime_v1", `{"runtime_ms": 3}`, since.Add(time.Hour)) + // Different type after since: excluded. + insertEvent("seats-new", "hb_ai_seats_v1", `{"count": 1}`, since.Add(time.Hour)) + + createdAts, err := db.ListUsageEventCreatedAtsByTypeSince(ctx, database.ListUsageEventCreatedAtsByTypeSinceParams{ + EventType: "hb_agent_runtime_v1", + Since: since, + }) + require.NoError(t, err) + require.Len(t, createdAts, 2) + normalized := make([]time.Time, len(createdAts)) + for i, ts := range createdAts { + normalized[i] = ts.UTC() + } + require.ElementsMatch(t, []time.Time{since, since.Add(time.Hour)}, normalized) +} + func TestListTasks(t *testing.T) { t.Parallel() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 3fb67221f0f..e9df5bb7977 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -10090,6 +10090,28 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time return items, nil } +const getTotalChatMessageRuntimeMsInRange = `-- name: GetTotalChatMessageRuntimeMsInRange :one +SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms +FROM chat_messages cm +WHERE cm.created_at >= $1::timestamptz + AND cm.created_at < $2::timestamptz + AND cm.runtime_ms IS NOT NULL +` + +type GetTotalChatMessageRuntimeMsInRangeParams struct { + StartTime time.Time `db:"start_time" json:"start_time"` + EndTime time.Time `db:"end_time" json:"end_time"` +} + +// Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes +// soft-deleted messages and messages from all chats. +func (q *sqlQuerier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) { + row := q.db.QueryRowContext(ctx, getTotalChatMessageRuntimeMsInRange, arg.StartTime, arg.EndTime) + var total_runtime_ms int64 + err := row.Scan(&total_runtime_ms) + return total_runtime_ms, err +} + const getUserChatSpendInPeriod = `-- name: GetUserChatSpendInPeriod :one SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros FROM chat_messages cm @@ -29088,6 +29110,42 @@ func (q *sqlQuerier) InsertUsageEvent(ctx context.Context, arg InsertUsageEventP return err } +const listUsageEventCreatedAtsByTypeSince = `-- name: ListUsageEventCreatedAtsByTypeSince :many +SELECT created_at +FROM usage_events +WHERE event_type = $1 + AND created_at >= $2::timestamptz +` + +type ListUsageEventCreatedAtsByTypeSinceParams struct { + EventType string `db:"event_type" json:"event_type"` + Since time.Time `db:"since" json:"since"` +} + +// Used by the usage generator to find missing heartbeat buckets. +func (q *sqlQuerier) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) { + rows, err := q.db.QueryContext(ctx, listUsageEventCreatedAtsByTypeSince, arg.EventType, arg.Since) + if err != nil { + return nil, err + } + defer rows.Close() + var items []time.Time + for rows.Next() { + var created_at time.Time + if err := rows.Scan(&created_at); err != nil { + return nil, err + } + items = append(items, created_at) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const selectUsageEventsForPublishing = `-- name: SelectUsageEventsForPublishing :many WITH usage_events AS ( UPDATE diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index da5d96f856c..9c3ed9e65a5 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -2483,6 +2483,15 @@ WHERE c.owner_id = @user_id::uuid AND cm.created_at < @end_time::timestamptz AND cm.total_cost_micros IS NOT NULL; +-- name: GetTotalChatMessageRuntimeMsInRange :one +-- Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes +-- soft-deleted messages and messages from all chats. +SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms +FROM chat_messages cm +WHERE cm.created_at >= @start_time::timestamptz + AND cm.created_at < @end_time::timestamptz + AND cm.runtime_ms IS NOT NULL; + -- name: CountEnabledModelsWithoutPricing :one -- Counts enabled, non-deleted model configs that lack both input and -- output pricing in their JSONB options.cost configuration. diff --git a/coderd/database/queries/usageevents.sql b/coderd/database/queries/usageevents.sql index 7ffcb1173b5..8ba706b0c8a 100644 --- a/coderd/database/queries/usageevents.sql +++ b/coderd/database/queries/usageevents.sql @@ -20,6 +20,13 @@ SELECT EXISTS( SELECT 1 FROM usage_events WHERE id = @id )::bool; +-- name: ListUsageEventCreatedAtsByTypeSince :many +-- Used by the usage generator to find missing heartbeat buckets. +SELECT created_at +FROM usage_events +WHERE event_type = @event_type + AND created_at >= @since::timestamptz; + -- name: SelectUsageEventsForPublishing :many WITH usage_events AS ( UPDATE diff --git a/coderd/pproflabel/pproflabel.go b/coderd/pproflabel/pproflabel.go index 5204a0681df..dcc17c5c0cc 100644 --- a/coderd/pproflabel/pproflabel.go +++ b/coderd/pproflabel/pproflabel.go @@ -38,6 +38,7 @@ const ( // ServiceWorkspaceBuildOrchestrator fulfills workspace build // orchestrations once their parent build reaches a terminal state. ServiceWorkspaceBuildOrchestrator = "workspace-build-orchestrator" + ServiceUsageEventGenerator = "usage-event-generator" RequestTypeTag = "coder_request_type" ) diff --git a/coderd/usage/inserter.go b/coderd/usage/inserter.go index 891f5c73873..acaaf4bbeb7 100644 --- a/coderd/usage/inserter.go +++ b/coderd/usage/inserter.go @@ -2,6 +2,7 @@ package usage import ( "context" + "time" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/usage/usagetypes" @@ -28,7 +29,11 @@ type Inserter interface { // // Inserts with the same `id` must be idempotent. The database enforces this by // ignoring duplicate records. - InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error + // + // Generators that backfill historical buckets pass the bucket start as + // `createdAt` rather than the insertion time. `createdAt` must be + // non-zero; implementations reject the zero time rather than storing it. + InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error } // AGPLInserter is a no-op implementation of Inserter. @@ -48,6 +53,6 @@ func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store // InsertHeartbeatUsageEvent is a no-op implementation of // InsertHeartbeatUsageEvent. -func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ usagetypes.HeartbeatEvent) error { +func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ time.Time, _ usagetypes.HeartbeatEvent) error { return nil } diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index 6c8fde416eb..9004fb2514c 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -30,6 +30,7 @@ type UsageEventType string const ( UsageEventTypeDCManagedAgentsV1 UsageEventType = "dc_managed_agents_v1" UsageEventTypeHBAISeatsV1 UsageEventType = "hb_ai_seats_v1" + UsageEventTypeHBAgentRuntimeV1 UsageEventType = "hb_agent_runtime_v1" ) func (e UsageEventType) Valid() bool { @@ -38,6 +39,8 @@ func (e UsageEventType) Valid() bool { return true case UsageEventTypeHBAISeatsV1: return true + case UsageEventTypeHBAgentRuntimeV1: + return true default: return false } @@ -105,6 +108,12 @@ func ParseEventWithType(eventType UsageEventType, data json.RawMessage) (Event, return nil, err } return event, nil + case UsageEventTypeHBAgentRuntimeV1: + var event HBAgentRuntime + if err := ParseEvent(data, &event); err != nil { + return nil, err + } + return event, nil default: return nil, UnknownEventTypeError{EventType: string(eventType)} } @@ -192,3 +201,38 @@ func (e HBAISeats) Fields() map[string]any { "count": e.Count, } } + +// HBAgentRuntime is the event associated with hb_agent_runtime_v1. RuntimeMs +// is the total agent-loop runtime in milliseconds consumed by Coder Agents +// (chats) in one UTC hour. Each measured step spans model streaming (including +// provider-executed tools) and stream retries, and ends when the model stream +// finishes. Time spent executing local tools between steps, including +// sub-agents that bill their own model calls, is excluded. +// +// This measures the new Coder Agents (the `chats` tables), not the deprecated +// Tasks counted by dc_managed_agents_v1. +type HBAgentRuntime struct { + RuntimeMs int64 `json:"runtime_ms"` +} + +var _ HeartbeatEvent = HBAgentRuntime{} + +func (HBAgentRuntime) usageEvent() {} +func (HBAgentRuntime) heartbeatUsageEvent() {} +func (HBAgentRuntime) EventType() UsageEventType { + return UsageEventTypeHBAgentRuntimeV1 +} + +func (e HBAgentRuntime) Valid() error { + if e.RuntimeMs < 0 { + return xerrors.New("runtime_ms cannot be negative") + } + // The runtime can be 0 (idle hour). + return nil +} + +func (e HBAgentRuntime) Fields() map[string]any { + return map[string]any{ + "runtime_ms": e.RuntimeMs, + } +} diff --git a/coderd/usage/usagetypes/events_test.go b/coderd/usage/usagetypes/events_test.go index fcfd076fc0e..ad017b0e073 100644 --- a/coderd/usage/usagetypes/events_test.go +++ b/coderd/usage/usagetypes/events_test.go @@ -76,4 +76,25 @@ func TestParseEventWithType(t *testing.T) { require.Equal(t, eventType, event.EventType()) require.Equal(t, map[string]any{"count": int64(1)}, event.Fields()) }) + + t.Run("HBAgentRuntimeV1", func(t *testing.T) { + t.Parallel() + + eventType := usagetypes.UsageEventTypeHBAgentRuntimeV1 + event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1234}`)) + require.NoError(t, err) + require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 1234}, event) + require.Equal(t, eventType, event.EventType()) + require.Equal(t, map[string]any{"runtime_ms": int64(1234)}, event.Fields()) + + event, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 0}`)) + require.NoError(t, err) + require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 0}, event) + + _, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": -1}`)) + require.ErrorContains(t, err, "runtime_ms cannot be negative") + + _, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1, "extra": "field"}`)) + require.ErrorContains(t, err, "unmarshal *usagetypes.HBAgentRuntime event") + }) } diff --git a/docs/ai-coder/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index 9ee528f1da1..7483599a827 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -4,6 +4,7 @@ - number of agent workspace builds consumed - number of AI Governance seats consumed +- total Coder Agent runtime consumed per hour, in milliseconds No user-identifiable information or additional metrics are sent to Tallyman. This information is also shared with [Metronome](https://metronome.com), a Stripe product and Coder partner for usage-based billing and reporting. @@ -33,6 +34,14 @@ Coder-Deployment-ID: 8a4e92f1-3b7c-4d5e-9f12-abc123def456 # your deployment ID "count": 1 }, "created_at": "2025-01-15T14:30:00Z" + }, + { + "id": "hb_agent_runtime_v1:2025-01-15_14:00:00", # unique event ID generated by Coder + "event_type": "hb_agent_runtime_v1", # aka. Coder Agent runtime + "event_data": { + "runtime_ms": 3600000 + }, + "created_at": "2025-01-15T14:00:00Z" # start of the hour the runtime was consumed in } ] } diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index ec3bc52d393..6a334d1b8e4 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -154,6 +154,12 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { usageCron.Start(ctx) closers.Add(usageCron) + // Usage generation is deliberately not license-gated; the + // publish_usage_data license flag only gates publishing to Tallyman. + usageGenerator := usage.NewGenerator(quartz.NewReal(), options.Logger.Named("usage-event-generator"), options.Database, *options.UsageInserter.Load()) + usageGenerator.Start(ctx) + closers.Add(usageGenerator) + // In-memory AI Bridge Proxy daemon. The bridge daemon itself is // started unconditionally by AGPL cli/server.go (chatd uses its // in-memory roundtripper regardless of license); only the proxy diff --git a/enterprise/coderd/usage/cron.go b/enterprise/coderd/usage/cron.go index 13ccbb927c4..47b85f6d377 100644 --- a/enterprise/coderd/usage/cron.go +++ b/enterprise/coderd/usage/cron.go @@ -23,7 +23,10 @@ import ( var epoch = time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC) const ( - cronDateFormat = "2006-01-02_15:04:05" + // usageEventIDTimeFormat is the timestamp layout used in every + // deterministic usage event ID, both the cron's boundary IDs and the + // generator's bucket IDs. + usageEventIDTimeFormat = "2006-01-02_15:04:05" ) // HeartbeatFunc generates a heartbeat event and its stable ID. @@ -145,7 +148,7 @@ func (c *Cron) run(ctx context.Context, job CronJob) { // Use the boundary (not wall-clock "now") for the stable ID // so all replicas targeting the same boundary produce the // same key. - stableID := string(job.EventType) + ":" + boundary.UTC().Format(cronDateFormat) + stableID := string(job.EventType) + ":" + boundary.UTC().Format(usageEventIDTimeFormat) // Skip if this bucket was already recorded — avoids running // the potentially expensive heartbeat function for a @@ -184,7 +187,7 @@ func (c *Cron) run(ctx context.Context, job CronJob) { continue } - if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, event); err != nil { + if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, c.clock.Now(), event); err != nil { c.log.Warn(ctx, "cron heartbeat insert failed", slog.F("job", job.Name), slog.Error(err), diff --git a/enterprise/coderd/usage/generator.go b/enterprise/coderd/usage/generator.go new file mode 100644 index 00000000000..91d911b57a7 --- /dev/null +++ b/enterprise/coderd/usage/generator.go @@ -0,0 +1,242 @@ +package usage + +import ( + "context" + "math/rand" + "sync" + "time" + + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/pproflabel" + agplusage "github.com/coder/coder/v2/coderd/usage" + "github.com/coder/coder/v2/coderd/usage/usagetypes" + "github.com/coder/quartz" +) + +const ( + // AgentRuntimeInterval is the bucket size of hb_agent_runtime_v1 events. + AgentRuntimeInterval = time.Hour + // AgentRuntimeWindow is the trailing window scanned for missing buckets. + // Buckets still missing beyond this window (e.g. because the deployment + // was down for longer) are forfeited, which can only ever undercount + // usage. + AgentRuntimeWindow = 7 * 24 * time.Hour + // AgentRuntimeEligibilityLag is how long after a bucket closes before it + // becomes eligible for generation, giving replicas time to commit + // in-flight chat messages with timestamps inside the bucket. This + // assumes chat message inserts commit within the lag of their + // statement-time created_at; a message committing later than the lag + // after its bucket closes lands in an already-sealed bucket and its + // runtime is dropped (undercount-only). + AgentRuntimeEligibilityLag = 5 * time.Minute + // agentRuntimeJitter staggers replicas after each hour boundary so one + // is likely to complete the work before others attempt it. + agentRuntimeJitter = 4 * time.Minute + // agentRuntimeStartupDelay is the floor on the first pass after start, + // giving the deployment time to finish booting before the generator + // competes for database work. Jitter is added on top of it. + agentRuntimeStartupDelay = time.Minute + // generatorTimerName tags the quartz timer so tests can trap it. + generatorTimerName = "agent-runtime-generator" +) + +// Generator reconciles hb_agent_runtime_v1 heartbeat usage events. Unlike +// 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. +// +// Events are generated unconditionally in enterprise builds; the +// publish_usage_data license flag only gates publishing to Tallyman. +type Generator struct { + clock quartz.Clock + log slog.Logger + db database.Store + ins agplusage.Inserter + + cancel context.CancelFunc + wg sync.WaitGroup + startOnce sync.Once +} + +// NewGenerator creates an unstarted Generator. +func NewGenerator(clock quartz.Clock, log slog.Logger, db database.Store, ins agplusage.Inserter) *Generator { + return &Generator{ + clock: clock, + log: log, + db: db, + ins: ins, + } +} + +// Start launches the reconciliation goroutine. Subsequent calls are no-ops; +// a closed Generator cannot be restarted. +func (g *Generator) Start(ctx context.Context) { + g.startOnce.Do(func() { + ctx, g.cancel = context.WithCancel(ctx) + g.wg.Add(1) + pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceUsageEventGenerator), func(ctx context.Context) { + g.run(ctx) + }) + }) +} + +// Close stops the Generator and waits for its goroutine to exit. +// It always returns nil; the error return exists to satisfy io.Closer, as the +// Generator is registered with the server's closer list. +func (g *Generator) Close() error { + if g.cancel != nil { + g.cancel() + } + g.wg.Wait() + return nil +} + +func (g *Generator) run(ctx context.Context) { + //nolint:gocritic // We are a publisher in this function. + ctx = dbauthz.AsUsagePublisher(ctx) + defer g.wg.Done() + + // The random initial delay staggers replicas that start simultaneously. + //nolint:gosec // Jitter does not need cryptographic randomness. + delay := agentRuntimeStartupDelay + time.Duration(rand.Int63n(int64(agentRuntimeJitter))) + for { + timer := g.clock.NewTimer(delay, generatorTimerName) + + select { + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return + case <-timer.C: + } + + err := g.generateAgentRuntimeEvents(ctx) + if ctx.Err() != nil { + return + } + if err != nil { + g.log.Warn(ctx, "generate agent runtime usage events", slog.Error(err)) + } + + // Wake at the next eligibility instant (hour boundary + lag), not + // the next hour boundary. Computing the tick against the + // lag-shifted clock keeps a bucket whose eligibility is still + // pending in this hour (e.g. a pass that ran just before HH:05) + // from waiting a whole extra hour. + _, delay = nextTick(g.clock.Now().Add(-AgentRuntimeEligibilityLag), AgentRuntimeInterval, agentRuntimeJitter) + } +} + +// generateAgentRuntimeEvents inserts one hb_agent_runtime_v1 event per +// missing hourly bucket in the trailing window. Per-bucket errors skip only +// that bucket; the next tick rescans the whole window, so transient failures +// self-heal. +func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error { + now := g.clock.Now().UTC() + // Bucket [H, H+1) becomes eligible at H + interval + lag. + latestEligible := now.Add(-AgentRuntimeInterval - AgentRuntimeEligibilityLag).Truncate(AgentRuntimeInterval) + earliest := now.Truncate(AgentRuntimeInterval).Add(-AgentRuntimeWindow) + if latestEligible.Before(earliest) { + return nil + } + + existingTimes, err := g.db.ListUsageEventCreatedAtsByTypeSince(ctx, database.ListUsageEventCreatedAtsByTypeSinceParams{ + EventType: string(usagetypes.UsageEventTypeHBAgentRuntimeV1), + Since: earliest, + }) + if err != nil { + return xerrors.Errorf("list existing agent runtime events: %w", err) + } + // 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 runtime is not lost locally: the row still holds it, and the + // event can be re-queued for publishing with + // + // UPDATE usage_events + // SET published_at = NULL, publish_started_at = NULL, failure_message = NULL + // WHERE id = 'hb_agent_runtime_v1:'; + // + // That re-arm only has an effect while the bucket is inside the + // publisher's 30-day cutoff: SelectUsageEventsForPublishing also + // filters created_at > now - INTERVAL '30 days', and created_at is the + // bucket start, so past that the UPDATE reports success but the row is + // never picked up again. The release gate (Tallyman must accept this + // event type before coderd ships it) is what keeps permanent + // rejections exceptional. + existing := make(map[time.Time]struct{}, len(existingTimes)) + for _, ts := range existingTimes { + // created_at is always the exact bucket start for this event type; + // truncation just normalizes timezone and precision. + existing[ts.UTC().Truncate(AgentRuntimeInterval)] = struct{}{} + } + + var filled, failed int + for bucket := earliest; !bucket.After(latestEligible); bucket = bucket.Add(AgentRuntimeInterval) { + if _, ok := existing[bucket]; ok { + continue + } + if ctx.Err() != nil { + return ctx.Err() + } + + err := g.generateBucket(ctx, bucket) + if err != nil { + if ctx.Err() != nil { + // A cancel landing mid-query surfaces as a bucket error. + // Report the cancellation instead of logging shutdown as + // a bucket failure. + return ctx.Err() + } + // Skip only the failed bucket so one bad bucket (e.g. an + // invalid runtime sum) cannot stall every later bucket until + // it ages out of the window. + g.log.Warn(ctx, "generate agent runtime usage event for bucket", + slog.F("bucket", bucket), + slog.Error(err), + ) + failed++ + continue + } + filled++ + } + if filled > 0 || failed > 0 { + g.log.Info(ctx, "generated agent runtime usage events", + slog.F("buckets_filled", filled), + slog.F("buckets_failed", failed), + slog.F("window_start", earliest), + slog.F("latest_eligible", latestEligible), + ) + } + return nil +} + +// generateBucket computes and inserts the event for a single hourly bucket. +func (g *Generator) generateBucket(ctx context.Context, bucket time.Time) error { + runtimeMs, err := g.db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ + StartTime: bucket, + EndTime: bucket.Add(AgentRuntimeInterval), + }) + if err != nil { + return xerrors.Errorf("sum chat message runtime: %w", err) + } + + // The deterministic ID makes concurrent inserts of the same bucket + // idempotent, and created_at is the bucket start (not the insertion + // 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 err != nil { + return xerrors.Errorf("insert usage event: %w", err) + } + return nil +} diff --git a/enterprise/coderd/usage/generator_test.go b/enterprise/coderd/usage/generator_test.go new file mode 100644 index 00000000000..aeff3c9c309 --- /dev/null +++ b/enterprise/coderd/usage/generator_test.go @@ -0,0 +1,481 @@ +package usage_test + +import ( + "context" + "database/sql" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/require" + + "cdr.dev/slog/v3" + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/usage/usagetypes" + "github.com/coder/coder/v2/enterprise/coderd/usage" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +// generatorTimerName must match the tag the Generator passes to +// clock.NewTimer so tests can trap its timers. +const generatorTimerName = "agent-runtime-generator" + +// warnSink counts log entries at Warn or above. The generator downgrades +// per-bucket failures to Warn logs (which slogtest tolerates), so tests that +// must prove the error paths stayed quiet assert on this counter instead. +type warnSink struct{ count atomic.Int64 } + +func (s *warnSink) LogEntry(_ context.Context, e slog.SinkEntry) { + if e.Level >= slog.LevelWarn { + s.count.Add(1) + } +} + +func (*warnSink) Sync() {} + +// generatorHarness runs the generator against a dbauthz-wrapped store so the +// tests also verify that the usage publisher subject holds the permissions +// the generator's queries require. +type generatorHarness struct { + db database.Store + authzDB database.Store + rawDB *sql.DB + + user database.User + modelConfig database.ChatModelConfig + chat database.Chat + chat2 database.Chat +} + +func newGeneratorHarness(t *testing.T) *generatorHarness { + t.Helper() + db, _, rawDB := dbtestutil.NewDBWithSQLDB(t) + log := slogtest.Make(t, nil) + authzDB := dbauthz.New(db, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), log, coderdtest.AccessControlStorePointer()) + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", + }) + mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "test-model", + ContextLimit: 8192, + }) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + chat2 := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: mc.ID, + }) + return &generatorHarness{ + db: db, + authzDB: authzDB, + rawDB: rawDB, + user: user, + modelConfig: mc, + chat: chat, + chat2: chat2, + } +} + +func (h *generatorHarness) insertRuntimeMessage(ctx context.Context, t *testing.T, chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) { + t.Helper() + msg := dbgen.ChatMessage(t, h.db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: h.user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: h.modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + RuntimeMs: sql.NullInt64{Int64: runtimeMs, Valid: true}, + }) + _, err := h.rawDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1, deleted = $2 WHERE id = $3", createdAt, deleted, msg.ID) + require.NoError(t, err) +} + +// fetchRuntimeEvents fails the test if any bucket has more than one event. +func (h *generatorHarness) fetchRuntimeEvents(ctx context.Context, t *testing.T) (map[time.Time]int64, map[time.Time]string) { + t.Helper() + rows, err := h.rawDB.QueryContext(ctx, ` + SELECT id, (event_data->>'runtime_ms')::bigint, created_at + FROM usage_events + WHERE event_type = 'hb_agent_runtime_v1' + `) + require.NoError(t, err) + defer rows.Close() + + runtimes := make(map[time.Time]int64) + ids := make(map[time.Time]string) + for rows.Next() { + var ( + id string + runtimeMs int64 + createdAt time.Time + ) + require.NoError(t, rows.Scan(&id, &runtimeMs, &createdAt)) + bucket := createdAt.UTC() + _, ok := runtimes[bucket] + require.False(t, ok, "duplicate event for bucket %s", bucket) + runtimes[bucket] = runtimeMs + ids[bucket] = id + } + require.NoError(t, rows.Err()) + return runtimes, ids +} + +func expectedBuckets(first, last time.Time, overrides map[time.Time]int64) map[time.Time]int64 { + expected := make(map[time.Time]int64) + for bucket := first; !bucket.After(last); bucket = bucket.Add(time.Hour) { + expected[bucket] = 0 + } + for bucket, runtimeMs := range overrides { + expected[bucket] = runtimeMs + } + return expected +} + +func TestGenerator(t *testing.T) { + t.Parallel() + + // startTime is exactly on an hour boundary, so the first tick (which + // fires 1-5 minutes later) always lands before the just-closed bucket + // [13:00, 14:00) becomes eligible at 14:05. + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + log := slogtest.Make(t, nil) + h := newGeneratorHarness(t) + clock := quartz.NewMock(t) + clock.Set(startTime) + + var ( + bucketA = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + bucketB = time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC) + // The most recent closed bucket; not eligible at the first tick. + bucketC = time.Date(2025, 3, 10, 13, 0, 0, 0, time.UTC) + // Window bounds at the first tick. + windowFirst = startTime.Add(-usage.AgentRuntimeWindow) // 2025-03-03 14:00 + windowLast = startTime.Add(-2 * time.Hour) // 2025-03-10 12:00 + ) + + // Bucket A: a message on the bucket start boundary, a message from a + // second chat, and a soft-deleted message. All must be counted. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA, false) + h.insertRuntimeMessage(ctx, t, h.chat2.ID, 2000, bucketA.Add(15*time.Minute), false) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 4000, bucketA.Add(30*time.Minute), true) + // Bucket B: a message exactly on the A/B boundary belongs to B. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 8000, bucketB, false) + // Older than the window: must never be generated. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 16000, windowFirst.Add(-30*time.Minute), false) + // Bucket C: only becomes eligible at 14:05, after the first tick. + h.insertRuntimeMessage(ctx, t, h.chat.ID, 32000, bucketC.Add(30*time.Minute), false) + + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + defer gen.Close() + + call := trap.MustWait(ctx) + call.MustRelease(ctx) + require.GreaterOrEqual(t, call.Duration, time.Minute) + require.Less(t, call.Duration, 5*time.Minute) + clock.Advance(call.Duration).MustWait(ctx) + + // The generator creates the next timer only after the tick completes, + // so trapping it synchronizes with the end of the pass. + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // The first pass fills every bucket in [windowFirst, windowLast]: bucket + // C is not yet eligible, the pre-window message is excluded, and idle + // hours are zero-filled. + runtimes, ids := h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets(windowFirst, windowLast, map[time.Time]int64{ + bucketA: 7000, + bucketB: 8000, + }), runtimes) + require.Equal(t, "hb_agent_runtime_v1:2025-03-10_10:00:00", ids[bucketA]) + + // The next tick fires at bucket C's eligibility instant (14:05 plus + // jitter), in the same hour the first pass ran, rather than waiting for + // the next hour boundary. + fireTime := clock.Now().Add(call.Duration) + require.Equal(t, startTime, fireTime.Truncate(usage.AgentRuntimeInterval)) + require.GreaterOrEqual(t, fireTime.Sub(fireTime.Truncate(usage.AgentRuntimeInterval)), usage.AgentRuntimeEligibilityLag) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // The second tick fills only the newly-eligible bucket C (13:00). The + // window's start advanced by an hour, but the bucket at the old + // windowFirst is kept (rows are never deleted). + runtimes, _ = h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets(windowFirst, windowLast.Add(time.Hour), map[time.Time]int64{ + bucketA: 7000, + bucketB: 8000, + bucketC: 32000, + }), runtimes) + + // The third tick fires after the next hour boundary (15:05 plus jitter) + // and fills the idle 14:00 bucket. + fireTime = clock.Now().Add(call.Duration) + require.Equal(t, startTime.Add(time.Hour), fireTime.Truncate(usage.AgentRuntimeInterval)) + require.GreaterOrEqual(t, fireTime.Sub(fireTime.Truncate(usage.AgentRuntimeInterval)), usage.AgentRuntimeEligibilityLag) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + runtimes, _ = h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets(windowFirst, windowLast.Add(2*time.Hour), map[time.Time]int64{ + bucketA: 7000, + bucketB: 8000, + bucketC: 32000, + }), runtimes) + + // A separate generator started later (e.g. another replica restarting) + // finds nothing to do: all buckets in its window already exist. + gen2 := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen2.Start(ctx) + defer gen2.Close() + call = trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + runtimes2, _ := h.fetchRuntimeEvents(ctx, t) + require.Equal(t, runtimes, runtimes2) +} + +// TestGeneratorBackfillAfterDowntime simulates a deployment that was down +// for several hours: a fresh generator's first pass backfills exactly the +// missing buckets. +func TestGeneratorBackfillAfterDowntime(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + log := slogtest.Make(t, nil) + h := newGeneratorHarness(t) + + windowFirst := startTime.Add(-usage.AgentRuntimeWindow) + gapBucket := time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 5000, gapBucket.Add(45*time.Minute), false) + + // Simulate events generated before the deployment went down at 09:00: + // every bucket in [windowFirst, 08:00] exists with runtime 1. + inserter := usage.NewDBInserter() + for bucket := windowFirst; !bucket.After(time.Date(2025, 3, 10, 8, 0, 0, 0, time.UTC)); bucket = bucket.Add(time.Hour) { + err := inserter.InsertHeartbeatUsageEvent(ctx, h.db, "hb_agent_runtime_v1:"+bucket.Format("2006-01-02_15:04:05"), bucket, usagetypes.HBAgentRuntime{RuntimeMs: 1}) + require.NoError(t, err) + } + + clock := quartz.NewMock(t) + clock.Set(startTime) + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + defer gen.Close() + + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // The pass must fill exactly the gap [09:00, 12:00] and leave the + // pre-existing rows untouched. + runtimes, _ := h.fetchRuntimeEvents(ctx, t) + expected := expectedBuckets(windowFirst, startTime.Add(-2*time.Hour), map[time.Time]int64{ + gapBucket: 5000, + }) + for bucket := windowFirst; !bucket.After(time.Date(2025, 3, 10, 8, 0, 0, 0, time.UTC)); bucket = bucket.Add(time.Hour) { + expected[bucket] = 1 + } + require.Equal(t, expected, runtimes) +} + +func TestGeneratorInserterArguments(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + log := slogtest.Make(t, nil) + h := newGeneratorHarness(t) + clock := quartz.NewMock(t) + clock.Set(startTime) + + var ( + bucketA = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + windowFirst = startTime.Add(-usage.AgentRuntimeWindow) + // The first pass stops at 12:00; the second pass fires after 14:05 + // and adds the newly eligible 13:00 bucket. + windowLast = startTime.Add(-time.Hour) + ) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA.Add(10*time.Minute), false) + + ins := coderdtest.NewUsageInserter() + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, ins) + gen.Start(ctx) + defer gen.Close() + + // Each pass requests its next timer only after finishing, so trapping + // that request synchronizes with the end of the pass. + for range 2 { + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + } + call := trap.MustWait(ctx) + call.MustRelease(ctx) + + var expected []coderdtest.HeartbeatEvent + for bucket := windowFirst; !bucket.After(windowLast); bucket = bucket.Add(usage.AgentRuntimeInterval) { + var runtimeMs int64 + if bucket.Equal(bucketA) { + runtimeMs = 1000 + } + expected = append(expected, coderdtest.HeartbeatEvent{ + ID: "hb_agent_runtime_v1:" + bucket.Format("2006-01-02_15:04:05"), + CreatedAt: bucket, + Event: usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs}, + }) + } + require.Equal(t, expected, ins.GetHeartbeatEvents()) +} + +// TestGeneratorConcurrentReplicas runs two generators against the same +// database concurrently and verifies exactly one event is produced per +// bucket. Each replica gets its own mock clock (as real replicas have their +// own wall clocks) so their first passes can be fired independently and run +// at the same time. +func TestGeneratorConcurrentReplicas(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + // If insert idempotency regressed (e.g. ON CONFLICT DO NOTHING was + // dropped), the losing replica's inserts would error and surface as + // Warn logs, which slogtest alone would not catch. + sink := &warnSink{} + log := slogtest.Make(t, nil).AppendSinks(sink) + h := newGeneratorHarness(t) + + bucketA := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA.Add(10*time.Minute), false) + + // Both replicas' first passes fire between 14:01 and 14:05, so they + // compute identical windows regardless of their random startup jitter. + var traps []*quartz.Trap + for range 2 { + clock := quartz.NewMock(t) + clock.Set(startTime) + trap := clock.Trap().NewTimer(generatorTimerName) + t.Cleanup(trap.Close) + traps = append(traps, trap) + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + // Cleanups run LIFO, so each generator is closed before its trap. + t.Cleanup(func() { _ = gen.Close() }) + + call := trap.MustWait(ctx) + call.MustRelease(ctx) + // Fire the initial timer without waiting for the pass to complete so + // both replicas' passes overlap. + clock.Advance(call.Duration).MustWait(ctx) + } + + // Wait for both passes to complete (each requests its next timer only + // after the pass finishes). + for _, trap := range traps { + call := trap.MustWait(ctx) + call.MustRelease(ctx) + } + + // fetchRuntimeEvents fails on duplicate buckets; the expected map proves + // both replicas raced without double-inserting, and the warn counter + // proves the duplicate inserts were deduplicated rather than rejected + // with errors. + runtimes, _ := h.fetchRuntimeEvents(ctx, t) + require.Equal(t, expectedBuckets( + startTime.Add(-usage.AgentRuntimeWindow), + startTime.Add(-2*time.Hour), + map[time.Time]int64{bucketA: 1000}, + ), runtimes) + require.Zero(t, sink.count.Load(), "no replica may log Warn or above during the race") +} + +// TestGeneratorPoisonBucket verifies that a bucket which fails +// deterministically on every tick does not stall generation of later +// buckets. +func TestGeneratorPoisonBucket(t *testing.T) { + t.Parallel() + + startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC) + + ctx := testutil.Context(t, testutil.WaitLong) + sink := &warnSink{} + log := slogtest.Make(t, nil).AppendSinks(sink) + h := newGeneratorHarness(t) + clock := quartz.NewMock(t) + clock.Set(startTime) + + var ( + poisonBucket = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC) + laterBucket = time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC) + ) + // Nothing prevents a negative runtime_ms row in the database, and a + // negative sum fails event validation on every tick. + h.insertRuntimeMessage(ctx, t, h.chat.ID, -5, poisonBucket.Add(10*time.Minute), false) + h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, laterBucket.Add(10*time.Minute), false) + + trap := clock.Trap().NewTimer(generatorTimerName) + defer trap.Close() + + gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter()) + gen.Start(ctx) + defer gen.Close() + + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(call.Duration).MustWait(ctx) + call = trap.MustWait(ctx) + call.MustRelease(ctx) + + // Every bucket except the poison bucket is generated, including buckets + // after it, and the failure is logged. + runtimes, _ := h.fetchRuntimeEvents(ctx, t) + expected := expectedBuckets( + startTime.Add(-usage.AgentRuntimeWindow), + startTime.Add(-2*time.Hour), + map[time.Time]int64{laterBucket: 1000}, + ) + delete(expected, poisonBucket) + require.Equal(t, expected, runtimes) + require.NotZero(t, sink.count.Load(), "the poison bucket failure must be logged") +} diff --git a/enterprise/coderd/usage/inserter.go b/enterprise/coderd/usage/inserter.go index 90fb6ab4ca8..80b209fd906 100644 --- a/enterprise/coderd/usage/inserter.go +++ b/enterprise/coderd/usage/inserter.go @@ -3,6 +3,7 @@ package usage import ( "context" "encoding/json" + "time" "github.com/google/uuid" "golang.org/x/xerrors" @@ -68,10 +69,16 @@ func (i *dbInserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.S } // InsertHeartbeatUsageEvent implements agplusage.Inserter. -func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error { +func (*dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error { if !event.EventType().IsHeartbeat() { return xerrors.Errorf("event type %q is not a heartbeat event", event.EventType()) } + // A zero createdAt stores the row at year 1, where bucket reconciliation + // can never match it again while the deterministic id turns every retry + // into a no-op, silently forfeiting the bucket's usage. + if createdAt.IsZero() { + return xerrors.Errorf("createdAt must be set for %q event", event.EventType()) + } if err := event.Valid(); err != nil { return xerrors.Errorf("invalid %q event: %w", event.EventType(), err) } @@ -87,6 +94,6 @@ func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database. ID: id, EventType: string(event.EventType()), EventData: jsonData, - CreatedAt: dbtime.Time(i.clock.Now()), + CreatedAt: dbtime.Time(createdAt), }) } diff --git a/enterprise/coderd/usage/inserter_test.go b/enterprise/coderd/usage/inserter_test.go index 7ac915be7a5..af48341f816 100644 --- a/enterprise/coderd/usage/inserter_test.go +++ b/enterprise/coderd/usage/inserter_test.go @@ -68,6 +68,35 @@ func TestInserter(t *testing.T) { } }) + t.Run("Heartbeat", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + inserter := usage.NewDBInserter() + + // Heartbeat inserts must store the provided id and createdAt + // verbatim. + event := usagetypes.HBAgentRuntime{RuntimeMs: 1234} + eventJSON := jsoninate(t, event) + id := "hb_agent_runtime_v1:2025-01-02_03:00:00" + createdAt := time.Date(2025, 1, 2, 3, 0, 0, 0, time.UTC) + + db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx interface{}, params database.InsertUsageEventParams) error { + assert.Equal(t, id, params.ID) + assert.Equal(t, event.EventType(), usagetypes.UsageEventType(params.EventType)) + assert.JSONEq(t, eventJSON, string(params.EventData)) + assert.Equal(t, dbtime.Time(createdAt), params.CreatedAt) + return nil + }, + ).Times(1) + + err := inserter.InsertHeartbeatUsageEvent(ctx, db, id, createdAt, event) + require.NoError(t, err) + }) + t.Run("InvalidEvent", func(t *testing.T) { t.Parallel() @@ -81,5 +110,26 @@ func TestInserter(t *testing.T) { Count: 0, // invalid }) assert.ErrorContains(t, err, `invalid "dc_managed_agents_v1" event: count must be greater than 0`) + + err = inserter.InsertHeartbeatUsageEvent(ctx, db, "some-id", time.Now(), usagetypes.HBAgentRuntime{ + RuntimeMs: -1, // invalid + }) + assert.ErrorContains(t, err, `invalid "hb_agent_runtime_v1" event: runtime_ms cannot be negative`) + }) + + t.Run("ZeroCreatedAt", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + // The mock store fails the test on any unexpected call, so no insert + // may reach the database. + db := dbmock.NewMockStore(ctrl) + + inserter := usage.NewDBInserter() + err := inserter.InsertHeartbeatUsageEvent(ctx, db, "some-id", time.Time{}, usagetypes.HBAgentRuntime{ + RuntimeMs: 1, + }) + assert.ErrorContains(t, err, `createdAt must be set for "hb_agent_runtime_v1" event`) }) }