Skip to content

Commit 5789ea5

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

File tree

19 files changed

+314
-346
lines changed

19 files changed

+314
-346
lines changed

‎coderd/agentapi/api.go

Lines changed: 3 additions & 5 deletions
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

Lines changed: 11 additions & 82 deletions
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

Lines changed: 52 additions & 39 deletions
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

Lines changed: 12 additions & 4 deletions
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 11 additions & 4 deletions
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)
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