From d8ba1892b671f4d4337889730de2f34ec82b6460 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Mon, 25 Nov 2024 11:38:00 +0000 Subject: [PATCH 1/3] refactor(provisionerdserver): use quartz.Clock instead of TimeNowFn --- .../provisionerdserver/provisionerdserver.go | 48 ++++++++++--------- .../provisionerdserver_test.go | 18 +++---- 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go index 8387f97ea21cb..71847b0562d0b 100644 --- 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,7 +408,7 @@ 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{ @@ -414,7 +416,7 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo 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,7 +1010,7 @@ 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 { @@ -1016,7 +1018,7 @@ func (s *server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto. } err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ ID: input.WorkspaceBuildID, - UpdatedAt: dbtime.Now(), + UpdatedAt: s.timeNow(), Deadline: build.Deadline, MaxDeadline: build.MaxDeadline, }) @@ -1382,7 +1384,7 @@ 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) @@ -1390,9 +1392,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: 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 index 84ae07ce7e5fd..526b8fdfde686 100644 --- 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.Clock 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, From 0ad16bbadc4ca027dcce767cf9b636a7b2280c72 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Mon, 25 Nov 2024 15:51:19 +0000 Subject: [PATCH 2/3] chore: pass through clock --- coderd/coderd.go | 1 + enterprise/coderd/provisionerdaemons.go | 1 + 2 files changed, 2 insertions(+) diff --git a/coderd/coderd.go b/coderd/coderd.go index b5450b1387dfe..d64727567720d 100644 --- 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/enterprise/coderd/provisionerdaemons.go b/enterprise/coderd/provisionerdaemons.go index 0eb3a51db57dd..615c225e38422 100644 --- 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, ) From 4ca5353d4a2385bdab6d150ee281471b4a536b2a Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Mon, 25 Nov 2024 16:12:35 +0000 Subject: [PATCH 3/3] chore: replace quartz.Clock with *quartz.Mock --- coderd/provisionerdserver/provisionerdserver_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go index 526b8fdfde686..325e639947f86 100644 --- a/coderd/provisionerdserver/provisionerdserver_test.go +++ b/coderd/provisionerdserver/provisionerdserver_test.go @@ -2189,7 +2189,7 @@ type overrides struct { externalAuthConfigs []*externalauth.Config templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore] userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore] - clock quartz.Clock + clock *quartz.Mock acquireJobLongPollDuration time.Duration heartbeatFn func(ctx context.Context) error heartbeatInterval time.Duration