Skip to content

chore: add usage tracking package #19095

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion CODEOWNERS
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ tailnet/proto/ @spikecurtis @johnstcn
vpn/vpn.proto @spikecurtis @johnstcn
vpn/version.go @spikecurtis @johnstcn


# This caching code is particularly tricky, and one must be very careful when
# altering it.
coderd/files/ @aslilac
Expand All @@ -34,3 +33,8 @@ site/CLAUDE.md
# requires elite ball knowledge of most of the scheduling code to make changes
# without inadvertently affecting other parts of the codebase.
coderd/schedule/autostop.go @deansheather @DanielleMaywood

# Usage tracking code requires intimate knowledge of Tallyman and Metronome, as
# well as guidance from revenue.
coderd/usage/ @deansheather @spikecurtis
enterprise/coderd/usage/ @deansheather @spikecurtis
2 changes: 2 additions & 0 deletions coderd/apidoc/docs.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions coderd/apidoc/swagger.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions coderd/database/check_constraint.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

49 changes: 49 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
}
Expand Down Expand Up @@ -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)...)
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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)
}
Expand Down
31 changes: 31 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mocked! Nice! Been slowly converting them all 👍

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)
}))
}
21 changes: 21 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

43 changes: 43 additions & 0 deletions coderd/database/dbmock/dbmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

30 changes: 30 additions & 0 deletions coderd/database/dump.sql

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
DROP TABLE usage_events;
25 changes: 25 additions & 0 deletions coderd/database/migrations/000359_create_usage_events_table.up.sql
Original file line number Diff line number Diff line change
@@ -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')),
Comment on lines +3 to +6
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding enum values is actually easy now:

ALTER TYPE notification_method ADD VALUE IF NOT EXISTS 'inbox';

Removing them though is a problem:

-- The migration is about an enum value change
-- As we can not remove a value from an enum, we can let the down migration empty
-- In order to avoid any failure, we use ADD VALUE IF NOT EXISTS to add the value

(Although we never run down migrations so....)


Making it an enum generates the proper Golang enum

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Spike asked me to not do this in a prior comment because you can't add a value and then reference it in the same transaction unless you're on pg 17. We run all migrations in a single transaction (for good reason, to avoid partial upgrades), which means we can never reference enum types in future transactions.

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);
Loading
Loading