Skip to content

Commit 5b9a65e

Browse files
f0sself0ssel
authored
chore: move Batcher and Tracker to workspacestats (#13418)
1 parent c7e7312 commit 5b9a65e

File tree

13 files changed

+105
-108
lines changed

13 files changed

+105
-108
lines changed

‎cli/server.go

+6-7
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,6 @@ import (
6262
"github.com/coder/coder/v2/cli/config"
6363
"github.com/coder/coder/v2/coderd"
6464
"github.com/coder/coder/v2/coderd/autobuild"
65-
"github.com/coder/coder/v2/coderd/batchstats"
6665
"github.com/coder/coder/v2/coderd/database"
6766
"github.com/coder/coder/v2/coderd/database/awsiamrds"
6867
"github.com/coder/coder/v2/coderd/database/dbmem"
@@ -87,7 +86,7 @@ import (
8786
stringutil "github.com/coder/coder/v2/coderd/util/strings"
8887
"github.com/coder/coder/v2/coderd/workspaceapps"
8988
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
90-
"github.com/coder/coder/v2/coderd/workspaceusage"
89+
"github.com/coder/coder/v2/coderd/workspacestats"
9190
"github.com/coder/coder/v2/codersdk"
9291
"github.com/coder/coder/v2/codersdk/drpc"
9392
"github.com/coder/coder/v2/cryptorand"
@@ -870,9 +869,9 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
870869
options.SwaggerEndpoint = vals.Swagger.Enable.Value()
871870
}
872871

873-
batcher, closeBatcher, err := batchstats.New(ctx,
874-
batchstats.WithLogger(options.Logger.Named("batchstats")),
875-
batchstats.WithStore(options.Database),
872+
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
873+
workspacestats.BatcherWithLogger(options.Logger.Named("batchstats")),
874+
workspacestats.BatcherWithStore(options.Database),
876875
)
877876
if err != nil {
878877
return xerrors.Errorf("failed to create agent stats batcher: %w", err)
@@ -977,8 +976,8 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
977976
defer purger.Close()
978977

979978
// Updates workspace usage
980-
tracker := workspaceusage.New(options.Database,
981-
workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")),
979+
tracker := workspacestats.NewTracker(options.Database,
980+
workspacestats.TrackerWithLogger(logger.Named("workspace_usage_tracker")),
982981
)
983982
options.WorkspaceUsageTracker = tracker
984983
defer tracker.Close()

‎coderd/agentapi/stats.go

-6
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,13 @@ import (
77
"golang.org/x/xerrors"
88
"google.golang.org/protobuf/types/known/durationpb"
99

10-
"github.com/google/uuid"
11-
1210
"cdr.dev/slog"
1311
agentproto "github.com/coder/coder/v2/agent/proto"
1412
"github.com/coder/coder/v2/coderd/database"
1513
"github.com/coder/coder/v2/coderd/database/dbtime"
1614
"github.com/coder/coder/v2/coderd/workspacestats"
1715
)
1816

19-
type StatsBatcher interface {
20-
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error
21-
}
22-
2317
type StatsAPI struct {
2418
AgentFn func(context.Context) (database.WorkspaceAgent, error)
2519
Database database.Store

‎coderd/agentapi/stats_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ type statsBatcher struct {
3939
lastStats *agentproto.Stats
4040
}
4141

42-
var _ agentapi.StatsBatcher = &statsBatcher{}
42+
var _ workspacestats.Batcher = &statsBatcher{}
4343

4444
func (b *statsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error {
4545
b.mu.Lock()

‎coderd/coderd.go

+8-11
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@ import (
4343
"github.com/coder/coder/v2/coderd/appearance"
4444
"github.com/coder/coder/v2/coderd/audit"
4545
"github.com/coder/coder/v2/coderd/awsidentity"
46-
"github.com/coder/coder/v2/coderd/batchstats"
4746
"github.com/coder/coder/v2/coderd/database"
4847
"github.com/coder/coder/v2/coderd/database/dbauthz"
4948
"github.com/coder/coder/v2/coderd/database/dbrollup"
@@ -69,7 +68,6 @@ import (
6968
"github.com/coder/coder/v2/coderd/util/slice"
7069
"github.com/coder/coder/v2/coderd/workspaceapps"
7170
"github.com/coder/coder/v2/coderd/workspacestats"
72-
"github.com/coder/coder/v2/coderd/workspaceusage"
7371
"github.com/coder/coder/v2/codersdk"
7472
"github.com/coder/coder/v2/codersdk/drpc"
7573
"github.com/coder/coder/v2/codersdk/healthsdk"
@@ -189,7 +187,7 @@ type Options struct {
189187
HTTPClient *http.Client
190188

191189
UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
192-
StatsBatcher *batchstats.Batcher
190+
StatsBatcher *workspacestats.DBBatcher
193191

194192
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
195193

@@ -206,7 +204,7 @@ type Options struct {
206204
// stats. This is used to provide insights in the WebUI.
207205
DatabaseRolluper *dbrollup.Rolluper
208206
// WorkspaceUsageTracker tracks workspace usage by the CLI.
209-
WorkspaceUsageTracker *workspaceusage.Tracker
207+
WorkspaceUsageTracker *workspacestats.UsageTracker
210208
}
211209

212210
// @title Coder API
@@ -384,8 +382,8 @@ func New(options *Options) *API {
384382
}
385383

386384
if options.WorkspaceUsageTracker == nil {
387-
options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
388-
workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
385+
options.WorkspaceUsageTracker = workspacestats.NewTracker(options.Database,
386+
workspacestats.TrackerWithLogger(options.Logger.Named("workspace_usage_tracker")),
389387
)
390388
}
391389

@@ -434,8 +432,7 @@ func New(options *Options) *API {
434432
options.Database,
435433
options.Pubsub,
436434
),
437-
dbRolluper: options.DatabaseRolluper,
438-
workspaceUsageTracker: options.WorkspaceUsageTracker,
435+
dbRolluper: options.DatabaseRolluper,
439436
}
440437

441438
var customRoleHandler CustomRoleHandler = &agplCustomRoleHandler{}
@@ -557,6 +554,7 @@ func New(options *Options) *API {
557554
Pubsub: options.Pubsub,
558555
TemplateScheduleStore: options.TemplateScheduleStore,
559556
StatsBatcher: options.StatsBatcher,
557+
UsageTracker: options.WorkspaceUsageTracker,
560558
UpdateAgentMetricsFn: options.UpdateAgentMetrics,
561559
AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize,
562560
})
@@ -1301,8 +1299,7 @@ type API struct {
13011299
Acquirer *provisionerdserver.Acquirer
13021300
// dbRolluper rolls up template usage stats from raw agent and app
13031301
// stats. This is used to provide insights in the WebUI.
1304-
dbRolluper *dbrollup.Rolluper
1305-
workspaceUsageTracker *workspaceusage.Tracker
1302+
dbRolluper *dbrollup.Rolluper
13061303
}
13071304

13081305
// Close waits for all WebSocket connections to drain before returning.
@@ -1341,7 +1338,7 @@ func (api *API) Close() error {
13411338
_ = (*coordinator).Close()
13421339
}
13431340
_ = api.agentProvider.Close()
1344-
api.workspaceUsageTracker.Close()
1341+
_ = api.statsReporter.Close()
13451342
return nil
13461343
}
13471344

‎coderd/coderdtest/coderdtest.go

+8-9
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@ import (
5454
"github.com/coder/coder/v2/coderd/audit"
5555
"github.com/coder/coder/v2/coderd/autobuild"
5656
"github.com/coder/coder/v2/coderd/awsidentity"
57-
"github.com/coder/coder/v2/coderd/batchstats"
5857
"github.com/coder/coder/v2/coderd/database"
5958
"github.com/coder/coder/v2/coderd/database/dbauthz"
6059
"github.com/coder/coder/v2/coderd/database/dbrollup"
@@ -71,7 +70,7 @@ import (
7170
"github.com/coder/coder/v2/coderd/util/ptr"
7271
"github.com/coder/coder/v2/coderd/workspaceapps"
7372
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
74-
"github.com/coder/coder/v2/coderd/workspaceusage"
73+
"github.com/coder/coder/v2/coderd/workspacestats"
7574
"github.com/coder/coder/v2/codersdk"
7675
"github.com/coder/coder/v2/codersdk/agentsdk"
7776
"github.com/coder/coder/v2/codersdk/drpc"
@@ -145,7 +144,7 @@ type Options struct {
145144
// Logger should only be overridden if you expect errors
146145
// as part of your test.
147146
Logger *slog.Logger
148-
StatsBatcher *batchstats.Batcher
147+
StatsBatcher *workspacestats.DBBatcher
149148

150149
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
151150
AllowWorkspaceRenames bool
@@ -272,10 +271,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
272271
if options.StatsBatcher == nil {
273272
ctx, cancel := context.WithCancel(context.Background())
274273
t.Cleanup(cancel)
275-
batcher, closeBatcher, err := batchstats.New(ctx,
276-
batchstats.WithStore(options.Database),
274+
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
275+
workspacestats.BatcherWithStore(options.Database),
277276
// Avoid cluttering up test output.
278-
batchstats.WithLogger(slog.Make(sloghuman.Sink(io.Discard))),
277+
workspacestats.BatcherWithLogger(slog.Make(sloghuman.Sink(io.Discard))),
279278
)
280279
require.NoError(t, err, "create stats batcher")
281280
options.StatsBatcher = batcher
@@ -337,10 +336,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
337336
options.WorkspaceUsageTrackerTick = make(chan time.Time, 1) // buffering just in case
338337
}
339338
// Close is called by API.Close()
340-
wuTracker := workspaceusage.New(
339+
wuTracker := workspacestats.NewTracker(
341340
options.Database,
342-
workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
343-
workspaceusage.WithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
341+
workspacestats.TrackerWithLogger(options.Logger.Named("workspace_usage_tracker")),
342+
workspacestats.TrackerWithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
344343
)
345344

346345
var mutex sync.RWMutex

‎coderd/insights_test.go

+8-9
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ import (
2121
"cdr.dev/slog/sloggers/slogtest"
2222
"github.com/coder/coder/v2/agent/agenttest"
2323
agentproto "github.com/coder/coder/v2/agent/proto"
24-
"github.com/coder/coder/v2/coderd/batchstats"
2524
"github.com/coder/coder/v2/coderd/coderdtest"
2625
"github.com/coder/coder/v2/coderd/database"
2726
"github.com/coder/coder/v2/coderd/database/dbauthz"
@@ -684,11 +683,11 @@ func TestTemplateInsights_Golden(t *testing.T) {
684683
// NOTE(mafredri): Ideally we would pass batcher as a coderd option and
685684
// insert using the agentClient, but we have a circular dependency on
686685
// the database.
687-
batcher, batcherCloser, err := batchstats.New(
686+
batcher, batcherCloser, err := workspacestats.NewBatcher(
688687
ctx,
689-
batchstats.WithStore(db),
690-
batchstats.WithLogger(logger.Named("batchstats")),
691-
batchstats.WithInterval(time.Hour),
688+
workspacestats.BatcherWithStore(db),
689+
workspacestats.BatcherWithLogger(logger.Named("batchstats")),
690+
workspacestats.BatcherWithInterval(time.Hour),
692691
)
693692
require.NoError(t, err)
694693
defer batcherCloser() // Flushes the stats, this is to ensure they're written.
@@ -1583,11 +1582,11 @@ func TestUserActivityInsights_Golden(t *testing.T) {
15831582
// NOTE(mafredri): Ideally we would pass batcher as a coderd option and
15841583
// insert using the agentClient, but we have a circular dependency on
15851584
// the database.
1586-
batcher, batcherCloser, err := batchstats.New(
1585+
batcher, batcherCloser, err := workspacestats.NewBatcher(
15871586
ctx,
1588-
batchstats.WithStore(db),
1589-
batchstats.WithLogger(logger.Named("batchstats")),
1590-
batchstats.WithInterval(time.Hour),
1587+
workspacestats.BatcherWithStore(db),
1588+
workspacestats.BatcherWithLogger(logger.Named("batchstats")),
1589+
workspacestats.BatcherWithInterval(time.Hour),
15911590
)
15921591
require.NoError(t, err)
15931592
defer batcherCloser() // Flushes the stats, this is to ensure they're written.

‎coderd/prometheusmetrics/prometheusmetrics_test.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,14 @@ import (
2121
"cdr.dev/slog/sloggers/slogtest"
2222

2323
"github.com/coder/coder/v2/coderd/agentmetrics"
24-
"github.com/coder/coder/v2/coderd/batchstats"
2524
"github.com/coder/coder/v2/coderd/coderdtest"
2625
"github.com/coder/coder/v2/coderd/database"
2726
"github.com/coder/coder/v2/coderd/database/dbgen"
2827
"github.com/coder/coder/v2/coderd/database/dbmem"
2928
"github.com/coder/coder/v2/coderd/database/dbtestutil"
3029
"github.com/coder/coder/v2/coderd/database/dbtime"
3130
"github.com/coder/coder/v2/coderd/prometheusmetrics"
31+
"github.com/coder/coder/v2/coderd/workspacestats"
3232
"github.com/coder/coder/v2/codersdk"
3333
"github.com/coder/coder/v2/codersdk/agentsdk"
3434
"github.com/coder/coder/v2/cryptorand"
@@ -391,14 +391,14 @@ func TestAgentStats(t *testing.T) {
391391
db, pubsub := dbtestutil.NewDB(t)
392392
log := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
393393

394-
batcher, closeBatcher, err := batchstats.New(ctx,
394+
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
395395
// We had previously set the batch size to 1 here, but that caused
396396
// intermittent test flakes due to a race between the batcher completing
397397
// its flush and the test asserting that the metrics were collected.
398398
// Instead, we close the batcher after all stats have been posted, which
399399
// forces a flush.
400-
batchstats.WithStore(db),
401-
batchstats.WithLogger(log),
400+
workspacestats.BatcherWithStore(db),
401+
workspacestats.BatcherWithLogger(log),
402402
)
403403
require.NoError(t, err, "create stats batcher failed")
404404
t.Cleanup(closeBatcher)

‎coderd/workspaces.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -1115,7 +1115,7 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
11151115
return
11161116
}
11171117

1118-
api.workspaceUsageTracker.Add(workspace.ID)
1118+
api.statsReporter.TrackUsage(workspace.ID)
11191119
rw.WriteHeader(http.StatusNoContent)
11201120
}
11211121

0 commit comments

Comments
 (0)