diff --git a/cli/server.go b/cli/server.go index 855c5e6c547bf..11c3ec50ba833 100644 --- a/cli/server.go +++ b/cli/server.go @@ -62,7 +62,6 @@ import ( "github.com/coder/coder/v2/cli/config" "github.com/coder/coder/v2/coderd" "github.com/coder/coder/v2/coderd/autobuild" - "github.com/coder/coder/v2/coderd/batchstats" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/awsiamrds" "github.com/coder/coder/v2/coderd/database/dbmem" @@ -87,7 +86,7 @@ import ( stringutil "github.com/coder/coder/v2/coderd/util/strings" "github.com/coder/coder/v2/coderd/workspaceapps" "github.com/coder/coder/v2/coderd/workspaceapps/appurl" - "github.com/coder/coder/v2/coderd/workspaceusage" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/drpc" "github.com/coder/coder/v2/cryptorand" @@ -870,9 +869,9 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd. options.SwaggerEndpoint = vals.Swagger.Enable.Value() } - batcher, closeBatcher, err := batchstats.New(ctx, - batchstats.WithLogger(options.Logger.Named("batchstats")), - batchstats.WithStore(options.Database), + batcher, closeBatcher, err := workspacestats.NewBatcher(ctx, + workspacestats.BatcherWithLogger(options.Logger.Named("batchstats")), + workspacestats.BatcherWithStore(options.Database), ) if err != nil { 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. defer purger.Close() // Updates workspace usage - tracker := workspaceusage.New(options.Database, - workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")), + tracker := workspacestats.NewTracker(options.Database, + workspacestats.TrackerWithLogger(logger.Named("workspace_usage_tracker")), ) options.WorkspaceUsageTracker = tracker defer tracker.Close() diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go index ee17897572f3d..a167fb5d6f275 100644 --- a/coderd/agentapi/stats.go +++ b/coderd/agentapi/stats.go @@ -7,8 +7,6 @@ import ( "golang.org/x/xerrors" "google.golang.org/protobuf/types/known/durationpb" - "github.com/google/uuid" - "cdr.dev/slog" agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/coderd/database" @@ -16,10 +14,6 @@ import ( "github.com/coder/coder/v2/coderd/workspacestats" ) -type StatsBatcher interface { - Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error -} - type StatsAPI struct { AgentFn func(context.Context) (database.WorkspaceAgent, error) Database database.Store diff --git a/coderd/agentapi/stats_test.go b/coderd/agentapi/stats_test.go index c304dea93ecc9..8b4d72fc1d579 100644 --- a/coderd/agentapi/stats_test.go +++ b/coderd/agentapi/stats_test.go @@ -39,7 +39,7 @@ type statsBatcher struct { lastStats *agentproto.Stats } -var _ agentapi.StatsBatcher = &statsBatcher{} +var _ workspacestats.Batcher = &statsBatcher{} func (b *statsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error { b.mu.Lock() diff --git a/coderd/coderd.go b/coderd/coderd.go index 60bfe9813c559..22154b97d963d 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -43,7 +43,6 @@ import ( "github.com/coder/coder/v2/coderd/appearance" "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/awsidentity" - "github.com/coder/coder/v2/coderd/batchstats" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbrollup" @@ -69,7 +68,6 @@ import ( "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/coderd/workspaceapps" "github.com/coder/coder/v2/coderd/workspacestats" - "github.com/coder/coder/v2/coderd/workspaceusage" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/drpc" "github.com/coder/coder/v2/codersdk/healthsdk" @@ -189,7 +187,7 @@ type Options struct { HTTPClient *http.Client UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) - StatsBatcher *batchstats.Batcher + StatsBatcher *workspacestats.DBBatcher WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions @@ -206,7 +204,7 @@ type Options struct { // stats. This is used to provide insights in the WebUI. DatabaseRolluper *dbrollup.Rolluper // WorkspaceUsageTracker tracks workspace usage by the CLI. - WorkspaceUsageTracker *workspaceusage.Tracker + WorkspaceUsageTracker *workspacestats.UsageTracker } // @title Coder API @@ -384,8 +382,8 @@ func New(options *Options) *API { } if options.WorkspaceUsageTracker == nil { - options.WorkspaceUsageTracker = workspaceusage.New(options.Database, - workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")), + options.WorkspaceUsageTracker = workspacestats.NewTracker(options.Database, + workspacestats.TrackerWithLogger(options.Logger.Named("workspace_usage_tracker")), ) } @@ -434,8 +432,7 @@ func New(options *Options) *API { options.Database, options.Pubsub, ), - dbRolluper: options.DatabaseRolluper, - workspaceUsageTracker: options.WorkspaceUsageTracker, + dbRolluper: options.DatabaseRolluper, } var customRoleHandler CustomRoleHandler = &agplCustomRoleHandler{} @@ -557,6 +554,7 @@ func New(options *Options) *API { Pubsub: options.Pubsub, TemplateScheduleStore: options.TemplateScheduleStore, StatsBatcher: options.StatsBatcher, + UsageTracker: options.WorkspaceUsageTracker, UpdateAgentMetricsFn: options.UpdateAgentMetrics, AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize, }) @@ -1301,8 +1299,7 @@ type API struct { Acquirer *provisionerdserver.Acquirer // dbRolluper rolls up template usage stats from raw agent and app // stats. This is used to provide insights in the WebUI. - dbRolluper *dbrollup.Rolluper - workspaceUsageTracker *workspaceusage.Tracker + dbRolluper *dbrollup.Rolluper } // Close waits for all WebSocket connections to drain before returning. @@ -1341,7 +1338,7 @@ func (api *API) Close() error { _ = (*coordinator).Close() } _ = api.agentProvider.Close() - api.workspaceUsageTracker.Close() + _ = api.statsReporter.Close() return nil } diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 3ebca07686d0e..9ca2f551978f1 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -54,7 +54,6 @@ import ( "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/autobuild" "github.com/coder/coder/v2/coderd/awsidentity" - "github.com/coder/coder/v2/coderd/batchstats" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbrollup" @@ -71,7 +70,7 @@ import ( "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/coderd/workspaceapps" "github.com/coder/coder/v2/coderd/workspaceapps/appurl" - "github.com/coder/coder/v2/coderd/workspaceusage" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/codersdk/drpc" @@ -145,7 +144,7 @@ type Options struct { // Logger should only be overridden if you expect errors // as part of your test. Logger *slog.Logger - StatsBatcher *batchstats.Batcher + StatsBatcher *workspacestats.DBBatcher WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions AllowWorkspaceRenames bool @@ -272,10 +271,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can if options.StatsBatcher == nil { ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) - batcher, closeBatcher, err := batchstats.New(ctx, - batchstats.WithStore(options.Database), + batcher, closeBatcher, err := workspacestats.NewBatcher(ctx, + workspacestats.BatcherWithStore(options.Database), // Avoid cluttering up test output. - batchstats.WithLogger(slog.Make(sloghuman.Sink(io.Discard))), + workspacestats.BatcherWithLogger(slog.Make(sloghuman.Sink(io.Discard))), ) require.NoError(t, err, "create stats batcher") options.StatsBatcher = batcher @@ -337,10 +336,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can options.WorkspaceUsageTrackerTick = make(chan time.Time, 1) // buffering just in case } // Close is called by API.Close() - wuTracker := workspaceusage.New( + wuTracker := workspacestats.NewTracker( options.Database, - workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")), - workspaceusage.WithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush), + workspacestats.TrackerWithLogger(options.Logger.Named("workspace_usage_tracker")), + workspacestats.TrackerWithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush), ) var mutex sync.RWMutex diff --git a/coderd/insights_test.go b/coderd/insights_test.go index 22e7ed6947bac..2447ec37f3516 100644 --- a/coderd/insights_test.go +++ b/coderd/insights_test.go @@ -21,7 +21,6 @@ import ( "cdr.dev/slog/sloggers/slogtest" "github.com/coder/coder/v2/agent/agenttest" agentproto "github.com/coder/coder/v2/agent/proto" - "github.com/coder/coder/v2/coderd/batchstats" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" @@ -684,11 +683,11 @@ func TestTemplateInsights_Golden(t *testing.T) { // NOTE(mafredri): Ideally we would pass batcher as a coderd option and // insert using the agentClient, but we have a circular dependency on // the database. - batcher, batcherCloser, err := batchstats.New( + batcher, batcherCloser, err := workspacestats.NewBatcher( ctx, - batchstats.WithStore(db), - batchstats.WithLogger(logger.Named("batchstats")), - batchstats.WithInterval(time.Hour), + workspacestats.BatcherWithStore(db), + workspacestats.BatcherWithLogger(logger.Named("batchstats")), + workspacestats.BatcherWithInterval(time.Hour), ) require.NoError(t, err) defer batcherCloser() // Flushes the stats, this is to ensure they're written. @@ -1583,11 +1582,11 @@ func TestUserActivityInsights_Golden(t *testing.T) { // NOTE(mafredri): Ideally we would pass batcher as a coderd option and // insert using the agentClient, but we have a circular dependency on // the database. - batcher, batcherCloser, err := batchstats.New( + batcher, batcherCloser, err := workspacestats.NewBatcher( ctx, - batchstats.WithStore(db), - batchstats.WithLogger(logger.Named("batchstats")), - batchstats.WithInterval(time.Hour), + workspacestats.BatcherWithStore(db), + workspacestats.BatcherWithLogger(logger.Named("batchstats")), + workspacestats.BatcherWithInterval(time.Hour), ) require.NoError(t, err) defer batcherCloser() // Flushes the stats, this is to ensure they're written. diff --git a/coderd/prometheusmetrics/prometheusmetrics_test.go b/coderd/prometheusmetrics/prometheusmetrics_test.go index 9c4c9fca0b66f..3e42b1ea783c6 100644 --- a/coderd/prometheusmetrics/prometheusmetrics_test.go +++ b/coderd/prometheusmetrics/prometheusmetrics_test.go @@ -21,7 +21,6 @@ import ( "cdr.dev/slog/sloggers/slogtest" "github.com/coder/coder/v2/coderd/agentmetrics" - "github.com/coder/coder/v2/coderd/batchstats" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" @@ -29,6 +28,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/prometheusmetrics" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/cryptorand" @@ -391,14 +391,14 @@ func TestAgentStats(t *testing.T) { db, pubsub := dbtestutil.NewDB(t) log := slogtest.Make(t, nil).Leveled(slog.LevelDebug) - batcher, closeBatcher, err := batchstats.New(ctx, + batcher, closeBatcher, err := workspacestats.NewBatcher(ctx, // We had previously set the batch size to 1 here, but that caused // intermittent test flakes due to a race between the batcher completing // its flush and the test asserting that the metrics were collected. // Instead, we close the batcher after all stats have been posted, which // forces a flush. - batchstats.WithStore(db), - batchstats.WithLogger(log), + workspacestats.BatcherWithStore(db), + workspacestats.BatcherWithLogger(log), ) require.NoError(t, err, "create stats batcher failed") t.Cleanup(closeBatcher) diff --git a/coderd/workspaces.go b/coderd/workspaces.go index 7d0344be4e321..1b3f076e8f4bf 100644 --- a/coderd/workspaces.go +++ b/coderd/workspaces.go @@ -1115,7 +1115,7 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) { return } - api.workspaceUsageTracker.Add(workspace.ID) + api.statsReporter.TrackUsage(workspace.ID) rw.WriteHeader(http.StatusNoContent) } diff --git a/coderd/batchstats/batcher.go b/coderd/workspacestats/batcher.go similarity index 86% rename from coderd/batchstats/batcher.go rename to coderd/workspacestats/batcher.go index bbff38b0413c0..2872c368dc61c 100644 --- a/coderd/batchstats/batcher.go +++ b/coderd/workspacestats/batcher.go @@ -1,4 +1,4 @@ -package batchstats +package workspacestats import ( "context" @@ -24,9 +24,13 @@ const ( defaultFlushInterval = time.Second ) -// Batcher holds a buffer of agent stats and periodically flushes them to -// its configured store. It also updates the workspace's last used time. -type Batcher struct { +type Batcher interface { + Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error +} + +// DBBatcher holds a buffer of agent stats and periodically flushes them to +// its configured store. +type DBBatcher struct { store database.Store log slog.Logger @@ -50,39 +54,39 @@ type Batcher struct { } // Option is a functional option for configuring a Batcher. -type Option func(b *Batcher) +type BatcherOption func(b *DBBatcher) -// WithStore sets the store to use for storing stats. -func WithStore(store database.Store) Option { - return func(b *Batcher) { +// BatcherWithStore sets the store to use for storing stats. +func BatcherWithStore(store database.Store) BatcherOption { + return func(b *DBBatcher) { b.store = store } } -// WithBatchSize sets the number of stats to store in a batch. -func WithBatchSize(size int) Option { - return func(b *Batcher) { +// BatcherWithBatchSize sets the number of stats to store in a batch. +func BatcherWithBatchSize(size int) BatcherOption { + return func(b *DBBatcher) { b.batchSize = size } } -// WithInterval sets the interval for flushes. -func WithInterval(d time.Duration) Option { - return func(b *Batcher) { +// BatcherWithInterval sets the interval for flushes. +func BatcherWithInterval(d time.Duration) BatcherOption { + return func(b *DBBatcher) { b.interval = d } } -// WithLogger sets the logger to use for logging. -func WithLogger(log slog.Logger) Option { - return func(b *Batcher) { +// BatcherWithLogger sets the logger to use for logging. +func BatcherWithLogger(log slog.Logger) BatcherOption { + return func(b *DBBatcher) { b.log = log } } -// New creates a new Batcher and starts it. -func New(ctx context.Context, opts ...Option) (*Batcher, func(), error) { - b := &Batcher{} +// NewBatcher creates a new Batcher and starts it. +func NewBatcher(ctx context.Context, opts ...BatcherOption) (*DBBatcher, func(), error) { + b := &DBBatcher{} b.log = slog.Make(sloghuman.Sink(os.Stderr)) b.flushLever = make(chan struct{}, 1) // Buffered so that it doesn't block. for _, opt := range opts { @@ -127,7 +131,7 @@ func New(ctx context.Context, opts ...Option) (*Batcher, func(), error) { } // Add adds a stat to the batcher for the given workspace and agent. -func (b *Batcher) Add( +func (b *DBBatcher) Add( now time.Time, agentID uuid.UUID, templateID uuid.UUID, @@ -174,7 +178,7 @@ func (b *Batcher) Add( } // Run runs the batcher. -func (b *Batcher) run(ctx context.Context) { +func (b *DBBatcher) run(ctx context.Context) { // nolint:gocritic // This is only ever used for one thing - inserting agent stats. authCtx := dbauthz.AsSystemRestricted(ctx) for { @@ -199,7 +203,7 @@ func (b *Batcher) run(ctx context.Context) { } // flush flushes the batcher's buffer. -func (b *Batcher) flush(ctx context.Context, forced bool, reason string) { +func (b *DBBatcher) flush(ctx context.Context, forced bool, reason string) { b.mu.Lock() b.flushForced.Store(true) start := time.Now() @@ -256,7 +260,7 @@ func (b *Batcher) flush(ctx context.Context, forced bool, reason string) { } // initBuf resets the buffer. b MUST be locked. -func (b *Batcher) initBuf(size int) { +func (b *DBBatcher) initBuf(size int) { b.buf = &database.InsertWorkspaceAgentStatsParams{ ID: make([]uuid.UUID, 0, b.batchSize), CreatedAt: make([]time.Time, 0, b.batchSize), @@ -280,7 +284,7 @@ func (b *Batcher) initBuf(size int) { b.connectionsByProto = make([]map[string]int64, 0, size) } -func (b *Batcher) resetBuf() { +func (b *DBBatcher) resetBuf() { b.buf.ID = b.buf.ID[:0] b.buf.CreatedAt = b.buf.CreatedAt[:0] b.buf.UserID = b.buf.UserID[:0] diff --git a/coderd/batchstats/batcher_internal_test.go b/coderd/workspacestats/batcher_internal_test.go similarity index 98% rename from coderd/batchstats/batcher_internal_test.go rename to coderd/workspacestats/batcher_internal_test.go index d153ac283b086..0e797986555e5 100644 --- a/coderd/batchstats/batcher_internal_test.go +++ b/coderd/workspacestats/batcher_internal_test.go @@ -1,4 +1,4 @@ -package batchstats +package workspacestats import ( "context" @@ -35,10 +35,10 @@ func TestBatchStats(t *testing.T) { tick := make(chan time.Time) flushed := make(chan int, 1) - b, closer, err := New(ctx, - WithStore(store), - WithLogger(log), - func(b *Batcher) { + b, closer, err := NewBatcher(ctx, + BatcherWithStore(store), + BatcherWithLogger(log), + func(b *DBBatcher) { b.tickCh = tick b.flushed = flushed }, diff --git a/coderd/workspacestats/reporter.go b/coderd/workspacestats/reporter.go index 8ae4bdd827ac3..c6b7afb3c68ad 100644 --- a/coderd/workspacestats/reporter.go +++ b/coderd/workspacestats/reporter.go @@ -22,16 +22,13 @@ import ( "github.com/coder/coder/v2/codersdk" ) -type StatsBatcher interface { - Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error -} - type ReporterOptions struct { Database database.Store Logger slog.Logger Pubsub pubsub.Pubsub TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore] - StatsBatcher StatsBatcher + StatsBatcher Batcher + UsageTracker *UsageTracker UpdateAgentMetricsFn func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) AppStatBatchSize int @@ -205,3 +202,11 @@ func UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, db database.Store, } return nil } + +func (r *Reporter) TrackUsage(workspaceID uuid.UUID) { + r.opts.UsageTracker.Add(workspaceID) +} + +func (r *Reporter) Close() error { + return r.opts.UsageTracker.Close() +} diff --git a/coderd/workspaceusage/tracker.go b/coderd/workspacestats/tracker.go similarity index 86% rename from coderd/workspaceusage/tracker.go rename to coderd/workspacestats/tracker.go index 118b021d71d52..33532247b36e0 100644 --- a/coderd/workspaceusage/tracker.go +++ b/coderd/workspacestats/tracker.go @@ -1,4 +1,4 @@ -package workspaceusage +package workspacestats import ( "bytes" @@ -25,10 +25,10 @@ type Store interface { BatchUpdateWorkspaceLastUsedAt(context.Context, database.BatchUpdateWorkspaceLastUsedAtParams) error } -// Tracker tracks and de-bounces updates to workspace usage activity. +// UsageTracker tracks and de-bounces updates to workspace usage activity. // It keeps an internal map of workspace IDs that have been used and // periodically flushes this to its configured Store. -type Tracker struct { +type UsageTracker struct { log slog.Logger // you know, for logs flushLock sync.Mutex // protects m flushErrors int // tracks the number of consecutive errors flushing @@ -42,10 +42,10 @@ type Tracker struct { flushCh chan int // used for testing. } -// New returns a new Tracker. It is the caller's responsibility +// NewTracker returns a new Tracker. It is the caller's responsibility // to call Close(). -func New(s Store, opts ...Option) *Tracker { - tr := &Tracker{ +func NewTracker(s Store, opts ...TrackerOption) *UsageTracker { + tr := &UsageTracker{ log: slog.Make(sloghuman.Sink(os.Stderr)), m: &uuidSet{}, s: s, @@ -67,33 +67,33 @@ func New(s Store, opts ...Option) *Tracker { return tr } -type Option func(*Tracker) +type TrackerOption func(*UsageTracker) -// WithLogger sets the logger to be used by Tracker. -func WithLogger(log slog.Logger) Option { - return func(h *Tracker) { +// TrackerWithLogger sets the logger to be used by Tracker. +func TrackerWithLogger(log slog.Logger) TrackerOption { + return func(h *UsageTracker) { h.log = log } } -// WithFlushInterval allows configuring the flush interval of Tracker. -func WithFlushInterval(d time.Duration) Option { - return func(h *Tracker) { +// TrackerWithFlushInterval allows configuring the flush interval of Tracker. +func TrackerWithFlushInterval(d time.Duration) TrackerOption { + return func(h *UsageTracker) { ticker := time.NewTicker(d) h.tickCh = ticker.C h.stopTick = ticker.Stop } } -// WithTickFlush allows passing two channels: one that reads +// TrackerWithTickFlush allows passing two channels: one that reads // a time.Time, and one that returns the number of marked workspaces // every time Tracker flushes. // For testing only and will panic if used outside of tests. -func WithTickFlush(tickCh <-chan time.Time, flushCh chan int) Option { +func TrackerWithTickFlush(tickCh <-chan time.Time, flushCh chan int) TrackerOption { if flag.Lookup("test.v") == nil { panic("developer error: WithTickFlush is not to be used outside of tests.") } - return func(h *Tracker) { + return func(h *UsageTracker) { h.tickCh = tickCh h.stopTick = func() {} h.flushCh = flushCh @@ -102,14 +102,14 @@ func WithTickFlush(tickCh <-chan time.Time, flushCh chan int) Option { // Add marks the workspace with the given ID as having been used recently. // Tracker will periodically flush this to its configured Store. -func (tr *Tracker) Add(workspaceID uuid.UUID) { +func (tr *UsageTracker) Add(workspaceID uuid.UUID) { tr.m.Add(workspaceID) } // flush updates last_used_at of all current workspace IDs. // If this is held while a previous flush is in progress, it will // deadlock until the previous flush has completed. -func (tr *Tracker) flush(now time.Time) { +func (tr *UsageTracker) flush(now time.Time) { // Copy our current set of IDs ids := tr.m.UniqueAndClear() count := len(ids) @@ -154,7 +154,7 @@ func (tr *Tracker) flush(now time.Time) { // loop periodically flushes every tick. // If loop is called after Close, it will exit immediately and log an error. -func (tr *Tracker) loop() { +func (tr *UsageTracker) loop() { select { case <-tr.doneCh: tr.log.Error(context.Background(), "developer error: Loop called after Close") @@ -186,7 +186,7 @@ func (tr *Tracker) loop() { // Close stops Tracker and returns once Loop has exited. // After calling Close(), Loop must not be called. -func (tr *Tracker) Close() error { +func (tr *UsageTracker) Close() error { tr.stopOnce.Do(func() { tr.stopCh <- struct{}{} tr.stopTick() diff --git a/coderd/workspaceusage/tracker_test.go b/coderd/workspacestats/tracker_test.go similarity index 96% rename from coderd/workspaceusage/tracker_test.go rename to coderd/workspacestats/tracker_test.go index ae9a9d2162d1c..99e9f9503b645 100644 --- a/coderd/workspaceusage/tracker_test.go +++ b/coderd/workspacestats/tracker_test.go @@ -1,4 +1,4 @@ -package workspaceusage_test +package workspacestats_test import ( "bytes" @@ -21,7 +21,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/database/pubsub" - "github.com/coder/coder/v2/coderd/workspaceusage" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) @@ -35,9 +35,9 @@ func TestTracker(t *testing.T) { tickCh := make(chan time.Time) flushCh := make(chan int, 1) - wut := workspaceusage.New(mDB, - workspaceusage.WithLogger(log), - workspaceusage.WithTickFlush(tickCh, flushCh), + wut := workspacestats.NewTracker(mDB, + workspacestats.TrackerWithLogger(log), + workspacestats.TrackerWithTickFlush(tickCh, flushCh), ) defer wut.Close()