Skip to content

Commit 3aea6d1

Browse files
deansheatherdeansheather
committed
chore: wire up usage tracking for managed agents
1 parent a25d856 commit 3aea6d1

File tree

10 files changed

+210
-66
lines changed

10 files changed

+210
-66
lines changed

‎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,
@@ -1687,6 +1697,9 @@ type API struct {
16871697
// BuildUsageChecker is a pointer as it's passed around to multiple
16881698
// components.
16891699
BuildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker]
1700+
// UsageInserter is a pointer to an atomic pointer because it is passed to
1701+
// multiple components.
1702+
UsageInserter *atomic.Pointer[usage.Inserter]
16901703

16911704
UpdatesProvider tailnet.WorkspaceUpdatesProvider
16921705

@@ -1902,6 +1915,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
19021915
&api.Auditor,
19031916
api.TemplateScheduleStore,
19041917
api.UserQuietHoursScheduleStore,
1918+
api.UsageInserter,
19051919
api.DeploymentValues,
19061920
provisionerdserver.Options{
19071921
OIDCConfig: api.OIDCConfig,

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929

3030
"cdr.dev/slog"
3131

32+
"github.com/coder/coder/v2/coderd/usage"
3233
"github.com/coder/coder/v2/coderd/util/slice"
3334

3435
"github.com/coder/coder/v2/codersdk/drpcsdk"
@@ -121,6 +122,7 @@ type server struct {
121122
DeploymentValues *codersdk.DeploymentValues
122123
NotificationsEnqueuer notifications.Enqueuer
123124
PrebuildsOrchestrator *atomic.Pointer[prebuilds.ReconciliationOrchestrator]
125+
UsageInserter *atomic.Pointer[usage.Inserter]
124126

125127
OIDCConfig promoauth.OAuth2Config
126128

@@ -174,6 +176,7 @@ func NewServer(
174176
auditor *atomic.Pointer[audit.Auditor],
175177
templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore],
176178
userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore],
179+
usageInserter *atomic.Pointer[usage.Inserter],
177180
deploymentValues *codersdk.DeploymentValues,
178181
options Options,
179182
enqueuer notifications.Enqueuer,
@@ -195,6 +198,9 @@ func NewServer(
195198
if userQuietHoursScheduleStore == nil {
196199
return nil, xerrors.New("userQuietHoursScheduleStore is nil")
197200
}
201+
if usageInserter == nil {
202+
return nil, xerrors.New("usageCollector is nil")
203+
}
198204
if deploymentValues == nil {
199205
return nil, xerrors.New("deploymentValues is nil")
200206
}
@@ -244,6 +250,7 @@ func NewServer(
244250
heartbeatInterval: options.HeartbeatInterval,
245251
heartbeatFn: options.HeartbeatFn,
246252
PrebuildsOrchestrator: prebuildsOrchestrator,
253+
UsageInserter: usageInserter,
247254
}
248255

249256
if s.heartbeatFn == nil {
@@ -2026,6 +2033,20 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
20262033
sidebarAppID = uuid.NullUUID{}
20272034
}
20282035

2036+
if hasAITask {
2037+
// Insert usage event for managed agents.
2038+
usageInserter := s.UsageInserter.Load()
2039+
if usageInserter != nil {
2040+
event := usage.DCManagedAgentsV1{
2041+
Count: 1,
2042+
}
2043+
err = (*usageInserter).InsertDiscreteUsageEvent(ctx, db, event)
2044+
if err != nil {
2045+
return xerrors.Errorf("insert %q event: %w", event.EventType(), err)
2046+
}
2047+
}
2048+
}
2049+
20292050
// Regardless of whether there is an AI task or not, update the field to indicate one way or the other since it
20302051
// always defaults to nil. ONLY if has_ai_task=true MUST ai_task_sidebar_app_id be set.
20312052
if err := db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{

‎coderd/provisionerdserver/provisionerdserver_test.go

Lines changed: 68 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import (
4444
"github.com/coder/coder/v2/coderd/schedule"
4545
"github.com/coder/coder/v2/coderd/schedule/cron"
4646
"github.com/coder/coder/v2/coderd/telemetry"
47+
"github.com/coder/coder/v2/coderd/usage"
4748
"github.com/coder/coder/v2/coderd/wspubsub"
4849
"github.com/coder/coder/v2/codersdk"
4950
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -67,6 +68,13 @@ func testUserQuietHoursScheduleStore() *atomic.Pointer[schedule.UserQuietHoursSc
6768
return ptr
6869
}
6970

71+
func testUsageInserter() *atomic.Pointer[usage.Inserter] {
72+
ptr := &atomic.Pointer[usage.Inserter]{}
73+
inserter := usage.NewAGPLInserter()
74+
ptr.Store(&inserter)
75+
return ptr
76+
}
77+
7078
func TestAcquireJob_LongPoll(t *testing.T) {
7179
t.Parallel()
7280
//nolint:dogsled
@@ -2672,7 +2680,10 @@ func TestCompleteJob(t *testing.T) {
26722680
t.Run(tc.name, func(t *testing.T) {
26732681
t.Parallel()
26742682

2675-
srv, db, _, pd := setup(t, false, &overrides{})
2683+
fakeUsageInserter, usageInserterPtr := newFakeUsageInserter()
2684+
srv, db, _, pd := setup(t, false, &overrides{
2685+
usageInserter: usageInserterPtr,
2686+
})
26762687

26772688
importJobID := uuid.New()
26782689
tvID := uuid.New()
@@ -2741,6 +2752,10 @@ func TestCompleteJob(t *testing.T) {
27412752
require.NoError(t, err)
27422753
require.True(t, version.HasAITask.Valid) // We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true.
27432754
require.Equal(t, tc.expected, version.HasAITask.Bool)
2755+
2756+
// We never expect a usage event to be collected for
2757+
// template imports.
2758+
require.Empty(t, fakeUsageInserter.collectedEvents)
27442759
})
27452760
}
27462761
})
@@ -2750,9 +2765,9 @@ func TestCompleteJob(t *testing.T) {
27502765
// will be set as well in that case.
27512766
t.Run("WorkspaceBuild", func(t *testing.T) {
27522767
type testcase struct {
2753-
name string
2754-
input *proto.CompletedJob_WorkspaceBuild
2755-
expected bool
2768+
name string
2769+
input *proto.CompletedJob_WorkspaceBuild
2770+
expectedHasAiTask bool
27562771
}
27572772

27582773
sidebarAppID := uuid.NewString()
@@ -2762,7 +2777,7 @@ func TestCompleteJob(t *testing.T) {
27622777
input: &proto.CompletedJob_WorkspaceBuild{
27632778
// No AiTasks defined.
27642779
},
2765-
expected: false,
2780+
expectedHasAiTask: false,
27662781
},
27672782
{
27682783
name: "has_ai_task is set to true",
@@ -2792,7 +2807,7 @@ func TestCompleteJob(t *testing.T) {
27922807
},
27932808
},
27942809
},
2795-
expected: true,
2810+
expectedHasAiTask: true,
27962811
},
27972812
// Checks regression for https://github.com/coder/coder/issues/18776
27982813
{
@@ -2808,13 +2823,16 @@ func TestCompleteJob(t *testing.T) {
28082823
},
28092824
},
28102825
},
2811-
expected: false,
2826+
expectedHasAiTask: false,
28122827
},
28132828
} {
28142829
t.Run(tc.name, func(t *testing.T) {
28152830
t.Parallel()
28162831

2817-
srv, db, _, pd := setup(t, false, &overrides{})
2832+
fakeUsageInserter, usageInserterPtr := newFakeUsageInserter()
2833+
srv, db, _, pd := setup(t, false, &overrides{
2834+
usageInserter: usageInserterPtr,
2835+
})
28182836

28192837
importJobID := uuid.New()
28202838
tvID := uuid.New()
@@ -2899,10 +2917,19 @@ func TestCompleteJob(t *testing.T) {
28992917
build, err = db.GetWorkspaceBuildByID(ctx, build.ID)
29002918
require.NoError(t, err)
29012919
require.True(t, build.HasAITask.Valid) // We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true.
2902-
require.Equal(t, tc.expected, build.HasAITask.Bool)
2920+
require.Equal(t, tc.expectedHasAiTask, build.HasAITask.Bool)
29032921

2904-
if tc.expected {
2922+
if tc.expectedHasAiTask {
29052923
require.Equal(t, sidebarAppID, build.AITaskSidebarAppID.UUID.String())
2924+
2925+
// Check that a usage event was collected.
2926+
require.Len(t, fakeUsageInserter.collectedEvents, 1)
2927+
require.Equal(t, usage.DCManagedAgentsV1{
2928+
Count: 1,
2929+
}, fakeUsageInserter.collectedEvents[0])
2930+
} else {
2931+
// Check that no usage event was collected.
2932+
require.Empty(t, fakeUsageInserter.collectedEvents)
29062933
}
29072934
})
29082935
}
@@ -3835,6 +3862,7 @@ type overrides struct {
38353862
externalAuthConfigs []*externalauth.Config
38363863
templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
38373864
userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore]
3865+
usageInserter *atomic.Pointer[usage.Inserter]
38383866
clock *quartz.Mock
38393867
acquireJobLongPollDuration time.Duration
38403868
heartbeatFn func(ctx context.Context) error
@@ -3855,6 +3883,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
38553883
var externalAuthConfigs []*externalauth.Config
38563884
tss := testTemplateScheduleStore()
38573885
uqhss := testUserQuietHoursScheduleStore()
3886+
usageInserter := testUsageInserter()
38583887
clock := quartz.NewReal()
38593888
pollDur := time.Duration(0)
38603889
if ov == nil {
@@ -3892,6 +3921,15 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
38923921
require.True(t, swapped)
38933922
}
38943923
}
3924+
if ov.usageInserter != nil {
3925+
tUsageInserter := usageInserter.Load()
3926+
// keep the initial test value if the override hasn't set the atomic pointer.
3927+
usageInserter = ov.usageInserter
3928+
if usageInserter.Load() == nil {
3929+
swapped := usageInserter.CompareAndSwap(nil, tUsageInserter)
3930+
require.True(t, swapped)
3931+
}
3932+
}
38953933
if ov.clock != nil {
38963934
clock = ov.clock
38973935
}
@@ -3947,6 +3985,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
39473985
auditPtr,
39483986
tss,
39493987
uqhss,
3988+
usageInserter,
39503989
deploymentValues,
39513990
provisionerdserver.Options{
39523991
ExternalAuthConfigs: externalAuthConfigs,
@@ -4061,3 +4100,22 @@ func (s *fakeStream) cancel() {
40614100
s.canceled = true
40624101
s.c.Broadcast()
40634102
}
4103+
4104+
type fakeUsageInserter struct {
4105+
collectedEvents []usage.Event
4106+
}
4107+
4108+
var _ usage.Inserter = &fakeUsageInserter{}
4109+
4110+
func newFakeUsageInserter() (*fakeUsageInserter, *atomic.Pointer[usage.Inserter]) {
4111+
ptr := &atomic.Pointer[usage.Inserter]{}
4112+
fake := &fakeUsageInserter{}
4113+
var inserter usage.Inserter = fake
4114+
ptr.Store(&inserter)
4115+
return fake, ptr
4116+
}
4117+
4118+
func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usage.DiscreteEvent) error {
4119+
f.collectedEvents = append(f.collectedEvents, event)
4120+
return nil
4121+
}

‎enterprise/cli/server.go

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"github.com/coder/coder/v2/enterprise/audit/backends"
2121
"github.com/coder/coder/v2/enterprise/coderd"
2222
"github.com/coder/coder/v2/enterprise/coderd/dormancy"
23+
"github.com/coder/coder/v2/enterprise/coderd/usage"
2324
"github.com/coder/coder/v2/enterprise/dbcrypt"
2425
"github.com/coder/coder/v2/enterprise/trialer"
2526
"github.com/coder/coder/v2/tailnet"
@@ -116,15 +117,57 @@ func (r *RootCmd) Server(_ func()) *serpent.Command {
116117
o.ExternalTokenEncryption = cs
117118
}
118119

120+
if o.LicenseKeys == nil {
121+
o.LicenseKeys = coderd.Keys
122+
}
123+
124+
multiCloser := &multiCloser{}
125+
126+
// Create the enterprise API.
119127
api, err := coderd.New(ctx, o)
120128
if err != nil {
121129
return nil, nil, err
122130
}
123-
return api.AGPL, api, nil
131+
multiCloser.Add(api)
132+
133+
// Start the enterprise usage publisher routine. This won't do anything
134+
// unless the deployment is licensed and one of the licenses has usage
135+
// publishing enabled.
136+
publisher := usage.NewTallymanPublisher(ctx, options.Logger, options.Database, o.LicenseKeys,
137+
usage.PublisherWithHTTPClient(api.HTTPClient),
138+
)
139+
err = publisher.Start()
140+
if err != nil {
141+
_ = multiCloser.Close()
142+
return nil, nil, xerrors.Errorf("start usage publisher: %w", err)
143+
}
144+
multiCloser.Add(publisher)
145+
146+
return api.AGPL, multiCloser, nil
124147
})
125148

126149
cmd.AddSubcommands(
127150
r.dbcryptCmd(),
128151
)
129152
return cmd
130153
}
154+
155+
type multiCloser struct {
156+
closers []io.Closer
157+
}
158+
159+
var _ io.Closer = &multiCloser{}
160+
161+
func (m *multiCloser) Add(closer io.Closer) {
162+
m.closers = append(m.closers, closer)
163+
}
164+
165+
func (m *multiCloser) Close() error {
166+
var mErr error
167+
for _, closer := range m.closers {
168+
if err := closer.Close(); err != nil {
169+
mErr = xerrors.Errorf("close %T: %w", closer, err)
170+
}
171+
}
172+
return mErr
173+
}

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