chore(coderd/database/dbpurge): replace usage of time.* with quartz (#14480)

Related to #10576

This PR introduces quartz to coderd/database/dbpurge and updates the following unit tests to make use of Quartz's functionality:

- TestPurge
- TestDeleteOldWorkspaceAgentLogs

Additionally, updates DeleteOldWorkspaceAgentLogs to replace the hard-coded interval with a parameter passed into the query. This aids in testing and brings us a step towards allowing operators to configure the cutoff interval for workspace agent logs.

cli/server.go

@@ -985,7 +985,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 			defer shutdownConns()
 
 			// Ensures that old database entries are cleaned up over time!
-			purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database)
+			purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database, quartz.NewReal())
 			defer purger.Close()
 
 			// Updates workspace usage

coderd/database/dbauthz/dbauthz.go

@@ -1144,11 +1144,11 @@ func (q *querier) DeleteOldProvisionerDaemons(ctx context.Context) error {
 	return q.db.DeleteOldProvisionerDaemons(ctx)
 }
 
-func (q *querier) DeleteOldWorkspaceAgentLogs(ctx context.Context) error {
+func (q *querier) DeleteOldWorkspaceAgentLogs(ctx context.Context, threshold time.Time) error {
 	if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceSystem); err != nil {
 		return err
 	}
-	return q.db.DeleteOldWorkspaceAgentLogs(ctx)
+	return q.db.DeleteOldWorkspaceAgentLogs(ctx, threshold)
 }
 
 func (q *querier) DeleteOldWorkspaceAgentStats(ctx context.Context) error {

coderd/database/dbauthz/dbauthz_test.go

@@ -2517,7 +2517,7 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 		}).Asserts(rbac.ResourceSystem, policy.ActionCreate)
 	}))
 	s.Run("DeleteOldWorkspaceAgentLogs", s.Subtest(func(db database.Store, check *expects) {
-		check.Args().Asserts(rbac.ResourceSystem, policy.ActionDelete)
+		check.Args(time.Time{}).Asserts(rbac.ResourceSystem, policy.ActionDelete)
 	}))
 	s.Run("InsertWorkspaceAgentStats", s.Subtest(func(db database.Store, check *expects) {
 		check.Args(database.InsertWorkspaceAgentStatsParams{}).Asserts(rbac.ResourceSystem, policy.ActionCreate).Errors(errMatchAny)

coderd/database/dbmem/dbmem.go

@@ -1706,19 +1706,15 @@ func (q *FakeQuerier) DeleteOldProvisionerDaemons(_ context.Context) error {
 	return nil
 }
 
-func (q *FakeQuerier) DeleteOldWorkspaceAgentLogs(_ context.Context) error {
+func (q *FakeQuerier) DeleteOldWorkspaceAgentLogs(_ context.Context, threshold time.Time) error {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 
-	now := dbtime.Now()
-	weekInterval := 7 * 24 * time.Hour
-	weekAgo := now.Add(-weekInterval)
-
 	var validLogs []database.WorkspaceAgentLog
 	for _, log := range q.workspaceAgentLogs {
 		var toBeDeleted bool
 		for _, agent := range q.workspaceAgents {
-			if agent.ID == log.AgentID && agent.LastConnectedAt.Valid && agent.LastConnectedAt.Time.Before(weekAgo) {
+			if agent.ID == log.AgentID && agent.LastConnectedAt.Valid && agent.LastConnectedAt.Time.Before(threshold) {
 				toBeDeleted = true
 				break
 			}

coderd/database/dbmetrics/dbmetrics.go

@@ -305,9 +305,9 @@ func (m metricsStore) DeleteOldProvisionerDaemons(ctx context.Context) error {
 	return r0
 }
 
-func (m metricsStore) DeleteOldWorkspaceAgentLogs(ctx context.Context) error {
+func (m metricsStore) DeleteOldWorkspaceAgentLogs(ctx context.Context, arg time.Time) error {
 	start := time.Now()
-	r0 := m.s.DeleteOldWorkspaceAgentLogs(ctx)
+	r0 := m.s.DeleteOldWorkspaceAgentLogs(ctx, arg)
 	m.queryLatencies.WithLabelValues("DeleteOldWorkspaceAgentLogs").Observe(time.Since(start).Seconds())
 	return r0
 }

coderd/database/dbmock/dbmock.go

@@ -501,17 +501,17 @@ func (mr *MockStoreMockRecorder) DeleteOldProvisionerDaemons(arg0 any) *gomock.C
 }
 
 // DeleteOldWorkspaceAgentLogs mocks base method.
-func (m *MockStore) DeleteOldWorkspaceAgentLogs(arg0 context.Context) error {
+func (m *MockStore) DeleteOldWorkspaceAgentLogs(arg0 context.Context, arg1 time.Time) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "DeleteOldWorkspaceAgentLogs", arg0)
+	ret := m.ctrl.Call(m, "DeleteOldWorkspaceAgentLogs", arg0, arg1)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // DeleteOldWorkspaceAgentLogs indicates an expected call of DeleteOldWorkspaceAgentLogs.
-func (mr *MockStoreMockRecorder) DeleteOldWorkspaceAgentLogs(arg0 any) *gomock.Call {
+func (mr *MockStoreMockRecorder) DeleteOldWorkspaceAgentLogs(arg0, arg1 any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldWorkspaceAgentLogs", reflect.TypeOf((*MockStore)(nil).DeleteOldWorkspaceAgentLogs), arg0)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldWorkspaceAgentLogs", reflect.TypeOf((*MockStore)(nil).DeleteOldWorkspaceAgentLogs), arg0, arg1)
 }
 
 // DeleteOldWorkspaceAgentStats mocks base method.

coderd/database/dbpurge/dbpurge.go

@@ -11,30 +11,28 @@ import (
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/quartz"
 )
 
 const (
-	delay = 10 * time.Minute
+	delay          = 10 * time.Minute
+	maxAgentLogAge = 7 * 24 * time.Hour
 )
 
 // New creates a new periodically purging database instance.
 // It is the caller's responsibility to call Close on the returned instance.
 //
 // This is for cleaning up old, unused resources from the database that take up space.
-func New(ctx context.Context, logger slog.Logger, db database.Store) io.Closer {
+func New(ctx context.Context, logger slog.Logger, db database.Store, clk quartz.Clock) io.Closer {
 	closed := make(chan struct{})
 
 	ctx, cancelFunc := context.WithCancel(ctx)
 	//nolint:gocritic // The system purges old db records without user input.
 	ctx = dbauthz.AsSystemRestricted(ctx)
 
-	// Use time.Nanosecond to force an initial tick. It will be reset to the
-	// correct duration after executing once.
-	ticker := time.NewTicker(time.Nanosecond)
-	doTick := func() {
+	ticker := clk.NewTicker(time.Nanosecond)
+	doTick := func(start time.Time) {
 		defer ticker.Reset(delay)
-
-		start := time.Now()
 		// Start a transaction to grab advisory lock, we don't want to run
 		// multiple purges at the same time (multiple replicas).
 		if err := db.InTx(func(tx database.Store) error {
@@ -49,7 +47,7 @@ func New(ctx context.Context, logger slog.Logger, db database.Store) io.Closer {
 				return nil
 			}
 
-			if err := tx.DeleteOldWorkspaceAgentLogs(ctx); err != nil {
+			if err := tx.DeleteOldWorkspaceAgentLogs(ctx, start.Add(-maxAgentLogAge)); err != nil {
 				return xerrors.Errorf("failed to delete old workspace agent logs: %w", err)
 			}
 			if err := tx.DeleteOldWorkspaceAgentStats(ctx); err != nil {
@@ -62,7 +60,7 @@ func New(ctx context.Context, logger slog.Logger, db database.Store) io.Closer {
 				return xerrors.Errorf("failed to delete old notification messages: %w", err)
 			}
 
-			logger.Info(ctx, "purged old database entries", slog.F("duration", time.Since(start)))
+			logger.Info(ctx, "purged old database entries", slog.F("duration", clk.Since(start)))
 
 			return nil
 		}, nil); err != nil {
@@ -78,9 +76,9 @@ func New(ctx context.Context, logger slog.Logger, db database.Store) io.Closer {
 			select {
 			case <-ctx.Done():
 				return
-			case <-ticker.C:
+			case tick := <-ticker.C:
 				ticker.Stop()
-				doTick()
+				doTick(tick)
 			}
 		}
 	}()

coderd/database/dbpurge/dbpurge_test.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/google/uuid"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 	"golang.org/x/exp/slices"
@@ -29,6 +28,7 @@ import (
 	"github.com/coder/coder/v2/provisionerd/proto"
 	"github.com/coder/coder/v2/provisionersdk"
 	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/quartz"
 )
 
 func TestMain(m *testing.M) {
@@ -36,19 +36,40 @@ func TestMain(m *testing.M) {
 }
 
 // Ensures no goroutines leak.
+//
+//nolint:paralleltest // It uses LockIDDBPurge.
 func TestPurge(t *testing.T) {
-	t.Parallel()
-	purger := dbpurge.New(context.Background(), slogtest.Make(t, nil), dbmem.New())
-	err := purger.Close()
-	require.NoError(t, err)
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	clk := quartz.NewMock(t)
+
+	// We want to make sure dbpurge is actually started so that this test is meaningful.
+	trapStop := clk.Trap().TickerStop()
+
+	purger := dbpurge.New(context.Background(), slogtest.Make(t, nil), dbmem.New(), clk)
+
+	// Wait for the initial nanosecond tick.
+	clk.Advance(time.Nanosecond).MustWait(ctx)
+	// Wait for ticker.Stop call that happens in the goroutine.
+	trapStop.MustWait(ctx).Release()
+	// Stop the trap now to avoid blocking further.
+	trapStop.Close()
+
+	require.NoError(t, purger.Close())
 }
 
 //nolint:paralleltest // It uses LockIDDBPurge.
 func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
-	db, _ := dbtestutil.NewDB(t)
-	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
 
 	now := dbtime.Now()
+	// TODO: must refactor DeleteOldWorkspaceAgentStats to allow passing in cutoff
+	//       before using quarts.NewMock()
+	clk := quartz.NewReal()
+	db, _ := dbtestutil.NewDB(t)
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
 
 	defer func() {
 		if t.Failed() {
@@ -78,9 +99,6 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
 		}
 	}()
 
-	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
-	defer cancel()
-
 	// given
 	// Note: We use increments of 2 hours to ensure we avoid any DST
 	// conflicts, verifying DST behavior is beyond the scope of this
@@ -114,7 +132,7 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
 	})
 
 	// when
-	closer := dbpurge.New(ctx, logger, db)
+	closer := dbpurge.New(ctx, logger, db, clk)
 	defer closer.Close()
 
 	// then
@@ -139,7 +157,7 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
 
 	// Start a new purger to immediately trigger delete after rollup.
 	_ = closer.Close()
-	closer = dbpurge.New(ctx, logger, db)
+	closer = dbpurge.New(ctx, logger, db, clk)
 	defer closer.Close()
 
 	// then
@@ -177,57 +195,75 @@ func TestDeleteOldWorkspaceAgentLogs(t *testing.T) {
 	t.Run("AgentHasNotConnectedSinceWeek_LogsExpired", func(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
 		defer cancel()
+		clk := quartz.NewMock(t)
+		clk.Set(now).MustWait(ctx)
 
-		// given
+		// After dbpurge completes, the ticker is reset. Trap this call.
+		trapReset := clk.Trap().TickerReset()
+		defer trapReset.Close()
+
+		// given: an agent with logs older than threshold
 		agent := mustCreateAgentWithLogs(ctx, t, db, user, org, tmpl, tv, now.Add(-8*24*time.Hour), t.Name())
 
-		// Make sure that agent logs have been collected.
+		// when dbpurge runs
+		closer := dbpurge.New(ctx, logger, db, clk)
+		defer closer.Close()
+		// Wait for the initial nanosecond tick.
+		clk.Advance(time.Nanosecond).MustWait(ctx)
+
+		trapReset.MustWait(ctx).Release() // Wait for ticker.Reset()
+		d, w := clk.AdvanceNext()
+		require.Equal(t, 10*time.Minute, d)
+
+		closer.Close() // doTick() has now run.
+		w.MustWait(ctx)
+
+		// then the logs should be gone
 		agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
-			AgentID: agent,
+			AgentID:      agent,
+			CreatedAfter: 0,
 		})
 		require.NoError(t, err)
-		require.NotZero(t, agentLogs, "agent logs must be present")
-
-		// when
-		closer := dbpurge.New(ctx, logger, db)
-		defer closer.Close()
-
-		// then
-		assert.Eventually(t, func() bool {
-			agentLogs, err = db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
-				AgentID: agent,
-			})
-			if err != nil {
-				return false
-			}
-			return !containsAgentLog(agentLogs, t.Name())
-		}, testutil.WaitShort, testutil.IntervalFast)
-		require.NoError(t, err)
-		require.NotContains(t, agentLogs, t.Name())
+		require.Empty(t, agentLogs, "expected agent logs to be empty")
 	})
 
 	//nolint:paralleltest // It uses LockIDDBPurge.
 	t.Run("AgentConnectedSixDaysAgo_LogsValid", func(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
 		defer cancel()
+		clk := quartz.NewMock(t)
+		clk.Set(now).MustWait(ctx)
 
-		// given
+		// After dbpurge completes, the ticker is reset. Trap this call.
+		trapReset := clk.Trap().TickerReset()
+		defer trapReset.Close()
+
+		// given: an agent with logs newer than threshold
 		agent := mustCreateAgentWithLogs(ctx, t, db, user, org, tmpl, tv, now.Add(-6*24*time.Hour), t.Name())
 
-		// when
-		closer := dbpurge.New(ctx, logger, db)
+		// when dbpurge runs
+		closer := dbpurge.New(ctx, logger, db, clk)
 		defer closer.Close()
 
-		// then
-		require.Eventually(t, func() bool {
-			agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
-				AgentID: agent,
-			})
-			if err != nil {
-				return false
-			}
-			return containsAgentLog(agentLogs, t.Name())
-		}, testutil.WaitShort, testutil.IntervalFast)
+		// Wait for the initial nanosecond tick.
+		clk.Advance(time.Nanosecond).MustWait(ctx)
+
+		trapReset.MustWait(ctx).Release() // Wait for ticker.Reset()
+		d, w := clk.AdvanceNext()
+		require.Equal(t, 10*time.Minute, d)
+
+		closer.Close() // doTick() has now run.
+		w.MustWait(ctx)
+
+		// then the logs should still be there
+		agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
+			AgentID: agent,
+		})
+		require.NoError(t, err)
+		require.NotEmpty(t, agentLogs)
+		for _, al := range agentLogs {
+			require.Equal(t, t.Name(), al.Output)
+		}
 	})
 }
 
@@ -246,6 +282,12 @@ func mustCreateAgentWithLogs(ctx context.Context, t *testing.T, db database.Stor
 		Level:     []database.LogLevel{database.LogLevelDebug},
 	})
 	require.NoError(t, err)
+	// Make sure that agent logs have been collected.
+	agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
+		AgentID: agent.ID,
+	})
+	require.NoError(t, err)
+	require.NotZero(t, agentLogs, "agent logs must be present")
 	return agent.ID
 }
 
@@ -268,19 +310,17 @@ func mustCreateAgent(t *testing.T, db database.Store, user database.User, org da
 		JobID:      job.ID,
 		Transition: database.WorkspaceTransitionStart,
 	})
