Skip to content

Commit 6eb02d1

Browse files
deansheatherdeansheather
authored
chore: wire up usage tracking for managed agents (#19096)
Wires up the usage collector and publisher to coderd. Relates to coder/internal#814
1 parent dd867bd commit 6eb02d1

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+539
-345
lines changed

‎cli/delete_test.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,6 @@ func TestDelete(t *testing.T) {
111111
// The API checks if the user has any workspaces, so we cannot delete a user
112112
// this way.
113113
ctx := testutil.Context(t, testutil.WaitShort)
114-
// nolint:gocritic // Unit test
115114
err := api.Database.UpdateUserDeletedByID(dbauthz.AsSystemRestricted(ctx), deleteMeUser.ID)
116115
require.NoError(t, err)
117116

‎cli/provisioners_test.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ func TestProvisioners_Golden(t *testing.T) {
3131
// Replace UUIDs with predictable values for golden files.
3232
replace := make(map[string]string)
3333
updateReplaceUUIDs := func(coderdAPI *coderd.API) {
34-
//nolint:gocritic // This is a test.
3534
systemCtx := dbauthz.AsSystemRestricted(context.Background())
3635
provisioners, err := coderdAPI.Database.GetProvisionerDaemons(systemCtx)
3736
require.NoError(t, err)

‎coderd/agentapi/subagent_test.go

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,7 @@ func TestSubAgentAPI(t *testing.T) {
163163
agentID, err := uuid.FromBytes(createResp.Agent.Id)
164164
require.NoError(t, err)
165165

166-
agent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test.
166+
agent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID)
167167
require.NoError(t, err)
168168

169169
assert.Equal(t, tt.agentName, agent.Name)
@@ -621,7 +621,7 @@ func TestSubAgentAPI(t *testing.T) {
621621
agentID, err := uuid.FromBytes(createResp.Agent.Id)
622622
require.NoError(t, err)
623623

624-
apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test.
624+
apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID)
625625
require.NoError(t, err)
626626

627627
// Sort the apps for determinism
@@ -751,7 +751,7 @@ func TestSubAgentAPI(t *testing.T) {
751751
agentID, err := uuid.FromBytes(createResp.Agent.Id)
752752
require.NoError(t, err)
753753

754-
apps, err := db.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test.
754+
apps, err := db.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID)
755755
require.NoError(t, err)
756756
require.Len(t, apps, 1)
757757
require.Equal(t, "k5jd7a99-duplicate-slug", apps[0].Slug)
@@ -789,7 +789,7 @@ func TestSubAgentAPI(t *testing.T) {
789789
require.NoError(t, err)
790790

791791
// Then: It is deleted.
792-
_, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgent.ID) //nolint:gocritic // this is a test.
792+
_, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgent.ID)
793793
require.ErrorIs(t, err, sql.ErrNoRows)
794794
})
795795

@@ -830,10 +830,10 @@ func TestSubAgentAPI(t *testing.T) {
830830
require.NoError(t, err)
831831

832832
// Then: The correct one is deleted.
833-
_, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID) //nolint:gocritic // this is a test.
833+
_, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID)
834834
require.ErrorIs(t, err, sql.ErrNoRows)
835835

836-
_, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentTwo.ID) //nolint:gocritic // this is a test.
836+
_, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentTwo.ID)
837837
require.NoError(t, err)
838838
})
839839

@@ -871,7 +871,7 @@ func TestSubAgentAPI(t *testing.T) {
871871
var notAuthorizedError dbauthz.NotAuthorizedError
872872
require.ErrorAs(t, err, &notAuthorizedError)
873873

874-
_, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID) //nolint:gocritic // this is a test.
874+
_, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), childAgentOne.ID)
875875
require.NoError(t, err)
876876
})
877877

@@ -912,7 +912,7 @@ func TestSubAgentAPI(t *testing.T) {
912912
require.NoError(t, err)
913913

914914
// Verify that the apps were created
915-
apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), subAgentID) //nolint:gocritic // this is a test.
915+
apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), subAgentID)
916916
require.NoError(t, err)
917917
require.Len(t, apps, 2)
918918

@@ -923,7 +923,7 @@ func TestSubAgentAPI(t *testing.T) {
923923
require.NoError(t, err)
924924

925925
// Then: The agent is deleted
926-
_, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), subAgentID) //nolint:gocritic // this is a test.
926+
_, err = api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), subAgentID)
927927
require.ErrorIs(t, err, sql.ErrNoRows)
928928

