refactor(coderd/provisionerdserver): use quartz.Clock instead of TimeNowFn (#15642)

DanielleMaywood · web-flow · commit 0896f339c403 · 2024-11-25T16:25:36.000Z
Replace `TimeNowFn` in `provisionerdserver` with `quartz.Clock` as
well as pass `coderd`'s `Clock` to `provisionerdserver`.
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -1648,6 +1648,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
 		provisionerdserver.Options{
 			OIDCConfig:          api.OIDCConfig,
 			ExternalAuthConfigs: api.ExternalAuthConfigs,
+			Clock:               api.Clock,
 		},
 		api.NotificationsEnqueuer,
 	)
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -46,6 +46,7 @@ import (
 	"github.com/coder/coder/v2/provisionerd/proto"
 	"github.com/coder/coder/v2/provisionersdk"
 	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+	"github.com/coder/quartz"
 )
 
 const (
@@ -61,8 +62,9 @@ const (
 type Options struct {
 	OIDCConfig          promoauth.OAuth2Config
 	ExternalAuthConfigs []*externalauth.Config
-	// TimeNowFn is only used in tests
-	TimeNowFn func() time.Time
+
+	// Clock for testing
+	Clock quartz.Clock
 
 	// AcquireJobLongPollDur is used in tests
 	AcquireJobLongPollDur time.Duration
@@ -104,7 +106,7 @@ type server struct {
 
 	OIDCConfig promoauth.OAuth2Config
 
-	TimeNowFn func() time.Time
+	Clock quartz.Clock
 
 	acquireJobLongPollDur time.Duration
 
@@ -191,6 +193,9 @@ func NewServer(
 	if options.HeartbeatInterval == 0 {
 		options.HeartbeatInterval = DefaultHeartbeatInterval
 	}
+	if options.Clock == nil {
+		options.Clock = quartz.NewReal()
+	}
 
 	s := &server{
 		lifecycleCtx:                lifecycleCtx,
@@ -213,7 +218,7 @@ func NewServer(
 		UserQuietHoursScheduleStore: userQuietHoursScheduleStore,
 		DeploymentValues:            deploymentValues,
 		OIDCConfig:                  options.OIDCConfig,
-		TimeNowFn:                   options.TimeNowFn,
+		Clock:                       options.Clock,
 		acquireJobLongPollDur:       options.AcquireJobLongPollDur,
 		heartbeatInterval:           options.HeartbeatInterval,
 		heartbeatFn:                 options.HeartbeatFn,
@@ -229,11 +234,8 @@ func NewServer(
 
 // timeNow should be used when trying to get the current time for math
 // calculations regarding workspace start and stop time.
-func (s *server) timeNow() time.Time {
-	if s.TimeNowFn != nil {
-		return dbtime.Time(s.TimeNowFn())
-	}
-	return dbtime.Now()
+func (s *server) timeNow(tags ...string) time.Time {
+	return dbtime.Time(s.Clock.Now(tags...))
 }
 
 // heartbeatLoop runs heartbeatOnce at the interval specified by HeartbeatInterval
@@ -365,7 +367,7 @@ func (s *server) AcquireJobWithCancel(stream proto.DRPCProvisionerDaemon_Acquire
 		logger.Error(streamCtx, "recv error and failed to cancel acquire job", slog.Error(recvErr))
 		// Well, this is awkward.  We hit an error receiving from the stream, but didn't cancel before we locked a job
 		// in the database.  We need to mark this job as failed so the end user can retry if they want to.
-		now := dbtime.Now()
+		now := s.timeNow()
 		err := s.Database.UpdateProvisionerJobWithCompleteByID(
 			//nolint:gocritic // Provisionerd has specific authz rules.
 			dbauthz.AsProvisionerd(context.Background()),
@@ -406,15 +408,15 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
 		err := s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
 			ID: job.ID,
 			CompletedAt: sql.NullTime{
-				Time:  dbtime.Now(),
+				Time:  s.timeNow(),
 				Valid: true,
 			},
 			Error: sql.NullString{
 				String: errorMessage,
 				Valid:  true,
 			},
 			ErrorCode: job.ErrorCode,
-			UpdatedAt: dbtime.Now(),
+			UpdatedAt: s.timeNow(),
 		})
 		if err != nil {
 			return xerrors.Errorf("update provisioner job: %w", err)
@@ -792,7 +794,7 @@ func (s *server) UpdateJob(ctx context.Context, request *proto.UpdateJobRequest)
 	}
 	err = s.Database.UpdateProvisionerJobByID(ctx, database.UpdateProvisionerJobByIDParams{
 		ID:        parsedID,
-		UpdatedAt: dbtime.Now(),
+		UpdatedAt: s.timeNow(),
 	})
 	if err != nil {
 		return nil, xerrors.Errorf("update job: %w", err)
@@ -869,7 +871,7 @@ func (s *server) UpdateJob(ctx context.Context, request *proto.UpdateJobRequest)
 		err := s.Database.UpdateTemplateVersionDescriptionByJobID(ctx, database.UpdateTemplateVersionDescriptionByJobIDParams{
 			JobID:     job.ID,
 			Readme:    string(request.Readme),
-			UpdatedAt: dbtime.Now(),
+			UpdatedAt: s.timeNow(),
 		})
 		if err != nil {
 			return nil, xerrors.Errorf("update template version description: %w", err)
@@ -958,7 +960,7 @@ func (s *server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto.
 		return nil, xerrors.Errorf("job already completed")
 	}
 	job.CompletedAt = sql.NullTime{
-		Time:  dbtime.Now(),
+		Time:  s.timeNow(),
 		Valid: true,
 	}
 	job.Error = sql.NullString{
@@ -973,7 +975,7 @@ func (s *server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto.
 	err = s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
 		ID:          jobID,
 		CompletedAt: job.CompletedAt,
-		UpdatedAt:   dbtime.Now(),
+		UpdatedAt:   s.timeNow(),
 		Error:       job.Error,
 		ErrorCode:   job.ErrorCode,
 	})
@@ -1008,15 +1010,15 @@ func (s *server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto.
 			if jobType.WorkspaceBuild.State != nil {
 				err = db.UpdateWorkspaceBuildProvisionerStateByID(ctx, database.UpdateWorkspaceBuildProvisionerStateByIDParams{
 					ID:               input.WorkspaceBuildID,
-					UpdatedAt:        dbtime.Now(),
+					UpdatedAt:        s.timeNow(),
 					ProvisionerState: jobType.WorkspaceBuild.State,
 				})
 				if err != nil {
 					return xerrors.Errorf("update workspace build state: %w", err)
 				}
 				err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
 					ID:          input.WorkspaceBuildID,
-					UpdatedAt:   dbtime.Now(),
+					UpdatedAt:   s.timeNow(),
 					Deadline:    build.Deadline,
 					MaxDeadline: build.MaxDeadline,
 				})
@@ -1382,17 +1384,17 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
 		err = s.Database.UpdateTemplateVersionExternalAuthProvidersByJobID(ctx, database.UpdateTemplateVersionExternalAuthProvidersByJobIDParams{
 			JobID:                 jobID,
 			ExternalAuthProviders: json.RawMessage(externalAuthProvidersMessage),
-			UpdatedAt:             dbtime.Now(),
+			UpdatedAt:             s.timeNow(),
 		})
 		if err != nil {
 			return nil, xerrors.Errorf("update template version external auth providers: %w", err)
 		}
 
 		err = s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
 			ID:        jobID,
-			UpdatedAt: dbtime.Now(),
+			UpdatedAt: s.timeNow(),
 			CompletedAt: sql.NullTime{
-				Time:  dbtime.Now(),
+				Time:  s.timeNow(),
 				Valid: true,
 			},
 			Error:     completedError,
@@ -1687,9 +1689,9 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
 
 		err = s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
 			ID:        jobID,
-			UpdatedAt: dbtime.Now(),
+			UpdatedAt: s.timeNow(),
 			CompletedAt: sql.NullTime{
-				Time:  dbtime.Now(),
+				Time:  s.timeNow(),
 				Valid: true,
 			},
 			Error:     sql.NullString{},
diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -22,6 +22,7 @@ import (
 	"storj.io/drpc"
 
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/quartz"
 	"github.com/coder/serpent"
 
 	"github.com/coder/coder/v2/buildinfo"
@@ -1211,14 +1212,13 @@ func TestCompleteJob(t *testing.T) {
 
 				// Simulate the given time starting from now.
 				require.False(t, c.now.IsZero())
-				start := time.Now()
+				clock := quartz.NewMock(t)
+				clock.Set(c.now)
 				tss := &atomic.Pointer[schedule.TemplateScheduleStore]{}
 				uqhss := &atomic.Pointer[schedule.UserQuietHoursScheduleStore]{}
 				auditor := audit.NewMock()
 				srv, db, ps, pd := setup(t, false, &overrides{
-					timeNowFn: func() time.Time {
-						return c.now.Add(time.Since(start))
-					},
+					clock:                       clock,
 					templateScheduleStore:       tss,
 					userQuietHoursScheduleStore: uqhss,
 					auditor:                     auditor,
@@ -2189,7 +2189,7 @@ type overrides struct {
 	externalAuthConfigs         []*externalauth.Config
 	templateScheduleStore       *atomic.Pointer[schedule.TemplateScheduleStore]
 	userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore]
-	timeNowFn                   func() time.Time
+	clock                       *quartz.Mock
 	acquireJobLongPollDuration  time.Duration
 	heartbeatFn                 func(ctx context.Context) error
 	heartbeatInterval           time.Duration
@@ -2209,7 +2209,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 	var externalAuthConfigs []*externalauth.Config
 	tss := testTemplateScheduleStore()
 	uqhss := testUserQuietHoursScheduleStore()
-	var timeNowFn func() time.Time
+	clock := quartz.NewReal()
 	pollDur := time.Duration(0)
 	if ov == nil {
 		ov = &overrides{}
@@ -2246,8 +2246,8 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 			require.True(t, swapped)
 		}
 	}
-	if ov.timeNowFn != nil {
-		timeNowFn = ov.timeNowFn
+	if ov.clock != nil {
+		clock = ov.clock
 	}
 	auditPtr := &atomic.Pointer[audit.Auditor]{}
 	var auditor audit.Auditor = audit.NewMock()
@@ -2296,7 +2296,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 		deploymentValues,
 		provisionerdserver.Options{
 			ExternalAuthConfigs:   externalAuthConfigs,
-			TimeNowFn:             timeNowFn,
+			Clock:                 clock,
 			OIDCConfig:            &oauth2.Config{},
 			AcquireJobLongPollDur: pollDur,
 			HeartbeatInterval:     ov.heartbeatInterval,
diff --git a/enterprise/coderd/provisionerdaemons.go b/enterprise/coderd/provisionerdaemons.go
@@ -403,6 +403,7 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request)
 		provisionerdserver.Options{
 			ExternalAuthConfigs: api.ExternalAuthConfigs,
 			OIDCConfig:          api.OIDCConfig,
+			Clock:               api.Clock,
 		},
 		api.NotificationsEnqueuer,
 	)

Original file line number	Diff line number	Diff line change
`@@ -1648,6 +1648,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n`
`1648`	`1648`	`provisionerdserver.Options{`
`1649`	`1649`	`OIDCConfig: api.OIDCConfig,`
`1650`	`1650`	`ExternalAuthConfigs: api.ExternalAuthConfigs,`
	`1651`	`+ Clock: api.Clock,`
`1651`	`1652`	`},`
`1652`	`1653`	`api.NotificationsEnqueuer,`
`1653`	`1654`	`)`
Original file line number	Diff line number	Diff line change
`@@ -403,6 +403,7 @@ func (api API) provisionerDaemonServe(rw http.ResponseWriter, r http.Request)`
`403`	`403`	`provisionerdserver.Options{`
`404`	`404`	`ExternalAuthConfigs: api.ExternalAuthConfigs,`
`405`	`405`	`OIDCConfig: api.OIDCConfig,`
	`406`	`+ Clock: api.Clock,`
`406`	`407`	`},`
`407`	`408`	`api.NotificationsEnqueuer,`
`408`	`409`	`)`