Skip to content

Commit 5789ea5

Browse files
f0sself0ssel
authored
chore: move stat reporting into workspacestats package (coder#13386)
1 parent afd9d3b commit 5789ea5

19 files changed

+314
-346
lines changed

‎coderd/agentapi/api.go

+3-5
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"github.com/coder/coder/v2/coderd/prometheusmetrics"
2525
"github.com/coder/coder/v2/coderd/schedule"
2626
"github.com/coder/coder/v2/coderd/tracing"
27+
"github.com/coder/coder/v2/coderd/workspacestats"
2728
"github.com/coder/coder/v2/codersdk/agentsdk"
2829
"github.com/coder/coder/v2/tailnet"
2930
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
@@ -59,7 +60,7 @@ type Options struct {
5960
DerpMapFn func() *tailcfg.DERPMap
6061
TailnetCoordinator *atomic.Pointer[tailnet.Coordinator]
6162
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
62-
StatsBatcher StatsBatcher
63+
StatsReporter *workspacestats.Reporter
6364
AppearanceFetcher *atomic.Pointer[appearance.Fetcher]
6465
PublishWorkspaceUpdateFn func(ctx context.Context, workspaceID uuid.UUID)
6566
PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
@@ -114,12 +115,9 @@ func New(opts Options) *API {
114115
api.StatsAPI = &StatsAPI{
115116
AgentFn: api.agent,
116117
Database: opts.Database,
117-
Pubsub: opts.Pubsub,
118118
Log: opts.Log,
119-
StatsBatcher: opts.StatsBatcher,
120-
TemplateScheduleStore: opts.TemplateScheduleStore,
119+
StatsReporter: opts.StatsReporter,
121120
AgentStatsRefreshInterval: opts.AgentStatsRefreshInterval,
122-
UpdateAgentMetricsFn: opts.UpdateAgentMetricsFn,
123121
}
124122

125123
api.LifecycleAPI = &LifecycleAPI{

‎coderd/agentapi/stats.go

+11-82
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,8 @@ package agentapi
22

33
import (
44
"context"
5-
"sync/atomic"
65
"time"
76

8-
"golang.org/x/sync/errgroup"
97
"golang.org/x/xerrors"
108
"google.golang.org/protobuf/types/known/durationpb"
119

@@ -15,10 +13,7 @@ import (
1513
agentproto "github.com/coder/coder/v2/agent/proto"
1614
"github.com/coder/coder/v2/coderd/database"
1715
"github.com/coder/coder/v2/coderd/database/dbtime"
18-
"github.com/coder/coder/v2/coderd/database/pubsub"
19-
"github.com/coder/coder/v2/coderd/prometheusmetrics"
20-
"github.com/coder/coder/v2/coderd/schedule"
21-
"github.com/coder/coder/v2/codersdk"
16+
"github.com/coder/coder/v2/coderd/workspacestats"
2217
)
2318

2419
type StatsBatcher interface {
@@ -28,12 +23,9 @@ type StatsBatcher interface {
2823
type StatsAPI struct {
2924
AgentFn func(context.Context) (database.WorkspaceAgent, error)
3025
Database database.Store
31-
Pubsub pubsub.Pubsub
3226
Log slog.Logger
33-
StatsBatcher StatsBatcher
34-
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
27+
StatsReporter *workspacestats.Reporter
3528
AgentStatsRefreshInterval time.Duration
36-
UpdateAgentMetricsFn func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
3729

3830
TimeNowFn func() time.Time // defaults to dbtime.Now()
3931
}
@@ -69,80 +61,17 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
6961
slog.F("payload", req),
7062
)
7163

72-
now := a.now()
73-
if req.Stats.ConnectionCount > 0 {
74-
var nextAutostart time.Time
75-
if workspace.AutostartSchedule.String != "" {
76-
templateSchedule, err := (*(a.TemplateScheduleStore.Load())).Get(ctx, a.Database, workspace.TemplateID)
77-
// If the template schedule fails to load, just default to bumping
78-
// without the next transition and log it.
79-
if err != nil {
80-
a.Log.Error(ctx, "failed to load template schedule bumping activity, defaulting to bumping by 60min",
81-
slog.F("workspace_id", workspace.ID),
82-
slog.F("template_id", workspace.TemplateID),
83-
slog.Error(err),
84-
)
85-
} else {
86-
next, allowed := schedule.NextAutostart(now, workspace.AutostartSchedule.String, templateSchedule)
87-
if allowed {
88-
nextAutostart = next
89-
}
90-
}
91-
}
92-
ActivityBumpWorkspace(ctx, a.Log.Named("activity_bump"), a.Database, workspace.ID, nextAutostart)
93-
}
94-
95-
var errGroup errgroup.Group
96-
errGroup.Go(func() error {
97-
err := a.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, req.Stats)
98-
if err != nil {
99-
a.Log.Error(ctx, "add agent stats to batcher", slog.Error(err))
100-
return xerrors.Errorf("insert workspace agent stats batch: %w", err)
101-
}
102-
return nil
103-
})
104-
errGroup.Go(func() error {
105-
// nolint:gocritic // (#13146) Will be moved soon as part of refactor.
106-
err := a.Database.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{
107-
ID: workspace.ID,
108-
LastUsedAt: now,
109-
})
110-
if err != nil {
111-
return xerrors.Errorf("update workspace LastUsedAt: %w", err)
112-
}
113-
return nil
114-
})
115-
if a.UpdateAgentMetricsFn != nil {
116-
errGroup.Go(func() error {
117-
user, err := a.Database.GetUserByID(ctx, workspace.OwnerID)
118-
if err != nil {
119-
return xerrors.Errorf("get user: %w", err)
120-
}
121-
122-
a.UpdateAgentMetricsFn(ctx, prometheusmetrics.AgentMetricLabels{
123-
Username: user.Username,
124-
WorkspaceName: workspace.Name,
125-
AgentName: workspaceAgent.Name,
126-
TemplateName: getWorkspaceAgentByIDRow.TemplateName,
127-
}, req.Stats.Metrics)
128-
return nil
129-
})
130-
}
131-
err = errGroup.Wait()
64+
err = a.StatsReporter.ReportAgentStats(
65+
ctx,
66+
a.now(),
67+
workspace,
68+
workspaceAgent,
69+
getWorkspaceAgentByIDRow.TemplateName,
70+
req.Stats,
71+
)
13272
if err != nil {
133-
return nil, xerrors.Errorf("update stats in database: %w", err)
73+
return nil, xerrors.Errorf("report agent stats: %w", err)
13474
}
13575

136-
// Tell the frontend about the new agent report, now that everything is updated
137-
a.publishWorkspaceAgentStats(ctx, workspace.ID)
138-
13976
return res, nil
14077
}
141-
142-
func (a *StatsAPI) publishWorkspaceAgentStats(ctx context.Context, workspaceID uuid.UUID) {
143-
err := a.Pubsub.Publish(codersdk.WorkspaceNotifyChannel(workspaceID), []byte{})
144-
if err != nil {
145-
a.Log.Warn(ctx, "failed to publish workspace agent stats",
146-
slog.F("workspace_id", workspaceID), slog.Error(err))
147-
}
148-
}

‎coderd/agentapi/stats_test.go

+52-39
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"github.com/coder/coder/v2/coderd/database/pubsub"
2323
"github.com/coder/coder/v2/coderd/prometheusmetrics"
2424
"github.com/coder/coder/v2/coderd/schedule"
25+
"github.com/coder/coder/v2/coderd/workspacestats"
2526
"github.com/coder/coder/v2/codersdk"
2627
"github.com/coder/coder/v2/testutil"
2728
)
@@ -129,21 +130,24 @@ func TestUpdateStates(t *testing.T) {
129130
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
130131
return agent, nil
131132
},
132-
Database: dbM,
133-
Pubsub: ps,
134-
StatsBatcher: batcher,
135-
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
133+
Database: dbM,
134+
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
135+
Database: dbM,
136+
Pubsub: ps,
137+
StatsBatcher: batcher,
138+
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
139+
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
140+
updateAgentMetricsFnCalled = true
141+
assert.Equal(t, prometheusmetrics.AgentMetricLabels{
142+
Username: user.Username,
143+
WorkspaceName: workspace.Name,
144+
AgentName: agent.Name,
145+
TemplateName: template.Name,
146+
}, labels)
147+
assert.Equal(t, req.Stats.Metrics, metrics)
148+
},
149+
}),
136150
AgentStatsRefreshInterval: 10 * time.Second,
137-
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
138-
updateAgentMetricsFnCalled = true
139-
assert.Equal(t, prometheusmetrics.AgentMetricLabels{
140-
Username: user.Username,
141-
WorkspaceName: workspace.Name,
142-
AgentName: agent.Name,
143-
TemplateName: template.Name,
144-
}, labels)
145-
assert.Equal(t, req.Stats.Metrics, metrics)
146-
},
147151
TimeNowFn: func() time.Time {
148152
return now
149153
},
@@ -232,13 +236,16 @@ func TestUpdateStates(t *testing.T) {
232236
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
233237
return agent, nil
234238
},
235-
Database: dbM,
236-
Pubsub: ps,
237-
StatsBatcher: batcher,
238-
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
239+
Database: dbM,
240+
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
241+
Database: dbM,
242+
Pubsub: ps,
243+
StatsBatcher: batcher,
244+
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
245+
// Ignored when nil.
246+
UpdateAgentMetricsFn: nil,
247+
}),
239248
AgentStatsRefreshInterval: 10 * time.Second,
240-
// Ignored when nil.
241-
UpdateAgentMetricsFn: nil,
242249
TimeNowFn: func() time.Time {
243250
return now
244251
},
@@ -274,12 +281,15 @@ func TestUpdateStates(t *testing.T) {
274281
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
275282
return agent, nil
276283
},
277-
Database: dbM,
278-
Pubsub: ps,
279-
StatsBatcher: nil, // should not be called
280-
TemplateScheduleStore: nil, // should not be called
284+
Database: dbM,
285+
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
286+
Database: dbM,
287+
Pubsub: ps,
288+
StatsBatcher: nil, // should not be called
289+
TemplateScheduleStore: nil, // should not be called
290+
UpdateAgentMetricsFn: nil, // should not be called
291+
}),
281292
AgentStatsRefreshInterval: 10 * time.Second,
282-
UpdateAgentMetricsFn: nil, // should not be called
283293
TimeNowFn: func() time.Time {
284294
panic("should not be called")
285295
},
@@ -343,21 +353,24 @@ func TestUpdateStates(t *testing.T) {
343353
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
344354
return agent, nil
345355
},
346-
Database: dbM,
347-
Pubsub: ps,
348-
StatsBatcher: batcher,
349-
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
356+
Database: dbM,
357+
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
358+
Database: dbM,
359+
Pubsub: ps,
360+
StatsBatcher: batcher,
361+
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
362+
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
363+
updateAgentMetricsFnCalled = true
364+
assert.Equal(t, prometheusmetrics.AgentMetricLabels{
365+
Username: user.Username,
366+
WorkspaceName: workspace.Name,
367+
AgentName: agent.Name,
368+
TemplateName: template.Name,
369+
}, labels)
370+
assert.Equal(t, req.Stats.Metrics, metrics)
371+
},
372+
}),
350373
AgentStatsRefreshInterval: 15 * time.Second,
351-
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
352-
updateAgentMetricsFnCalled = true
353-
assert.Equal(t, prometheusmetrics.AgentMetricLabels{
354-
Username: user.Username,
355-
WorkspaceName: workspace.Name,
356-
AgentName: agent.Name,
357-
TemplateName: template.Name,
358-
}, labels)
359-
assert.Equal(t, req.Stats.Metrics, metrics)
360-
},
361374
TimeNowFn: func() time.Time {
362375
return now
363376
},

