Skip to content

Commit d773056

Browse files
deansheatherdeansheather
committed
chore: add usage tracking package
1 parent d7bdb3c commit d773056

22 files changed

+1867
-16
lines changed

‎CODEOWNERS

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ tailnet/proto/ @spikecurtis @johnstcn
77
vpn/vpn.proto @spikecurtis @johnstcn
88
vpn/version.go @spikecurtis @johnstcn
99

10-
1110
# This caching code is particularly tricky, and one must be very careful when
1211
# altering it.
1312
coderd/files/ @aslilac
@@ -34,3 +33,8 @@ site/CLAUDE.md
3433
# requires elite ball knowledge of most of the scheduling code to make changes
3534
# without inadvertently affecting other parts of the codebase.
3635
coderd/schedule/autostop.go @deansheather @DanielleMaywood
36+
37+
# Usage tracking code requires intimate knowledge of Tallyman and Metronome, as
38+
# well as guidance from revenue.
39+
coderd/usage/ @deansheather
40+
enterprise/coderd/usage/ @deansheather

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3951,6 +3951,13 @@ func (q *querier) InsertTemplateVersionWorkspaceTag(ctx context.Context, arg dat
39513951
return q.db.InsertTemplateVersionWorkspaceTag(ctx, arg)
39523952
}
39533953

3954+
func (q *querier) InsertUsageEvent(ctx context.Context, arg database.InsertUsageEventParams) error {
3955+
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceSystem); err != nil {
3956+
return err
3957+
}
3958+
return q.db.InsertUsageEvent(ctx, arg)
3959+
}
3960+
39543961
func (q *querier) InsertUser(ctx context.Context, arg database.InsertUserParams) (database.User, error) {
39553962
// Always check if the assigned roles can actually be assigned by this actor.
39563963
impliedRoles := append([]rbac.RoleIdentifier{rbac.RoleMember()}, q.convertToDeploymentRoles(arg.RBACRoles)...)
@@ -4306,6 +4313,14 @@ func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string)
43064313
return q.db.RevokeDBCryptKey(ctx, activeKeyDigest)
43074314
}
43084315

4316+
func (q *querier) SelectUsageEventsForPublishing(ctx context.Context, arg time.Time) ([]database.UsageEvent, error) {
4317+
// ActionUpdate because we're updating the publish_started_at column.
4318+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
4319+
return nil, err
4320+
}
4321+
return q.db.SelectUsageEventsForPublishing(ctx, arg)
4322+
}
4323+
43094324
func (q *querier) TryAcquireLock(ctx context.Context, id int64) (bool, error) {
43104325
return q.db.TryAcquireLock(ctx, id)
43114326
}
@@ -4787,6 +4802,13 @@ func (q *querier) UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, arg da
47874802
return fetchAndExec(q.log, q.auth, policy.ActionUpdate, fetch, q.db.UpdateTemplateWorkspacesLastUsedAt)(ctx, arg)
47884803
}
47894804

4805+
func (q *querier) UpdateUsageEventsPostPublish(ctx context.Context, arg database.UpdateUsageEventsPostPublishParams) error {
4806+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
4807+
return err
4808+
}
4809+
return q.db.UpdateUsageEventsPostPublish(ctx, arg)
4810+
}
4811+
47904812
func (q *querier) UpdateUserDeletedByID(ctx context.Context, id uuid.UUID) error {
47914813
return deleteQ(q.log, q.auth, q.db.GetUserByID, q.db.UpdateUserDeletedByID)(ctx, id)
47924814
}

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5666,3 +5666,27 @@ func (s *MethodTestSuite) TestUserSecrets() {
56665666
Asserts(userSecret, policy.ActionRead, userSecret, policy.ActionDelete)
56675667
}))
56685668
}
5669+
5670+
func (s *MethodTestSuite) TestUsageEvents() {
5671+
s.Run("InsertUsageEvent", s.Subtest(func(db database.Store, check *expects) {
5672+
check.Args(database.InsertUsageEventParams{
5673+
ID: "1",
5674+
EventType: database.UsageEventTypeDcManagedAgentsV1,
5675+
EventData: []byte("{}"),
5676+
CreatedAt: dbtime.Now(),
5677+
}).Asserts(rbac.ResourceSystem, policy.ActionCreate)
5678+
}))
5679+
5680+
s.Run("SelectUsageEventsForPublishing", s.Subtest(func(db database.Store, check *expects) {
5681+
check.Args(dbtime.Now()).Asserts(rbac.ResourceSystem, policy.ActionUpdate)
5682+
}))
5683+
5684+
s.Run("UpdateUsageEventsPostPublish", s.Subtest(func(db database.Store, check *expects) {
5685+
check.Args(database.UpdateUsageEventsPostPublishParams{
5686+
Now: dbtime.Now(),
5687+
IDs: []string{"1", "2"},
5688+
FailureMessages: []string{"error", "error"},
5689+
SetPublishedAts: []bool{false, false},
5690+
}).Asserts(rbac.ResourceSystem, policy.ActionUpdate)
5691+
}))
5692+
}

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go

Lines changed: 43 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dump.sql

Lines changed: 37 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000358_create_usage_events_table.down.sql

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
DROP TABLE usage_events;
2+
DROP TYPE usage_event_type;

‎coderd/database/migrations/000358_create_usage_events_table.up.sql

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
CREATE TYPE usage_event_type AS ENUM (
2+
'dc_managed_agents_v1'
3+
);
4+
5+
COMMENT ON TYPE usage_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).';
6+
7+
CREATE TABLE usage_events (
8+
id TEXT PRIMARY KEY,
9+
event_type usage_event_type NOT NULL,
10+
event_data JSONB NOT NULL,
11+
created_at TIMESTAMP WITH TIME ZONE NOT NULL,
12+
publish_started_at TIMESTAMP WITH TIME ZONE DEFAULT NULL,
13+
published_at TIMESTAMP WITH TIME ZONE DEFAULT NULL,
14+
failure_message TEXT DEFAULT NULL
15+
);
16+
17+
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.';
18+
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.';
19+
COMMENT ON COLUMN usage_events.event_data IS 'Event payload. Determined by the matching usage struct for this event type.';
20+
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.';
21+
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.';
22+
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.';
23+
24+
CREATE INDEX idx_usage_events_created_at ON usage_events (created_at);
25+
CREATE INDEX idx_usage_events_publish_started_at ON usage_events (publish_started_at);
26+
CREATE INDEX idx_usage_events_published_at ON usage_events (published_at);

‎coderd/database/modelmethods.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"encoding/hex"
55
"sort"
66
"strconv"
7+
"strings"
78
"time"
89

910
"github.com/google/uuid"
@@ -636,3 +637,11 @@ func (m WorkspaceAgentVolumeResourceMonitor) Debounce(
636637
func (s UserSecret) RBACObject() rbac.Object {
637638
return rbac.ResourceUserSecret.WithID(s.ID).WithOwner(s.UserID.String())
638639
}
640+
641+
func (e UsageEventType) IsDiscrete() bool {
642+
return e.Valid() && strings.HasPrefix(string(e), "dc_")
643+
}
644+
645+
func (e UsageEventType) IsHeartbeat() bool {
646+
return e.Valid() && strings.HasPrefix(string(e), "hb_")
647+
}

‎coderd/database/models.go

Lines changed: 72 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)