refactor and add test for locking

coder · mafredri · Mar 22, 2024 · Mar 18, 2024 · Mar 22, 2024 · Mar 22, 2024
commit 2040b875d7ff7d0732bf972cefbe1d985d2ec2d8
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -181,7 +181,6 @@ type Options struct {
 
 	UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
 	StatsBatcher       *batchstats.Batcher
-	DBRollupInterval   time.Duration
 
 	WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
 
@@ -194,6 +193,9 @@ type Options struct {
 	// NewTicker is used for unit tests to replace "time.NewTicker".
 	NewTicker func(duration time.Duration) (tick <-chan time.Time, done func())
 
+	// DatabaseRolluper rolls up template usage stats from raw agent and app
+	// stats. This is used to provide insights in the WebUI.
+	DatabaseRolluper *dbrollup.Rolluper
 	// WorkspaceUsageTracker tracks workspace usage by the CLI.
 	WorkspaceUsageTracker *workspaceusage.Tracker
 }
@@ -344,9 +346,6 @@ func New(options *Options) *API {
 	if options.StatsBatcher == nil {
 		panic("developer error: options.StatsBatcher is nil")
 	}
-	if options.DBRollupInterval == 0 {
-		options.DBRollupInterval = dbrollup.DefaultInterval
-	}
 
 	siteCacheDir := options.CacheDir
 	if siteCacheDir != "" {
@@ -371,6 +370,10 @@ func New(options *Options) *API {
 		OIDC:   options.OIDCConfig,
 	}
 
+	if options.DatabaseRolluper == nil {
+		options.DatabaseRolluper = dbrollup.New(options.Logger.Named("dbrollup"), options.Database)
+	}
+
 	if options.WorkspaceUsageTracker == nil {
 		options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
 			workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
@@ -421,11 +424,7 @@ func New(options *Options) *API {
 			options.Database,
 			options.Pubsub,
 		),
-		rolluper: dbrollup.New(
-			options.Logger,
-			options.Database,
-			options.DBRollupInterval,
-		),
+		dbRolluper:            options.DatabaseRolluper,
 		workspaceUsageTracker: options.WorkspaceUsageTracker,
 	}
 
@@ -1208,9 +1207,9 @@ type API struct {
 	statsBatcher *batchstats.Batcher
 
 	Acquirer *provisionerdserver.Acquirer
-	// rolluper rolls up template usage stats from raw agent and app
+	// dbRolluper rolls up template usage stats from raw agent and app
 	// stats. This is used to provide insights in the WebUI.
-	rolluper              *dbrollup.Rolluper
+	dbRolluper            *dbrollup.Rolluper
 	workspaceUsageTracker *workspaceusage.Tracker
 }
 
@@ -1225,7 +1224,7 @@ func (api *API) Close() error {
 	api.WebsocketWaitGroup.Wait()
 	api.WebsocketWaitMutex.Unlock()
 
-	api.rolluper.Close()
+	api.dbRolluper.Close()
 	api.metricsCache.Close()
 	if api.updateChecker != nil {
 		api.updateChecker.Close()

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -57,6 +57,7 @@ import (
 	"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"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/externalauth"
@@ -147,7 +148,7 @@ type Options struct {
 	WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
 	AllowWorkspaceRenames              bool
 	NewTicker                          func(duration time.Duration) (<-chan time.Time, func())
-	DBRollupInterval                   time.Duration
+	DatabaseRolluper                   *dbrollup.Rolluper
 	WorkspaceUsageTrackerFlush         chan int
 	WorkspaceUsageTrackerTick          chan time.Time
 }
@@ -492,7 +493,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
 			AllowWorkspaceRenames:              options.AllowWorkspaceRenames,
 			NewTicker:                          options.NewTicker,
-			DBRollupInterval:                   options.DBRollupInterval,
+			DatabaseRolluper:                   options.DatabaseRolluper,
 			WorkspaceUsageTracker:              wuTracker,
 		}
 }

diff --git a/coderd/database/dbrollup/dbrollup.go b/coderd/database/dbrollup/dbrollup.go
@@ -2,6 +2,7 @@ package dbrollup
 
 import (
 	"context"
+	"flag"
 	"time"
 
 	"golang.org/x/sync/errgroup"
@@ -19,36 +20,68 @@ const (
 	DefaultInterval = 5 * time.Minute
 )
 
+type Event struct {
+	TemplateUsageStats bool
+}
+
 type Rolluper struct {
-	cancel context.CancelFunc
-	closed chan struct{}
-	db     database.Store
-	logger slog.Logger
+	cancel   context.CancelFunc
+	closed   chan struct{}
+	db       database.Store
+	logger   slog.Logger
+	interval time.Duration
+	event    chan<- Event
+}
+
+type Option func(*Rolluper)
+
+// WithInterval sets the interval between rollups.
+func WithInterval(interval time.Duration) Option {
+	return func(r *Rolluper) {
+		r.interval = interval
+	}
+}
+
+// WithEventChannel sets the event channel to use for rollup events.
+//
+// This is only used for testing.
+func WithEventChannel(ch chan<- Event) Option {
+	if flag.Lookup("test.v") == nil {
+		panic("developer error: WithEventChannel is not to be used outside of tests")
+	}
+	return func(r *Rolluper) {
+		r.event = ch
+	}
 }
 
 // New creates a new DB rollup service that periodically runs rollup queries.
 // It is the caller's responsibility to call Close on the returned instance.
 //
 // This is for e.g. generating insights data (template_usage_stats) from
 // raw data (workspace_agent_stats, workspace_app_stats).
-func New(logger slog.Logger, db database.Store, interval time.Duration) *Rolluper {
+func New(logger slog.Logger, db database.Store, opts ...Option) *Rolluper {
 	ctx, cancel := context.WithCancel(context.Background())
 
 	r := &Rolluper{
-		cancel: cancel,
-		closed: make(chan struct{}),
-		db:     db,
-		logger: logger.Named("dbrollup"),
+		cancel:   cancel,
+		closed:   make(chan struct{}),
+		db:       db,
+		logger:   logger,
+		interval: DefaultInterval,
+	}
+
+	for _, opt := range opts {
+		opt(r)
 	}
 
 	//nolint:gocritic // The system rolls up database tables without user input.
 	ctx = dbauthz.AsSystemRestricted(ctx)
-	go r.start(ctx, interval)
+	go r.start(ctx)
 
 	return r
 }
 
-func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
+func (r *Rolluper) start(ctx context.Context) {
 	defer close(r.closed)
 
 	do := func() {
@@ -58,7 +91,7 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
 		now := time.Now()
 
 		// Track whether or not we performed a rollup (we got the advisory lock).
-		templateUsageStats := false
+		var ev Event
 
 		eg.Go(func() error {
 			return r.db.InTx(func(tx database.Store) error {
@@ -72,7 +105,7 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
 					return nil
 				}
 
-				templateUsageStats = true
+				ev.TemplateUsageStats = true
 				return tx.UpsertTemplateUsageStats(ctx)
 			}, nil)
 		})
@@ -86,12 +119,22 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
 			if ctx.Err() == nil {
 				r.logger.Error(ctx, "failed to rollup data", slog.Error(err))
 			}
-		} else {
-			r.logger.Debug(ctx,
-				"rolled up data",
-				slog.F("took", time.Since(now)),
-				slog.F("template_usage_stats", templateUsageStats),
-			)
+			return
+		}
+
+		r.logger.Debug(ctx,
+			"rolled up data",
+			slog.F("took", time.Since(now)),
+			slog.F("event", ev),
+		)
+
+		// For testing.
+		if r.event != nil {
+			select {
+			case <-ctx.Done():
+				return
+			case r.event <- ev:
+			}
 		}
 	}
 
@@ -108,11 +151,11 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
 		case <-t.C:
 			// Ensure we're on the interval.
 			now := time.Now()
-			next := now.Add(interval).Truncate(interval) // Ensure we're on the interval and synced with the clock.
+			next := now.Add(r.interval).Truncate(r.interval) // Ensure we're on the interval and synced with the clock.
 			d := next.Sub(now)
 			// Safety check (shouldn't be possible).
 			if d <= 0 {
-				d = interval
+				d = r.interval
 			}
 			t.Reset(d)