Skip to content

Commit 530e334

Browse files
f0sself0ssel
committed
chore: move Batcher and Tracker to workspacestats
1 parent c7e7312 commit 530e334

File tree

10 files changed

+89
-86
lines changed

10 files changed

+89
-86
lines changed

‎cli/server.go

Lines changed: 6 additions & 7 deletions
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/coderd.go

Lines changed: 8 additions & 11 deletions
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.Batcher
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

Lines changed: 8 additions & 9 deletions
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.Batcher
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

Lines changed: 8 additions & 9 deletions
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

Lines changed: 4 additions & 4 deletions
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/batchstats/batcher.go renamed to ‎coderd/workspacestats/batcher.go

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package batchstats
1+
package workspacestats
22

33
import (
44
"context"
@@ -24,8 +24,12 @@ const (
2424
defaultFlushInterval = time.Second
2525
)
2626

27+
type StatsBatcher interface {
28+
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error
29+
}
30+
2731
// Batcher holds a buffer of agent stats and periodically flushes them to
28-
// its configured store. It also updates the workspace's last used time.
32+
// its configured store.
2933
type Batcher struct {
3034
store database.Store
3135
log slog.Logger
@@ -50,38 +54,38 @@ type Batcher struct {
5054
}
5155

5256
// Option is a functional option for configuring a Batcher.
53-
type Option func(b *Batcher)
57+
type BatcherOption func(b *Batcher)
5458

55-
// WithStore sets the store to use for storing stats.
56-
func WithStore(store database.Store) Option {
59+
// BatcherWithStore sets the store to use for storing stats.
60+
func BatcherWithStore(store database.Store) BatcherOption {
5761
return func(b *Batcher) {
5862
b.store = store
5963
}
6064
}
6165

62-
// WithBatchSize sets the number of stats to store in a batch.
63-
func WithBatchSize(size int) Option {
66+
// BatcherWithBatchSize sets the number of stats to store in a batch.
67+
func BatcherWithBatchSize(size int) BatcherOption {
6468
return func(b *Batcher) {
6569
b.batchSize = size
6670
}
6771
}
6872

69-
// WithInterval sets the interval for flushes.
70-
func WithInterval(d time.Duration) Option {
73+
// BatcherWithInterval sets the interval for flushes.
74+
func BatcherWithInterval(d time.Duration) BatcherOption {
7175
return func(b *Batcher) {
7276
b.interval = d
7377
}
7478
}
7579

76-
// WithLogger sets the logger to use for logging.
77-
func WithLogger(log slog.Logger) Option {
80+
// BatcherWithLogger sets the logger to use for logging.
81+
func BatcherWithLogger(log slog.Logger) BatcherOption {
7882
return func(b *Batcher) {
7983
b.log = log
8084
}
8185
}
8286

83-
// New creates a new Batcher and starts it.
84-
func New(ctx context.Context, opts ...Option) (*Batcher, func(), error) {
87+
// NewBatcher creates a new Batcher and starts it.
88+
func NewBatcher(ctx context.Context, opts ...BatcherOption) (*Batcher, func(), error) {
8589
b := &Batcher{}
8690
b.log = slog.Make(sloghuman.Sink(os.Stderr))
8791
b.flushLever = make(chan struct{}, 1) // Buffered so that it doesn't block.

‎coderd/batchstats/batcher_internal_test.go renamed to ‎coderd/workspacestats/batcher_internal_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package batchstats
1+
package workspacestats
22

33
import (
44
"context"
@@ -35,9 +35,9 @@ func TestBatchStats(t *testing.T) {
3535
tick := make(chan time.Time)
3636
flushed := make(chan int, 1)
3737

38-
b, closer, err := New(ctx,
39-
WithStore(store),
40-
WithLogger(log),
38+
b, closer, err := NewBatcher(ctx,
39+
BatcherWithStore(store),
40+
BatcherWithLogger(log),
4141
func(b *Batcher) {
4242
b.tickCh = tick
4343
b.flushed = flushed

‎coderd/workspacestats/reporter.go

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,13 @@ import (
2222
"github.com/coder/coder/v2/codersdk"
2323
)
2424

25-
type StatsBatcher interface {
26-
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error
27-
}
28-
2925
type ReporterOptions struct {
3026
Database database.Store
3127
Logger slog.Logger
3228
Pubsub pubsub.Pubsub
3329
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
3430
StatsBatcher StatsBatcher
31+
UsageTracker *UsageTracker
3532
UpdateAgentMetricsFn func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
3633

3734
AppStatBatchSize int
@@ -205,3 +202,11 @@ func UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, db database.Store,
205202
}
206203
return nil
207204
}
205+
206+
func (r *Reporter) TrackUsage(workspaceID uuid.UUID) {
207+
r.opts.UsageTracker.Add(workspaceID)
208+
}
209+
210+
func (r *Reporter) Close() error {
211+
return r.opts.UsageTracker.Close()
212+
}

0 commit comments

Comments
 (0)