chore: add usage tracking package (#19095)

Not used in coderd yet, see stack.

Adds two new packages:
- `coderd/usage`: provides an interface for the "Collector" as well as a stub implementation for AGPL
- `enterprise/coderd/usage`: provides an interface for the "Publisher" as well as a Tallyman implementation

Relates to https://github.com/coder/internal/issues/814

coderd/database/check_constraint.go

@@ -9,6 +9,7 @@ const (
 	CheckOneTimePasscodeSet                        CheckConstraint = "one_time_passcode_set"                            // users
 	CheckMaxProvisionerLogsLength                  CheckConstraint = "max_provisioner_logs_length"                      // provisioner_jobs
 	CheckValidationMonotonicOrder                  CheckConstraint = "validation_monotonic_order"                       // template_version_parameters
+	CheckUsageEventTypeCheck                       CheckConstraint = "usage_event_type_check"                           // usage_events
 	CheckMaxLogsLength                             CheckConstraint = "max_logs_length"                                  // workspace_agents
 	CheckSubsystemsNotNone                         CheckConstraint = "subsystems_not_none"                              // workspace_agents
 	CheckWorkspaceBuildsAiTaskSidebarAppIDRequired CheckConstraint = "workspace_builds_ai_task_sidebar_app_id_required" // workspace_builds

coderd/database/dbauthz/dbauthz.go

@@ -509,6 +509,25 @@ var (
 		}),
 		Scope: rbac.ScopeAll,
 	}.WithCachedASTValue()
+
+	subjectUsageTracker = rbac.Subject{
+		Type:         rbac.SubjectTypeUsageTracker,
+		FriendlyName: "Usage Tracker",
+		ID:           uuid.Nil.String(),
+		Roles: rbac.Roles([]rbac.Role{
+			{
+				Identifier:  rbac.RoleIdentifier{Name: "usage-tracker"},
+				DisplayName: "Usage Tracker",
+				Site: rbac.Permissions(map[string][]policy.Action{
+					rbac.ResourceLicense.Type:    {policy.ActionRead},
+					rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
+				}),
+				Org:  map[string][]rbac.Permission{},
+				User: []rbac.Permission{},
+			},
+		}),
+		Scope: rbac.ScopeAll,
+	}.WithCachedASTValue()
 )
 
 // AsProvisionerd returns a context with an actor that has permissions required
@@ -579,10 +598,18 @@ func AsPrebuildsOrchestrator(ctx context.Context) context.Context {
 	return As(ctx, subjectPrebuildsOrchestrator)
 }
 
+// AsFileReader returns a context with an actor that has permissions required
+// for reading all files.
 func AsFileReader(ctx context.Context) context.Context {
 	return As(ctx, subjectFileReader)
 }
 
+// AsUsageTracker returns a context with an actor that has permissions required
+// for creating, reading, and updating usage events.
+func AsUsageTracker(ctx context.Context) context.Context {
+	return As(ctx, subjectUsageTracker)
+}
+
 var AsRemoveActor = rbac.Subject{
 	ID: "remove-actor",
 }
@@ -3951,6 +3978,13 @@ func (q *querier) InsertTemplateVersionWorkspaceTag(ctx context.Context, arg dat
 	return q.db.InsertTemplateVersionWorkspaceTag(ctx, arg)
 }
 
+func (q *querier) InsertUsageEvent(ctx context.Context, arg database.InsertUsageEventParams) error {
+	if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceUsageEvent); err != nil {
+		return err
+	}
+	return q.db.InsertUsageEvent(ctx, arg)
+}
+
 func (q *querier) InsertUser(ctx context.Context, arg database.InsertUserParams) (database.User, error) {
 	// Always check if the assigned roles can actually be assigned by this actor.
 	impliedRoles := append([]rbac.RoleIdentifier{rbac.RoleMember()}, q.convertToDeploymentRoles(arg.RBACRoles)...)
@@ -4306,6 +4340,14 @@ func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string)
 	return q.db.RevokeDBCryptKey(ctx, activeKeyDigest)
 }
 
