Skip to content

Commit 6109d9c

Browse files
f0sself0ssel
committed
chore: move Batcher and Tracker to workspacestats
1 parent 4758952 commit 6109d9c

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"
@@ -869,9 +868,9 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
869868
options.SwaggerEndpoint = vals.Swagger.Enable.Value()
870869
}
871870

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

978977
// Updates workspace usage
979-
tracker := workspaceusage.New(options.Database,
980-
workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")),
978+
tracker := workspacestats.NewTracker(options.Database,
979+
workspacestats.TrackerWithLogger(logger.Named("workspace_usage_tracker")),
981980
)
982981
options.WorkspaceUsageTracker = tracker
983982
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
})
@@ -1300,8 +1298,7 @@ type API struct {
13001298
Acquirer *provisionerdserver.Acquirer
13011299
// dbRolluper rolls up template usage stats from raw agent and app
13021300
// stats. This is used to provide insights in the WebUI.
1303-
dbRolluper *dbrollup.Rolluper
1304-
workspaceUsageTracker *workspaceusage.Tracker
1301+
dbRolluper *dbrollup.Rolluper
13051302
}
13061303

13071304
// Close waits for all WebSocket connections to drain before returning.
@@ -1340,7 +1337,7 @@ func (api *API) Close() error {
13401337
_ = (*coordinator).Close()
13411338
}
13421339
_ = api.agentProvider.Close()
1343-
api.workspaceUsageTracker.Close()
1340+
_ = api.statsReporter.Close()
13441341
return nil
13451342
}
13461343

‎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"
@@ -144,7 +143,7 @@ type Options struct {
144143
// Logger should only be overridden if you expect errors
145144
// as part of your test.
146145
Logger *slog.Logger
147-
StatsBatcher *batchstats.Batcher
146+
StatsBatcher *workspacestats.Batcher
148147

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

345344
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
@@ -192,3 +189,11 @@ func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspac
192189

193190
return nil
194191
}
192+
193+
func (r *Reporter) TrackUsage(workspaceID uuid.UUID) {
194+
r.opts.UsageTracker.Add(workspaceID)
195+
}
196+
197+
func (r *Reporter) Close() error {
198+
return r.opts.UsageTracker.Close()
199+
}

0 commit comments

Comments
 (0)