progress

coder · deansheather · Jul 20, 2023 · Jun 20, 2023 · Jun 20, 2023 · Jun 21, 2023
commit 06272e26fc419a20d8ef682bbba17cd560d89d67
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2558,13 +2558,6 @@ func (q *querier) UpdateWorkspaceTTL(ctx context.Context, arg database.UpdateWor
 	return update(q.log, q.auth, fetch, q.db.UpdateWorkspaceTTL)(ctx, arg)
 }
 
-func (q *querier) UpdateWorkspaceTTLToBeWithinTemplateMax(ctx context.Context, arg database.UpdateWorkspaceTTLToBeWithinTemplateMaxParams) error {
-	fetch := func(ctx context.Context, arg database.UpdateWorkspaceTTLToBeWithinTemplateMaxParams) (database.Template, error) {
-		return q.db.GetTemplateByID(ctx, arg.TemplateID)
-	}
-	return fetchAndExec(q.log, q.auth, rbac.ActionUpdate, fetch, q.db.UpdateWorkspaceTTLToBeWithinTemplateMax)(ctx, arg)
-}
-
 func (q *querier) UpsertAppSecurityKey(ctx context.Context, data string) error {
 	// No authz checks as this is done during startup
 	return q.db.UpsertAppSecurityKey(ctx, data)

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -5276,26 +5276,6 @@ func (q *fakeQuerier) UpdateWorkspaceTTL(_ context.Context, arg database.UpdateW
 	return sql.ErrNoRows
 }
 
-func (q *fakeQuerier) UpdateWorkspaceTTLToBeWithinTemplateMax(_ context.Context, arg database.UpdateWorkspaceTTLToBeWithinTemplateMaxParams) error {
-	if err := validateDatabaseType(arg); err != nil {
-		return err
-	}
-
-	q.mutex.Lock()
-	defer q.mutex.Unlock()
-
-	for index, workspace := range q.workspaces {
-		if workspace.TemplateID != arg.TemplateID || !workspace.Ttl.Valid || workspace.Ttl.Int64 < arg.TemplateMaxTTL {
-			continue
-		}
-
-		workspace.Ttl = sql.NullInt64{Int64: arg.TemplateMaxTTL, Valid: true}
-		q.workspaces[index] = workspace
-	}
-
-	return nil
-}
-
 func (q *fakeQuerier) UpsertAppSecurityKey(_ context.Context, data string) error {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()

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/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -329,20 +329,6 @@ SET
 WHERE
 	id = $1;
 
--- name: UpdateWorkspaceTTLToBeWithinTemplateMax :exec
-UPDATE
-	workspaces
-SET
-	ttl = LEAST(ttl, @template_max_ttl::bigint)
-WHERE
-	template_id = @template_id
-	-- LEAST() does not pick NULL, so filter it out as we don't want to set a
-	-- TTL on the workspace if it's unset.
-	--
-	-- During build time, the template max TTL will still be used if the
-	-- workspace TTL is NULL.
-	AND ttl IS NOT NULL;
-
 -- name: GetDeploymentWorkspaceStats :one
 WITH workspaces_with_jobs AS (
 	SELECT

diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -927,41 +927,68 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
 					deadline = now.Add(templateSchedule.DefaultTTL)
 				}
 			}
-			if templateSchedule.MaxTTL > 0 {
-				maxDeadline = now.Add(templateSchedule.MaxTTL)
+			if templateSchedule.RestartRequirement.DaysOfWeek != 0 {
+				// The template has a restart requirement, so determine the max
+				// deadline of this workspace build.
 
-				if deadline.IsZero() || maxDeadline.Before(deadline) {
-					// If the workspace doesn't have a deadline or the max
-					// deadline is sooner than the workspace deadline, use the
-					// max deadline as the actual deadline.
+				// First, get the user's quiet hours schedule (this will return
+				// the default if the user has not set their own schedule).
+				userQuietHoursSchedule, err := (*server.UserQuietHoursScheduleStore.Load()).GetUserQuietHoursScheduleOptions(ctx, db, workspace.OwnerID)
+				if err != nil {
+					return xerrors.Errorf("get user quiet hours schedule options: %w", err)
+				}
+
+				// If the schedule is nil, that means the deployment isn't
+				// entitled to use quiet hours or the default schedule has not
+				// been set. In this case, do not set a max deadline on the
+				// workspace.
+				if userQuietHoursSchedule.Schedule != nil {
+					loc := userQuietHoursSchedule.Schedule.Location()
+					now := time.Now().In(loc)
+					startOfDay := now.Truncate(24 * time.Hour)
+
+					// First, determine if we can restart today or if the
+					// schedule is too near/already passed.
 					//
-					// Notably, this isn't affected by the user's quiet hours
-					// schedule below because we'd still like to use the max TTL
-					// as the TTL for the workspace if it's not set.
-					deadline = maxDeadline
+					// Allow an hour of leeway (i.e. any workspaces started
+					// within an hour of the scheduled stop time will always
+					// bounce to the next stop window).
+					todaySchedule := userQuietHoursSchedule.Schedule.Next(startOfDay)
+					if todaySchedule.Before(now.Add(-time.Hour)) {
+						// Set the first stop day we try to tomorrow because
+						// today's schedule is too close to now or has already
+						// passed.
+						startOfDay = startOfDay.Add(24 * time.Hour)
+					}
+
+					// Get the current day of week and iterate through the days
+					// of week until we wrap or find a day present in the
+					// restart requirement.
+					requirementDays := templateSchedule.RestartRequirement.DaysMap()
+					for i := 0; i < len(schedule.DaysOfWeek)+1; i++ {
+						if i == len(schedule.DaysOfWeek) {
+							// We've wrapped, so somehow we couldn't find a day
+							// in the restart requirement. This shouldn't happen
+							// because the restart requirement has a day set.
+							return xerrors.Errorf("could not find suitable day for template restart requirement in the next 7 days")
+						}
+						if requirementDays[startOfDay.Weekday()] {
+							break
+						}
+						startOfDay = startOfDay.Add(24 * time.Hour)
+					}
+
+					// Get the next occurrence of the restart schedule after
+					// the start of today.
+					maxDeadline = userQuietHoursSchedule.Schedule.Next(startOfDay)
 				}
-			}
 
-			userQuietHoursSchedule, err := (*server.UserQuietHoursScheduleStore.Load()).GetUserQuietHoursScheduleOptions(ctx, db, workspace.OwnerID)
-			if err != nil {
-				return xerrors.Errorf("get user quiet hours schedule options: %w", err)
-			}
-			if userQuietHoursSchedule.Schedule != nil {
-				// Round the max deadline up to the nearest occurrence of the
-				// user's quiet hours schedule. This ensures that workspaces
-				// can't be force-stopped due to max TTL during business hours.
-
-				// Get the schedule occurrence that happens right before, during
-				// or after the max deadline.
-				// TODO: change to the quiet hours window BEFORE max TTL
-				scheduleDur := userQuietHoursSchedule.Duration
-				if scheduleDur > 1*time.Hour {
-					// Allow a 15 minute buffer when possible so we're not too
-					// constrained with the autostop time.
-					scheduleDur -= 15 * time.Minute
+				// If the workspace doesn't have a deadline or the max deadline
+				// is sooner than the workspace deadline, use the max deadline
+				// as the actual deadline.
+				if deadline.IsZero() || maxDeadline.Before(deadline) {
+					deadline = maxDeadline
 				}
-				windowStart := userQuietHoursSchedule.Schedule.Next(maxDeadline.Add(scheduleDur))
-				maxDeadline = windowStart
 			}
 
 			err = db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{

diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -47,7 +47,7 @@ func testTemplateScheduleStore() *atomic.Pointer[schedule.TemplateScheduleStore]
 	return ptr
 }
 
-func testUserMaintenanceScheduleStore() *atomic.Pointer[schedule.UserQuietHoursScheduleStore] {
+func testUserQuietHoursScheduleStore() *atomic.Pointer[schedule.UserQuietHoursScheduleStore] {
 	ptr := &atomic.Pointer[schedule.UserQuietHoursScheduleStore]{}
 	store := schedule.NewAGPLUserQuietHoursScheduleStore()
 	ptr.Store(&store)
@@ -71,7 +71,7 @@ func TestAcquireJob(t *testing.T) {
 			AcquireJobDebounce:          time.Hour,
 			Auditor:                     mockAuditor(),
 			TemplateScheduleStore:       testTemplateScheduleStore(),
-			UserQuietHoursScheduleStore: testUserMaintenanceScheduleStore(),
+			UserQuietHoursScheduleStore: testUserQuietHoursScheduleStore(),
 			Tracer:                      trace.NewNoopTracerProvider().Tracer("noop"),
 			DeploymentValues:            &codersdk.DeploymentValues{},
 		}
@@ -1278,7 +1278,7 @@ func setup(t *testing.T, ignoreLogErrors bool) *provisionerdserver.Server {
 		Telemetry:                   telemetry.NewNoop(),
 		Auditor:                     mockAuditor(),
 		TemplateScheduleStore:       testTemplateScheduleStore(),
-		UserQuietHoursScheduleStore: testUserMaintenanceScheduleStore(),
+		UserQuietHoursScheduleStore: testUserQuietHoursScheduleStore(),
 		Tracer:                      trace.NewNoopTracerProvider().Tracer("noop"),
 		DeploymentValues:            &codersdk.DeploymentValues{},
 	}

diff --git a/coderd/schedule/template.go b/coderd/schedule/template.go
@@ -9,15 +9,54 @@ import (
 	"github.com/coder/coder/coderd/database"
 )
 
+var DaysOfWeek = []time.Weekday{
+	time.Sunday,
+	time.Monday,
+	time.Tuesday,
+	time.Wednesday,
+	time.Thursday,
+	time.Friday,
+	time.Saturday,
+}
+
+type TemplateRestartRequirement struct {
+	// DaysOfWeek is a bitmap of which days of the week the workspace must be
+	// restarted. If fully zero, the workspace is not required to be restarted
+	// ever.
+	//
+	// First bit is Sunday, second bit is Monday, ..., seventh bit is Saturday,
+	// eighth bit is unused.
+	DaysOfWeek uint8
+}
+
+// Days returns the days of the week that the workspace must be restarted.
+func (r TemplateRestartRequirement) Days() []time.Weekday {
+	days := make([]time.Weekday, 0, 7)
+	for i, day := range DaysOfWeek {
+		if r.DaysOfWeek&(1<<uint(i)) != 0 {
+			days = append(days, day)
+		}
+	}
+	return days
+}
+
+// DaysMap returns a map of the days of the week that the workspace must be
+// restarted.
+func (r TemplateRestartRequirement) DaysMap() map[time.Weekday]bool {
+	days := make(map[time.Weekday]bool)
+	for i, day := range DaysOfWeek {
+		days[day] = r.DaysOfWeek&(1<<uint(i)) != 0
+	}
+	return days
+}
+
 type TemplateScheduleOptions struct {
 	UserAutostartEnabled bool          `json:"user_autostart_enabled"`
 	UserAutostopEnabled  bool          `json:"user_autostop_enabled"`
 	DefaultTTL           time.Duration `json:"default_ttl"`
-	// If MaxTTL is set, the workspace must be stopped before this time or it
-	// will be stopped automatically.
-	//
-	// If set, users cannot disable automatic workspace shutdown.
-	MaxTTL time.Duration `json:"max_ttl"`
+	// RestartRequirement dictates when the workspace must be restarted. This
+	// used to be handled by MaxTTL.
+	RestartRequirement TemplateRestartRequirement `json:"restart_requirement"`
 	// FailureTTL dictates the duration after which failed workspaces will be
 	// stopped automatically.
 	FailureTTL time.Duration `json:"failure_ttl"`
@@ -56,9 +95,11 @@ func (*agplTemplateScheduleStore) GetTemplateScheduleOptions(ctx context.Context
 		UserAutostartEnabled: true,
 		UserAutostopEnabled:  true,
 		DefaultTTL:           time.Duration(tpl.DefaultTTL),
-		// Disregard the values in the database, since MaxTTL, FailureTTL,
-		// InactivityTTL, and LockedTTL are enterprise features.
-		MaxTTL:        0,
+		// Disregard the values in the database, since RestartRequirement,
+		// FailureTTL, InactivityTTL, and LockedTTL are enterprise features.
+		RestartRequirement: TemplateRestartRequirement{
+			DaysOfWeek: 0,
+		},
 		FailureTTL:    0,
 		InactivityTTL: 0,
 		LockedTTL:     0,
@@ -71,6 +112,7 @@ func (*agplTemplateScheduleStore) SetTemplateScheduleOptions(ctx context.Context
 		return tpl, nil
 	}
 
+	// TODO: fix storage to use new restart requirement
 	return db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{
 		ID:         tpl.ID,
 		UpdatedAt:  database.Now(),