929929
// And: The apps are *retained* to avoid causing issues
@@ -1068,7 +1068,7 @@ func TestSubAgentAPI(t *testing.T) {
10681068
agentID, err := uuid.FromBytes(createResp.Agent.Id)
10691069
require.NoError(t, err)
10701070

1071-
subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test.
1071+
subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID)
10721072
require.NoError(t, err)
10731073

10741074
require.Equal(t, len(tt.expectedApps), len(subAgent.DisplayApps), "display apps count mismatch")
@@ -1118,14 +1118,14 @@ func TestSubAgentAPI(t *testing.T) {
11181118
require.NoError(t, err)
11191119

11201120
// Verify display apps
1121-
subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test.
1121+
subAgent, err := api.Database.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), agentID)
11221122
require.NoError(t, err)
11231123
require.Len(t, subAgent.DisplayApps, 2)
11241124
require.Equal(t, database.DisplayAppVscode, subAgent.DisplayApps[0])
11251125
require.Equal(t, database.DisplayAppWebTerminal, subAgent.DisplayApps[1])
11261126

11271127
// Verify regular apps
1128-
apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // this is a test.
1128+
apps, err := api.Database.GetWorkspaceAppsByAgentID(dbauthz.AsSystemRestricted(ctx), agentID)
11291129
require.NoError(t, err)
11301130
require.Len(t, apps, 1)
11311131
require.Equal(t, "v4qhkq17-custom-app", apps[0].Slug)
@@ -1190,7 +1190,7 @@ func TestSubAgentAPI(t *testing.T) {
11901190
})
11911191

11921192
// When: We list the sub agents.
1193-
listResp, err := api.ListSubAgents(ctx, &proto.ListSubAgentsRequest{}) //nolint:gocritic // this is a test.
1193+
listResp, err := api.ListSubAgents(ctx, &proto.ListSubAgentsRequest{})
11941194
require.NoError(t, err)
11951195

11961196
listedChildAgents := listResp.Agents

‎coderd/coderd.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"github.com/coder/coder/v2/coderd/oauth2provider"
2424
"github.com/coder/coder/v2/coderd/pproflabel"
2525
"github.com/coder/coder/v2/coderd/prebuilds"
26+
"github.com/coder/coder/v2/coderd/usage"
2627
"github.com/coder/coder/v2/coderd/wsbuilder"
2728

