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

Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions coderd/coderd.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
48 changes: 25 additions & 23 deletions coderd/provisionerdserver/provisionerdserver.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand Down Expand Up @@ -104,7 +106,7 @@ type server struct {

OIDCConfig promoauth.OAuth2Config

TimeNowFn func() time.Time
Clock quartz.Clock

acquireJobLongPollDur time.Duration

Expand Down Expand Up @@ -191,6 +193,9 @@ func NewServer(
if options.HeartbeatInterval == 0 {
options.HeartbeatInterval = DefaultHeartbeatInterval
}
if options.Clock == nil {
options.Clock = quartz.NewReal()
}
Comment on lines +196 to +198
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍 Even if api.Clock is nil for some reason this should keep us covered.


s := &server{
lifecycleCtx: lifecycleCtx,
Expand All @@ -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,
Expand All @@ -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...))
Comment on lines +237 to +238
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice 👍

}

// heartbeatLoop runs heartbeatOnce at the interval specified by HeartbeatInterval
Expand Down Expand Up @@ -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()),
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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{
Expand All @@ -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,
})
Expand Down Expand Up @@ -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,
})
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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{},
Expand Down
18 changes: 9 additions & 9 deletions coderd/provisionerdserver/provisionerdserver_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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{}
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions enterprise/coderd/provisionerdaemons.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
Loading