refactor: unexport Loop, remove panic, simplify external API

coder · johnstcn · Mar 20, 2024 · Mar 18, 2024 · Mar 18, 2024 · Mar 18, 2024
commit c9ac9d23b1deefad63c407cd448a8c7d6bbff7a4
diff --git a/cli/portforward_test.go b/cli/portforward_test.go
@@ -21,9 +21,7 @@ import (
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
-	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
-	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/pty/ptytest"
 	"github.com/coder/coder/v2/testutil"
@@ -99,14 +97,11 @@ func TestPortForward(t *testing.T) {
 	// Setup agent once to be shared between test-cases (avoid expensive
 	// non-parallel setup).
 	var (
-		db, ps  = dbtestutil.NewDB(t)
-		wuTick  = make(chan time.Time)
-		wuFlush = make(chan int, 1)
-		wut     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlush), workspaceusage.WithTickChannel(wuTick))
-		client  = coderdtest.New(t, &coderdtest.Options{
-			WorkspaceUsageTracker: wut,
-			Database:              db,
-			Pubsub:                ps,
+		wuTick     = make(chan time.Time)
+		wuFlush    = make(chan int, 1)
+		client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{
+			WorkspaceUsageTrackerTick:  wuTick,
+			WorkspaceUsageTrackerFlush: wuFlush,
 		})
 		admin              = coderdtest.CreateFirstUser(t, client)
 		member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)

diff --git a/cli/server.go b/cli/server.go
@@ -86,6 +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/codersdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/cryptorand"
@@ -968,6 +969,13 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			purger := dbpurge.New(ctx, logger, options.Database)
 			defer purger.Close()
 
+			// Updates workspace usage
+			tracker := workspaceusage.New(options.Database,
+				workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")),
+			)
+			options.WorkspaceUsageTracker = tracker
+			defer tracker.Close()
+
 			// Wrap the server in middleware that redirects to the access URL if
 			// the request is not to a local IP.
 			var handler http.Handler = coderAPI.RootHandler

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -371,7 +371,6 @@ func New(options *Options) *API {
 			workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
 		)
 	}
-	go options.WorkspaceUsageTracker.Loop()
 
 	ctx, cancel := context.WithCancel(context.Background())
 	r := chi.NewRouter()

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -147,7 +147,8 @@ type Options struct {
 	WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
 	AllowWorkspaceRenames              bool
 	NewTicker                          func(duration time.Duration) (<-chan time.Time, func())
-	WorkspaceUsageTracker              *workspaceusage.Tracker
+	WorkspaceUsageTrackerFlush         chan int
+	WorkspaceUsageTrackerTick          chan time.Time
 }
 
 // New constructs a codersdk client connected to an in-memory API instance.
@@ -308,35 +309,35 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	hangDetector.Start()
 	t.Cleanup(hangDetector.Close)
 