+
 	return dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
 		ResourceID: resource.ID,
 	})
 }
 
-func containsAgentLog(daemons []database.WorkspaceAgentLog, output string) bool {
-	return slices.ContainsFunc(daemons, func(d database.WorkspaceAgentLog) bool {
-		return d.Output == output
-	})
-}
-
 //nolint:paralleltest // It uses LockIDDBPurge.
 func TestDeleteOldProvisionerDaemons(t *testing.T) {
+	// TODO: must refactor DeleteOldProvisionerDaemons to allow passing in cutoff
+	//       before using quartz.NewMock
+	clk := quartz.NewReal()
 	db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
 	defaultOrg := dbgen.Organization(t, db, database.Organization{})
 	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
@@ -347,7 +387,7 @@ func TestDeleteOldProvisionerDaemons(t *testing.T) {
 	require.NoError(t, err)
 
 	// when
-	closer := dbpurge.New(ctx, logger, db)
+	closer := dbpurge.New(ctx, logger, db, clk)
 	defer closer.Close()
 
 	// then

coderd/database/querier.go

@@ -89,7 +89,7 @@ type sqlcQuerier interface {
 	DeleteOldProvisionerDaemons(ctx context.Context) error
 	// If an agent hasn't connected in the last 7 days, we purge it's logs.
 	// Logs can take up a lot of space, so it's important we clean up frequently.
-	DeleteOldWorkspaceAgentLogs(ctx context.Context) error
+	DeleteOldWorkspaceAgentLogs(ctx context.Context, threshold time.Time) error
 	DeleteOldWorkspaceAgentStats(ctx context.Context) error
 	DeleteOrganization(ctx context.Context, id uuid.UUID) error
 	DeleteOrganizationMember(ctx context.Context, arg DeleteOrganizationMemberParams) error

coderd/database/queries.sql.go

@@ -10484,13 +10484,13 @@ func (q *sqlQuerier) UpsertWorkspaceAgentPortShare(ctx context.Context, arg Upse
 const deleteOldWorkspaceAgentLogs = `-- name: DeleteOldWorkspaceAgentLogs :exec
 DELETE FROM workspace_agent_logs WHERE agent_id IN
 	(SELECT id FROM workspace_agents WHERE last_connected_at IS NOT NULL
-		AND last_connected_at < NOW() - INTERVAL '7 day')
+		AND last_connected_at < $1 :: timestamptz)
 `
 
 // If an agent hasn't connected in the last 7 days, we purge it's logs.
 // Logs can take up a lot of space, so it's important we clean up frequently.
-func (q *sqlQuerier) DeleteOldWorkspaceAgentLogs(ctx context.Context) error {
-	_, err := q.db.ExecContext(ctx, deleteOldWorkspaceAgentLogs)
+func (q *sqlQuerier) DeleteOldWorkspaceAgentLogs(ctx context.Context, threshold time.Time) error {
+	_, err := q.db.ExecContext(ctx, deleteOldWorkspaceAgentLogs, threshold)
 	return err
 }
 

coderd/database/queries/workspaceagents.sql

@@ -192,7 +192,7 @@ SELECT * FROM workspace_agent_log_sources WHERE workspace_agent_id = ANY(@ids ::
 -- name: DeleteOldWorkspaceAgentLogs :exec
 DELETE FROM workspace_agent_logs WHERE agent_id IN
 	(SELECT id FROM workspace_agents WHERE last_connected_at IS NOT NULL
-		AND last_connected_at < NOW() - INTERVAL '7 day');
+		AND last_connected_at < @threshold :: timestamptz);
 
 -- name: GetWorkspaceAgentsInLatestBuildByWorkspaceID :many
 SELECT