From cd35f879fe7023ee0006f3244b1d9541d714cfc4 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Tue, 8 Aug 2023 09:04:56 +0000 Subject: [PATCH 1/5] feat: update workspace deadline when template policy changes --- coderd/database/dbauthz/dbauthz.go | 4 + coderd/database/dbfake/dbfake.go | 28 +++++++ coderd/database/dbmetrics/dbmetrics.go | 7 ++ coderd/database/dbmock/dbmock.go | 15 ++++ coderd/database/querier.go | 1 + coderd/database/queries.sql.go | 66 ++++++++++++++++ coderd/database/queries/workspacebuilds.sql | 24 ++++++ enterprise/coderd/coderd.go | 2 +- enterprise/coderd/schedule/template.go | 84 ++++++++++++++++++++- enterprise/coderd/schedule/template_test.go | 20 +++++ 10 files changed, 246 insertions(+), 5 deletions(-) create mode 100644 enterprise/coderd/schedule/template_test.go diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 488842dcaf351..675a89dd4be4d 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -784,6 +784,10 @@ func (q *querier) GetActiveUserCount(ctx context.Context) (int64, error) { return q.db.GetActiveUserCount(ctx) } +func (q *querier) GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]database.WorkspaceBuild, error) { + panic("not implemented") +} + func (q *querier) GetAllTailnetAgents(ctx context.Context) ([]database.TailnetAgent, error) { if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceTailnetCoordinator); err != nil { return []database.TailnetAgent{}, err diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go index cc0d05aaee9c3..aa33e9b8748f7 100644 --- a/coderd/database/dbfake/dbfake.go +++ b/coderd/database/dbfake/dbfake.go @@ -918,6 +918,34 @@ func (q *FakeQuerier) GetActiveUserCount(_ context.Context) (int64, error) { return active, nil } +func (q *FakeQuerier) GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]database.WorkspaceBuild, error) { + workspaceIDs := func() []uuid.UUID { + q.mutex.RLock() + defer q.mutex.RUnlock() + + ids := []uuid.UUID{} + for _, workspace := range q.workspaces { + if workspace.TemplateID == templateID { + ids = append(ids, workspace.ID) + } + } + return ids + }() + + builds, err := q.GetLatestWorkspaceBuildsByWorkspaceIDs(ctx, workspaceIDs) + if err != nil { + return nil, err + } + + filteredBuilds := []database.WorkspaceBuild{} + for _, build := range builds { + if build.Transition == database.WorkspaceTransitionStart { + filteredBuilds = append(filteredBuilds, build) + } + } + return filteredBuilds, nil +} + func (*FakeQuerier) GetAllTailnetAgents(_ context.Context) ([]database.TailnetAgent, error) { return nil, ErrUnimplemented } diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go index 85ee8c26a0d51..85a5f29f54071 100644 --- a/coderd/database/dbmetrics/dbmetrics.go +++ b/coderd/database/dbmetrics/dbmetrics.go @@ -237,6 +237,13 @@ func (m metricsStore) GetActiveUserCount(ctx context.Context) (int64, error) { return count, err } +func (m metricsStore) GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]database.WorkspaceBuild, error) { + start := time.Now() + r0, r1 := m.s.GetActiveWorkspaceBuildsByTemplateID(ctx, templateID) + m.queryLatencies.WithLabelValues("GetActiveWorkspaceBuildsByTemplateID").Observe(time.Since(start).Seconds()) + return r0, r1 +} + func (m metricsStore) GetAllTailnetAgents(ctx context.Context) ([]database.TailnetAgent, error) { start := time.Now() r0, r1 := m.s.GetAllTailnetAgents(ctx) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index cb7278369884b..921a6f50e0fdf 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -371,6 +371,21 @@ func (mr *MockStoreMockRecorder) GetActiveUserCount(arg0 interface{}) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetActiveUserCount", reflect.TypeOf((*MockStore)(nil).GetActiveUserCount), arg0) } +// GetActiveWorkspaceBuildsByTemplateID mocks base method. +func (m *MockStore) GetActiveWorkspaceBuildsByTemplateID(arg0 context.Context, arg1 uuid.UUID) ([]database.WorkspaceBuild, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetActiveWorkspaceBuildsByTemplateID", arg0, arg1) + ret0, _ := ret[0].([]database.WorkspaceBuild) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetActiveWorkspaceBuildsByTemplateID indicates an expected call of GetActiveWorkspaceBuildsByTemplateID. +func (mr *MockStoreMockRecorder) GetActiveWorkspaceBuildsByTemplateID(arg0, arg1 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetActiveWorkspaceBuildsByTemplateID", reflect.TypeOf((*MockStore)(nil).GetActiveWorkspaceBuildsByTemplateID), arg0, arg1) +} + // GetAllTailnetAgents mocks base method. func (m *MockStore) GetAllTailnetAgents(arg0 context.Context) ([]database.TailnetAgent, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index b308589ffc350..0dc123ea5e557 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -48,6 +48,7 @@ type sqlcQuerier interface { GetAPIKeysByUserID(ctx context.Context, arg GetAPIKeysByUserIDParams) ([]APIKey, error) GetAPIKeysLastUsedAfter(ctx context.Context, lastUsed time.Time) ([]APIKey, error) GetActiveUserCount(ctx context.Context) (int64, error) + GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceBuild, error) GetAllTailnetAgents(ctx context.Context) ([]TailnetAgent, error) GetAllTailnetClients(ctx context.Context) ([]TailnetClient, error) GetAppSecurityKey(ctx context.Context) (string, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 3e07fc22f1aa3..d97080f81d735 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -7815,6 +7815,72 @@ func (q *sqlQuerier) InsertWorkspaceBuildParameters(ctx context.Context, arg Ins return err } +const getActiveWorkspaceBuildsByTemplateID = `-- name: GetActiveWorkspaceBuildsByTemplateID :many +SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason, wb.daily_cost, wb.max_deadline, wb.initiator_by_avatar_url, wb.initiator_by_username +FROM ( + SELECT + workspace_id, MAX(build_number) as max_build_number + FROM + workspace_build_with_user AS workspace_builds + WHERE + workspace_id IN ( + SELECT + id + FROM + workspaces + WHERE + template_id = $1 + ) + GROUP BY + workspace_id +) m +JOIN + workspace_build_with_user AS wb + ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number +WHERE + wb.transition = 'start'::workspace_transition +` + +func (q *sqlQuerier) GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceBuild, error) { + rows, err := q.db.QueryContext(ctx, getActiveWorkspaceBuildsByTemplateID, templateID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []WorkspaceBuild + for rows.Next() { + var i WorkspaceBuild + if err := rows.Scan( + &i.ID, + &i.CreatedAt, + &i.UpdatedAt, + &i.WorkspaceID, + &i.TemplateVersionID, + &i.BuildNumber, + &i.Transition, + &i.InitiatorID, + &i.ProvisionerState, + &i.JobID, + &i.Deadline, + &i.Reason, + &i.DailyCost, + &i.MaxDeadline, + &i.InitiatorByAvatarUrl, + &i.InitiatorByUsername, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getLatestWorkspaceBuildByWorkspaceID = `-- name: GetLatestWorkspaceBuildByWorkspaceID :one SELECT id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, initiator_by_avatar_url, initiator_by_username diff --git a/coderd/database/queries/workspacebuilds.sql b/coderd/database/queries/workspacebuilds.sql index ea2ccdb8d08ce..1020b729c4f27 100644 --- a/coderd/database/queries/workspacebuilds.sql +++ b/coderd/database/queries/workspacebuilds.sql @@ -144,3 +144,27 @@ SET WHERE id = $1; +-- name: GetActiveWorkspaceBuildsByTemplateID :many +SELECT wb.* +FROM ( + SELECT + workspace_id, MAX(build_number) as max_build_number + FROM + workspace_build_with_user AS workspace_builds + WHERE + workspace_id IN ( + SELECT + id + FROM + workspaces + WHERE + template_id = $1 + ) + GROUP BY + workspace_id +) m +JOIN + workspace_build_with_user AS wb + ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number +WHERE + wb.transition = 'start'::workspace_transition; diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 71d975e3ef1d6..d5439057aed30 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -507,7 +507,7 @@ func (api *API) updateEntitlements(ctx context.Context) error { if initial, changed, enabled := featureChanged(codersdk.FeatureAdvancedTemplateScheduling); shouldUpdate(initial, changed, enabled) { if enabled { - templateStore := schedule.NewEnterpriseTemplateScheduleStore() + templateStore := schedule.NewEnterpriseTemplateScheduleStore(api.AGPL.UserQuietHoursScheduleStore) templateStoreInterface := agplschedule.TemplateScheduleStore(templateStore) api.AGPL.TemplateScheduleStore.Store(&templateStoreInterface) } else { diff --git a/enterprise/coderd/schedule/template.go b/enterprise/coderd/schedule/template.go index 278e315dda3af..c03ce45964699 100644 --- a/enterprise/coderd/schedule/template.go +++ b/enterprise/coderd/schedule/template.go @@ -9,7 +9,9 @@ import ( "golang.org/x/xerrors" "github.com/coder/coder/coderd/database" + "github.com/coder/coder/coderd/database/db2sdk" agpl "github.com/coder/coder/coderd/schedule" + "github.com/coder/coder/codersdk" ) // EnterpriseTemplateScheduleStore provides an agpl.TemplateScheduleStore that @@ -20,12 +22,18 @@ type EnterpriseTemplateScheduleStore struct { // workspace build. This value is determined by a feature flag, licensing, // and whether a default user quiet hours schedule is set. UseRestartRequirement atomic.Bool + + // UserQuietHoursScheduleStore is used when recalculating build deadlines on + // update. + UserQuietHoursScheduleStore *atomic.Pointer[agpl.UserQuietHoursScheduleStore] } var _ agpl.TemplateScheduleStore = &EnterpriseTemplateScheduleStore{} -func NewEnterpriseTemplateScheduleStore() *EnterpriseTemplateScheduleStore { - return &EnterpriseTemplateScheduleStore{} +func NewEnterpriseTemplateScheduleStore(userQuietHoursStore *atomic.Pointer[agpl.UserQuietHoursScheduleStore]) *EnterpriseTemplateScheduleStore { + return &EnterpriseTemplateScheduleStore{ + UserQuietHoursScheduleStore: userQuietHoursStore, + } } // Get implements agpl.TemplateScheduleStore. @@ -65,7 +73,7 @@ func (s *EnterpriseTemplateScheduleStore) Get(ctx context.Context, db database.S } // Set implements agpl.TemplateScheduleStore. -func (*EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.Store, tpl database.Template, opts agpl.TemplateScheduleOptions) (database.Template, error) { +func (s *EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.Store, tpl database.Template, opts agpl.TemplateScheduleOptions) (database.Template, error) { if int64(opts.DefaultTTL) == tpl.DefaultTTL && int64(opts.MaxTTL) == tpl.MaxTTL && int16(opts.RestartRequirement.DaysOfWeek) == tpl.RestartRequirementDaysOfWeek && @@ -115,7 +123,75 @@ func (*EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.Sto return xerrors.Errorf("update deleting_at of all workspaces for new locked_ttl %q: %w", opts.LockedTTL, err) } - // TODO: update all workspace max_deadlines to be within new bounds + // Recalculate max_deadline and deadline for all running workspace + // builds on this template. + builds, err := db.GetActiveWorkspaceBuildsByTemplateID(ctx, template.ID) + if err != nil { + return xerrors.Errorf("get active workspace builds: %w", err) + } + for _, build := range builds { + if !build.MaxDeadline.IsZero() && build.MaxDeadline.Before(time.Now().Add(time.Hour)) { + // Skip this since it's already too close to the max_deadline. + continue + } + + workspace, err := db.GetWorkspaceByID(ctx, build.WorkspaceID) + if err != nil { + return xerrors.Errorf("get workspace %q: %w", build.WorkspaceID, err) + } + + job, err := db.GetProvisionerJobByID(ctx, build.JobID) + if err != nil { + return xerrors.Errorf("get provisioner job %q: %w", build.JobID, err) + } + if db2sdk.ProvisionerJobStatus(job) != codersdk.ProvisionerJobSucceeded { + continue + } + + // If the job completed before the autostop epoch, then it must be + // skipped to avoid failures below. + if job.CompletedAt.Time.Before(agpl.TemplateRestartRequirementEpoch(time.UTC).Add(time.Hour * 7 * 24)) { + continue + } + + autostop, err := agpl.CalculateAutostop(ctx, agpl.CalculateAutostopParams{ + Database: db, + TemplateScheduleStore: s, + UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(), + // Use the job completion time as the time we calculate autostop + // from. + Now: job.CompletedAt.Time, + Workspace: workspace, + }) + if err != nil { + return xerrors.Errorf("calculate new autostop for workspace %q: %w", workspace.ID, err) + } + + // If max deadline is before now, then set it to two hours from now. + now := database.Now() + if autostop.MaxDeadline.Before(now) { + autostop.MaxDeadline = now.Add(time.Hour * 2) + } + + // If the current deadline on the build is after the new + // max_deadline, then set it to the max_deadline. + autostop.Deadline = build.Deadline + if build.Deadline.After(build.MaxDeadline) { + autostop.Deadline = build.MaxDeadline + } + + // Update the workspace build. + err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{ + ID: build.ID, + UpdatedAt: database.Now(), + Deadline: autostop.Deadline, + MaxDeadline: autostop.MaxDeadline, + }) + if err != nil { + return xerrors.Errorf("update workspace build %q: %w", build.ID, err) + } + } + template, err = db.GetTemplateByID(ctx, tpl.ID) if err != nil { return xerrors.Errorf("get updated template schedule: %w", err) diff --git a/enterprise/coderd/schedule/template_test.go b/enterprise/coderd/schedule/template_test.go new file mode 100644 index 0000000000000..0fa186252dda1 --- /dev/null +++ b/enterprise/coderd/schedule/template_test.go @@ -0,0 +1,20 @@ +package schedule + +import ( + "github.com/coder/coder/coderd/database" + "github.com/coder/coder/coderd/database/dbgen" + "github.com/coder/coder/coderd/database/dbtestutil" +) + +func TestTemplateUpdateBuildDeadlines(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + + var ( + org = dbgen.Organization(t, db, database.Organization{}) + template = dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + + ) +} \ No newline at end of file From 359d4a5072025c35f3db42834652b152c71f8013 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Tue, 8 Aug 2023 12:10:45 +0000 Subject: [PATCH 2/5] tests --- coderd/database/dbauthz/dbauthz.go | 6 +- enterprise/coderd/schedule/template.go | 156 +++--- enterprise/coderd/schedule/template_test.go | 511 +++++++++++++++++++- enterprise/coderd/workspaces_test.go | 31 +- 4 files changed, 625 insertions(+), 79 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 675a89dd4be4d..be78ed250edd2 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -785,7 +785,11 @@ func (q *querier) GetActiveUserCount(ctx context.Context) (int64, error) { } func (q *querier) GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]database.WorkspaceBuild, error) { - panic("not implemented") + // This is a system-only function. + if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil { + return []database.WorkspaceBuild{}, err + } + return q.db.GetActiveWorkspaceBuildsByTemplateID(ctx, templateID) } func (q *querier) GetAllTailnetAgents(ctx context.Context) ([]database.TailnetAgent, error) { diff --git a/enterprise/coderd/schedule/template.go b/enterprise/coderd/schedule/template.go index c03ce45964699..9d315ad38669e 100644 --- a/enterprise/coderd/schedule/template.go +++ b/enterprise/coderd/schedule/template.go @@ -10,6 +10,7 @@ import ( "github.com/coder/coder/coderd/database" "github.com/coder/coder/coderd/database/db2sdk" + "github.com/coder/coder/coderd/database/dbauthz" agpl "github.com/coder/coder/coderd/schedule" "github.com/coder/coder/codersdk" ) @@ -26,6 +27,9 @@ type EnterpriseTemplateScheduleStore struct { // UserQuietHoursScheduleStore is used when recalculating build deadlines on // update. UserQuietHoursScheduleStore *atomic.Pointer[agpl.UserQuietHoursScheduleStore] + + // Custom time.Now() function to use in tests. Defaults to database.Now(). + TimeNowFn func() time.Time } var _ agpl.TemplateScheduleStore = &EnterpriseTemplateScheduleStore{} @@ -36,6 +40,13 @@ func NewEnterpriseTemplateScheduleStore(userQuietHoursStore *atomic.Pointer[agpl } } +func (s *EnterpriseTemplateScheduleStore) now() time.Time { + if s.TimeNowFn != nil { + return s.TimeNowFn() + } + return database.Now() +} + // Get implements agpl.TemplateScheduleStore. func (s *EnterpriseTemplateScheduleStore) Get(ctx context.Context, db database.Store, templateID uuid.UUID) (agpl.TemplateScheduleOptions, error) { tpl, err := db.GetTemplateByID(ctx, templateID) @@ -96,7 +107,7 @@ func (s *EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.S err = db.InTx(func(db database.Store) error { err := db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{ ID: tpl.ID, - UpdatedAt: database.Now(), + UpdatedAt: s.now(), AllowUserAutostart: opts.UserAutostartEnabled, AllowUserAutostop: opts.UserAutostopEnabled, DefaultTTL: int64(opts.DefaultTTL), @@ -123,85 +134,102 @@ func (s *EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.S return xerrors.Errorf("update deleting_at of all workspaces for new locked_ttl %q: %w", opts.LockedTTL, err) } - // Recalculate max_deadline and deadline for all running workspace - // builds on this template. - builds, err := db.GetActiveWorkspaceBuildsByTemplateID(ctx, template.ID) + template, err = db.GetTemplateByID(ctx, tpl.ID) if err != nil { - return xerrors.Errorf("get active workspace builds: %w", err) + return xerrors.Errorf("get updated template schedule: %w", err) } - for _, build := range builds { - if !build.MaxDeadline.IsZero() && build.MaxDeadline.Before(time.Now().Add(time.Hour)) { - // Skip this since it's already too close to the max_deadline. - continue - } - workspace, err := db.GetWorkspaceByID(ctx, build.WorkspaceID) + // Recalculate max_deadline and deadline for all running workspace + // builds on this template. + if s.UseRestartRequirement.Load() { + err = s.updateWorkspaceBuilds(ctx, db, template) if err != nil { - return xerrors.Errorf("get workspace %q: %w", build.WorkspaceID, err) + return xerrors.Errorf("update workspace builds: %w", err) } + } - job, err := db.GetProvisionerJobByID(ctx, build.JobID) - if err != nil { - return xerrors.Errorf("get provisioner job %q: %w", build.JobID, err) - } - if db2sdk.ProvisionerJobStatus(job) != codersdk.ProvisionerJobSucceeded { - continue - } + return nil + }, nil) + if err != nil { + return database.Template{}, err + } - // If the job completed before the autostop epoch, then it must be - // skipped to avoid failures below. - if job.CompletedAt.Time.Before(agpl.TemplateRestartRequirementEpoch(time.UTC).Add(time.Hour * 7 * 24)) { - continue - } + return template, nil +} - autostop, err := agpl.CalculateAutostop(ctx, agpl.CalculateAutostopParams{ - Database: db, - TemplateScheduleStore: s, - UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(), - // Use the job completion time as the time we calculate autostop - // from. - Now: job.CompletedAt.Time, - Workspace: workspace, - }) - if err != nil { - return xerrors.Errorf("calculate new autostop for workspace %q: %w", workspace.ID, err) - } +func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuilds(ctx context.Context, db database.Store, template database.Template) error { + //nolint:gocritic // This function will retrieve all workspace builds on + // the template and update their max deadline to be within the new + // policy parameters. + ctx = dbauthz.AsSystemRestricted(ctx) - // If max deadline is before now, then set it to two hours from now. - now := database.Now() - if autostop.MaxDeadline.Before(now) { - autostop.MaxDeadline = now.Add(time.Hour * 2) - } + builds, err := db.GetActiveWorkspaceBuildsByTemplateID(ctx, template.ID) + if err != nil { + return xerrors.Errorf("get active workspace builds: %w", err) + } - // If the current deadline on the build is after the new - // max_deadline, then set it to the max_deadline. - autostop.Deadline = build.Deadline - if build.Deadline.After(build.MaxDeadline) { - autostop.Deadline = build.MaxDeadline - } + for _, build := range builds { + if !build.MaxDeadline.IsZero() && build.MaxDeadline.Before(s.now().Add(2*time.Hour)) { + // Skip this since it's already too close to the max_deadline. + continue + } - // Update the workspace build. - err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{ - ID: build.ID, - UpdatedAt: database.Now(), - Deadline: autostop.Deadline, - MaxDeadline: autostop.MaxDeadline, - }) - if err != nil { - return xerrors.Errorf("update workspace build %q: %w", build.ID, err) - } + workspace, err := db.GetWorkspaceByID(ctx, build.WorkspaceID) + if err != nil { + return xerrors.Errorf("get workspace %q: %w", build.WorkspaceID, err) } - template, err = db.GetTemplateByID(ctx, tpl.ID) + job, err := db.GetProvisionerJobByID(ctx, build.JobID) if err != nil { - return xerrors.Errorf("get updated template schedule: %w", err) + return xerrors.Errorf("get provisioner job %q: %w", build.JobID, err) + } + if db2sdk.ProvisionerJobStatus(job) != codersdk.ProvisionerJobSucceeded { + continue } - return nil - }, nil) - if err != nil { - return database.Template{}, err + // If the job completed before the autostop epoch, then it must be + // skipped to avoid failures below. Add a week to account for timezones. + if job.CompletedAt.Time.Before(agpl.TemplateRestartRequirementEpoch(time.UTC).Add(time.Hour * 7 * 24)) { + continue + } + + autostop, err := agpl.CalculateAutostop(ctx, agpl.CalculateAutostopParams{ + Database: db, + TemplateScheduleStore: s, + UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(), + // Use the job completion time as the time we calculate autostop + // from. + Now: job.CompletedAt.Time, + Workspace: workspace, + }) + if err != nil { + return xerrors.Errorf("calculate new autostop for workspace %q: %w", workspace.ID, err) + } + + // If max deadline is before now()+2h, then set it to that. + now := s.now() + if autostop.MaxDeadline.Before(now.Add(2 * time.Hour)) { + autostop.MaxDeadline = now.Add(time.Hour * 2) + } + + // If the current deadline on the build is after the new + // max_deadline, then set it to the max_deadline. + autostop.Deadline = build.Deadline + if autostop.Deadline.After(autostop.MaxDeadline) { + autostop.Deadline = autostop.MaxDeadline + } + + // Update the workspace build. + err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{ + ID: build.ID, + UpdatedAt: now, + Deadline: autostop.Deadline, + MaxDeadline: autostop.MaxDeadline, + }) + if err != nil { + return xerrors.Errorf("update workspace build %q: %w", build.ID, err) + } } - return template, nil + return nil } diff --git a/enterprise/coderd/schedule/template_test.go b/enterprise/coderd/schedule/template_test.go index 0fa186252dda1..1defaa0664895 100644 --- a/enterprise/coderd/schedule/template_test.go +++ b/enterprise/coderd/schedule/template_test.go @@ -1,9 +1,23 @@ -package schedule +package schedule_test import ( + "database/sql" + "encoding/json" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/coder/coder/coderd/database" "github.com/coder/coder/coderd/database/dbgen" "github.com/coder/coder/coderd/database/dbtestutil" + agplschedule "github.com/coder/coder/coderd/schedule" + "github.com/coder/coder/enterprise/coderd/schedule" + "github.com/coder/coder/testutil" ) func TestTemplateUpdateBuildDeadlines(t *testing.T) { @@ -12,9 +26,498 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { db, _ := dbtestutil.NewDB(t) var ( - org = dbgen.Organization(t, db, database.Organization{}) - template = dbgen.Template(t, db, database.Template{ + org = dbgen.Organization(t, db, database.Organization{}) + user = dbgen.User(t, db, database.User{}) + file = dbgen.File(t, db, database.File{ + CreatedBy: user.ID, + }) + templateJob = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{ OrganizationID: org.ID, + FileID: file.ID, + InitiatorID: user.ID, + Tags: database.StringMap{ + "foo": "bar", + }, + }) + templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + JobID: templateJob.ID, + }) + ) + + const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC + ctx := testutil.Context(t, testutil.WaitLong) + user, err := db.UpdateUserQuietHoursSchedule(ctx, database.UpdateUserQuietHoursScheduleParams{ + ID: user.ID, + QuietHoursSchedule: userQuietHoursSchedule, + }) + require.NoError(t, err) + + realNow := time.Now().UTC() + nowY, nowM, nowD := realNow.Date() + buildTime := time.Date(nowY, nowM, nowD, 12, 0, 0, 0, time.UTC) // noon today UTC + nextQuietHours := time.Date(nowY, nowM, nowD+1, 0, 0, 0, 0, time.UTC) // midnight tomorrow UTC + + // Workspace old max_deadline too soon + cases := []struct { + name string + now time.Time + deadline time.Time + maxDeadline time.Time + newDeadline time.Time // 0 for no change + newMaxDeadline time.Time + }{ + { + name: "SkippedWorkspaceMaxDeadlineTooSoon", + now: buildTime, + deadline: buildTime, + maxDeadline: buildTime.Add(1 * time.Hour), + // Unchanged since the max deadline is too soon. + newDeadline: time.Time{}, + newMaxDeadline: buildTime.Add(1 * time.Hour), + }, + { + name: "NewWorkspaceMaxDeadlineBeforeNow", + // After the new max deadline... + now: nextQuietHours.Add(6 * time.Hour), + deadline: buildTime, + // Far into the future... + maxDeadline: nextQuietHours.Add(24 * time.Hour), + newDeadline: time.Time{}, + // We will use now() + 2 hours if the newly calculated max deadline + // from the workspace build time is before now. + newMaxDeadline: nextQuietHours.Add(8 * time.Hour), + }, + { + name: "NewWorkspaceMaxDeadlineSoon", + // Right before the new max deadline... + now: nextQuietHours.Add(-1 * time.Hour), + deadline: buildTime, + // Far into the future... + maxDeadline: nextQuietHours.Add(24 * time.Hour), + newDeadline: time.Time{}, + // We will use now() + 2 hours if the newly calculated max deadline + // from the workspace build time is within the next 2 hours. + newMaxDeadline: nextQuietHours.Add(1 * time.Hour), + }, + { + name: "NewWorkspaceMaxDeadlineFuture", + // Well before the new max deadline... + now: nextQuietHours.Add(-6 * time.Hour), + deadline: buildTime, + // Far into the future... + maxDeadline: nextQuietHours.Add(24 * time.Hour), + newDeadline: time.Time{}, + newMaxDeadline: nextQuietHours, + }, + { + name: "DeadlineAfterNewWorkspaceMaxDeadline", + // Well before the new max deadline... + now: nextQuietHours.Add(-6 * time.Hour), + // Far into the future... + deadline: nextQuietHours.Add(24 * time.Hour), + maxDeadline: nextQuietHours.Add(24 * time.Hour), + // The deadline should match since it is after the new max deadline. + newDeadline: nextQuietHours, + newMaxDeadline: nextQuietHours, + }, + } + + for _, c := range cases { + c := c + + t.Run(c.name, func(t *testing.T) { + t.Parallel() + + t.Log("buildTime", buildTime) + t.Log("nextQuietHours", nextQuietHours) + t.Log("now", c.now) + t.Log("deadline", c.deadline) + t.Log("maxDeadline", c.maxDeadline) + t.Log("newDeadline", c.newDeadline) + t.Log("newMaxDeadline", c.newMaxDeadline) + + var ( + template = dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + ActiveVersionID: templateVersion.ID, + CreatedBy: user.ID, + }) + ws = dbgen.Workspace(t, db, database.Workspace{ + OrganizationID: org.ID, + OwnerID: user.ID, + TemplateID: template.ID, + }) + job = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{ + OrganizationID: org.ID, + FileID: file.ID, + InitiatorID: user.ID, + Provisioner: database.ProvisionerTypeEcho, + Tags: database.StringMap{ + c.name: "yeah", + }, + }) + wsBuild = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + BuildNumber: 1, + JobID: job.ID, + InitiatorID: user.ID, + TemplateVersionID: templateVersion.ID, + }) + ) + acquiredJob, err := db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ + StartedAt: sql.NullTime{ + Time: buildTime, + Valid: true, + }, + WorkerID: uuid.NullUUID{ + UUID: uuid.New(), + Valid: true, + }, + Types: []database.ProvisionerType{database.ProvisionerTypeEcho}, + Tags: json.RawMessage(fmt.Sprintf(`{%q: "yeah"}`, c.name)), + }) + require.NoError(t, err) + require.Equal(t, job.ID, acquiredJob.ID) + err = db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{ + ID: job.ID, + CompletedAt: sql.NullTime{ + Time: buildTime, + Valid: true, + }, + UpdatedAt: buildTime, + }) + require.NoError(t, err) + + err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{ + ID: wsBuild.ID, + UpdatedAt: buildTime, + ProvisionerState: []byte{}, + Deadline: c.deadline, + MaxDeadline: c.maxDeadline, + }) + require.NoError(t, err) + + wsBuild, err = db.GetWorkspaceBuildByID(ctx, wsBuild.ID) + require.NoError(t, err) + + userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule) + require.NoError(t, err) + userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{} + userQuietHoursStorePtr.Store(&userQuietHoursStore) + + // Set the template policy. + templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr) + templateScheduleStore.UseRestartRequirement.Store(true) + templateScheduleStore.TimeNowFn = func() time.Time { + return c.now + } + _, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{ + UserAutostartEnabled: false, + UserAutostopEnabled: false, + DefaultTTL: 0, + MaxTTL: 0, + UseRestartRequirement: true, + RestartRequirement: agplschedule.TemplateRestartRequirement{ + // Every day + DaysOfWeek: 0b01111111, + Weeks: 0, + }, + FailureTTL: 0, + InactivityTTL: 0, + LockedTTL: 0, + }) + require.NoError(t, err) + + // Check that the workspace build has the expected deadlines. + newBuild, err := db.GetWorkspaceBuildByID(ctx, wsBuild.ID) + require.NoError(t, err) + + if c.newDeadline.IsZero() { + c.newDeadline = wsBuild.Deadline + } + require.WithinDuration(t, c.newDeadline, newBuild.Deadline, time.Second) + require.WithinDuration(t, c.newMaxDeadline, newBuild.MaxDeadline, time.Second) + }) + } +} + +func TestTemplateUpdateBuildDeadlinesSkip(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + + var ( + org = dbgen.Organization(t, db, database.Organization{}) + user = dbgen.User(t, db, database.User{}) + file = dbgen.File(t, db, database.File{ + CreatedBy: user.ID, + }) + templateJob = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{ + OrganizationID: org.ID, + FileID: file.ID, + InitiatorID: user.ID, + Tags: database.StringMap{ + "foo": "bar", + }, + }) + templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + JobID: templateJob.ID, + }) + template = dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + ActiveVersionID: templateVersion.ID, + CreatedBy: user.ID, + }) + otherTemplate = dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + ActiveVersionID: templateVersion.ID, + CreatedBy: user.ID, + }) ) -} \ No newline at end of file + + // Create a workspace that will be shared by two builds. + ws := dbgen.Workspace(t, db, database.Workspace{ + OrganizationID: org.ID, + OwnerID: user.ID, + TemplateID: template.ID, + }) + + const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC + ctx := testutil.Context(t, testutil.WaitLong) + user, err := db.UpdateUserQuietHoursSchedule(ctx, database.UpdateUserQuietHoursScheduleParams{ + ID: user.ID, + QuietHoursSchedule: userQuietHoursSchedule, + }) + require.NoError(t, err) + + realNow := time.Now().UTC() + nowY, nowM, nowD := realNow.Date() + buildTime := time.Date(nowY, nowM, nowD, 12, 0, 0, 0, time.UTC) // noon today UTC + now := time.Date(nowY, nowM, nowD, 18, 0, 0, 0, time.UTC) // 6pm today UTC + nextQuietHours := time.Date(nowY, nowM, nowD+1, 0, 0, 0, 0, time.UTC) // midnight tomorrow UTC + + // A date very far in the future which would definitely be updated. + originalMaxDeadline := time.Date(nowY+1, nowM, nowD, 0, 0, 0, 0, time.UTC) + + _ = otherTemplate + + builds := []struct { + name string + templateID uuid.UUID + // Nil workspaceID means create a new workspace. + workspaceID uuid.UUID + buildNumber int32 + buildStarted bool + buildCompleted bool + buildError bool + + shouldBeUpdated bool + + // Set below: + wsBuild database.WorkspaceBuild + }{ + { + name: "DifferentTemplate", + templateID: otherTemplate.ID, + workspaceID: uuid.Nil, + buildNumber: 1, + buildStarted: true, + buildCompleted: true, + buildError: false, + shouldBeUpdated: false, + }, + { + name: "NonStartedBuild", + templateID: template.ID, + workspaceID: uuid.Nil, + buildNumber: 1, + buildStarted: false, + buildCompleted: false, + buildError: false, + shouldBeUpdated: false, + }, + { + name: "InProgressBuild", + templateID: template.ID, + workspaceID: uuid.Nil, + buildNumber: 1, + buildStarted: true, + buildCompleted: false, + buildError: false, + shouldBeUpdated: false, + }, + { + name: "FailedBuild", + templateID: template.ID, + workspaceID: uuid.Nil, + buildNumber: 1, + buildStarted: true, + buildCompleted: true, + buildError: true, + shouldBeUpdated: false, + }, + { + name: "NonLatestBuild", + templateID: template.ID, + workspaceID: ws.ID, + buildNumber: 1, + buildStarted: true, + buildCompleted: true, + buildError: false, + // This build was successful but is not the latest build for this + // workspace, see the next build. + shouldBeUpdated: false, + }, + { + name: "LatestBuild", + templateID: template.ID, + workspaceID: ws.ID, + buildNumber: 2, + buildStarted: true, + buildCompleted: true, + buildError: false, + shouldBeUpdated: true, + }, + { + name: "LatestBuildOtherWorkspace", + templateID: template.ID, + workspaceID: uuid.Nil, + buildNumber: 1, + buildStarted: true, + buildCompleted: true, + buildError: false, + shouldBeUpdated: true, + }, + } + + for i, b := range builds { + wsID := b.workspaceID + if wsID == uuid.Nil { + ws := dbgen.Workspace(t, db, database.Workspace{ + OrganizationID: org.ID, + OwnerID: user.ID, + TemplateID: b.templateID, + }) + wsID = ws.ID + } + job := dbgen.ProvisionerJob(t, db, database.ProvisionerJob{ + OrganizationID: org.ID, + FileID: file.ID, + InitiatorID: user.ID, + Provisioner: database.ProvisionerTypeEcho, + Tags: database.StringMap{ + wsID.String(): "yeah", + }, + }) + wsBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: wsID, + BuildNumber: b.buildNumber, + JobID: job.ID, + InitiatorID: user.ID, + TemplateVersionID: templateVersion.ID, + }) + + err := db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{ + ID: wsBuild.ID, + UpdatedAt: buildTime, + ProvisionerState: []byte{}, + Deadline: originalMaxDeadline, + MaxDeadline: originalMaxDeadline, + }) + require.NoError(t, err) + + wsBuild, err = db.GetWorkspaceBuildByID(ctx, wsBuild.ID) + require.NoError(t, err) + + builds[i].wsBuild = wsBuild + + if !b.buildStarted { + continue + } + + acquiredJob, err := db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{ + StartedAt: sql.NullTime{ + Time: buildTime, + Valid: true, + }, + WorkerID: uuid.NullUUID{ + UUID: uuid.New(), + Valid: true, + }, + Types: []database.ProvisionerType{database.ProvisionerTypeEcho}, + Tags: json.RawMessage(fmt.Sprintf(`{%q: "yeah"}`, wsID)), + }) + require.NoError(t, err) + require.Equal(t, job.ID, acquiredJob.ID) + + if !b.buildCompleted { + continue + } + + buildError := "" + if b.buildError { + buildError = "error" + } + err = db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{ + ID: job.ID, + CompletedAt: sql.NullTime{ + Time: buildTime, + Valid: true, + }, + Error: sql.NullString{ + String: buildError, + Valid: b.buildError, + }, + UpdatedAt: buildTime, + }) + require.NoError(t, err) + } + + userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule) + require.NoError(t, err) + userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{} + userQuietHoursStorePtr.Store(&userQuietHoursStore) + + // Set the template policy. + templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr) + templateScheduleStore.UseRestartRequirement.Store(true) + templateScheduleStore.TimeNowFn = func() time.Time { + return now + } + _, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{ + UserAutostartEnabled: false, + UserAutostopEnabled: false, + DefaultTTL: 0, + MaxTTL: 0, + UseRestartRequirement: true, + RestartRequirement: agplschedule.TemplateRestartRequirement{ + // Every day + DaysOfWeek: 0b01111111, + Weeks: 0, + }, + FailureTTL: 0, + InactivityTTL: 0, + LockedTTL: 0, + }) + require.NoError(t, err) + + // Check each build. + for i, b := range builds { + msg := fmt.Sprintf("build %d: %s", i, b.name) + newBuild, err := db.GetWorkspaceBuildByID(ctx, b.wsBuild.ID) + require.NoError(t, err) + + if b.shouldBeUpdated { + assert.Equal(t, nextQuietHours, newBuild.Deadline, msg) + assert.Equal(t, nextQuietHours, newBuild.MaxDeadline, msg) + } else { + assert.Equal(t, originalMaxDeadline, newBuild.Deadline, msg) + assert.Equal(t, originalMaxDeadline, newBuild.MaxDeadline, msg) + } + } +} diff --git a/enterprise/coderd/workspaces_test.go b/enterprise/coderd/workspaces_test.go index 07e501032f9e0..10140bd747957 100644 --- a/enterprise/coderd/workspaces_test.go +++ b/enterprise/coderd/workspaces_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "sync/atomic" "testing" "time" @@ -24,6 +25,16 @@ import ( "github.com/coder/coder/testutil" ) +// agplUserQuietHoursScheduleStore is passed to +// NewEnterpriseTemplateScheduleStore as we don't care about updating the +// schedule and having it recalculate the build deadline in these tests. +func agplUserQuietHoursScheduleStore() *atomic.Pointer[agplschedule.UserQuietHoursScheduleStore] { + store := agplschedule.NewAGPLUserQuietHoursScheduleStore() + p := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{} + p.Store(&store) + return p +} + func TestCreateWorkspace(t *testing.T) { t.Parallel() @@ -100,7 +111,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -147,7 +158,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -193,7 +204,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -235,7 +246,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -292,7 +303,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -334,7 +345,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -376,7 +387,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -427,7 +438,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -497,7 +508,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: ticker, IncludeProvisionerDaemon: true, AutobuildStats: statCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, @@ -558,7 +569,7 @@ func TestWorkspaceAutobuild(t *testing.T) { AutobuildTicker: tickCh, IncludeProvisionerDaemon: true, AutobuildStats: statsCh, - TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(), + TemplateScheduleStore: schedule.NewEnterpriseTemplateScheduleStore(agplUserQuietHoursScheduleStore()), }, LicenseOptions: &coderdenttest.LicenseOptions{ Features: license.Features{codersdk.FeatureAdvancedTemplateScheduling: 1}, From 3ea43481cce63b64653c037724b6ff47cec49986 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Fri, 11 Aug 2023 13:22:49 +0000 Subject: [PATCH 3/5] fixup! tests --- enterprise/coderd/schedule/template_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/enterprise/coderd/schedule/template_test.go b/enterprise/coderd/schedule/template_test.go index 1defaa0664895..bc97014e0ea0e 100644 --- a/enterprise/coderd/schedule/template_test.go +++ b/enterprise/coderd/schedule/template_test.go @@ -513,11 +513,11 @@ func TestTemplateUpdateBuildDeadlinesSkip(t *testing.T) { require.NoError(t, err) if b.shouldBeUpdated { - assert.Equal(t, nextQuietHours, newBuild.Deadline, msg) - assert.Equal(t, nextQuietHours, newBuild.MaxDeadline, msg) + assert.WithinDuration(t, nextQuietHours, newBuild.Deadline, time.Second, msg) + assert.WithinDuration(t, nextQuietHours, newBuild.MaxDeadline, time.Second, msg) } else { - assert.Equal(t, originalMaxDeadline, newBuild.Deadline, msg) - assert.Equal(t, originalMaxDeadline, newBuild.MaxDeadline, msg) + assert.WithinDuration(t, originalMaxDeadline, newBuild.Deadline, time.Second, msg) + assert.WithinDuration(t, originalMaxDeadline, newBuild.MaxDeadline, time.Second, msg) } } } From 49162053b856e08f06f1e5e5232f58528f8641e1 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Mon, 14 Aug 2023 20:37:27 +0000 Subject: [PATCH 4/5] chore: add instrumentation to everything --- coderd/schedule/autostop.go | 10 ++ coderd/schedule/template.go | 7 ++ enterprise/coderd/schedule/template.go | 134 +++++++++++++++---------- enterprise/coderd/schedule/user.go | 27 +++-- 4 files changed, 120 insertions(+), 58 deletions(-) diff --git a/coderd/schedule/autostop.go b/coderd/schedule/autostop.go index 6934640045506..64016677e1204 100644 --- a/coderd/schedule/autostop.go +++ b/coderd/schedule/autostop.go @@ -4,9 +4,12 @@ import ( "context" "time" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" "golang.org/x/xerrors" "github.com/coder/coder/coderd/database" + "github.com/coder/coder/coderd/tracing" ) const ( @@ -72,6 +75,13 @@ type AutostopTime struct { // Deadline is a cost saving measure, while max deadline is a // compliance/updating measure. func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (AutostopTime, error) { + ctx, span := tracing.StartSpan(ctx, + trace.WithAttributes(attribute.String("coder.workspace_id", params.Workspace.ID.String())), + trace.WithAttributes(attribute.String("coder.template_id", params.Workspace.TemplateID.String())), + ) + defer span.End() + defer span.End() + var ( db = params.Database workspace = params.Workspace diff --git a/coderd/schedule/template.go b/coderd/schedule/template.go index de97dffc9ac2a..29809fab854ba 100644 --- a/coderd/schedule/template.go +++ b/coderd/schedule/template.go @@ -8,6 +8,7 @@ import ( "golang.org/x/xerrors" "github.com/coder/coder/coderd/database" + "github.com/coder/coder/coderd/tracing" ) const MaxTemplateRestartRequirementWeeks = 16 @@ -122,6 +123,9 @@ func NewAGPLTemplateScheduleStore() TemplateScheduleStore { } func (*agplTemplateScheduleStore) Get(ctx context.Context, db database.Store, templateID uuid.UUID) (TemplateScheduleOptions, error) { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + tpl, err := db.GetTemplateByID(ctx, templateID) if err != nil { return TemplateScheduleOptions{}, err @@ -148,6 +152,9 @@ func (*agplTemplateScheduleStore) Get(ctx context.Context, db database.Store, te } func (*agplTemplateScheduleStore) Set(ctx context.Context, db database.Store, tpl database.Template, opts TemplateScheduleOptions) (database.Template, error) { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + if int64(opts.DefaultTTL) == tpl.DefaultTTL { // Avoid updating the UpdatedAt timestamp if nothing will be changed. return tpl, nil diff --git a/enterprise/coderd/schedule/template.go b/enterprise/coderd/schedule/template.go index 9d315ad38669e..7a7481c7438ba 100644 --- a/enterprise/coderd/schedule/template.go +++ b/enterprise/coderd/schedule/template.go @@ -6,12 +6,15 @@ import ( "time" "github.com/google/uuid" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" "golang.org/x/xerrors" "github.com/coder/coder/coderd/database" "github.com/coder/coder/coderd/database/db2sdk" "github.com/coder/coder/coderd/database/dbauthz" agpl "github.com/coder/coder/coderd/schedule" + "github.com/coder/coder/coderd/tracing" "github.com/coder/coder/codersdk" ) @@ -49,6 +52,9 @@ func (s *EnterpriseTemplateScheduleStore) now() time.Time { // Get implements agpl.TemplateScheduleStore. func (s *EnterpriseTemplateScheduleStore) Get(ctx context.Context, db database.Store, templateID uuid.UUID) (agpl.TemplateScheduleOptions, error) { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + tpl, err := db.GetTemplateByID(ctx, templateID) if err != nil { return agpl.TemplateScheduleOptions{}, err @@ -85,6 +91,9 @@ func (s *EnterpriseTemplateScheduleStore) Get(ctx context.Context, db database.S // Set implements agpl.TemplateScheduleStore. func (s *EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.Store, tpl database.Template, opts agpl.TemplateScheduleOptions) (database.Template, error) { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + if int64(opts.DefaultTTL) == tpl.DefaultTTL && int64(opts.MaxTTL) == tpl.MaxTTL && int16(opts.RestartRequirement.DaysOfWeek) == tpl.RestartRequirementDaysOfWeek && @@ -105,6 +114,9 @@ func (s *EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.S var template database.Template err = db.InTx(func(db database.Store) error { + ctx, span := tracing.StartSpanWithName(ctx, "(*schedule.EnterpriseTemplateScheduleStore).Set()-InTx()") + defer span.End() + err := db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{ ID: tpl.ID, UpdatedAt: s.now(), @@ -158,6 +170,9 @@ func (s *EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.S } func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuilds(ctx context.Context, db database.Store, template database.Template) error { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + //nolint:gocritic // This function will retrieve all workspace builds on // the template and update their max deadline to be within the new // policy parameters. @@ -169,66 +184,81 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuilds(ctx context.Cont } for _, build := range builds { - if !build.MaxDeadline.IsZero() && build.MaxDeadline.Before(s.now().Add(2*time.Hour)) { - // Skip this since it's already too close to the max_deadline. - continue - } - - workspace, err := db.GetWorkspaceByID(ctx, build.WorkspaceID) + err := s.updateWorkspaceBuild(ctx, db, template, build) if err != nil { - return xerrors.Errorf("get workspace %q: %w", build.WorkspaceID, err) + return xerrors.Errorf("update workspace build %q: %w", build.ID, err) } + } - job, err := db.GetProvisionerJobByID(ctx, build.JobID) - if err != nil { - return xerrors.Errorf("get provisioner job %q: %w", build.JobID, err) - } - if db2sdk.ProvisionerJobStatus(job) != codersdk.ProvisionerJobSucceeded { - continue - } + return nil +} - // If the job completed before the autostop epoch, then it must be - // skipped to avoid failures below. Add a week to account for timezones. - if job.CompletedAt.Time.Before(agpl.TemplateRestartRequirementEpoch(time.UTC).Add(time.Hour * 7 * 24)) { - continue - } +func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Context, db database.Store, template database.Template, build database.WorkspaceBuild) error { + ctx, span := tracing.StartSpan(ctx, + trace.WithAttributes(attribute.String("coder.workspace_id", build.WorkspaceID.String())), + trace.WithAttributes(attribute.String("coder.workspace_build_id", build.ID.String())), + ) + defer span.End() - autostop, err := agpl.CalculateAutostop(ctx, agpl.CalculateAutostopParams{ - Database: db, - TemplateScheduleStore: s, - UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(), - // Use the job completion time as the time we calculate autostop - // from. - Now: job.CompletedAt.Time, - Workspace: workspace, - }) - if err != nil { - return xerrors.Errorf("calculate new autostop for workspace %q: %w", workspace.ID, err) - } + if !build.MaxDeadline.IsZero() && build.MaxDeadline.Before(s.now().Add(2*time.Hour)) { + // Skip this since it's already too close to the max_deadline. + return nil + } - // If max deadline is before now()+2h, then set it to that. - now := s.now() - if autostop.MaxDeadline.Before(now.Add(2 * time.Hour)) { - autostop.MaxDeadline = now.Add(time.Hour * 2) - } + workspace, err := db.GetWorkspaceByID(ctx, build.WorkspaceID) + if err != nil { + return xerrors.Errorf("get workspace %q: %w", build.WorkspaceID, err) + } - // If the current deadline on the build is after the new - // max_deadline, then set it to the max_deadline. - autostop.Deadline = build.Deadline - if autostop.Deadline.After(autostop.MaxDeadline) { - autostop.Deadline = autostop.MaxDeadline - } + job, err := db.GetProvisionerJobByID(ctx, build.JobID) + if err != nil { + return xerrors.Errorf("get provisioner job %q: %w", build.JobID, err) + } + if db2sdk.ProvisionerJobStatus(job) != codersdk.ProvisionerJobSucceeded { + // Only touch builds that are completed. + return nil + } - // Update the workspace build. - err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{ - ID: build.ID, - UpdatedAt: now, - Deadline: autostop.Deadline, - MaxDeadline: autostop.MaxDeadline, - }) - if err != nil { - return xerrors.Errorf("update workspace build %q: %w", build.ID, err) - } + // If the job completed before the autostop epoch, then it must be skipped + // to avoid failures below. Add a week to account for timezones. + if job.CompletedAt.Time.Before(agpl.TemplateRestartRequirementEpoch(time.UTC).Add(time.Hour * 7 * 24)) { + return nil + } + + autostop, err := agpl.CalculateAutostop(ctx, agpl.CalculateAutostopParams{ + Database: db, + TemplateScheduleStore: s, + UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(), + // Use the job completion time as the time we calculate autostop from. + Now: job.CompletedAt.Time, + Workspace: workspace, + }) + if err != nil { + return xerrors.Errorf("calculate new autostop for workspace %q: %w", workspace.ID, err) + } + + // If max deadline is before now()+2h, then set it to that. + now := s.now() + if autostop.MaxDeadline.Before(now.Add(2 * time.Hour)) { + autostop.MaxDeadline = now.Add(time.Hour * 2) + } + + // If the current deadline on the build is after the new max_deadline, then + // set it to the max_deadline. + autostop.Deadline = build.Deadline + if autostop.Deadline.After(autostop.MaxDeadline) { + autostop.Deadline = autostop.MaxDeadline + } + + // Update the workspace build. + err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{ + ID: build.ID, + UpdatedAt: now, + Deadline: autostop.Deadline, + MaxDeadline: autostop.MaxDeadline, + }) + if err != nil { + return xerrors.Errorf("update workspace build %q: %w", build.ID, err) } return nil diff --git a/enterprise/coderd/schedule/user.go b/enterprise/coderd/schedule/user.go index c7d76f86119c8..64de37cee4e29 100644 --- a/enterprise/coderd/schedule/user.go +++ b/enterprise/coderd/schedule/user.go @@ -9,6 +9,7 @@ import ( "github.com/coder/coder/coderd/database" agpl "github.com/coder/coder/coderd/schedule" + "github.com/coder/coder/coderd/tracing" ) // enterpriseUserQuietHoursScheduleStore provides an @@ -29,7 +30,8 @@ func NewEnterpriseUserQuietHoursScheduleStore(defaultSchedule string) (agpl.User defaultSchedule: defaultSchedule, } - _, err := s.parseSchedule(defaultSchedule) + // The context is only used for tracing so using a background ctx is fine. + _, err := s.parseSchedule(context.Background(), defaultSchedule) if err != nil { return nil, xerrors.Errorf("parse default schedule: %w", err) } @@ -37,7 +39,10 @@ func NewEnterpriseUserQuietHoursScheduleStore(defaultSchedule string) (agpl.User return s, nil } -func (s *enterpriseUserQuietHoursScheduleStore) parseSchedule(rawSchedule string) (agpl.UserQuietHoursScheduleOptions, error) { +func (s *enterpriseUserQuietHoursScheduleStore) parseSchedule(ctx context.Context, rawSchedule string) (agpl.UserQuietHoursScheduleOptions, error) { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + userSet := true if strings.TrimSpace(rawSchedule) == "" { userSet = false @@ -64,16 +69,22 @@ func (s *enterpriseUserQuietHoursScheduleStore) parseSchedule(rawSchedule string } func (s *enterpriseUserQuietHoursScheduleStore) Get(ctx context.Context, db database.Store, userID uuid.UUID) (agpl.UserQuietHoursScheduleOptions, error) { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + user, err := db.GetUserByID(ctx, userID) if err != nil { return agpl.UserQuietHoursScheduleOptions{}, xerrors.Errorf("get user by ID: %w", err) } - return s.parseSchedule(user.QuietHoursSchedule) + return s.parseSchedule(ctx, user.QuietHoursSchedule) } func (s *enterpriseUserQuietHoursScheduleStore) Set(ctx context.Context, db database.Store, userID uuid.UUID, rawSchedule string) (agpl.UserQuietHoursScheduleOptions, error) { - opts, err := s.parseSchedule(rawSchedule) + ctx, span := tracing.StartSpan(ctx) + defer span.End() + + opts, err := s.parseSchedule(ctx, rawSchedule) if err != nil { return opts, err } @@ -91,8 +102,12 @@ func (s *enterpriseUserQuietHoursScheduleStore) Set(ctx context.Context, db data return agpl.UserQuietHoursScheduleOptions{}, xerrors.Errorf("update user quiet hours schedule: %w", err) } - // TODO(@dean): update max_deadline for all active builds for this user to clamp to - // the new schedule. + // We don't update workspace build deadlines when the user changes their own + // quiet hours schedule, because they could potentially keep their workspace + // running forever. + // + // Workspace build deadlines are updated when the template admin changes the + // template's settings however. return opts, nil } From 71f302239c400903611176918274bdb56ded6fc8 Mon Sep 17 00:00:00 2001 From: Dean Sheather Date: Mon, 14 Aug 2023 20:59:59 +0000 Subject: [PATCH 5/5] fixup! chore: add instrumentation to everything --- enterprise/coderd/schedule/template.go | 4 ++-- enterprise/coderd/schedule/user.go | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/enterprise/coderd/schedule/template.go b/enterprise/coderd/schedule/template.go index 7a7481c7438ba..c0169fcfc66ec 100644 --- a/enterprise/coderd/schedule/template.go +++ b/enterprise/coderd/schedule/template.go @@ -184,7 +184,7 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuilds(ctx context.Cont } for _, build := range builds { - err := s.updateWorkspaceBuild(ctx, db, template, build) + err := s.updateWorkspaceBuild(ctx, db, build) if err != nil { return xerrors.Errorf("update workspace build %q: %w", build.ID, err) } @@ -193,7 +193,7 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuilds(ctx context.Cont return nil } -func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Context, db database.Store, template database.Template, build database.WorkspaceBuild) error { +func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Context, db database.Store, build database.WorkspaceBuild) error { ctx, span := tracing.StartSpan(ctx, trace.WithAttributes(attribute.String("coder.workspace_id", build.WorkspaceID.String())), trace.WithAttributes(attribute.String("coder.workspace_build_id", build.ID.String())), diff --git a/enterprise/coderd/schedule/user.go b/enterprise/coderd/schedule/user.go index 64de37cee4e29..002965c8e70db 100644 --- a/enterprise/coderd/schedule/user.go +++ b/enterprise/coderd/schedule/user.go @@ -40,7 +40,7 @@ func NewEnterpriseUserQuietHoursScheduleStore(defaultSchedule string) (agpl.User } func (s *enterpriseUserQuietHoursScheduleStore) parseSchedule(ctx context.Context, rawSchedule string) (agpl.UserQuietHoursScheduleOptions, error) { - ctx, span := tracing.StartSpan(ctx) + _, span := tracing.StartSpan(ctx) defer span.End() userSet := true