-	if options.WorkspaceUsageTracker == nil {
-		// Did last_used_at not update? Scratching your noggin? Here's why.
-		// Workspace usage tracking must be triggered manually in tests.
-		// The vast majority of existing tests do not depend on last_used_at
-		// and adding an extra background goroutine to all existing tests may
-		// lead to future flakes and goleak complaints.
-		// To do this, pass in your own WorkspaceUsageTracker like so:
-		//
-		// 	db, ps  = dbtestutil.NewDB(t)
-		//   wuTick  = make(chan time.Time)
-		//   wuFlush = make(chan int, 1)
-		//   wut     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlush), workspaceusage.WithTickChannel(wuTick))
-		//   client  = coderdtest.New(t, &coderdtest.Options{
-		//     WorkspaceUsageTracker: wut,
-		//     Database:              db,
-		//     Pubsub:                ps,
-		//   })
-		//
-		// See TestPortForward for how this works in practice.
-		wutFlush := make(chan int, 1)      // buffering just in case
-		wutTick := make(chan time.Time, 1) // buffering just in case
-		options.WorkspaceUsageTracker = workspaceusage.New(
-			options.Database,
-			workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
-			workspaceusage.WithFlushChannel(wutFlush),
-			workspaceusage.WithTickChannel(wutTick),
-		)
+	// Did last_used_at not update? Scratching your noggin? Here's why.
+	// Workspace usage tracking must be triggered manually in tests.
+	// The vast majority of existing tests do not depend on last_used_at
+	// and adding an extra time-based background goroutine to all existing
+	// tests may lead to future flakes and goleak complaints.
+	// Instead, pass in your own flush and ticker like so:
+	//
+	//   tickCh = make(chan time.Time)
+	//   flushCh = make(chan int, 1)
+	//   client  = coderdtest.New(t, &coderdtest.Options{
+	//     WorkspaceUsageTrackerFlush: flushCh,
+	//     WorkspaceUsageTrackerTick: tickCh
+	//   })
+	//
+	// Now to trigger a tick, just write to `tickCh`.
+	// Reading from `flushCh` will ensure that workspaceusage.Tracker flushed.
+	// See TestPortForward or TestTracker_MultipleInstances for how this works in practice.
+	if options.WorkspaceUsageTrackerFlush == nil {
+		options.WorkspaceUsageTrackerFlush = make(chan int, 1) // buffering just in case
+	}
+	if options.WorkspaceUsageTrackerTick == nil {
+		options.WorkspaceUsageTrackerTick = make(chan time.Time, 1) // buffering just in case
 	}
-	t.Cleanup(options.WorkspaceUsageTracker.Close)
+	// Close is called by API.Close()
+	wuTracker := workspaceusage.New(
+		options.Database,
+		workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
+		workspaceusage.WithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
+	)
 
 	var mutex sync.RWMutex
 	var handler http.Handler
@@ -486,7 +487,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
 			AllowWorkspaceRenames:              options.AllowWorkspaceRenames,
 			NewTicker:                          options.NewTicker,
-			WorkspaceUsageTracker:              options.WorkspaceUsageTracker,
+			WorkspaceUsageTracker:              wuTracker,
 		}
 }
 

