feat: update workspace deadline when template policy changes

coder · deansheather · Aug 14, 2023 · Aug 8, 2023 · Aug 8, 2023 · Aug 11, 2023
commit cd35f879fe7023ee0006f3244b1d9541d714cfc4
diff --git 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
@@ -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
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git 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
@@ -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
@@ -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
@@ -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,
+
+	)
+}