‎coderd/coderd.go

+12-4
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ import (
6868
"github.com/coder/coder/v2/coderd/updatecheck"
6969
"github.com/coder/coder/v2/coderd/util/slice"
7070
"github.com/coder/coder/v2/coderd/workspaceapps"
71+
"github.com/coder/coder/v2/coderd/workspacestats"
7172
"github.com/coder/coder/v2/coderd/workspaceusage"
7273
"github.com/coder/coder/v2/codersdk"
7374
"github.com/coder/coder/v2/codersdk/drpc"
@@ -550,13 +551,22 @@ func New(options *Options) *API {
550551
api.Logger.Fatal(api.ctx, "failed to initialize tailnet client service", slog.Error(err))
551552
}
552553

554+
api.statsReporter = workspacestats.NewReporter(workspacestats.ReporterOptions{
555+
Database: options.Database,
556+
Logger: options.Logger.Named("workspacestats"),
557+
Pubsub: options.Pubsub,
558+
TemplateScheduleStore: options.TemplateScheduleStore,
559+
StatsBatcher: options.StatsBatcher,
560+
UpdateAgentMetricsFn: options.UpdateAgentMetrics,
561+
AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize,
562+
})
553563
workspaceAppsLogger := options.Logger.Named("workspaceapps")
554564
if options.WorkspaceAppsStatsCollectorOptions.Logger == nil {
555565
named := workspaceAppsLogger.Named("stats_collector")
556566
options.WorkspaceAppsStatsCollectorOptions.Logger = &named
557567
}
558568
if options.WorkspaceAppsStatsCollectorOptions.Reporter == nil {
559-
options.WorkspaceAppsStatsCollectorOptions.Reporter = workspaceapps.NewStatsDBReporter(options.Database, workspaceapps.DefaultStatsDBReporterBatchSize)
569+
options.WorkspaceAppsStatsCollectorOptions.Reporter = api.statsReporter
560570
}
561571

562572
api.workspaceAppServer = &workspaceapps.Server{
@@ -626,8 +636,6 @@ func New(options *Options) *API {
626636
cors := httpmw.Cors(options.DeploymentValues.Dangerous.AllowAllCors.Value())
627637
prometheusMW := httpmw.Prometheus(options.PrometheusRegistry)
628638

629-
api.statsBatcher = options.StatsBatcher
630-
631639
r.Use(
632640
httpmw.Recover(api.Logger),
633641
tracing.StatusWriterMiddleware,
@@ -1287,7 +1295,7 @@ type API struct {
12871295
healthCheckGroup *singleflight.Group[string, *healthsdk.HealthcheckReport]
12881296
healthCheckCache atomic.Pointer[healthsdk.HealthcheckReport]
12891297

1290-
statsBatcher *batchstats.Batcher
1298+
statsReporter *workspacestats.Reporter
12911299

12921300
Acquirer *provisionerdserver.Acquirer
12931301
// dbRolluper rolls up template usage stats from raw agent and app

‎coderd/database/dbauthz/setup_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,7 @@ func (s *MethodTestSuite) NotAuthorizedErrorTest(ctx context.Context, az *coderd
264264
// any case where the error is nil and the response is an empty slice.
265265
if err != nil || !hasEmptySliceResponse(resp) {
266266
s.Errorf(err, "method should an error with cancellation")
267-
s.ErrorIsf(err, context.Canceled, "error should match context.Cancelled")
267+
s.ErrorIsf(err, context.Canceled, "error should match context.Canceled")
268268
}
269269
})
270270
}

‎coderd/insights_test.go

+11-4
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"github.com/coder/coder/v2/coderd/rbac"
3232
"github.com/coder/coder/v2/coderd/rbac/policy"
3333
"github.com/coder/coder/v2/coderd/workspaceapps"
34+
"github.com/coder/coder/v2/coderd/workspacestats"
3435
"github.com/coder/coder/v2/codersdk"
3536
"github.com/coder/coder/v2/codersdk/agentsdk"
3637
"github.com/coder/coder/v2/codersdk/workspacesdk"
@@ -736,9 +737,12 @@ func TestTemplateInsights_Golden(t *testing.T) {
736737
})
737738
}
738739
}
739-
reporter := workspaceapps.NewStatsDBReporter(db, workspaceapps.DefaultStatsDBReporterBatchSize)
740+
reporter := workspacestats.NewReporter(workspacestats.ReporterOptions{
741+
Database: db,
742+
AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize,
743+
})
740744
//nolint:gocritic // This is a test.
741-
err = reporter.Report(dbauthz.AsSystemRestricted(ctx), stats)
745+
err = reporter.ReportAppStats(dbauthz.AsSystemRestricted(ctx), stats)
742746
require.NoError(t, err, "want no error inserting app stats")
743747

744748
return client, events
@@ -1632,9 +1636,12 @@ func TestUserActivityInsights_Golden(t *testing.T) {
16321636
})
16331637
}
16341638
}
1635-
reporter := workspaceapps.NewStatsDBReporter(db, workspaceapps.DefaultStatsDBReporterBatchSize)
1639+
reporter := workspacestats.NewReporter(workspacestats.ReporterOptions{
1640+
Database: db,
1641+
AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize,
1642+
})
16361643
//nolint:gocritic // This is a test.
1637-
err = reporter.Report(dbauthz.AsSystemRestricted(ctx), stats)
1644+
err = reporter.ReportAppStats(dbauthz.AsSystemRestricted(ctx), stats)
16381645
require.NoError(t, err, "want no error inserting app stats")
16391646

16401647
return client, events

0 commit comments

Comments
 (0)