diff --git a/coderd/workspaceusage/tracker.go b/coderd/workspaceusage/tracker.go
@@ -45,7 +45,7 @@ type Tracker struct {
 // New returns a new Tracker. It is the caller's responsibility
 // to call Close().
 func New(s Store, opts ...Option) *Tracker {
-	hb := &Tracker{
+	tr := &Tracker{
 		log:      slog.Make(sloghuman.Sink(os.Stderr)),
 		m:        &uuidSet{},
 		s:        s,
@@ -56,14 +56,15 @@ func New(s Store, opts ...Option) *Tracker {
 		flushCh:  nil,
 	}
 	for _, opt := range opts {
-		opt(hb)
+		opt(tr)
 	}
-	if hb.tickCh == nil && hb.stopTick == nil {
-		ticker := time.NewTicker(DefaultFlushInterval)
-		hb.tickCh = ticker.C
-		hb.stopTick = ticker.Stop
+	if tr.tickCh == nil && tr.stopTick == nil {
+		tick := time.NewTicker(DefaultFlushInterval)
+		tr.tickCh = tick.C
+		tr.stopTick = tick.Stop
 	}
-	return hb
+	go tr.loop()
+	return tr
 }
 
 type Option func(*Tracker)
@@ -84,27 +85,18 @@ func WithFlushInterval(d time.Duration) Option {
 	}
 }
 
-// WithFlushChannel allows passing a channel that receives
-// the number of marked workspaces every time Tracker flushes.
+// WithTickFlush 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 WithFlushChannel(c chan int) Option {
+func WithTickFlush(tickCh <-chan time.Time, flushCh chan int) Option {
 	if flag.Lookup("test.v") == nil {
-		panic("developer error: WithFlushChannel is not to be used outside of tests.")
+		panic("developer error: WithTickFlush is not to be used outside of tests.")
 	}
 	return func(h *Tracker) {
-		h.flushCh = c
-	}
-}
-
-// WithTickChannel allows passing a channel to replace a ticker.
-// For testing only and will panic if used outside of tests.
-func WithTickChannel(c chan time.Time) Option {
-	if flag.Lookup("test.v") == nil {
-		panic("developer error: WithTickChannel is not to be used outside of tests.")
-	}
-	return func(h *Tracker) {
-		h.tickCh = c
+		h.tickCh = tickCh
 		h.stopTick = func() {}
+		h.flushCh = flushCh
 	}
 }
 
@@ -159,13 +151,13 @@ func (tr *Tracker) flush(now time.Time) {
 	tr.log.Info(ctx, "updated workspaces last_used_at", slog.F("count", count), slog.F("now", now))
 }
 
-// Loop periodically flushes every tick.
-// If Loop is called after Close, it will panic. Don't do this.
-func (tr *Tracker) Loop() {
-	// Calling Loop after Close() is an error.
+// loop periodically flushes every tick.
+// If loop is called after Close, it will exit immediately and log an error.
+func (tr *Tracker) loop() {
 	select {
 	case <-tr.doneCh:
-		panic("developer error: Loop called after Close")
+		tr.log.Error(context.Background(), "developer error: Loop called after Close")
+		return
 	default:
 	}
 	defer func() {
@@ -193,12 +185,13 @@ 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() {
+func (tr *Tracker) Close() error {
 	tr.stopOnce.Do(func() {
 		tr.stopCh <- struct{}{}
 		tr.stopTick()
 		<-tr.doneCh
 	})
+	return nil
 }
 
 // uuidSet is a set of UUIDs. Safe for concurrent usage.

diff --git a/coderd/workspaceusage/tracker_test.go b/coderd/workspaceusage/tracker_test.go
@@ -35,11 +35,12 @@ func TestTracker(t *testing.T) {
 
 	tickCh := make(chan time.Time)
 	flushCh := make(chan int, 1)
-	wut := workspaceusage.New(mDB, workspaceusage.WithLogger(log), workspaceusage.WithTickChannel(tickCh), workspaceusage.WithFlushChannel(flushCh))
+	wut := workspaceusage.New(mDB,
+		workspaceusage.WithLogger(log),
+		workspaceusage.WithTickFlush(tickCh, flushCh),
+	)
 	defer wut.Close()
 
-	go wut.Loop()
-
 	// 1. No marked workspaces should imply no flush.
 	now := dbtime.Now()
 	tickCh <- now
@@ -106,9 +107,6 @@ func TestTracker(t *testing.T) {
 	// 5. Closing multiple times should not be a problem.
 	wut.Close()
 	wut.Close()
-
-	// 6. Running Loop() again should panic.
-	require.Panics(t, wut.Loop)
 }
 
 // This test performs a more 'integration-style' test with multiple instances.
@@ -127,25 +125,23 @@ func TestTracker_MultipleInstances(t *testing.T) {
 		ps       = pubsub.NewInMemory()
 		wuTickA  = make(chan time.Time)
 		wuFlushA = make(chan int, 1)
-		wutA     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlushA), workspaceusage.WithTickChannel(wuTickA))
 		wuTickB  = make(chan time.Time)
 		wuFlushB = make(chan int, 1)
-		wutB     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlushB), workspaceusage.WithTickChannel(wuTickB))
 		clientA  = coderdtest.New(t, &coderdtest.Options{
-			WorkspaceUsageTracker: wutA,
-			Database:              db,
-			Pubsub:                ps,
+			WorkspaceUsageTrackerTick:  wuTickA,
+			WorkspaceUsageTrackerFlush: wuFlushA,
+			Database:                   db,
+			Pubsub:                     ps,
 		})
 		clientB = coderdtest.New(t, &coderdtest.Options{
-			WorkspaceUsageTracker: wutB,
-			Database:              db,
-			Pubsub:                ps,
+			WorkspaceUsageTrackerTick:  wuTickB,
+			WorkspaceUsageTrackerFlush: wuFlushB,
+			Database:                   db,
+			Pubsub:                     ps,
 		})
 		owner = coderdtest.CreateFirstUser(t, clientA)
 		now   = dbtime.Now()
 	)
-	defer wutA.Close()
-	defer wutB.Close()
 
 	clientB.SetSessionToken(clientA.SessionToken())