+func (q *querier) SelectUsageEventsForPublishing(ctx context.Context, arg time.Time) ([]database.UsageEvent, error) {
+	// ActionUpdate because we're updating the publish_started_at column.
+	if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceUsageEvent); err != nil {
+		return nil, err
+	}
+	return q.db.SelectUsageEventsForPublishing(ctx, arg)
+}
+
 func (q *querier) TryAcquireLock(ctx context.Context, id int64) (bool, error) {
 	return q.db.TryAcquireLock(ctx, id)
 }
@@ -4787,6 +4829,13 @@ func (q *querier) UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, arg da
 	return fetchAndExec(q.log, q.auth, policy.ActionUpdate, fetch, q.db.UpdateTemplateWorkspacesLastUsedAt)(ctx, arg)
 }
 
+func (q *querier) UpdateUsageEventsPostPublish(ctx context.Context, arg database.UpdateUsageEventsPostPublishParams) error {
+	if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceUsageEvent); err != nil {
+		return err
+	}
+	return q.db.UpdateUsageEventsPostPublish(ctx, arg)
+}
+
 func (q *querier) UpdateUserDeletedByID(ctx context.Context, id uuid.UUID) error {
 	return deleteQ(q.log, q.auth, q.db.GetUserByID, q.db.UpdateUserDeletedByID)(ctx, id)
 }

coderd/database/dbauthz/dbauthz_test.go

@@ -5666,3 +5666,34 @@ func (s *MethodTestSuite) TestUserSecrets() {
 			Asserts(userSecret, policy.ActionRead, userSecret, policy.ActionDelete)
 	}))
 }
+
+func (s *MethodTestSuite) TestUsageEvents() {
+	s.Run("InsertUsageEvent", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		params := database.InsertUsageEventParams{
+			ID:        "1",
+			EventType: "dc_managed_agents_v1",
+			EventData: []byte("{}"),
+			CreatedAt: dbtime.Now(),
+		}
+		db.EXPECT().InsertUsageEvent(gomock.Any(), params).Return(nil)
+		check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionCreate)
+	}))
+
+	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)
+		check.Args(now).Asserts(rbac.ResourceUsageEvent, policy.ActionUpdate)
+	}))
+
+	s.Run("UpdateUsageEventsPostPublish", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		now := dbtime.Now()
+		params := database.UpdateUsageEventsPostPublishParams{
+			Now:             now,
+			IDs:             []string{"1", "2"},
+			FailureMessages: []string{"error", "error"},
+			SetPublishedAts: []bool{false, false},
+		}
+		db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), params).Return(nil)
+		check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionUpdate)
+	}))
+}

coderd/database/dbmetrics/querymetrics.go

@@ -2392,6 +2392,13 @@ func (m queryMetricsStore) InsertTemplateVersionWorkspaceTag(ctx context.Context
 	return r0, r1
 }
 