2829
"github.com/andybalholm/brotli"
@@ -200,6 +201,7 @@ type Options struct {
200201
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
201202
UserQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore]
202203
AccessControlStore *atomic.Pointer[dbauthz.AccessControlStore]
204+
UsageInserter *atomic.Pointer[usage.Inserter]
203205
// CoordinatorResumeTokenProvider is used to provide and validate resume
204206
// tokens issued by and passed to the coordinator DRPC API.
205207
CoordinatorResumeTokenProvider tailnet.ResumeTokenProvider
@@ -428,6 +430,13 @@ func New(options *Options) *API {
428430
v := schedule.NewAGPLUserQuietHoursScheduleStore()
429431
options.UserQuietHoursScheduleStore.Store(&v)
430432
}
433+
if options.UsageInserter == nil {
434+
options.UsageInserter = &atomic.Pointer[usage.Inserter]{}
435+
}
436+
if options.UsageInserter.Load() == nil {
437+
inserter := usage.NewAGPLInserter()
438+
options.UsageInserter.Store(&inserter)
439+
}
431440
if options.OneTimePasscodeValidityPeriod == 0 {
432441
options.OneTimePasscodeValidityPeriod = 20 * time.Minute
433442
}
@@ -590,6 +599,7 @@ func New(options *Options) *API {
590599
UserQuietHoursScheduleStore: options.UserQuietHoursScheduleStore,
591600
AccessControlStore: options.AccessControlStore,
592601
BuildUsageChecker: &buildUsageChecker,
602+
UsageInserter: options.UsageInserter,
593603
FileCache: files.New(options.PrometheusRegistry, options.Authorizer),
594604
Experiments: experiments,
595605
WebpushDispatcher: options.WebPushDispatcher,
@@ -1690,6 +1700,9 @@ type API struct {
16901700
// BuildUsageChecker is a pointer as it's passed around to multiple
16911701
// components.
16921702
BuildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker]
1703+
// UsageInserter is a pointer to an atomic pointer because it is passed to
1704+
// multiple components.
1705+
UsageInserter *atomic.Pointer[usage.Inserter]
16931706

16941707
UpdatesProvider tailnet.WorkspaceUpdatesProvider
16951708

@@ -1905,6 +1918,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
19051918
&api.Auditor,
19061919
api.TemplateScheduleStore,
19071920
api.UserQuietHoursScheduleStore,
1921+
api.UsageInserter,
19081922
api.DeploymentValues,
19091923
provisionerdserver.Options{
19101924
OIDCConfig: api.OIDCConfig,

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,8 @@ var (
213213
// Provisionerd creates workspaces resources monitor
214214
rbac.ResourceWorkspaceAgentResourceMonitor.Type: {policy.ActionCreate},
215215
rbac.ResourceWorkspaceAgentDevcontainers.Type: {policy.ActionCreate},
216+
// Provisionerd creates usage events
217+
rbac.ResourceUsageEvent.Type: {policy.ActionCreate},
216218
}),
217219
Org: map[string][]rbac.Permission{},
218220
User: []rbac.Permission{},
@@ -510,17 +512,19 @@ var (
510512
Scope: rbac.ScopeAll,
511513
}.WithCachedASTValue()
512514

513-
subjectUsageTracker = rbac.Subject{
514-
Type: rbac.SubjectTypeUsageTracker,
515-
FriendlyName: "Usage Tracker",
515+
subjectUsagePublisher = rbac.Subject{
516+
Type: rbac.SubjectTypeUsagePublisher,
517+
FriendlyName: "Usage Publisher",
516518
ID: uuid.Nil.String(),
517519
Roles: rbac.Roles([]rbac.Role{
518520
{
519-
Identifier: rbac.RoleIdentifier{Name: "usage-tracker"},
520-
DisplayName: "Usage Tracker",
521+
Identifier: rbac.RoleIdentifier{Name: "usage-publisher"},
522+
DisplayName: "Usage Publisher",
521523
Site: rbac.Permissions(map[string][]policy.Action{
522-
rbac.ResourceLicense.Type: {policy.ActionRead},
523-
rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
524+
rbac.ResourceLicense.Type: {policy.ActionRead},
525+
// The usage publisher doesn't create events, just
526+
// reads/processes them.
527+
rbac.ResourceUsageEvent.Type: {policy.ActionRead, policy.ActionUpdate},
524528
}),
525529
Org: map[string][]rbac.Permission{},
526530
User: []rbac.Permission{},
@@ -604,10 +608,10 @@ func AsFileReader(ctx context.Context) context.Context {
604608
return As(ctx, subjectFileReader)
605609
}
606610

607-
// AsUsageTracker returns a context with an actor that has permissions required
608-
// for creating, reading, and updating usage events.
609-
func AsUsageTracker(ctx context.Context) context.Context {
610-
return As(ctx, subjectUsageTracker)
611+
// AsUsagePublisher returns a context with an actor that has permissions
612+
// required for creating, reading, and updating usage events.
613+
func AsUsagePublisher(ctx context.Context) context.Context {
614+
return As(ctx, subjectUsagePublisher)
611615
}
612616

613617
var AsRemoveActor = rbac.Subject{
@@ -3038,7 +3042,7 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe
30383042
}
30393043

30403044
func (q *querier) GetUnexpiredLicenses(ctx context.Context) ([]database.License, error) {
3041-
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
3045+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceLicense); err != nil {
30423046
return nil, err
30433047
}
30443048
return q.db.GetUnexpiredLicenses(ctx)

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -758,6 +758,18 @@ func (s *MethodTestSuite) TestLicense() {
758758
check.Args().Asserts(l, policy.ActionRead).
759759
Returns([]database.License{l})
760760
}))
761+
s.Run("GetUnexpiredLicenses", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
762+
l := database.License{
763+
ID: 1,
764+
Exp: time.Now().Add(time.Hour * 24 * 30),
765+
UUID: uuid.New(),
766+
}
767+
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).
768+
Return([]database.License{l}, nil).
769+
AnyTimes()
770+
check.Args().Asserts(rbac.ResourceLicense, policy.ActionRead).
771+
Returns([]database.License{l})
772+
}))
761773
s.Run("InsertLicense", s.Subtest(func(db database.Store, check *expects) {
762774
check.Args(database.InsertLicenseParams{}).
763775
Asserts(rbac.ResourceLicense, policy.ActionCreate)
@@ -3770,9 +3782,6 @@ func (s *MethodTestSuite) TestSystemFunctions() {
37703782
s.Run("GetActiveUserCount", s.Subtest(func(db database.Store, check *expects) {
37713783
check.Args(false).Asserts(rbac.ResourceSystem, policy.ActionRead).Returns(int64(0))
37723784
}))
3773-
s.Run("GetUnexpiredLicenses", s.Subtest(func(db database.Store, check *expects) {
3774-
check.Args().Asserts(rbac.ResourceSystem, policy.ActionRead)
3775-
}))
37763785
s.Run("GetAuthorizationUserRoles", s.Subtest(func(db database.Store, check *expects) {
37773786
u := dbgen.User(s.T(), db, database.User{})
37783787
check.Args(u.ID).Asserts(rbac.ResourceSystem, policy.ActionRead)

‎coderd/externalauth/externalauth_test.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -337,7 +337,6 @@ func TestRefreshToken(t *testing.T) {
337337
require.Equal(t, 1, validateCalls, "token is validated")
338338
require.Equal(t, 1, refreshCalls, "token is refreshed")
339339
require.NotEqualf(t, link.OAuthAccessToken, updated.OAuthAccessToken, "token is updated")
340-
//nolint:gocritic // testing
341340
dbLink, err := db.GetExternalAuthLink(dbauthz.AsSystemRestricted(context.Background()), database.GetExternalAuthLinkParams{
342341
ProviderID: link.ProviderID,
343342
UserID: link.UserID,

‎coderd/files/cache_test.go

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,6 @@ func TestCancelledFetch(t *testing.T) {
4545
cache := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
4646

4747
// Cancel the context for the first call; should fail.
48-
//nolint:gocritic // Unit testing
4948
ctx, cancel := context.WithCancel(dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort)))
5049
cancel()
5150
_, err := cache.Acquire(ctx, dbM, fileID)
@@ -71,7 +70,6 @@ func TestCancelledConcurrentFetch(t *testing.T) {
7170

7271
cache := files.LeakCache{Cache: files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})}
7372

74-
//nolint:gocritic // Unit testing
7573
ctx := dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort))
7674

7775
// Cancel the context for the first call; should fail.
@@ -99,7 +97,6 @@ func TestConcurrentFetch(t *testing.T) {
9997
})
10098

10199
cache := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
102-
//nolint:gocritic // Unit testing
103100
ctx := dbauthz.AsFileReader(testutil.Context(t, testutil.WaitShort))
104101

105102
// Expect 2 calls to Acquire before we continue the test
@@ -151,7 +148,6 @@ func TestCacheRBAC(t *testing.T) {
151148
Scope: rbac.ScopeAll,
152149
})
153150

154-
//nolint:gocritic // Unit testing
155151
cacheReader := dbauthz.AsFileReader(ctx)
156152

157153
t.Run("NoRolesOpen", func(t *testing.T) {
@@ -207,7 +203,6 @@ func cachePromMetricName(metric string) string {
207203

208204
func TestConcurrency(t *testing.T) {
209205
t.Parallel()
210-
//nolint:gocritic // Unit testing
211206
ctx := dbauthz.AsFileReader(t.Context())
212207

213208
const fileSize = 10
@@ -268,7 +263,6 @@ func TestConcurrency(t *testing.T) {
268263

269264
func TestRelease(t *testing.T) {
270265
t.Parallel()
271-
//nolint:gocritic // Unit testing
272266
ctx := dbauthz.AsFileReader(t.Context())
273267

274268
const fileSize = 10

‎coderd/idpsync/group_test.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -328,7 +328,6 @@ func TestGroupSyncTable(t *testing.T) {
328328
},
329329
}
330330

331-
//nolint:gocritic // testing
332331
defOrg, err := db.GetDefaultOrganization(dbauthz.AsSystemRestricted(ctx))
333332
require.NoError(t, err)
334333
SetupOrganization(t, s, db, user, defOrg.ID, def)
@@ -527,7 +526,6 @@ func TestApplyGroupDifference(t *testing.T) {
527526
db, _ := dbtestutil.NewDB(t)
528527

529528
ctx := testutil.Context(t, testutil.WaitMedium)
530-
//nolint:gocritic // testing
531529
ctx = dbauthz.AsSystemRestricted(ctx)
532530

533531
org := dbgen.Organization(t, db, database.Organization{})

‎coderd/idpsync/role_test.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,6 @@ func TestRoleSyncTable(t *testing.T) {
273273
}
274274

275275
// Also assert site wide roles
276-
//nolint:gocritic // unit testing assertions
277276
allRoles, err := db.GetAuthorizationUserRoles(dbauthz.AsSystemRestricted(ctx), user.ID)
278277
require.NoError(t, err)
279278

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy