Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 0896f33

Browse files
DanielleMaywoodDanielleMaywood
authored
refactor(coderd/provisionerdserver): use quartz.Clock instead of TimeNowFn (#15642)
Replace `TimeNowFn` in `provisionerdserver` with `quartz.Clock` as well as pass `coderd`'s `Clock` to `provisionerdserver`.
1 parent bbc549d commit 0896f33

File tree

4 files changed

+36
-32
lines changed

4 files changed

+36
-32
lines changed

‎coderd/coderd.go

+1
Original file line numberDiff line numberDiff line change
@@ -1648,6 +1648,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
16481648
provisionerdserver.Options{
16491649
OIDCConfig: api.OIDCConfig,
16501650
ExternalAuthConfigs: api.ExternalAuthConfigs,
1651+
Clock: api.Clock,
16511652
},
16521653
api.NotificationsEnqueuer,
16531654
)

‎coderd/provisionerdserver/provisionerdserver.go

+25-23
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ import (
4646
"github.com/coder/coder/v2/provisionerd/proto"
4747
"github.com/coder/coder/v2/provisionersdk"
4848
sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
49+
"github.com/coder/quartz"
4950
)
5051

5152
const (
@@ -61,8 +62,9 @@ const (
6162
type Options struct {
6263
OIDCConfig promoauth.OAuth2Config
6364
ExternalAuthConfigs []*externalauth.Config
64-
// TimeNowFn is only used in tests
65-
TimeNowFn func() time.Time
65+
66+
// Clock for testing
67+
Clock quartz.Clock
6668

6769
// AcquireJobLongPollDur is used in tests
6870
AcquireJobLongPollDur time.Duration
@@ -104,7 +106,7 @@ type server struct {
104106

105107
OIDCConfig promoauth.OAuth2Config
106108

107-
TimeNowFn func() time.Time
109+
Clock quartz.Clock
108110

109111
acquireJobLongPollDur time.Duration
110112

@@ -191,6 +193,9 @@ func NewServer(
191193
if options.HeartbeatInterval == 0 {
192194
options.HeartbeatInterval = DefaultHeartbeatInterval
193195
}
196+
if options.Clock == nil {
197+
options.Clock = quartz.NewReal()
198+
}
194199

195200
s := &server{
196201
lifecycleCtx: lifecycleCtx,
@@ -213,7 +218,7 @@ func NewServer(
213218
UserQuietHoursScheduleStore: userQuietHoursScheduleStore,
214219
DeploymentValues: deploymentValues,
215220
OIDCConfig: options.OIDCConfig,
216-
TimeNowFn: options.TimeNowFn,
221+
Clock: options.Clock,
217222
acquireJobLongPollDur: options.AcquireJobLongPollDur,
218223
heartbeatInterval: options.HeartbeatInterval,
219224
heartbeatFn: options.HeartbeatFn,
@@ -229,11 +234,8 @@ func NewServer(
229234

230235
// timeNow should be used when trying to get the current time for math
231236
// calculations regarding workspace start and stop time.
232-
func (s *server) timeNow() time.Time {
233-
if s.TimeNowFn != nil {
234-
return dbtime.Time(s.TimeNowFn())
235-
}
236-
return dbtime.Now()
237+
func (s *server) timeNow(tags ...string) time.Time {
238+
return dbtime.Time(s.Clock.Now(tags...))
237239
}
238240

239241
// heartbeatLoop runs heartbeatOnce at the interval specified by HeartbeatInterval
@@ -365,7 +367,7 @@ func (s *server) AcquireJobWithCancel(stream proto.DRPCProvisionerDaemon_Acquire
365367
logger.Error(streamCtx, "recv error and failed to cancel acquire job", slog.Error(recvErr))
366368
// Well, this is awkward. We hit an error receiving from the stream, but didn't cancel before we locked a job
367369
// in the database. We need to mark this job as failed so the end user can retry if they want to.
368-
now := dbtime.Now()
370+
now := s.timeNow()
369371
err := s.Database.UpdateProvisionerJobWithCompleteByID(
370372
//nolint:gocritic // Provisionerd has specific authz rules.
371373
dbauthz.AsProvisionerd(context.Background()),
@@ -406,15 +408,15 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
406408
err := s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
407409
ID: job.ID,
408410
CompletedAt: sql.NullTime{
409-
Time: dbtime.Now(),
411+
Time: s.timeNow(),
410412
Valid: true,
411413
},
412414
Error: sql.NullString{
413415
String: errorMessage,
414416
Valid: true,
415417
},
416418
ErrorCode: job.ErrorCode,
417-
UpdatedAt: dbtime.Now(),
419+
UpdatedAt: s.timeNow(),
418420
})
419421
if err != nil {
420422
return xerrors.Errorf("update provisioner job: %w", err)
@@ -792,7 +794,7 @@ func (s *server) UpdateJob(ctx context.Context, request *proto.UpdateJobRequest)
792794
}
793795
err = s.Database.UpdateProvisionerJobByID(ctx, database.UpdateProvisionerJobByIDParams{
794796
ID: parsedID,
795-
UpdatedAt: dbtime.Now(),
797+
UpdatedAt: s.timeNow(),
796798
})
797799
if err != nil {
798800
return nil, xerrors.Errorf("update job: %w", err)
@@ -869,7 +871,7 @@ func (s *server) UpdateJob(ctx context.Context, request *proto.UpdateJobRequest)
869871
err := s.Database.UpdateTemplateVersionDescriptionByJobID(ctx, database.UpdateTemplateVersionDescriptionByJobIDParams{
870872
JobID: job.ID,
871873
Readme: string(request.Readme),
872-
UpdatedAt: dbtime.Now(),
874+
UpdatedAt: s.timeNow(),
873875
})
874876
if err != nil {
875877
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.
958960
return nil, xerrors.Errorf("job already completed")
959961
}
960962
job.CompletedAt = sql.NullTime{
961-
Time: dbtime.Now(),
963+
Time: s.timeNow(),
962964
Valid: true,
963965
}
964966
job.Error = sql.NullString{
@@ -973,7 +975,7 @@ func (s *server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto.
973975
err = s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
974976
ID: jobID,
975977
CompletedAt: job.CompletedAt,
976-
UpdatedAt: dbtime.Now(),
978+
UpdatedAt: s.timeNow(),
977979
Error: job.Error,
978980
ErrorCode: job.ErrorCode,
979981
})
@@ -1008,15 +1010,15 @@ func (s *server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto.
10081010
if jobType.WorkspaceBuild.State != nil {
10091011
err = db.UpdateWorkspaceBuildProvisionerStateByID(ctx, database.UpdateWorkspaceBuildProvisionerStateByIDParams{
10101012
ID: input.WorkspaceBuildID,
1011-
UpdatedAt: dbtime.Now(),
1013+
UpdatedAt: s.timeNow(),
10121014
ProvisionerState: jobType.WorkspaceBuild.State,
10131015
})
10141016
if err != nil {
10151017
return xerrors.Errorf("update workspace build state: %w", err)
10161018
}
10171019
err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
10181020
ID: input.WorkspaceBuildID,
1019-
UpdatedAt: dbtime.Now(),
1021+
UpdatedAt: s.timeNow(),
10201022
Deadline: build.Deadline,
10211023
MaxDeadline: build.MaxDeadline,
10221024
})
@@ -1382,17 +1384,17 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
13821384
err = s.Database.UpdateTemplateVersionExternalAuthProvidersByJobID(ctx, database.UpdateTemplateVersionExternalAuthProvidersByJobIDParams{
13831385
JobID: jobID,
13841386
ExternalAuthProviders: json.RawMessage(externalAuthProvidersMessage),
1385-
UpdatedAt: dbtime.Now(),
1387+
UpdatedAt: s.timeNow(),
13861388
})
13871389
if err != nil {
13881390
return nil, xerrors.Errorf("update template version external auth providers: %w", err)
13891391
}
13901392

13911393
err = s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
13921394
ID: jobID,
1393-
UpdatedAt: dbtime.Now(),
1395+
UpdatedAt: s.timeNow(),
13941396
CompletedAt: sql.NullTime{
1395-
Time: dbtime.Now(),
1397+
Time: s.timeNow(),
13961398
Valid: true,
13971399
},
13981400
Error: completedError,
@@ -1687,9 +1689,9 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
16871689

16881690
err = s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
16891691
ID: jobID,
1690-
UpdatedAt: dbtime.Now(),
1692+
UpdatedAt: s.timeNow(),
16911693
CompletedAt: sql.NullTime{
1692-
Time: dbtime.Now(),
1694+
Time: s.timeNow(),
16931695
Valid: true,
16941696
},
16951697
Error: sql.NullString{},

‎coderd/provisionerdserver/provisionerdserver_test.go

+9-9
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"storj.io/drpc"
2323

2424
"cdr.dev/slog/sloggers/slogtest"
25+
"github.com/coder/quartz"
2526
"github.com/coder/serpent"
2627

2728
"github.com/coder/coder/v2/buildinfo"
@@ -1211,14 +1212,13 @@ func TestCompleteJob(t *testing.T) {
12111212

12121213
// Simulate the given time starting from now.
12131214
require.False(t, c.now.IsZero())
1214-
start := time.Now()
1215+
clock := quartz.NewMock(t)
1216+
clock.Set(c.now)
12151217
tss := &atomic.Pointer[schedule.TemplateScheduleStore]{}
12161218
uqhss := &atomic.Pointer[schedule.UserQuietHoursScheduleStore]{}
12171219
auditor := audit.NewMock()
12181220
srv, db, ps, pd := setup(t, false, &overrides{
1219-
timeNowFn: func() time.Time {
1220-
return c.now.Add(time.Since(start))
1221-
},
1221+
clock: clock,
12221222
templateScheduleStore: tss,
12231223
userQuietHoursScheduleStore: uqhss,
12241224
auditor: auditor,
@@ -2189,7 +2189,7 @@ type overrides struct {
21892189
externalAuthConfigs []*externalauth.Config
21902190
templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
21912191
userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore]
2192-
timeNowFn func() time.Time
2192+
clock *quartz.Mock
21932193
acquireJobLongPollDuration time.Duration
21942194
heartbeatFn func(ctx context.Context) error
21952195
heartbeatInterval time.Duration
@@ -2209,7 +2209,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
22092209
var externalAuthConfigs []*externalauth.Config
22102210
tss := testTemplateScheduleStore()
22112211
uqhss := testUserQuietHoursScheduleStore()
2212-
var timeNowFn func() time.Time
2212+
clock := quartz.NewReal()
22132213
pollDur := time.Duration(0)
22142214
if ov == nil {
22152215
ov = &overrides{}
@@ -2246,8 +2246,8 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
22462246
require.True(t, swapped)
22472247
}
22482248
}
2249-
if ov.timeNowFn != nil {
2250-
timeNowFn = ov.timeNowFn
2249+
if ov.clock != nil {
2250+
clock = ov.clock
22512251
}
22522252
auditPtr := &atomic.Pointer[audit.Auditor]{}
22532253
var auditor audit.Auditor = audit.NewMock()
@@ -2296,7 +2296,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
22962296
deploymentValues,
22972297
provisionerdserver.Options{
22982298
ExternalAuthConfigs: externalAuthConfigs,
2299-
TimeNowFn: timeNowFn,
2299+
Clock: clock,
23002300
OIDCConfig: &oauth2.Config{},
23012301
AcquireJobLongPollDur: pollDur,
23022302
HeartbeatInterval: ov.heartbeatInterval,

‎enterprise/coderd/provisionerdaemons.go

+1
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,7 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request)
403403
provisionerdserver.Options{
404404
ExternalAuthConfigs: api.ExternalAuthConfigs,
405405
OIDCConfig: api.OIDCConfig,
406+
Clock: api.Clock,
406407
},
407408
api.NotificationsEnqueuer,
408409
)

0 commit comments

Comments
 (0)