+func (m queryMetricsStore) InsertUsageEvent(ctx context.Context, arg database.InsertUsageEventParams) error {
+	start := time.Now()
+	r0 := m.s.InsertUsageEvent(ctx, arg)
+	m.queryLatencies.WithLabelValues("InsertUsageEvent").Observe(time.Since(start).Seconds())
+	return r0
+}
+
 func (m queryMetricsStore) InsertUser(ctx context.Context, arg database.InsertUserParams) (database.User, error) {
 	start := time.Now()
 	user, err := m.s.InsertUser(ctx, arg)
@@ -2651,6 +2658,13 @@ func (m queryMetricsStore) RevokeDBCryptKey(ctx context.Context, activeKeyDigest
 	return r0
 }
 
+func (m queryMetricsStore) SelectUsageEventsForPublishing(ctx context.Context, arg time.Time) ([]database.UsageEvent, error) {
+	start := time.Now()
+	r0, r1 := m.s.SelectUsageEventsForPublishing(ctx, arg)
+	m.queryLatencies.WithLabelValues("SelectUsageEventsForPublishing").Observe(time.Since(start).Seconds())
+	return r0, r1
+}
+
 func (m queryMetricsStore) TryAcquireLock(ctx context.Context, pgTryAdvisoryXactLock int64) (bool, error) {
 	start := time.Now()
 	ok, err := m.s.TryAcquireLock(ctx, pgTryAdvisoryXactLock)
@@ -2938,6 +2952,13 @@ func (m queryMetricsStore) UpdateTemplateWorkspacesLastUsedAt(ctx context.Contex
 	return r0
 }
 
+func (m queryMetricsStore) UpdateUsageEventsPostPublish(ctx context.Context, arg database.UpdateUsageEventsPostPublishParams) error {
+	start := time.Now()
+	r0 := m.s.UpdateUsageEventsPostPublish(ctx, arg)
+	m.queryLatencies.WithLabelValues("UpdateUsageEventsPostPublish").Observe(time.Since(start).Seconds())
+	return r0
+}
+
 func (m queryMetricsStore) UpdateUserDeletedByID(ctx context.Context, id uuid.UUID) error {
 	start := time.Now()
 	r0 := m.s.UpdateUserDeletedByID(ctx, id)

coderd/database/dbmock/dbmock.go

@@ -5107,6 +5107,20 @@ func (mr *MockStoreMockRecorder) InsertTemplateVersionWorkspaceTag(ctx, arg any)
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertTemplateVersionWorkspaceTag", reflect.TypeOf((*MockStore)(nil).InsertTemplateVersionWorkspaceTag), ctx, arg)
 }
 
+// InsertUsageEvent mocks base method.
+func (m *MockStore) InsertUsageEvent(ctx context.Context, arg database.InsertUsageEventParams) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "InsertUsageEvent", ctx, arg)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// InsertUsageEvent indicates an expected call of InsertUsageEvent.
+func (mr *MockStoreMockRecorder) InsertUsageEvent(ctx, arg any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertUsageEvent", reflect.TypeOf((*MockStore)(nil).InsertUsageEvent), ctx, arg)
+}
+
 // InsertUser mocks base method.
 func (m *MockStore) InsertUser(ctx context.Context, arg database.InsertUserParams) (database.User, error) {
 	m.ctrl.T.Helper()
@@ -5682,6 +5696,21 @@ func (mr *MockStoreMockRecorder) RevokeDBCryptKey(ctx, activeKeyDigest any) *gom
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RevokeDBCryptKey", reflect.TypeOf((*MockStore)(nil).RevokeDBCryptKey), ctx, activeKeyDigest)
 }
 
+// SelectUsageEventsForPublishing mocks base method.
+func (m *MockStore) SelectUsageEventsForPublishing(ctx context.Context, now time.Time) ([]database.UsageEvent, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SelectUsageEventsForPublishing", ctx, now)
+	ret0, _ := ret[0].([]database.UsageEvent)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SelectUsageEventsForPublishing indicates an expected call of SelectUsageEventsForPublishing.
+func (mr *MockStoreMockRecorder) SelectUsageEventsForPublishing(ctx, now any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SelectUsageEventsForPublishing", reflect.TypeOf((*MockStore)(nil).SelectUsageEventsForPublishing), ctx, now)
+}
+
 // TryAcquireLock mocks base method.
 func (m *MockStore) TryAcquireLock(ctx context.Context, pgTryAdvisoryXactLock int64) (bool, error) {
 	m.ctrl.T.Helper()
@@ -6270,6 +6299,20 @@ func (mr *MockStoreMockRecorder) UpdateTemplateWorkspacesLastUsedAt(ctx, arg any
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTemplateWorkspacesLastUsedAt", reflect.TypeOf((*MockStore)(nil).UpdateTemplateWorkspacesLastUsedAt), ctx, arg)
 }
 
+// UpdateUsageEventsPostPublish mocks base method.
+func (m *MockStore) UpdateUsageEventsPostPublish(ctx context.Context, arg database.UpdateUsageEventsPostPublishParams) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateUsageEventsPostPublish", ctx, arg)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateUsageEventsPostPublish indicates an expected call of UpdateUsageEventsPostPublish.
+func (mr *MockStoreMockRecorder) UpdateUsageEventsPostPublish(ctx, arg any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateUsageEventsPostPublish", reflect.TypeOf((*MockStore)(nil).UpdateUsageEventsPostPublish), ctx, arg)
+}
+
 // UpdateUserDeletedByID mocks base method.
 func (m *MockStore) UpdateUserDeletedByID(ctx context.Context, id uuid.UUID) error {
 	m.ctrl.T.Helper()

coderd/database/dump.sql

@@ -1832,6 +1832,31 @@ CREATE VIEW template_with_names AS
 
 COMMENT ON VIEW template_with_names IS 'Joins in the display name information such as username, avatar, and organization name.';
 
+CREATE TABLE usage_events (
+    id text NOT NULL,
+    event_type text NOT NULL,
+    event_data jsonb NOT NULL,
+    created_at timestamp with time zone NOT NULL,
+    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))
+);
+
+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.';
+
+COMMENT ON COLUMN usage_events.id IS 'For "discrete" event types, this is a random UUID. For "heartbeat" event types, this is a combination of the event type and a truncated timestamp.';
+
+COMMENT ON COLUMN usage_events.event_type IS 'The usage event type with version. "dc" means "discrete" (e.g. a single event, for counters), "hb" means "heartbeat" (e.g. a recurring event that contains a total count of usage generated from the database, for gauges).';
+
+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.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.';
+
+COMMENT ON COLUMN usage_events.failure_message IS 'Set to an error message when the event is temporarily or permanently unsuccessfully published to the usage collector service.';
+
 CREATE TABLE user_configs (
     user_id uuid NOT NULL,
     key character varying(256) NOT NULL,
@@ -2681,6 +2706,9 @@ ALTER TABLE ONLY template_versions
 ALTER TABLE ONLY templates
     ADD CONSTRAINT templates_pkey PRIMARY KEY (id);
 
+ALTER TABLE ONLY usage_events
+    ADD CONSTRAINT usage_events_pkey PRIMARY KEY (id);
+
 ALTER TABLE ONLY user_configs
     ADD CONSTRAINT user_configs_pkey PRIMARY KEY (user_id, key);
 
@@ -2849,6 +2877,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_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);
 
 CREATE INDEX idx_user_status_changes_changed_at ON user_status_changes USING btree (changed_at);

coderd/database/migrations/000359_create_usage_events_table.down.sql

@@ -0,0 +1 @@
+DROP TABLE usage_events;

coderd/database/migrations/000359_create_usage_events_table.up.sql

@@ -0,0 +1,25 @@
+CREATE TABLE usage_events (
+  id TEXT PRIMARY KEY,
+  -- We use a TEXT column with a CHECK constraint rather than an enum because of
+  -- the limitations with adding new values to an enum and using them in the
+  -- same transaction.
+  event_type TEXT NOT NULL CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1')),
+  event_data JSONB NOT NULL,
+  created_at TIMESTAMP WITH TIME ZONE NOT NULL,
+  publish_started_at TIMESTAMP WITH TIME ZONE DEFAULT NULL,
+  published_at TIMESTAMP WITH TIME ZONE DEFAULT NULL,
+  failure_message TEXT DEFAULT NULL
+);
+
+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.';
+COMMENT ON COLUMN usage_events.id IS 'For "discrete" event types, this is a random UUID. For "heartbeat" event types, this is a combination of the event type and a truncated timestamp.';
+COMMENT ON COLUMN usage_events.event_type IS 'The usage event type with version. "dc" means "discrete" (e.g. a single event, for counters), "hb" means "heartbeat" (e.g. a recurring event that contains a total count of usage generated from the database, for gauges).';
+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.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.';
+COMMENT ON COLUMN usage_events.failure_message IS 'Set to an error message when the event is temporarily or permanently unsuccessfully published to the usage collector service.';
+
+-- Create an index with all three fields used by the
+-- SelectUsageEventsForPublishing query.
+CREATE INDEX idx_usage_events_select_for_publishing
+  ON usage_events (published_at, publish_started_at, created_at);

coderd/database/migrations/testdata/fixtures/000359_create_usage_events_table.up.sql

@@ -0,0 +1,60 @@
+INSERT INTO usage_events (
+    id,
+    event_type,
+    event_data,
+    created_at,
+    publish_started_at,
+    published_at,
+    failure_message
+)
+VALUES
+-- Unpublished dc_managed_agents_v1 event.
+(
+    'event1',
+    'dc_managed_agents_v1',
+    '{"count":1}',
+    '2023-01-01 00:00:00+00',
+    NULL,
+    NULL,
+    NULL
+),
+-- Successfully published dc_managed_agents_v1 event.
+(
+    'event2',
+    'dc_managed_agents_v1',
+    '{"count":2}',
+    '2023-01-01 00:00:00+00',
+    NULL,
+    '2023-01-01 00:00:02+00',
+    NULL
+),
+-- Publish in progress dc_managed_agents_v1 event.
+(
+    'event3',
+    'dc_managed_agents_v1',
+    '{"count":3}',
+    '2023-01-01 00:00:00+00',
+    '2023-01-01 00:00:01+00',
+    NULL,
+    NULL
+),
+-- Temporarily failed to publish dc_managed_agents_v1 event.
+(
+    'event4',
+    'dc_managed_agents_v1',
+    '{"count":4}',
+    '2023-01-01 00:00:00+00',
+    NULL,
+    NULL,
+    'publish failed temporarily'
+),
+-- Permanently failed to publish dc_managed_agents_v1 event.
+(
+    'event5',
+    'dc_managed_agents_v1',
+    '{"count":5}',
+    '2023-01-01 00:00:00+00',
+    NULL,
+    '2023-01-01 00:00:02+00',
+    'publish failed permanently'
+)

coderd/database/models.go

@@ -3759,6 +3759,23 @@ type TemplateVersionWorkspaceTag struct {
 	Value             string    `db:"value" json:"value"`
 }
 
+// usage_events contains usage data that is collected from the product and potentially shipped to the usage collector service.
+type UsageEvent struct {
+	// For "discrete" event types, this is a random UUID. For "heartbeat" event types, this is a combination of the event type and a truncated timestamp.
+	ID string `db:"id" json:"id"`
+	// The usage event type with version. "dc" means "discrete" (e.g. a single event, for counters), "hb" means "heartbeat" (e.g. a recurring event that contains a total count of usage generated from the database, for gauges).
+	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"`
+	// 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.
+	PublishedAt sql.NullTime `db:"published_at" json:"published_at"`
+	// Set to an error message when the event is temporarily or permanently unsuccessfully published to the usage collector service.
+	FailureMessage sql.NullString `db:"failure_message" json:"failure_message"`
+}
+
 type User struct {
 	ID             uuid.UUID      `db:"id" json:"id"`
 	Email          string         `db:"email" json:"email"`

coderd/database/querier.go

@@ -522,6 +522,9 @@ type sqlcQuerier interface {
 	InsertTemplateVersionTerraformValuesByJobID(ctx context.Context, arg InsertTemplateVersionTerraformValuesByJobIDParams) error
 	InsertTemplateVersionVariable(ctx context.Context, arg InsertTemplateVersionVariableParams) (TemplateVersionVariable, error)
 	InsertTemplateVersionWorkspaceTag(ctx context.Context, arg InsertTemplateVersionWorkspaceTagParams) (TemplateVersionWorkspaceTag, error)
+	// Duplicate events are ignored intentionally to allow for multiple replicas to
+	// publish heartbeat events.
+	InsertUsageEvent(ctx context.Context, arg InsertUsageEventParams) error
 	InsertUser(ctx context.Context, arg InsertUserParams) (User, error)
 	// InsertUserGroupsByID adds a user to all provided groups, if they exist.
 	// If there is a conflict, the user is already a member
@@ -568,6 +571,11 @@ type sqlcQuerier interface {
 	RemoveUserFromAllGroups(ctx context.Context, userID uuid.UUID) error
 	RemoveUserFromGroups(ctx context.Context, arg RemoveUserFromGroupsParams) ([]uuid.UUID, error)
 	RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error
+	// Note that this selects from the CTE, not the original table. The CTE is named
+	// the same as the original table to trick sqlc into reusing the existing struct
+	// for the table.
+	// The CTE and the reorder is required because UPDATE doesn't guarantee order.
+	SelectUsageEventsForPublishing(ctx context.Context, now time.Time) ([]UsageEvent, error)
 	// Non blocking lock. Returns true if the lock was acquired, false otherwise.
 	//
 	// This must be called from within a transaction. The lock will be automatically
@@ -614,6 +622,7 @@ type sqlcQuerier interface {
 	UpdateTemplateVersionDescriptionByJobID(ctx context.Context, arg UpdateTemplateVersionDescriptionByJobIDParams) error
 	UpdateTemplateVersionExternalAuthProvidersByJobID(ctx context.Context, arg UpdateTemplateVersionExternalAuthProvidersByJobIDParams) error
 	UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, arg UpdateTemplateWorkspacesLastUsedAtParams) error
+	UpdateUsageEventsPostPublish(ctx context.Context, arg UpdateUsageEventsPostPublishParams) error
 	UpdateUserDeletedByID(ctx context.Context, id uuid.UUID) error
 	UpdateUserGithubComUserID(ctx context.Context, arg UpdateUserGithubComUserIDParams) error
 	UpdateUserHashedOneTimePasscode(ctx context.Context, arg UpdateUserHashedOneTimePasscodeParams) error

coderd/database/queries.sql.go

@@ -13519,6 +13519,161 @@ func (q *sqlQuerier) DisableForeignKeysAndTriggers(ctx context.Context) error {
 	return err
 }
 
+const insertUsageEvent = `-- name: InsertUsageEvent :exec
+INSERT INTO
+    usage_events (
+        id,
+        event_type,
+        event_data,
+        created_at,
+        publish_started_at,
+        published_at,
+        failure_message
+    )
+VALUES
+    ($1, $2, $3, $4, NULL, NULL, NULL)
+ON CONFLICT (id) DO NOTHING
+`
+
+type InsertUsageEventParams struct {
+	ID        string          `db:"id" json:"id"`
+	EventType string          `db:"event_type" json:"event_type"`
+	EventData json.RawMessage `db:"event_data" json:"event_data"`
+	CreatedAt time.Time       `db:"created_at" json:"created_at"`
+}
+
+// Duplicate events are ignored intentionally to allow for multiple replicas to
+// publish heartbeat events.
+func (q *sqlQuerier) InsertUsageEvent(ctx context.Context, arg InsertUsageEventParams) error {
+	_, err := q.db.ExecContext(ctx, insertUsageEvent,
+		arg.ID,
+		arg.EventType,
+		arg.EventData,
+		arg.CreatedAt,
+	)
+	return err
+}
+
+const selectUsageEventsForPublishing = `-- name: SelectUsageEventsForPublishing :many
+WITH usage_events AS (
+    UPDATE
+        usage_events
+    SET
+        publish_started_at = $1::timestamptz
+    WHERE
+        id IN (
+            SELECT
+                potential_event.id
+            FROM
+                usage_events potential_event
+            WHERE
+                -- Do not publish events that have already been published or
+                -- have permanently failed to publish.
+                potential_event.published_at IS NULL
+                -- Do not publish events that are already being published by
+                -- another replica.
+                AND (
+                    potential_event.publish_started_at IS NULL
+                    -- If the event has publish_started_at set, it must be older
+                    -- than an hour ago. This is so we can retry publishing
+                    -- events where the replica exited or couldn't update the
+                    -- row.
+                    -- The parenthesis around @now::timestamptz are necessary to
+                    -- avoid sqlc from generating an extra argument.
+                    OR potential_event.publish_started_at < ($1::timestamptz) - INTERVAL '1 hour'
+                )
+                -- Do not publish events older than 30 days. Tallyman will
+                -- always permanently reject these events anyways. This is to
+                -- avoid duplicate events being billed to customers, as
+                -- Metronome will only deduplicate events within 34 days.
+                -- Also, the same parenthesis thing here as above.
+                AND potential_event.created_at > ($1::timestamptz) - INTERVAL '30 days'
+            ORDER BY potential_event.created_at ASC
+            FOR UPDATE SKIP LOCKED
+            LIMIT 100
+        )
+    RETURNING id, event_type, event_data, created_at, publish_started_at, published_at, failure_message
+)
+SELECT id, event_type, event_data, created_at, publish_started_at, published_at, failure_message
+FROM usage_events
+ORDER BY created_at ASC
+`
+
+// Note that this selects from the CTE, not the original table. The CTE is named
+// the same as the original table to trick sqlc into reusing the existing struct
+// for the table.
+// The CTE and the reorder is required because UPDATE doesn't guarantee order.
+func (q *sqlQuerier) SelectUsageEventsForPublishing(ctx context.Context, now time.Time) ([]UsageEvent, error) {
+	rows, err := q.db.QueryContext(ctx, selectUsageEventsForPublishing, now)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []UsageEvent
+	for rows.Next() {
+		var i UsageEvent
+		if err := rows.Scan(
+			&i.ID,
+			&i.EventType,
+			&i.EventData,
+			&i.CreatedAt,
+			&i.PublishStartedAt,
+			&i.PublishedAt,
+			&i.FailureMessage,
+		); err != nil {
+			return nil, err
+		}
+		items = append(items, i)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
+const updateUsageEventsPostPublish = `-- name: UpdateUsageEventsPostPublish :exec
+UPDATE
+    usage_events
+SET
+    publish_started_at = NULL,
+    published_at = CASE WHEN input.set_published_at THEN $1::timestamptz ELSE NULL END,
+    failure_message = NULLIF(input.failure_message, '')
+FROM (
+    SELECT
+        UNNEST($2::text[]) AS id,
+        UNNEST($3::text[]) AS failure_message,
+        UNNEST($4::boolean[]) AS set_published_at
+) input
+WHERE
+    input.id = usage_events.id
+    -- If the number of ids, failure messages, and set published ats are not the
+    -- same, do not do anything. Unfortunately you can't really throw from a
+    -- query without writing a function or doing some jank like dividing by
+    -- zero, so this is the best we can do.
+    AND cardinality($2::text[]) = cardinality($3::text[])
+    AND cardinality($2::text[]) = cardinality($4::boolean[])
+`
+
+type UpdateUsageEventsPostPublishParams struct {
+	Now             time.Time `db:"now" json:"now"`
+	IDs             []string  `db:"ids" json:"ids"`
+	FailureMessages []string  `db:"failure_messages" json:"failure_messages"`
+	SetPublishedAts []bool    `db:"set_published_ats" json:"set_published_ats"`
+}
+
+func (q *sqlQuerier) UpdateUsageEventsPostPublish(ctx context.Context, arg UpdateUsageEventsPostPublishParams) error {
+	_, err := q.db.ExecContext(ctx, updateUsageEventsPostPublish,
+		arg.Now,
+		pq.Array(arg.IDs),
+		pq.Array(arg.FailureMessages),
+		pq.Array(arg.SetPublishedAts),
+	)
+	return 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

coderd/database/queries/usageevents.sql

@@ -0,0 +1,86 @@
+-- name: InsertUsageEvent :exec
+-- Duplicate events are ignored intentionally to allow for multiple replicas to
+-- publish heartbeat events.
+INSERT INTO
+    usage_events (
+        id,
+        event_type,
+        event_data,
+        created_at,
+        publish_started_at,
+        published_at,
+        failure_message
+    )
+VALUES
+    (@id, @event_type, @event_data, @created_at, NULL, NULL, NULL)
+ON CONFLICT (id) DO NOTHING;
+
+-- name: SelectUsageEventsForPublishing :many
+WITH usage_events AS (
+    UPDATE
+        usage_events
+    SET
+        publish_started_at = @now::timestamptz
+    WHERE
+        id IN (
+            SELECT
+                potential_event.id
+            FROM
+                usage_events potential_event
+            WHERE
+                -- Do not publish events that have already been published or
+                -- have permanently failed to publish.
+                potential_event.published_at IS NULL
+                -- Do not publish events that are already being published by
+                -- another replica.
+                AND (
+                    potential_event.publish_started_at IS NULL
+                    -- If the event has publish_started_at set, it must be older
+                    -- than an hour ago. This is so we can retry publishing
+                    -- events where the replica exited or couldn't update the
+                    -- row.
+                    -- The parenthesis around @now::timestamptz are necessary to
+                    -- avoid sqlc from generating an extra argument.
+                    OR potential_event.publish_started_at < (@now::timestamptz) - INTERVAL '1 hour'
+                )
+                -- Do not publish events older than 30 days. Tallyman will
+                -- always permanently reject these events anyways. This is to
+                -- avoid duplicate events being billed to customers, as
+                -- Metronome will only deduplicate events within 34 days.
+                -- Also, the same parenthesis thing here as above.
+                AND potential_event.created_at > (@now::timestamptz) - INTERVAL '30 days'
+            ORDER BY potential_event.created_at ASC
+            FOR UPDATE SKIP LOCKED
+            LIMIT 100
+        )
+    RETURNING *
+)
+SELECT *
+-- Note that this selects from the CTE, not the original table. The CTE is named
+-- the same as the original table to trick sqlc into reusing the existing struct
+-- for the table.
+FROM usage_events
+-- The CTE and the reorder is required because UPDATE doesn't guarantee order.
+ORDER BY created_at ASC;
+
+-- name: UpdateUsageEventsPostPublish :exec
+UPDATE
+    usage_events
+SET
+    publish_started_at = NULL,
+    published_at = CASE WHEN input.set_published_at THEN @now::timestamptz ELSE NULL END,
+    failure_message = NULLIF(input.failure_message, '')
+FROM (
+    SELECT
+        UNNEST(@ids::text[]) AS id,
+        UNNEST(@failure_messages::text[]) AS failure_message,
+        UNNEST(@set_published_ats::boolean[]) AS set_published_at
+) input
+WHERE
+    input.id = usage_events.id
+    -- If the number of ids, failure messages, and set published ats are not the
+    -- same, do not do anything. Unfortunately you can't really throw from a
+    -- query without writing a function or doing some jank like dividing by
+    -- zero, so this is the best we can do.
+    AND cardinality(@ids::text[]) = cardinality(@failure_messages::text[])
+    AND cardinality(@ids::text[]) = cardinality(@set_published_ats::boolean[]);

coderd/database/unique_constraint.go

@@ -67,6 +67,7 @@ const (
 	UniqueTemplateVersionsPkey                                UniqueConstraint = "template_versions_pkey"                                          // ALTER TABLE ONLY template_versions ADD CONSTRAINT template_versions_pkey PRIMARY KEY (id);
 	UniqueTemplateVersionsTemplateIDNameKey                   UniqueConstraint = "template_versions_template_id_name_key"                          // ALTER TABLE ONLY template_versions ADD CONSTRAINT template_versions_template_id_name_key UNIQUE (template_id, name);
 	UniqueTemplatesPkey                                       UniqueConstraint = "templates_pkey"                                                  // ALTER TABLE ONLY templates ADD CONSTRAINT templates_pkey PRIMARY KEY (id);
+	UniqueUsageEventsPkey                                     UniqueConstraint = "usage_events_pkey"                                               // ALTER TABLE ONLY usage_events ADD CONSTRAINT usage_events_pkey PRIMARY KEY (id);
 	UniqueUserConfigsPkey                                     UniqueConstraint = "user_configs_pkey"                                               // ALTER TABLE ONLY user_configs ADD CONSTRAINT user_configs_pkey PRIMARY KEY (user_id, key);
 	UniqueUserDeletedPkey                                     UniqueConstraint = "user_deleted_pkey"                                               // ALTER TABLE ONLY user_deleted ADD CONSTRAINT user_deleted_pkey PRIMARY KEY (id);
 	UniqueUserLinksPkey                                       UniqueConstraint = "user_links_pkey"                                                 // ALTER TABLE ONLY user_links ADD CONSTRAINT user_links_pkey PRIMARY KEY (user_id, login_type);