Skip to content

Commit f2963cd

Browse files
deansheatherdeansheather
committed
chore: add usage tracking package
1 parent 4bced62 commit f2963cd

23 files changed

+3708
-1815
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
@@ -29,3 +28,8 @@ site/src/api/countriesGenerated.ts
2928
site/src/api/rbacresourcesGenerated.ts
3029
site/src/api/typesGenerated.ts
3130
site/CLAUDE.md
31+
32+
# Usage tracking code requires intimate knowledge of Tallyman and Metronome, as
33+
# well as guidance from revenue.
34+
coderd/usage/ @deansheather
35+
enterprise/coderd/usage/ @deansheather

‎coderd/database/dbauthz/dbauthz.go

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

3916+
func (q *querier) InsertUsageEvent(ctx context.Context, arg database.InsertUsageEventParams) error {
3917+
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceSystem); err != nil {
3918+
return err
3919+
}
3920+
return q.db.InsertUsageEvent(ctx, arg)
3921+
}
3922+
39163923
func (q *querier) InsertUser(ctx context.Context, arg database.InsertUserParams) (database.User, error) {
39173924
// Always check if the assigned roles can actually be assigned by this actor.
39183925
impliedRoles := append([]rbac.RoleIdentifier{rbac.RoleMember()}, q.convertToDeploymentRoles(arg.RBACRoles)...)
@@ -4260,6 +4267,14 @@ func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string)
42604267
return q.db.RevokeDBCryptKey(ctx, activeKeyDigest)
42614268
}
42624269

4270+
func (q *querier) SelectUsageEventsForPublishing(ctx context.Context, arg time.Time) ([]database.UsageEvent, error) {
4271+
// ActionUpdate because we're updating the publish_started_at column.
4272+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
4273+
return nil, err
4274+
}
4275+
return q.db.SelectUsageEventsForPublishing(ctx, arg)
4276+
}
4277+
42634278
func (q *querier) TryAcquireLock(ctx context.Context, id int64) (bool, error) {
42644279
return q.db.TryAcquireLock(ctx, id)
42654280
}
@@ -4725,6 +4740,13 @@ func (q *querier) UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, arg da
47254740
return fetchAndExec(q.log, q.auth, policy.ActionUpdate, fetch, q.db.UpdateTemplateWorkspacesLastUsedAt)(ctx, arg)
47264741
}
47274742

4743+
func (q *querier) UpdateUsageEventsPostPublish(ctx context.Context, arg database.UpdateUsageEventsPostPublishParams) error {
4744+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
4745+
return err
4746+
}
4747+
return q.db.UpdateUsageEventsPostPublish(ctx, arg)
4748+
}
4749+
47284750
func (q *querier) UpdateUserDeletedByID(ctx context.Context, id uuid.UUID) error {
47294751
return deleteQ(q.log, q.auth, q.db.GetUserByID, q.db.UpdateUserDeletedByID)(ctx, id)
47304752
}

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5845,3 +5845,27 @@ func (s *MethodTestSuite) TestAuthorizePrebuiltWorkspace() {
58455845
}).Asserts(w, policy.ActionUpdate, w.AsPrebuild(), policy.ActionUpdate)
58465846
}))
58475847
}
5848+
5849+
func (s *MethodTestSuite) TestUsageEvents() {
5850+
s.Run("InsertUsageEvent", s.Subtest(func(db database.Store, check *expects) {
5851+
check.Args(database.InsertUsageEventParams{
5852+
ID: "1",
5853+
EventType: database.UsageEventTypeDcManagedAgentsV1,
5854+
EventData: []byte("{}"),
5855+
CreatedAt: dbtime.Now(),
5856+
}).Asserts(rbac.ResourceSystem, policy.ActionCreate)
5857+
}))
5858+
5859+
s.Run("SelectUsageEventsForPublishing", s.Subtest(func(db database.Store, check *expects) {
5860+
check.Args(dbtime.Now()).Asserts(rbac.ResourceSystem, policy.ActionUpdate)
5861+
}))
5862+
5863+
s.Run("UpdateUsageEventsPostPublish", s.Subtest(func(db database.Store, check *expects) {
5864+
check.Args(database.UpdateUsageEventsPostPublishParams{
5865+
Now: dbtime.Now(),
5866+
IDs: []string{"1", "2"},
5867+
FailureMessages: []string{"error", "error"},
5868+
SetPublishedAts: []bool{false, false},
5869+
}).Asserts(rbac.ResourceSystem, policy.ActionUpdate)
5870+
}))
5871+
}

‎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/000353_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/000353_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"
@@ -628,3 +629,11 @@ func (m WorkspaceAgentVolumeResourceMonitor) Debounce(
628629

629630
return m.DebouncedUntil, false
630631
}
632+
633+
func (e UsageEventType) IsDiscrete() bool {
634+
return e.Valid() && strings.HasPrefix(string(e), "dc_")
635+
}
636+
637+
func (e UsageEventType) IsHeartbeat() bool {
638+
return e.Valid() && strings.HasPrefix(string(e), "hb_")
639+
}

‎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)