From 84de391f26869ee4a71a38610c8e835b46572d31 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Wed, 18 Mar 2026 09:30:22 -0500 Subject: [PATCH] chore: add tallyman events for ai seat tracking (#22689) AI seat tracking inserted as heartbeat into usage table. --- coderd/coderdtest/usage.go | 46 +++- coderd/database/dbauthz/dbauthz.go | 7 + coderd/database/dbauthz/dbauthz_test.go | 6 + coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 ++ coderd/database/dump.sql | 21 +- .../000444_usage_events_ai_seats.down.sql | 38 ++++ .../000444_usage_events_ai_seats.up.sql | 50 ++++ .../000444_usage_events_ai_seats.up.sql | 20 ++ coderd/database/querier.go | 1 + coderd/database/querier_test.go | 74 ++++++ coderd/database/queries.sql.go | 13 ++ coderd/database/queries/usageevents.sql | 5 + coderd/pproflabel/pproflabel.go | 1 + .../provisionerdserver_test.go | 23 +- coderd/usage/inserter.go | 21 ++ coderd/usage/usagetypes/events.go | 42 ++++ coderd/usage/usagetypes/events_test.go | 11 + enterprise/cli/server.go | 15 ++ enterprise/coderd/usage/cron.go | 215 ++++++++++++++++++ enterprise/coderd/usage/cron_internal_test.go | 101 ++++++++ enterprise/coderd/usage/cron_test.go | 98 ++++++++ enterprise/coderd/usage/heartbeats.go | 31 +++ enterprise/coderd/usage/inserter.go | 24 ++ enterprise/coderd/workspaces_test.go | 8 +- 25 files changed, 860 insertions(+), 34 deletions(-) create mode 100644 coderd/database/migrations/000444_usage_events_ai_seats.down.sql create mode 100644 coderd/database/migrations/000444_usage_events_ai_seats.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000444_usage_events_ai_seats.up.sql create mode 100644 enterprise/coderd/usage/cron.go create mode 100644 enterprise/coderd/usage/cron_internal_test.go create mode 100644 enterprise/coderd/usage/cron_test.go create mode 100644 enterprise/coderd/usage/heartbeats.go diff --git a/coderd/coderdtest/usage.go b/coderd/coderdtest/usage.go index 4da724b177..c713912867 100644 --- a/coderd/coderdtest/usage.go +++ b/coderd/coderdtest/usage.go @@ -13,32 +13,64 @@ var _ usage.Inserter = (*UsageInserter)(nil) type UsageInserter struct { sync.Mutex - events []usagetypes.DiscreteEvent + discreteEvents []usagetypes.DiscreteEvent + heartbeatEvents []usagetypes.HeartbeatEvent + seenHeartbeats map[string]struct{} } func NewUsageInserter() *UsageInserter { return &UsageInserter{ - events: []usagetypes.DiscreteEvent{}, + discreteEvents: []usagetypes.DiscreteEvent{}, + seenHeartbeats: map[string]struct{}{}, + heartbeatEvents: []usagetypes.HeartbeatEvent{}, } } func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error { u.Lock() defer u.Unlock() - u.events = append(u.events, event) + u.discreteEvents = append(u.discreteEvents, event) return nil } -func (u *UsageInserter) GetEvents() []usagetypes.DiscreteEvent { +func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, event usagetypes.HeartbeatEvent) error { u.Lock() defer u.Unlock() - eventsCopy := make([]usagetypes.DiscreteEvent, len(u.events)) - copy(eventsCopy, u.events) + if _, seen := u.seenHeartbeats[id]; seen { + return nil + } + + u.seenHeartbeats[id] = struct{}{} + u.heartbeatEvents = append(u.heartbeatEvents, event) + return nil +} + +func (u *UsageInserter) GetHeartbeatEvents() []usagetypes.HeartbeatEvent { + u.Lock() + defer u.Unlock() + eventsCopy := make([]usagetypes.HeartbeatEvent, len(u.heartbeatEvents)) + copy(eventsCopy, u.heartbeatEvents) return eventsCopy } +func (u *UsageInserter) GetDiscreteEvents() []usagetypes.DiscreteEvent { + u.Lock() + defer u.Unlock() + eventsCopy := make([]usagetypes.DiscreteEvent, len(u.discreteEvents)) + copy(eventsCopy, u.discreteEvents) + return eventsCopy +} + +func (u *UsageInserter) TotalEventCount() int { + u.Lock() + defer u.Unlock() + return len(u.discreteEvents) + len(u.heartbeatEvents) +} + func (u *UsageInserter) Reset() { u.Lock() defer u.Unlock() - u.events = []usagetypes.DiscreteEvent{} + u.seenHeartbeats = map[string]struct{}{} + u.discreteEvents = []usagetypes.DiscreteEvent{} + u.heartbeatEvents = []usagetypes.HeartbeatEvent{} } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 081982bacb..014a9efed1 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -6805,6 +6805,13 @@ func (q *querier) UpsertWorkspaceAppAuditSession(ctx context.Context, arg databa return q.db.UpsertWorkspaceAppAuditSession(ctx, arg) } +func (q *querier) UsageEventExistsByID(ctx context.Context, id string) (bool, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil { + return false, err + } + return q.db.UsageEventExistsByID(ctx, id) +} + func (q *querier) ValidateGroupIDs(ctx context.Context, groupIDs []uuid.UUID) (database.ValidateGroupIDsRow, error) { // This check is probably overly restrictive, but the "correct" check isn't // necessarily obvious. It's only used as a verification check for ACLs right diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index bb1659f054..8ac4a41874 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -5123,6 +5123,12 @@ func (s *MethodTestSuite) TestUsageEvents() { check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionCreate) })) + s.Run("UsageEventExistsByID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + id := uuid.NewString() + db.EXPECT().UsageEventExistsByID(gomock.Any(), id).Return(true, nil) + check.Args(id).Asserts(rbac.ResourceUsageEvent, policy.ActionRead) + })) + s.Run("SelectUsageEventsForPublishing", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { now := dbtime.Now() db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), now).Return([]database.UsageEvent{}, nil) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index c7c9246fe9..762e1d0974 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -4792,6 +4792,14 @@ func (m queryMetricsStore) UpsertWorkspaceAppAuditSession(ctx context.Context, a return r0, r1 } +func (m queryMetricsStore) UsageEventExistsByID(ctx context.Context, id string) (bool, error) { + start := time.Now() + r0, r1 := m.s.UsageEventExistsByID(ctx, id) + m.queryLatencies.WithLabelValues("UsageEventExistsByID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UsageEventExistsByID").Inc() + return r0, r1 +} + func (m queryMetricsStore) ValidateGroupIDs(ctx context.Context, groupIds []uuid.UUID) (database.ValidateGroupIDsRow, error) { start := time.Now() r0, r1 := m.s.ValidateGroupIDs(ctx, groupIds) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 8f950096e9..894898506c 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -8945,6 +8945,21 @@ func (mr *MockStoreMockRecorder) UpsertWorkspaceAppAuditSession(ctx, arg any) *g return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertWorkspaceAppAuditSession", reflect.TypeOf((*MockStore)(nil).UpsertWorkspaceAppAuditSession), ctx, arg) } +// UsageEventExistsByID mocks base method. +func (m *MockStore) UsageEventExistsByID(ctx context.Context, id string) (bool, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UsageEventExistsByID", ctx, id) + ret0, _ := ret[0].(bool) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UsageEventExistsByID indicates an expected call of UsageEventExistsByID. +func (mr *MockStoreMockRecorder) UsageEventExistsByID(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UsageEventExistsByID", reflect.TypeOf((*MockStore)(nil).UsageEventExistsByID), ctx, id) +} + // ValidateGroupIDs mocks base method. func (m *MockStore) ValidateGroupIDs(ctx context.Context, groupIds []uuid.UUID) (database.ValidateGroupIDsRow, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index ce24788cc4..11927e8553 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -620,28 +620,35 @@ CREATE FUNCTION aggregate_usage_event() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN - -- Check for supported event types and throw error for unknown types - IF NEW.event_type NOT IN ('dc_managed_agents_v1') THEN + -- 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 ( - -- Extract the date from the created_at timestamp, always using UTC for - -- consistency 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 + -- 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; @@ -2655,7 +2662,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 = 'dc_managed_agents_v1'::text)) + CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_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.'; @@ -3710,6 +3717,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_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); CREATE INDEX idx_user_deleted_deleted_at ON user_deleted USING btree (deleted_at); diff --git a/coderd/database/migrations/000444_usage_events_ai_seats.down.sql b/coderd/database/migrations/000444_usage_events_ai_seats.down.sql new file mode 100644 index 0000000000..e1bbf8ae3e --- /dev/null +++ b/coderd/database/migrations/000444_usage_events_ai_seats.down.sql @@ -0,0 +1,38 @@ +DROP INDEX IF EXISTS idx_usage_events_ai_seats; + +-- Remove hb_ai_seats_v1 rows so the original constraint can be restored. +DELETE FROM usage_events WHERE event_type = 'hb_ai_seats_v1'; +DELETE FROM usage_events_daily WHERE event_type = 'hb_ai_seats_v1'; + +-- Restore original constraint. +ALTER TABLE usage_events + DROP CONSTRAINT usage_event_type_check, + ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1')); + +-- Restore the original aggregate function without hb_ai_seats_v1 support. +CREATE OR REPLACE FUNCTION aggregate_usage_event() +RETURNS TRIGGER AS $$ +BEGIN + IF NEW.event_type NOT IN ('dc_managed_agents_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 + 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) + ) + END; + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; diff --git a/coderd/database/migrations/000444_usage_events_ai_seats.up.sql b/coderd/database/migrations/000444_usage_events_ai_seats.up.sql new file mode 100644 index 0000000000..9950915eef --- /dev/null +++ b/coderd/database/migrations/000444_usage_events_ai_seats.up.sql @@ -0,0 +1,50 @@ +-- Expand the CHECK constraint to allow hb_ai_seats_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')); + +-- Partial index for efficient lookups of AI seat heartbeat events by time. +-- This will be used for the admin dashboard to see seat count over time. +CREATE INDEX idx_usage_events_ai_seats + ON usage_events (event_type, created_at) + WHERE event_type = 'hb_ai_seats_v1'; + +-- Update the aggregate function to handle hb_ai_seats_v1 events. +-- Heartbeat events replace the previous value for the same time period. +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/testdata/fixtures/000444_usage_events_ai_seats.up.sql b/coderd/database/migrations/testdata/fixtures/000444_usage_events_ai_seats.up.sql new file mode 100644 index 0000000000..39d94c31d3 --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000444_usage_events_ai_seats.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_ai_seats_v1 event. +( + 'ai-seats-event1', + 'hb_ai_seats_v1', + '{"count":3}', + '2023-06-01 00:00:00+00', + NULL, + NULL, + NULL +); diff --git a/coderd/database/querier.go b/coderd/database/querier.go index e93c8a4c04..eeb24606ce 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -907,6 +907,7 @@ type sqlcQuerier interface { // was started. This means that a new row was inserted (no previous session) or // the updated_at is older than stale interval. UpsertWorkspaceAppAuditSession(ctx context.Context, arg UpsertWorkspaceAppAuditSessionParams) (bool, error) + UsageEventExistsByID(ctx context.Context, id string) (bool, error) ValidateGroupIDs(ctx context.Context, groupIds []uuid.UUID) (ValidateGroupIDsRow, error) ValidateUserIDs(ctx context.Context, userIds []uuid.UUID) (ValidateUserIDsRow, error) } diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index aaaa0f8022..37d12150bd 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -8315,6 +8315,80 @@ func TestUsageEventsTrigger(t *testing.T) { require.WithinDuration(t, time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC), rows[1].Day, time.Second) }) + t.Run("HeartbeatAISeats", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + + // Insert a heartbeat event. + err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: "hb-1", + EventType: "hb_ai_seats_v1", + EventData: []byte(`{"count": 10}`), + CreatedAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }) + require.NoError(t, err) + + rows := getDailyRows(ctx, sqlDB) + require.Len(t, rows, 1) + require.Equal(t, "hb_ai_seats_v1", rows[0].EventType) + require.JSONEq(t, `{"count": 10}`, string(rows[0].UsageData)) + + // Insert a higher count on the same day — should take the max. + err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: "hb-2", + EventType: "hb_ai_seats_v1", + EventData: []byte(`{"count": 50}`), + CreatedAt: time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC), + }) + require.NoError(t, err) + + rows = getDailyRows(ctx, sqlDB) + require.Len(t, rows, 1) + require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData)) + + // Insert a lower count on the same day — should keep the max (50). + err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: "hb-3", + EventType: "hb_ai_seats_v1", + EventData: []byte(`{"count": 25}`), + CreatedAt: time.Date(2025, 1, 1, 18, 0, 0, 0, time.UTC), + }) + require.NoError(t, err) + + rows = getDailyRows(ctx, sqlDB) + require.Len(t, rows, 1) + require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData)) + + // Insert on a different day. + err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: "hb-4", + EventType: "hb_ai_seats_v1", + EventData: []byte(`{"count": 5}`), + CreatedAt: time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC), + }) + require.NoError(t, err) + + rows = getDailyRows(ctx, sqlDB) + require.Len(t, rows, 2) + require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData)) + require.JSONEq(t, `{"count": 5}`, string(rows[1].UsageData)) + + // Also insert a dc_managed_agents_v1 on the same first day to + // verify different event types get separate daily rows. + err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: "dc-1", + EventType: "dc_managed_agents_v1", + EventData: []byte(`{"count": 7}`), + CreatedAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }) + require.NoError(t, err) + + rows = getDailyRows(ctx, sqlDB) + require.Len(t, rows, 3) + }) + t.Run("UnknownEventType", func(t *testing.T) { t.Parallel() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 08636649d2..7d4013a057 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -18845,6 +18845,19 @@ func (q *sqlQuerier) UpdateUsageEventsPostPublish(ctx context.Context, arg Updat return err } +const usageEventExistsByID = `-- name: UsageEventExistsByID :one +SELECT EXISTS( + SELECT 1 FROM usage_events WHERE id = $1 +)::bool +` + +func (q *sqlQuerier) UsageEventExistsByID(ctx context.Context, id string) (bool, error) { + row := q.db.QueryRowContext(ctx, usageEventExistsByID, id) + var column_1 bool + err := row.Scan(&column_1) + return column_1, err +} + const getUserLinkByLinkedID = `-- name: GetUserLinkByLinkedID :one SELECT user_links.user_id, user_links.login_type, user_links.linked_id, user_links.oauth_access_token, user_links.oauth_refresh_token, user_links.oauth_expiry, user_links.oauth_access_token_key_id, user_links.oauth_refresh_token_key_id, user_links.claims diff --git a/coderd/database/queries/usageevents.sql b/coderd/database/queries/usageevents.sql index 291e275c60..7ffcb1173b 100644 --- a/coderd/database/queries/usageevents.sql +++ b/coderd/database/queries/usageevents.sql @@ -15,6 +15,11 @@ VALUES (@id, @event_type, @event_data, @created_at, NULL, NULL, NULL) ON CONFLICT (id) DO NOTHING; +-- name: UsageEventExistsByID :one +SELECT EXISTS( + SELECT 1 FROM usage_events WHERE id = @id +)::bool; + -- name: SelectUsageEventsForPublishing :many WITH usage_events AS ( UPDATE diff --git a/coderd/pproflabel/pproflabel.go b/coderd/pproflabel/pproflabel.go index bde5be1b36..f686c1c428 100644 --- a/coderd/pproflabel/pproflabel.go +++ b/coderd/pproflabel/pproflabel.go @@ -34,6 +34,7 @@ const ( ServiceAgentMetricAggregator = "agent-metrics-aggregator" // ServiceTallymanPublisher publishes usage events to coder/tallyman. ServiceTallymanPublisher = "tallyman-publisher" + ServiceUsageEventCron = "usage-event-cron" RequestTypeTag = "coder_request_type" ) diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go index 336d91903a..267b453b41 100644 --- a/coderd/provisionerdserver/provisionerdserver_test.go +++ b/coderd/provisionerdserver/provisionerdserver_test.go @@ -3020,7 +3020,7 @@ func TestCompleteJob(t *testing.T) { // We never expect a usage event to be collected for // template imports. - require.Empty(t, fakeUsageInserter.collectedEvents) + require.Equal(t, 0, fakeUsageInserter.TotalEventCount()) }) } }) @@ -3371,13 +3371,13 @@ func TestCompleteJob(t *testing.T) { if tc.expectUsageEvent { // Check that a usage event was collected. - require.Len(t, fakeUsageInserter.collectedEvents, 1) + require.Len(t, fakeUsageInserter.GetDiscreteEvents(), 1) require.Equal(t, usagetypes.DCManagedAgentsV1{ Count: 1, - }, fakeUsageInserter.collectedEvents[0]) + }, fakeUsageInserter.GetDiscreteEvents()[0]) } else { // Check that no usage event was collected. - require.Empty(t, fakeUsageInserter.collectedEvents) + require.Equal(t, 0, fakeUsageInserter.TotalEventCount()) } }) } @@ -5032,21 +5032,10 @@ func (s *fakeStream) cancel() { s.c.Broadcast() } -type fakeUsageInserter struct { - collectedEvents []usagetypes.Event -} - -var _ usage.Inserter = &fakeUsageInserter{} - -func newFakeUsageInserter() (*fakeUsageInserter, *atomic.Pointer[usage.Inserter]) { +func newFakeUsageInserter() (*coderdtest.UsageInserter, *atomic.Pointer[usage.Inserter]) { poitr := &atomic.Pointer[usage.Inserter]{} - fake := &fakeUsageInserter{} + fake := coderdtest.NewUsageInserter() var inserter usage.Inserter = fake poitr.Store(&inserter) return fake, poitr } - -func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error { - f.collectedEvents = append(f.collectedEvents, event) - return nil -} diff --git a/coderd/usage/inserter.go b/coderd/usage/inserter.go index 7a0f42daf4..891f5c7387 100644 --- a/coderd/usage/inserter.go +++ b/coderd/usage/inserter.go @@ -14,6 +14,21 @@ type Inserter interface { // The caller context must be authorized to create usage events in the // database. InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event usagetypes.DiscreteEvent) error + + // InsertHeartbeatUsageEvent writes a heartbeat usage event to the database + // within the given transaction. + // + // The caller context must be authorized to create usage events in the database. + // + // The `id` should be a stable identifier for the event. Heartbeat events may be + // emitted by multiple replicas of the same daemon, so the same logical event + // may be submitted multiple times concurrently. For this reason the identifier + // must be deterministic and stateless, allowing duplicate submissions to be + // safely ignored. + // + // 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 } // AGPLInserter is a no-op implementation of Inserter. @@ -30,3 +45,9 @@ func NewAGPLInserter() Inserter { func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ usagetypes.DiscreteEvent) error { return nil } + +// InsertHeartbeatUsageEvent is a no-op implementation of +// InsertHeartbeatUsageEvent. +func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ usagetypes.HeartbeatEvent) error { + return nil +} diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index ef5ac79d45..6c8fde416e 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -29,12 +29,15 @@ type UsageEventType string // ParseEventWithType function. const ( UsageEventTypeDCManagedAgentsV1 UsageEventType = "dc_managed_agents_v1" + UsageEventTypeHBAISeatsV1 UsageEventType = "hb_ai_seats_v1" ) func (e UsageEventType) Valid() bool { switch e { case UsageEventTypeDCManagedAgentsV1: return true + case UsageEventTypeHBAISeatsV1: + return true default: return false } @@ -96,6 +99,12 @@ func ParseEventWithType(eventType UsageEventType, data json.RawMessage) (Event, return nil, err } return event, nil + case UsageEventTypeHBAISeatsV1: + var event HBAISeats + if err := ParseEvent(data, &event); err != nil { + return nil, err + } + return event, nil default: return nil, UnknownEventTypeError{EventType: string(eventType)} } @@ -121,6 +130,12 @@ type DiscreteEvent interface { discreteUsageEvent() // marker method, also prevents external types from implementing this interface } +// HeartbeatEvent is a usage event that is collected as a heartbeat. +type HeartbeatEvent interface { + Event + heartbeatUsageEvent() // marker method, also prevents external types from implementing this interface +} + // DCManagedAgentsV1 is a discrete usage event for the number of managed agents. // This event is sent in the following situations: // - Once on first startup after usage tracking is added to the product with @@ -150,3 +165,30 @@ func (e DCManagedAgentsV1) Fields() map[string]any { "count": e.Count, } } + +// HBAISeats is a heartbeat event for the total number of AI seats consumed. +type HBAISeats struct { + Count int64 `json:"count"` +} + +var _ HeartbeatEvent = HBAISeats{} + +func (HBAISeats) usageEvent() {} +func (HBAISeats) heartbeatUsageEvent() {} +func (HBAISeats) EventType() UsageEventType { + return UsageEventTypeHBAISeatsV1 +} + +func (e HBAISeats) Valid() error { + if e.Count < 0 { + return xerrors.New("count cannot be negative") + } + // The count can be 0 + return nil +} + +func (e HBAISeats) Fields() map[string]any { + return map[string]any{ + "count": e.Count, + } +} diff --git a/coderd/usage/usagetypes/events_test.go b/coderd/usage/usagetypes/events_test.go index a04e5d4df0..fcfd076fc0 100644 --- a/coderd/usage/usagetypes/events_test.go +++ b/coderd/usage/usagetypes/events_test.go @@ -65,4 +65,15 @@ func TestParseEventWithType(t *testing.T) { require.Equal(t, eventType, event.EventType()) require.Equal(t, map[string]any{"count": uint64(1)}, event.Fields()) }) + + t.Run("HBAISeatsV1", func(t *testing.T) { + t.Parallel() + + eventType := usagetypes.UsageEventTypeHBAISeatsV1 + event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"count": 1}`)) + require.NoError(t, err) + require.Equal(t, usagetypes.HBAISeats{Count: 1}, event) + require.Equal(t, eventType, event.EventType()) + require.Equal(t, map[string]any{"count": int64(1)}, event.Fields()) + }) } diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go index 5a2cb0a8fc..4a51912f69 100644 --- a/enterprise/cli/server.go +++ b/enterprise/cli/server.go @@ -9,6 +9,7 @@ import ( "errors" "io" "net/url" + "time" "golang.org/x/xerrors" "tailscale.com/derp" @@ -147,6 +148,20 @@ func (r *RootCmd) Server(_ func()) *serpent.Command { } closers.Add(publisher) + // usageCron are heartbeat events to the usage table. These events are eventually sent + // to Tallyman. + usageCron := usage.NewCron(quartz.NewReal(), options.Logger.Named("usage-cron"), options.Database, *options.UsageInserter.Load()) + // ai-seats heartbeats track the number of users that have used an AI feature. + // These users consume a seat for the AI addon to our License. + _ = usageCron.Register(usage.CronJob{ + Name: "ai-seats", + Interval: usage.AISeatsInterval, + Jitter: 10 * time.Minute, + Fn: usage.AISeatsHeartbeat(options.Database), + }) + usageCron.Start(ctx) + closers.Add(usageCron) + // In-memory aibridge daemon. // TODO(@deansheather): the lifecycle of the aibridged server is // probably better managed by the enterprise API type itself. Managing diff --git a/enterprise/coderd/usage/cron.go b/enterprise/coderd/usage/cron.go new file mode 100644 index 0000000000..13ccbb927c --- /dev/null +++ b/enterprise/coderd/usage/cron.go @@ -0,0 +1,215 @@ +package usage + +import ( + "context" + "math/rand" + "sync" + "sync/atomic" + "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" +) + +// epoch is a fixed reference point for aligning interval boundaries. +// All replicas use this same epoch so their buckets are identical. +var epoch = time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC) + +const ( + cronDateFormat = "2006-01-02_15:04:05" +) + +// HeartbeatFunc generates a heartbeat event and its stable ID. +// It is called periodically by the cron. Returning an error skips +// the insert for that tick and logs a warning. +type HeartbeatFunc func(ctx context.Context) (event usagetypes.HeartbeatEvent, err error) + +// CronJob defines a periodic heartbeat job. +type CronJob struct { + // Name is a human-readable label used in logs. + Name string + // Interval is the base duration between ticks. + Interval time.Duration + // EventType must match the events generated by the Fn. + EventType usagetypes.UsageEventType + // Jitter is the maximum random delay added after the boundary. + // The actual offset is uniformly distributed in [0, Jitter). + // This staggers replicas so one is likely to complete the work + // before others attempt it, allowing them to skip via the + // existence check (heartbeat inserts are idempotent). + Jitter time.Duration + // Fn produces the heartbeat event. + Fn HeartbeatFunc +} + +// Cron runs registered CronJobs on the dbInserter's clock. Stopping +// the context passed to Start cancels all jobs. Daemon restarts +// naturally restart the timers since Start() creates them fresh — +// there is no state to persist or recover. +type Cron struct { + clock quartz.Clock + log slog.Logger + db database.Store + ins agplusage.Inserter + jobs []CronJob + + // cancel cancels the context on all running jobs. If the ctx passed into `Start` + // is canceled, the jobs will also stop. + cancel context.CancelFunc + + // wg ensures all job goroutines have exited before Close returns. + wg sync.WaitGroup + + // startOnce ensures Start is idempotent. + startOnce sync.Once + started atomic.Bool +} + +// NewCron creates a Cron that periodically generates and inserts +// heartbeat events. The clock controls all timers so that tests can +// advance time deterministically via quartz.Mock. +func NewCron(clock quartz.Clock, log slog.Logger, db database.Store, ins agplusage.Inserter) *Cron { + return &Cron{ + clock: clock, + log: log, + db: db, + ins: ins, + } +} + +// Register adds a job. It must be called before Start; calling it +// after Start returns an error. +func (c *Cron) Register(job CronJob) error { + if !job.EventType.IsHeartbeat() { + return xerrors.New("event type must be a heartbeat type") + } + if c.started.Load() { + return xerrors.New("cannot register a job after Start has been called") + } + c.jobs = append(c.jobs, job) + return nil +} + +// Start launches a goroutine per job. Subsequent calls are no-ops. +// On daemon restart a new Cron should be created. +func (c *Cron) Start(ctx context.Context) { + c.startOnce.Do(func() { + c.started.Store(true) + ctx, c.cancel = context.WithCancel(ctx) + for _, job := range c.jobs { + c.wg.Add(1) + pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceUsageEventCron, "job", job.Name), func(ctx context.Context) { + c.run(ctx, job) + }) + } + }) +} + +// Close cancels all jobs and waits for goroutines to exit. +func (c *Cron) Close() error { + if c.cancel != nil { + c.cancel() + } + c.wg.Wait() + return nil +} + +func (c *Cron) run(ctx context.Context, job CronJob) { + //nolint:gocritic // We are a publisher in this function + ctx = dbauthz.AsUsagePublisher(ctx) + defer c.wg.Done() + for { + boundary, delay := nextTick(c.clock.Now(), job.Interval, job.Jitter) + + // Use a quartz timer so the wait honors ctx cancellation and + // tests can advance time deterministically. + timer := c.clock.NewTimer(delay, job.Name) + + select { + case <-ctx.Done(): + if !timer.Stop() { + // Drain the channel if the timer already fired. + <-timer.C + } + return + case <-timer.C: + } + + // 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) + + // Skip if this bucket was already recorded — avoids running + // the potentially expensive heartbeat function for a + // duplicate. + exists, err := c.db.UsageEventExistsByID(ctx, stableID) + if err != nil { + c.log.Warn(ctx, "cron heartbeat existence check failed", + slog.F("job", job.Name), + slog.Error(err), + ) + continue + } + if exists { + c.log.Debug(ctx, "cron heartbeat already recorded, skipping", + slog.F("job", job.Name), + slog.F("id", stableID), + ) + continue + } + + event, err := job.Fn(ctx) + if err != nil { + c.log.Error(ctx, "cron heartbeat func failed", + slog.F("job", job.Name), + slog.Error(err), + ) + continue + } + + if event.EventType() != job.EventType { + c.log.Error(ctx, "cron heartbeat func returned wrong event type", + slog.F("job", job.Name), + slog.F("expected", job.EventType), + slog.F("actual", event.EventType()), + ) + continue + } + + if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, event); err != nil { + c.log.Warn(ctx, "cron heartbeat insert failed", + slog.F("job", job.Name), + slog.Error(err), + ) + } + } +} + +// nextTick computes the delay until the next epoch-aligned boundary +// for the given interval, plus a random jitter in [0, jitter). It +// returns the target boundary and the total delay from now. +func nextTick(now time.Time, interval, jitter time.Duration) (boundary time.Time, delay time.Duration) { + boundary = nextBoundary(now, interval) + delay = boundary.Sub(now) + if jitter > 0 { + //nolint:gosec // Jitter does not need cryptographic randomness. + delay += time.Duration(rand.Int63n(int64(jitter))) + } + return boundary, delay +} + +// nextBoundary returns the first multiple of interval (relative to +// epoch) that is strictly after t. +func nextBoundary(t time.Time, interval time.Duration) time.Time { + since := t.Sub(epoch) + n := since / interval + return epoch.Add((n + 1) * interval) +} diff --git a/enterprise/coderd/usage/cron_internal_test.go b/enterprise/coderd/usage/cron_internal_test.go new file mode 100644 index 0000000000..b2d96cc1c7 --- /dev/null +++ b/enterprise/coderd/usage/cron_internal_test.go @@ -0,0 +1,101 @@ +package usage + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestNextBoundary(t *testing.T) { + t.Parallel() + + tcs := []struct { + name string + T time.Time + interval time.Duration + expected time.Time + }{ + { + name: "exactly_on_boundary", + T: time.Date(2023, 1, 1, 8, 0, 0, 0, time.UTC), + interval: 4 * time.Hour, + // On a boundary → returns the next one. + expected: time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC), + }, + { + name: "1ns_after_boundary", + T: time.Date(2023, 1, 1, 8, 0, 0, 1, time.UTC), + interval: 4 * time.Hour, + expected: time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC), + }, + { + name: "1ns_before_boundary", + T: time.Date(2023, 1, 1, 7, 59, 59, 999999999, time.UTC), + interval: 4 * time.Hour, + expected: time.Date(2023, 1, 1, 8, 0, 0, 0, time.UTC), + }, + { + name: "mid_interval", + T: time.Date(2023, 1, 1, 10, 0, 0, 0, time.UTC), + interval: 4 * time.Hour, + expected: time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC), + }, + { + name: "5min_interval", + T: time.Date(2026, 3, 13, 14, 2, 30, 0, time.UTC), + interval: 5 * time.Minute, + expected: time.Date(2026, 3, 13, 14, 5, 0, 0, time.UTC), + }, + { + name: "1hr_interval", + T: time.Date(2026, 6, 15, 9, 45, 0, 0, time.UTC), + interval: 1 * time.Hour, + expected: time.Date(2026, 6, 15, 10, 0, 0, 0, time.UTC), + }, + } + + for _, tc := range tcs { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := nextBoundary(tc.T, tc.interval) + require.Equal(t, tc.expected, got) + }) + } +} + +func TestNextTick(t *testing.T) { + t.Parallel() + + t.Run("NoJitter", func(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 3, 13, 14, 2, 30, 0, time.UTC) + interval := 4 * time.Hour + + boundary, delay := nextTick(now, interval, 0) + + expectedBoundary := time.Date(2026, 3, 13, 16, 0, 0, 0, time.UTC) + require.Equal(t, expectedBoundary, boundary) + require.Equal(t, boundary.Sub(now), delay) + }) + + t.Run("WithJitter", func(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 3, 13, 14, 2, 30, 0, time.UTC) + interval := 4 * time.Hour + jitter := 10 * time.Minute + + boundary, delay := nextTick(now, interval, jitter) + + expectedBoundary := time.Date(2026, 3, 13, 16, 0, 0, 0, time.UTC) + require.Equal(t, expectedBoundary, boundary) + + base := boundary.Sub(now) + require.GreaterOrEqual(t, delay, base, + "delay must be at least the base distance to boundary") + require.Less(t, delay, base+jitter, + "delay must be less than base + jitter") + }) +} diff --git a/enterprise/coderd/usage/cron_test.go b/enterprise/coderd/usage/cron_test.go new file mode 100644 index 0000000000..c2cf9e44d9 --- /dev/null +++ b/enterprise/coderd/usage/cron_test.go @@ -0,0 +1,98 @@ +package usage_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "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/coderd/usage/usagetypes" + "github.com/coder/coder/v2/enterprise/coderd/usage" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestCron(t *testing.T) { + t.Parallel() + + t.Run("BasicTick", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + clock := quartz.NewMock(t) + + // The existence check should return false so the event gets + // inserted. + db.EXPECT().UsageEventExistsByID(gomock.Any(), gomock.Any()). + Return(false, nil).AnyTimes() + + inserted := make(chan database.InsertUsageEventParams, 1) + db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, params database.InsertUsageEventParams) error { + inserted <- params + return nil + }).AnyTimes() + + inserter := usage.NewDBInserter(usage.InserterWithClock(clock)) + cron := usage.NewCron(clock, slogtest.Make(t, nil), db, inserter) + require.NoError(t, cron.Register(usage.CronJob{ + Name: "test-job", + Interval: 5 * time.Minute, + EventType: usagetypes.UsageEventTypeHBAISeatsV1, + Fn: func(_ context.Context) (usagetypes.HeartbeatEvent, error) { + return usagetypes.HBAISeats{Count: 42}, nil + }, + })) + + timerTrap := clock.Trap().NewTimer("test-job") + + cron.Start(ctx) + defer cron.Close() + defer timerTrap.Close() + + // Wait for timer creation, then fire it. The delay is the + // time until the next epoch-aligned boundary for the 5-minute + // interval — we don't assert the exact value since it depends + // on the mock clock's current time. + timerCall := timerTrap.MustWait(ctx) + timerCall.MustRelease(ctx) + clock.Advance(timerCall.Duration) + + // Verify the event was inserted with an epoch-aligned ID. + select { + case params := <-inserted: + assert.Contains(t, params.ID, "hb_ai_seats_v1:") + case <-ctx.Done(): + t.Fatal("timed out waiting for insert") + } + }) +} + +// TestAISeatsHeartbeat checks that AISeatsHeartbeat returns the +// correct event type and count. +func TestAISeatsHeartbeat(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + db.EXPECT().GetActiveAISeatCount(gomock.Any()).Return(int64(42), nil) + + fn := usage.AISeatsHeartbeat(db) + event, err := fn(ctx) + require.NoError(t, err) + + // Verify the event type and count. + hb, ok := event.(usagetypes.HBAISeats) + require.True(t, ok) + assert.Equal(t, int64(42), hb.Count) +} diff --git a/enterprise/coderd/usage/heartbeats.go b/enterprise/coderd/usage/heartbeats.go new file mode 100644 index 0000000000..c0171b4be9 --- /dev/null +++ b/enterprise/coderd/usage/heartbeats.go @@ -0,0 +1,31 @@ +package usage + +import ( + "context" + "time" + + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/usage/usagetypes" +) + +const ( + AISeatsInterval = 4 * time.Hour +) + +// AISeatsHeartbeat returns a HeartbeatFunc that queries the active +// AI seat count and emits it as an HBAISeats heartbeat event. +func AISeatsHeartbeat(db database.Store) HeartbeatFunc { + return func(ctx context.Context) (usagetypes.HeartbeatEvent, error) { + //nolint:gocritic // We are a publisher in this function + ctx = dbauthz.AsUsagePublisher(ctx) + count, err := db.GetActiveAISeatCount(ctx) + if err != nil { + return nil, xerrors.Errorf("get active AI seat count: %w", err) + } + + return usagetypes.HBAISeats{Count: count}, nil + } +} diff --git a/enterprise/coderd/usage/inserter.go b/enterprise/coderd/usage/inserter.go index f3566595a1..90fb6ab4ca 100644 --- a/enterprise/coderd/usage/inserter.go +++ b/enterprise/coderd/usage/inserter.go @@ -66,3 +66,27 @@ func (i *dbInserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.S CreatedAt: dbtime.Time(i.clock.Now()), }) } + +// InsertHeartbeatUsageEvent implements agplusage.Inserter. +func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error { + if !event.EventType().IsHeartbeat() { + return xerrors.Errorf("event type %q is not a heartbeat event", event.EventType()) + } + if err := event.Valid(); err != nil { + return xerrors.Errorf("invalid %q event: %w", event.EventType(), err) + } + + jsonData, err := json.Marshal(event.Fields()) + if err != nil { + return xerrors.Errorf("marshal event as JSON: %w", err) + } + + // Duplicate events are ignored by the query, so we don't need to check the + // error. + return tx.InsertUsageEvent(ctx, database.InsertUsageEventParams{ + ID: id, + EventType: string(event.EventType()), + EventData: jsonData, + CreatedAt: dbtime.Time(i.clock.Now()), + }) +} diff --git a/enterprise/coderd/workspaces_test.go b/enterprise/coderd/workspaces_test.go index 58ea5d78bf..129ac0df32 100644 --- a/enterprise/coderd/workspaces_test.go +++ b/enterprise/coderd/workspaces_test.go @@ -4788,7 +4788,7 @@ func TestWorkspaceAITask(t *testing.T) { wrk := coderdtest.CreateWorkspace(t, client, template.ID) build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, wrk.LatestBuild.ID) require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status) - require.Len(t, usage.GetEvents(), 0) + require.Len(t, usage.GetDiscreteEvents(), 0) }) t.Run("CreateTaskWorkspace", func(t *testing.T) { @@ -4815,7 +4815,7 @@ func TestWorkspaceAITask(t *testing.T) { build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, wrk.LatestBuild.ID) require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status) - require.Len(t, usage.GetEvents(), 1) + require.Len(t, usage.GetDiscreteEvents(), 1) usage.Reset() // Clean slate for easy checks // Stopping the workspace should not create additional usage. @@ -4825,7 +4825,7 @@ func TestWorkspaceAITask(t *testing.T) { }) require.NoError(t, err) coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID) - require.Len(t, usage.GetEvents(), 0) + require.Len(t, usage.GetDiscreteEvents(), 0) usage.Reset() // Clean slate for easy checks // Starting the workspace manually **WILL** create usage, as it's @@ -4836,6 +4836,6 @@ func TestWorkspaceAITask(t *testing.T) { }) require.NoError(t, err) coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID) - require.Len(t, usage.GetEvents(), 1) + require.Len(t, usage.GetDiscreteEvents(), 1) }) }