chore: improve build deadline code (#19203)

- Adds/improves a lot of comments to make the autostop calculation code
clearer
- Changes the behavior of the enterprise template schedule store to
match the behavior of the workspace TTL endpoint when the new TTL is
zero
- Fixes a bug in the workspace TTL endpoint where it could unset the
build deadline, even though a max_deadline was specified
- Adds a new constraint to the workspace_builds table that enforces the
deadline is non-zero and below the max_deadline if it is set
- Adds CHECK constraint enum generation to scripts/dbgen, used for
testing the above constraint
- Adds Dean and Danielle as CODEOWNERS for the autostop calculation code

CODEOWNERS

@@ -29,3 +29,8 @@ site/src/api/countriesGenerated.ts
 site/src/api/rbacresourcesGenerated.ts
 site/src/api/typesGenerated.ts
 site/CLAUDE.md
+
+# The blood and guts of the autostop algorithm, which is quite complex and
+# requires elite ball knowledge of most of the scheduling code to make changes
+# without inadvertently affecting other parts of the codebase.
+coderd/schedule/autostop.go @deansheather @DanielleMaywood

coderd/database/check_constraint.go

@@ -0,0 +1,16 @@
+// Code generated by scripts/dbgen/main.go. DO NOT EDIT.
+package database
+
+// CheckConstraint represents a named check constraint on a table.
+type CheckConstraint string
+
+// CheckConstraint enums.
+const (
+	CheckOneTimePasscodeSet                        CheckConstraint = "one_time_passcode_set"                            // users
+	CheckMaxProvisionerLogsLength                  CheckConstraint = "max_provisioner_logs_length"                      // provisioner_jobs
+	CheckValidationMonotonicOrder                  CheckConstraint = "validation_monotonic_order"                       // template_version_parameters
+	CheckMaxLogsLength                             CheckConstraint = "max_logs_length"                                  // workspace_agents
+	CheckSubsystemsNotNone                         CheckConstraint = "subsystems_not_none"                              // workspace_agents
+	CheckWorkspaceBuildsAiTaskSidebarAppIDRequired CheckConstraint = "workspace_builds_ai_task_sidebar_app_id_required" // workspace_builds
+	CheckWorkspaceBuildsDeadlineBelowMaxDeadline   CheckConstraint = "workspace_builds_deadline_below_max_deadline"     // workspace_builds
+)

coderd/database/dump.sql

@@ -2224,7 +2224,8 @@ CREATE TABLE workspace_builds (
     template_version_preset_id uuid,
     has_ai_task boolean,
     ai_task_sidebar_app_id uuid,
-    CONSTRAINT workspace_builds_ai_task_sidebar_app_id_required CHECK (((((has_ai_task IS NULL) OR (has_ai_task = false)) AND (ai_task_sidebar_app_id IS NULL)) OR ((has_ai_task = true) AND (ai_task_sidebar_app_id IS NOT NULL))))
+    CONSTRAINT workspace_builds_ai_task_sidebar_app_id_required CHECK (((((has_ai_task IS NULL) OR (has_ai_task = false)) AND (ai_task_sidebar_app_id IS NULL)) OR ((has_ai_task = true) AND (ai_task_sidebar_app_id IS NOT NULL)))),
+    CONSTRAINT workspace_builds_deadline_below_max_deadline CHECK ((((deadline <> '0001-01-01 00:00:00+00'::timestamp with time zone) AND (deadline <= max_deadline)) OR (max_deadline = '0001-01-01 00:00:00+00'::timestamp with time zone)))
 );
 
 CREATE VIEW workspace_build_with_user AS

coderd/database/errors.go

@@ -59,6 +59,28 @@ func IsForeignKeyViolation(err error, foreignKeyConstraints ...ForeignKeyConstra
 	return false
 }
 
+// IsCheckViolation checks if the error is due to a check violation. If one or
+// more specific check constraints are given as arguments, the error must be
+// caused by one of them. If no constraints are given, this function returns
+// true for any check violation.
+func IsCheckViolation(err error, checkConstraints ...CheckConstraint) bool {
+	var pqErr *pq.Error
+	if errors.As(err, &pqErr) {
+		if pqErr.Code.Name() == "check_violation" {
+			if len(checkConstraints) == 0 {
+				return true
+			}
+			for _, cc := range checkConstraints {
+				if pqErr.Constraint == string(cc) {
+					return true
+				}
+			}
+		}
+	}
+
+	return false
+}
+
 // IsQueryCanceledError checks if the error is due to a query being canceled.
 func IsQueryCanceledError(err error) bool {
 	var pqErr *pq.Error

coderd/database/migrations/000356_enforce_deadline_below_max_deadline.down.sql

@@ -0,0 +1,2 @@
+ALTER TABLE workspace_builds
+    DROP CONSTRAINT workspace_builds_deadline_below_max_deadline;

coderd/database/migrations/000356_enforce_deadline_below_max_deadline.up.sql

@@ -0,0 +1,20 @@
+-- New constraint: (deadline IS NOT zero AND deadline <= max_deadline) UNLESS max_deadline is zero.
+-- Unfortunately, "zero" here means `time.Time{}`...
+
+-- Update previous builds that would fail this new constraint. This matches the
+-- intended behaviour of the autostop algorithm.
+UPDATE
+    workspace_builds
+SET
+    deadline = max_deadline
+WHERE
+    deadline > max_deadline
+    AND max_deadline != '0001-01-01 00:00:00+00';
+
+-- Add the new constraint.
+ALTER TABLE workspace_builds
+    ADD CONSTRAINT workspace_builds_deadline_below_max_deadline
+        CHECK (
+            (deadline != '0001-01-01 00:00:00+00'::timestamptz AND deadline <= max_deadline)
+            OR max_deadline = '0001-01-01 00:00:00+00'::timestamptz
+        );

coderd/database/querier_test.go

@@ -6003,3 +6003,102 @@ func TestGetRunningPrebuiltWorkspaces(t *testing.T) {
 	require.Len(t, runningPrebuilds, 1, "expected only one running prebuilt workspace")
 	require.Equal(t, runningPrebuild.ID, runningPrebuilds[0].ID, "expected the running prebuilt workspace to be returned")
 }
+
+func TestWorkspaceBuildDeadlineConstraint(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	db, _ := dbtestutil.NewDB(t)
+	org := dbgen.Organization(t, db, database.Organization{})
+	user := dbgen.User(t, db, database.User{})
+	template := dbgen.Template(t, db, database.Template{
+		CreatedBy:      user.ID,
+		OrganizationID: org.ID,
+	})
+	templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+		TemplateID:     uuid.NullUUID{UUID: template.ID, Valid: true},
+		OrganizationID: org.ID,
+		CreatedBy:      user.ID,
+	})
+	workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+		OwnerID:    user.ID,
+		TemplateID: template.ID,
+		Name:       "test-workspace",
+		Deleted:    false,
+	})
+	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+		OrganizationID: org.ID,
+		InitiatorID:    database.PrebuildsSystemUserID,
+		Provisioner:    database.ProvisionerTypeEcho,
+		Type:           database.ProvisionerJobTypeWorkspaceBuild,
+		StartedAt:      sql.NullTime{Time: time.Now().Add(-time.Minute), Valid: true},
+		CompletedAt:    sql.NullTime{Time: time.Now(), Valid: true},
+	})
+	workspaceBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+		WorkspaceID:       workspace.ID,
+		TemplateVersionID: templateVersion.ID,
+		JobID:             job.ID,
+		BuildNumber:       1,
+	})
+
+	cases := []struct {
+		name        string
+		deadline    time.Time
+		maxDeadline time.Time
+		expectOK    bool
+	}{
+		{
+			name:        "no deadline or max_deadline",
+			deadline:    time.Time{},
+			maxDeadline: time.Time{},
+			expectOK:    true,
+		},
+		{
+			name:        "deadline set when max_deadline is not set",
+			deadline:    time.Now().Add(time.Hour),
+			maxDeadline: time.Time{},
+			expectOK:    true,
+		},
+		{
+			name:        "deadline before max_deadline",
+			deadline:    time.Now().Add(-time.Hour),
+			maxDeadline: time.Now().Add(time.Hour),
+			expectOK:    true,
+		},
+		{
+			name:        "deadline is max_deadline",
+			deadline:    time.Now().Add(time.Hour),
+			maxDeadline: time.Now().Add(time.Hour),
+			expectOK:    true,
+		},
+
+		{
+			name:        "deadline after max_deadline",
+			deadline:    time.Now().Add(time.Hour),
+			maxDeadline: time.Now().Add(-time.Hour),
+			expectOK:    false,
+		},
+		{
+			name:        "deadline is not set when max_deadline is set",
+			deadline:    time.Time{},
+			maxDeadline: time.Now().Add(time.Hour),
+			expectOK:    false,
+		},
+	}
+
+	for _, c := range cases {
+		err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+			ID:          workspaceBuild.ID,
+			Deadline:    c.deadline,
+			MaxDeadline: c.maxDeadline,
+			UpdatedAt:   time.Now(),
+		})
+		if c.expectOK {
+			require.NoError(t, err)
+		} else {
+			require.Error(t, err)
+			require.True(t, database.IsCheckViolation(err, database.CheckWorkspaceBuildsDeadlineBelowMaxDeadline))
+		}
+	}
+}

coderd/provisionerdserver/provisionerdserver.go

@@ -1866,8 +1866,9 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 			Database:                    db,
 			TemplateScheduleStore:       templateScheduleStore,
 			UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(),
-			Now:                         now,
-			Workspace:                   workspace.WorkspaceTable(),
+			// `now` is used below to set the build completion time.
+			WorkspaceBuildCompletedAt: now,
+			Workspace:                 workspace.WorkspaceTable(),
 			// Allowed to be the empty string.
 			WorkspaceAutostart: workspace.AutostartSchedule.String,
 		})

coderd/schedule/autostop.go

@@ -50,8 +50,19 @@ type CalculateAutostopParams struct {
 	// by autobuild.NextAutostart
 	WorkspaceAutostart string
 
-	Now       time.Time
-	Workspace database.WorkspaceTable
+	// WorkspaceBuildCompletedAt is the time when the workspace build was
+	// completed.
+	//
+	// We always want to calculate using the build completion time, and not just
+	// the current time, to avoid forcing a workspace build's max_deadline being
+	// pushed to the next potential cron instance.
+	//
+	// E.g. if this function is called for an existing workspace build, which
+	//      currently has a max_deadline within the next 2 hours (see leeway
+	//      above), and the current time is passed into this function, the
+	//      max_deadline will be updated to be much later than expected.
+	WorkspaceBuildCompletedAt time.Time
+	Workspace                 database.WorkspaceTable
 }
 
 type AutostopTime struct {
@@ -68,8 +79,8 @@ type AutostopTime struct {
 // Deadline is the time when the workspace will be stopped, as long as it
 // doesn't see any new activity (such as SSH, app requests, etc.). When activity
 // is detected the deadline is bumped by the workspace's TTL (this only happens
-// when activity is detected and more than 20% of the TTL has passed to save
-// database queries).
+// when activity is detected and more than 5% of the TTL has passed to save
+// database queries, see the ActivityBumpWorkspace query).
 //
 // MaxDeadline is the maximum value for deadline. The deadline cannot be bumped
 // past this value, so it denotes the absolute deadline that the workspace build
@@ -77,55 +88,45 @@ type AutostopTime struct {
 // requirement" settings and the user's "quiet hours" settings to pick a time
 // outside of working hours.
 //
-// Deadline is a cost saving measure, while max deadline is a
-// compliance/updating measure.
+// Note that the deadline is checked at the database level:
+//
+//	(deadline IS NOT zero AND deadline <= max_deadline) UNLESS max_deadline is zero.
+//
+// Deadline is intended as a cost saving measure, not as a hard policy. It is
+// derived from either the workspace's TTL or the template's TTL, depending on
+// the template's policy, to ensure workspaces are stopped when they are idle.
+//
+// MaxDeadline is intended as a compliance policy. It is derived from the
+// template's autostop requirement to cap workspace uptime and effectively force
+// people to update often.
+//
+// Note that only the build's CURRENT deadline property influences automation in
+// the autobuild package. As stated above, the MaxDeadline property is only used
+// to cap the value of a build's deadline.
 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
-		now       = params.Now
+		db               = params.Database
+		workspace        = params.Workspace
+		buildCompletedAt = params.WorkspaceBuildCompletedAt
 
 		autostop AutostopTime
 	)
 
-	var ttl time.Duration
-	if workspace.Ttl.Valid {
-		// When the workspace is made it copies the template's TTL, and the user
-		// can unset it to disable it (unless the template has
-		// UserAutoStopEnabled set to false, see below).
-		ttl = time.Duration(workspace.Ttl.Int64)
-	}
-
-	if workspace.Ttl.Valid {
-		// When the workspace is made it copies the template's TTL, and the user
-		// can unset it to disable it (unless the template has
-		// UserAutoStopEnabled set to false, see below).
-		autostop.Deadline = now.Add(time.Duration(workspace.Ttl.Int64))
-	}
-
 	templateSchedule, err := params.TemplateScheduleStore.Get(ctx, db, workspace.TemplateID)
 	if err != nil {
 		return autostop, xerrors.Errorf("get template schedule options: %w", err)
 	}
-	if !templateSchedule.UserAutostopEnabled {
-		// The user is not permitted to set their own TTL, so use the template
-		// default.
-		ttl = 0
-		if templateSchedule.DefaultTTL > 0 {
-			ttl = templateSchedule.DefaultTTL
-		}
-	}
 
+	ttl := workspaceTTL(workspace, templateSchedule)
 	if ttl > 0 {
 		// Only apply non-zero TTLs.
-		autostop.Deadline = now.Add(ttl)
+		autostop.Deadline = buildCompletedAt.Add(ttl)
 		if params.WorkspaceAutostart != "" {
 			// If the deadline passes the next autostart, we need to extend the deadline to
 			// autostart + deadline. ActivityBumpWorkspace already covers this case
@@ -137,14 +138,14 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut
 			// 3. User starts workspace at 9:45pm.
 			//	- The initial deadline is calculated to be 9:45am
 			//	- This crosses the autostart deadline, so the deadline is extended to 9pm
-			nextAutostart, ok := NextAutostart(params.Now, params.WorkspaceAutostart, templateSchedule)
+			nextAutostart, ok := NextAutostart(params.WorkspaceBuildCompletedAt, params.WorkspaceAutostart, templateSchedule)
 			if ok && autostop.Deadline.After(nextAutostart) {
 				autostop.Deadline = nextAutostart.Add(ttl)
 			}
 		}
 	}
 
-	// Otherwise, use the autostop_requirement algorithm.
+	// Enforce the template autostop requirement if it's configured correctly.
 	if templateSchedule.AutostopRequirement.DaysOfWeek != 0 {
 		// The template has a autostop requirement, so determine the max deadline
 		// of this workspace build.
@@ -161,10 +162,10 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut
 		// workspace.
 		if userQuietHoursSchedule.Schedule != nil {
 			loc := userQuietHoursSchedule.Schedule.Location()
-			now := now.In(loc)
+			buildCompletedAtInLoc := buildCompletedAt.In(loc)
 			// Add the leeway here so we avoid checking today's quiet hours if
 			// the workspace was started <1h before midnight.
-			startOfStopDay := truncateMidnight(now.Add(autostopRequirementLeeway))
+			startOfStopDay := truncateMidnight(buildCompletedAtInLoc.Add(autostopRequirementLeeway))
 
 			// If the template schedule wants to only autostop on n-th weeks
 			// then change the startOfDay to be the Monday of the next
@@ -183,7 +184,7 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut
 			// hour of the scheduled stop time will always bounce to the next
 			// stop window).
 			checkSchedule := userQuietHoursSchedule.Schedule.Next(startOfStopDay.Add(autostopRequirementBuffer))
-			if checkSchedule.Before(now.Add(autostopRequirementLeeway)) {
+			if checkSchedule.Before(buildCompletedAtInLoc.Add(autostopRequirementLeeway)) {
 				// Set the first stop day we try to tomorrow because today's
 				// schedule is too close to now or has already passed.
 				startOfStopDay = nextDayMidnight(startOfStopDay)
@@ -213,14 +214,17 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut
 				startOfStopDay = nextDayMidnight(startOfStopDay)
 			}
 
-			// If the startOfDay is within an hour of now, then we add an hour.
+			// If the startOfDay is within an hour of the build completion time,
+			// then we add an hour.
 			checkTime := startOfStopDay
-			if checkTime.Before(now.Add(time.Hour)) {
-				checkTime = now.Add(time.Hour)
+			if checkTime.Before(buildCompletedAtInLoc.Add(time.Hour)) {
+				checkTime = buildCompletedAtInLoc.Add(time.Hour)
 			} else {
-				// If it's not within an hour of now, subtract 15 minutes to
-				// give a little leeway. This prevents skipped stop events
-				// because autostart perfectly lines up with autostop.
+				// If it's not within an hour of the build completion time,
+				// subtract 15 minutes to give a little leeway. This prevents
+				// skipped stop events because the build time (e.g. autostart
+				// time) perfectly lines up with the max_deadline minus the
+				// leeway.
 				checkTime = checkTime.Add(autostopRequirementBuffer)
 			}
 
@@ -238,15 +242,35 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut
 		autostop.Deadline = autostop.MaxDeadline
 	}
 
-	if (!autostop.Deadline.IsZero() && autostop.Deadline.Before(now)) || (!autostop.MaxDeadline.IsZero() && autostop.MaxDeadline.Before(now)) {
+	if (!autostop.Deadline.IsZero() && autostop.Deadline.Before(buildCompletedAt)) || (!autostop.MaxDeadline.IsZero() && autostop.MaxDeadline.Before(buildCompletedAt)) {
 		// Something went wrong with the deadline calculation, so we should
 		// bail.
-		return autostop, xerrors.Errorf("deadline calculation error, computed deadline or max deadline is in the past for workspace build: deadline=%q maxDeadline=%q now=%q", autostop.Deadline, autostop.MaxDeadline, now)
+		return autostop, xerrors.Errorf("deadline calculation error, computed deadline or max deadline is in the past for workspace build: deadline=%q maxDeadline=%q now=%q", autostop.Deadline, autostop.MaxDeadline, buildCompletedAt)
 	}
 
 	return autostop, nil
 }
 
+// workspaceTTL returns the TTL to use for a workspace.
+//
+// If the template forbids custom workspace TTLs, then we always use the
+// template's configured TTL (or 0 if the template has no TTL configured).
+func workspaceTTL(workspace database.WorkspaceTable, templateSchedule TemplateScheduleOptions) time.Duration {
+	// If the template forbids custom workspace TTLs, then we always use the
+	// template's configured TTL (or 0 if the template has no TTL configured).
+	if !templateSchedule.UserAutostopEnabled {
+		// This is intentionally a nested if statement because of the else if.
+		if templateSchedule.DefaultTTL > 0 {
+			return templateSchedule.DefaultTTL
+		}
+		return 0
+	}
+	if workspace.Ttl.Valid {
+		return time.Duration(workspace.Ttl.Int64)
+	}
+	return 0
+}
+
 // truncateMidnight truncates a time to midnight in the time object's timezone.
 // t.Truncate(24 * time.Hour) truncates based on the internal time and doesn't
 // factor daylight savings properly.

coderd/schedule/autostop_test.go

@@ -76,8 +76,8 @@ func TestCalculateAutoStop(t *testing.T) {
 	t.Log("saturdayMidnightAfterDstOut", saturdayMidnightAfterDstOut)
 
 	cases := []struct {
-		name string
-		now  time.Time
+		name             string
+		buildCompletedAt time.Time
 
 		wsAutostart       string
 		templateAutoStart schedule.TemplateAutostartRequirement
@@ -98,7 +98,7 @@ func TestCalculateAutoStop(t *testing.T) {
 	}{
 		{
 			name:                        "OK",
-			now:                         now,
+			buildCompletedAt:            now,
 			templateAllowAutostop:       true,
 			templateDefaultTTL:          0,
 			templateAutostopRequirement: schedule.TemplateAutostopRequirement{},
@@ -108,7 +108,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                        "Delete",
-			now:                         now,
+			buildCompletedAt:            now,
 			templateAllowAutostop:       true,
 			templateDefaultTTL:          0,
 			templateAutostopRequirement: schedule.TemplateAutostopRequirement{},
@@ -118,7 +118,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                        "WorkspaceTTL",
-			now:                         now,
+			buildCompletedAt:            now,
 			templateAllowAutostop:       true,
 			templateDefaultTTL:          0,
 			templateAutostopRequirement: schedule.TemplateAutostopRequirement{},
@@ -128,7 +128,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                        "TemplateDefaultTTLIgnored",
-			now:                         now,
+			buildCompletedAt:            now,
 			templateAllowAutostop:       true,
 			templateDefaultTTL:          time.Hour,
 			templateAutostopRequirement: schedule.TemplateAutostopRequirement{},
@@ -138,7 +138,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                        "WorkspaceTTLOverridesTemplateDefaultTTL",
-			now:                         now,
+			buildCompletedAt:            now,
 			templateAllowAutostop:       true,
 			templateDefaultTTL:          2 * time.Hour,
 			templateAutostopRequirement: schedule.TemplateAutostopRequirement{},
@@ -148,7 +148,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                        "TemplateBlockWorkspaceTTL",
-			now:                         now,
+			buildCompletedAt:            now,
 			templateAllowAutostop:       false,
 			templateDefaultTTL:          3 * time.Hour,
 			templateAutostopRequirement: schedule.TemplateAutostopRequirement{},
@@ -158,7 +158,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "TemplateAutostopRequirement",
-			now:                    wednesdayMidnightUTC,
+			buildCompletedAt:       wednesdayMidnightUTC,
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -172,7 +172,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "TemplateAutostopRequirement1HourSkip",
-			now:                    saturdayMidnightSydney.Add(-59 * time.Minute),
+			buildCompletedAt:       saturdayMidnightSydney.Add(-59 * time.Minute),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -188,7 +188,7 @@ func TestCalculateAutoStop(t *testing.T) {
 			// The next autostop requirement should be skipped if the
 			// workspace is started within 1 hour of it.
 			name:                   "TemplateAutostopRequirementDaily",
-			now:                    fridayEveningSydney,
+			buildCompletedAt:       fridayEveningSydney,
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -202,7 +202,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "TemplateAutostopRequirementFortnightly/Skip",
-			now:                    wednesdayMidnightUTC,
+			buildCompletedAt:       wednesdayMidnightUTC,
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -216,7 +216,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "TemplateAutostopRequirementFortnightly/NoSkip",
-			now:                    wednesdayMidnightUTC.AddDate(0, 0, 7),
+			buildCompletedAt:       wednesdayMidnightUTC.AddDate(0, 0, 7),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -230,7 +230,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "TemplateAutostopRequirementTriweekly/Skip",
-			now:                    wednesdayMidnightUTC,
+			buildCompletedAt:       wednesdayMidnightUTC,
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -246,7 +246,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "TemplateAutostopRequirementTriweekly/NoSkip",
-			now:                    wednesdayMidnightUTC.AddDate(0, 0, 7),
+			buildCompletedAt:       wednesdayMidnightUTC.AddDate(0, 0, 7),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -262,7 +262,7 @@ func TestCalculateAutoStop(t *testing.T) {
 			name: "TemplateAutostopRequirementOverridesWorkspaceTTL",
 			// now doesn't have to be UTC, but it helps us ensure that
 			// timezones are compared correctly in this test.
-			now:                    fridayEveningSydney.In(time.UTC),
+			buildCompletedAt:       fridayEveningSydney.In(time.UTC),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -276,7 +276,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "TemplateAutostopRequirementOverridesTemplateDefaultTTL",
-			now:                    fridayEveningSydney.In(time.UTC),
+			buildCompletedAt:       fridayEveningSydney.In(time.UTC),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     3 * time.Hour,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -293,7 +293,7 @@ func TestCalculateAutoStop(t *testing.T) {
 			// The epoch is 2023-01-02 in each timezone. We set the time to
 			// 1 second before 11pm the previous day, as this is the latest time
 			// we allow due to our 2h leeway logic.
-			now:                    time.Date(2023, 1, 1, 21, 59, 59, 0, sydneyLoc),
+			buildCompletedAt:       time.Date(2023, 1, 1, 21, 59, 59, 0, sydneyLoc),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -306,7 +306,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "DaylightSavings/OK",
-			now:                    duringDst,
+			buildCompletedAt:       duringDst,
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -320,7 +320,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "DaylightSavings/SwitchMidWeek/In",
-			now:                    beforeDstIn,
+			buildCompletedAt:       beforeDstIn,
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -334,7 +334,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "DaylightSavings/SwitchMidWeek/Out",
-			now:                    beforeDstOut,
+			buildCompletedAt:       beforeDstOut,
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: sydneyQuietHours,
@@ -348,7 +348,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "DaylightSavings/QuietHoursFallsOnDstSwitch/In",
-			now:                    beforeDstIn.Add(-24 * time.Hour),
+			buildCompletedAt:       beforeDstIn.Add(-24 * time.Hour),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: dstInQuietHours,
@@ -362,7 +362,7 @@ func TestCalculateAutoStop(t *testing.T) {
 		},
 		{
 			name:                   "DaylightSavings/QuietHoursFallsOnDstSwitch/Out",
-			now:                    beforeDstOut.Add(-24 * time.Hour),
+			buildCompletedAt:       beforeDstOut.Add(-24 * time.Hour),
 			templateAllowAutostop:  true,
 			templateDefaultTTL:     0,
 			userQuietHoursSchedule: dstOutQuietHours,
@@ -382,7 +382,7 @@ func TestCalculateAutoStop(t *testing.T) {
 			// activity on the workspace.
 			name: "AutostopCrossAutostartBorder",
 			// Starting at 9:45pm, with the autostart at 9am.
-			now:                   pastDateNight,
+			buildCompletedAt:      pastDateNight,
 			templateAllowAutostop: false,
 			templateDefaultTTL:    time.Hour * 12,
 			workspaceTTL:          time.Hour * 12,
@@ -405,7 +405,7 @@ func TestCalculateAutoStop(t *testing.T) {
 			// Same as AutostopCrossAutostartBorder, but just misses the autostart.
 			name: "AutostopCrossMissAutostartBorder",
 			// Starting at 8:45pm, with the autostart at 9am.
-			now:                   time.Date(pastDateNight.Year(), pastDateNight.Month(), pastDateNight.Day(), 20, 30, 0, 0, chicago),
+			buildCompletedAt:      time.Date(pastDateNight.Year(), pastDateNight.Month(), pastDateNight.Day(), 20, 30, 0, 0, chicago),
 			templateAllowAutostop: false,
 			templateDefaultTTL:    time.Hour * 12,
 			workspaceTTL:          time.Hour * 12,
@@ -429,7 +429,7 @@ func TestCalculateAutoStop(t *testing.T) {
 			// The autostop deadline is before the autostart threshold.
 			name: "AutostopCrossAutostartBorderMaxEarlyDeadline",
 			// Starting at 9:45pm, with the autostart at 9am.
-			now:                   pastDateNight,
+			buildCompletedAt:      pastDateNight,
 			templateAllowAutostop: false,
 			templateDefaultTTL:    time.Hour * 12,
 			workspaceTTL:          time.Hour * 12,
@@ -459,7 +459,7 @@ func TestCalculateAutoStop(t *testing.T) {
 			// So the deadline is > 12 hours, but stops at the max deadline.
 			name: "AutostopCrossAutostartBorderMaxDeadline",
 			// Starting at 9:45pm, with the autostart at 9am.
-			now:                   pastDateNight,
+			buildCompletedAt:      pastDateNight,
 			templateAllowAutostop: false,
 			templateDefaultTTL:    time.Hour * 12,
 			workspaceTTL:          time.Hour * 12,
@@ -571,7 +571,7 @@ func TestCalculateAutoStop(t *testing.T) {
 				Database:                    db,
 				TemplateScheduleStore:       templateScheduleStore,
 				UserQuietHoursScheduleStore: userQuietHoursScheduleStore,
-				Now:                         c.now,
+				WorkspaceBuildCompletedAt:   c.buildCompletedAt,
 				Workspace:                   workspace,
 				WorkspaceAutostart:          c.wsAutostart,
 			})

coderd/workspaces.go

@@ -45,8 +45,8 @@ import (
 )
 
 var (
-	ttlMin = time.Minute //nolint:revive // min here means 'minimum' not 'minutes'
-	ttlMax = 30 * 24 * time.Hour
+	ttlMinimum = time.Minute
+	ttlMaximum = 30 * 24 * time.Hour
 
 	errTTLMin              = xerrors.New("time until shutdown must be at least one minute")
 	errTTLMax              = xerrors.New("time until shutdown must be less than 30 days")
@@ -1190,8 +1190,22 @@ func (api *API) putWorkspaceTTL(rw http.ResponseWriter, r *http.Request) {
 
 			if build.Transition == database.WorkspaceTransitionStart {
 				if err = s.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
-					ID:          build.ID,
-					Deadline:    time.Time{},
+					ID: build.ID,
+					// Use the max_deadline as the new build deadline. It will
+					// either be zero (our target), or a non-zero value that we
+					// need to abide by anyway due to template policy.
+					//
+					// Previously, we would always set the deadline to zero,
+					// which was incorrect behavior. When max_deadline is
+					// non-zero, deadline must be set to a non-zero value that
+					// is less than max_deadline.
+					//
+					// Disabling TTL autostop (at a workspace or template level)
+					// does not trump the template's autostop requirement.
+					//
+					// Refer to the comments on schedule.CalculateAutostop for
+					// more information.
+					Deadline:    build.MaxDeadline,
 					MaxDeadline: build.MaxDeadline,
 					UpdatedAt:   dbtime.Time(api.Clock.Now()),
 				}); err != nil {
@@ -2391,11 +2405,11 @@ func validWorkspaceTTLMillis(millis *int64, templateDefault time.Duration) (sql.
 
 	dur := time.Duration(*millis) * time.Millisecond
 	truncated := dur.Truncate(time.Minute)
-	if truncated < ttlMin {
+	if truncated < ttlMinimum {
 		return sql.NullInt64{}, errTTLMin
 	}
 
-	if truncated > ttlMax {
+	if truncated > ttlMaximum {
 		return sql.NullInt64{}, errTTLMax
 	}
 

coderd/workspaces_test.go

@@ -2897,6 +2897,55 @@ func TestWorkspaceUpdateTTL(t *testing.T) {
 		}
 	})
 
+	t.Run("RemoveAutostopWithRunningWorkspaceWithMaxDeadline", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			ctx         = testutil.Context(t, testutil.WaitLong)
+			client, db  = coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+			user        = coderdtest.CreateFirstUser(t, client)
+			version     = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
+			_           = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+			template    = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+			deadline    = 8 * time.Hour
+			maxDeadline = 10 * time.Hour
+			workspace   = coderdtest.CreateWorkspace(t, client, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
+				cwr.TTLMillis = ptr.Ref(deadline.Milliseconds())
+			})
+			build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+		)
+
+		// This is a hack, but the max_deadline isn't precisely configurable
+		// without a lot of unnecessary hassle.
+		dbBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), build.ID) //nolint:gocritic // test
+		require.NoError(t, err)
+		dbJob, err := db.GetProvisionerJobByID(dbauthz.AsSystemRestricted(ctx), dbBuild.JobID) //nolint:gocritic // test
+		require.NoError(t, err)
+		require.True(t, dbJob.CompletedAt.Valid)
+		expectedMaxDeadline := dbJob.CompletedAt.Time.Add(maxDeadline)
+		err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{ //nolint:gocritic // test
+			ID:          build.ID,
+			Deadline:    dbBuild.Deadline,
+			MaxDeadline: expectedMaxDeadline,
+			UpdatedAt:   dbtime.Now(),
+		})
+		require.NoError(t, err)
+
+		// Remove autostop.
+		err = client.UpdateWorkspaceTTL(ctx, workspace.ID, codersdk.UpdateWorkspaceTTLRequest{
+			TTLMillis: nil,
+		})
+		require.NoError(t, err)
+
+		// Expect that the deadline is set to the max_deadline.
+		build, err = client.WorkspaceBuild(ctx, build.ID)
+		require.NoError(t, err)
+		require.True(t, build.Deadline.Valid)
+		require.WithinDuration(t, build.Deadline.Time, expectedMaxDeadline, time.Second)
+		require.True(t, build.MaxDeadline.Valid)
+		require.WithinDuration(t, build.MaxDeadline.Time, expectedMaxDeadline, time.Second)
+	})
+
 	t.Run("CustomAutostopDisabledByTemplate", func(t *testing.T) {
 		t.Parallel()
 		var (

enterprise/coderd/schedule/template.go

@@ -350,14 +350,23 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Conte
 		return nil
 	}
 
+	// Calculate the new autostop max_deadline from the workspace. Since
+	// autostop is always calculated from the build completion time, we don't
+	// want to use the returned autostop.Deadline property as it will likely be
+	// in the distant past.
+	//
+	// The only exception is if the newly calculated workspace TTL is now zero,
+	// which means the workspace can now stay on indefinitely.
+	//
+	// This also matches the behavior of updating a workspace's TTL, where we
+	// don't apply the changes until the workspace is rebuilt.
 	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.WorkspaceTable(),
-		WorkspaceAutostart: workspace.AutostartSchedule.String,
+		WorkspaceBuildCompletedAt:   job.CompletedAt.Time,
+		Workspace:                   workspace.WorkspaceTable(),
+		WorkspaceAutostart:          workspace.AutostartSchedule.String,
 	})
 	if err != nil {
 		return xerrors.Errorf("calculate new autostop for workspace %q: %w", workspace.ID, err)
@@ -389,9 +398,24 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Conte
 		autostop.MaxDeadline = now.Add(time.Hour * 2)
 	}
 
+	// If the new deadline is zero, the workspace can now stay on indefinitely.
+	// Otherwise, we want to discard the new value as per the comment above the
+	// CalculateAutostop call.
+	//
+	// We could potentially calculate a new deadline based on the TTL setting
+	// (on either the workspace or the template based on the template's policy)
+	// against the current time, but doing nothing here matches the current
+	// behavior of the workspace TTL update endpoint.
+	//
+	// Per the documentation of CalculateAutostop, the deadline is not intended
+	// as a policy measure, so it's fine that we don't update it when the
+	// template schedule changes.
+	if !autostop.Deadline.IsZero() {
+		autostop.Deadline = build.Deadline
+	}
+
 	// 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.MaxDeadline.IsZero() && autostop.Deadline.After(autostop.MaxDeadline) {
 		autostop.Deadline = autostop.MaxDeadline
 	}

enterprise/coderd/schedule/template_test.go

@@ -23,7 +23,6 @@ import (
 	"github.com/coder/coder/v2/coderd/notifications"
 	"github.com/coder/coder/v2/coderd/notifications/notificationstest"
 	agplschedule "github.com/coder/coder/v2/coderd/schedule"
-	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/cryptorand"
 	"github.com/coder/coder/v2/enterprise/coderd/schedule"
 	"github.com/coder/coder/v2/testutil"
@@ -73,17 +72,23 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 	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
+	defaultTTL := 8 * time.Hour
+
 	cases := []struct {
-		name        string
-		now         time.Time
+		name string
+		now  time.Time
+		// Before:
 		deadline    time.Time
 		maxDeadline time.Time
-		// Set to nil for no change.
-		newDeadline    *time.Time
+		// After:
+		newDeadline    time.Time
 		newMaxDeadline time.Time
-		noQuietHours   bool
-		autostopReq    *agplschedule.TemplateAutostopRequirement
+		// Config:
+		noQuietHours bool
+		// Note that ttl will not influence the new build at all unless it's 0
+		// AND the build does not have a max deadline post recalculation.
+		ttl         time.Duration
+		autostopReq *agplschedule.TemplateAutostopRequirement
 	}{
 		{
 			name:        "SkippedWorkspaceMaxDeadlineTooSoon",
@@ -91,8 +96,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			deadline:    buildTime,
 			maxDeadline: buildTime.Add(1 * time.Hour),
 			// Unchanged since the max deadline is too soon.
-			newDeadline:    nil,
+			newDeadline:    buildTime,
 			newMaxDeadline: buildTime.Add(1 * time.Hour),
+			ttl:            defaultTTL, // no effect
 		},
 		{
 			name: "NewWorkspaceMaxDeadlineBeforeNow",
@@ -101,10 +107,11 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			deadline: buildTime,
 			// Far into the future...
 			maxDeadline: nextQuietHours.Add(24 * time.Hour),
-			newDeadline: nil,
+			newDeadline: buildTime,
 			// 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),
+			ttl:            defaultTTL, // no effect
 		},
 		{
 			name: "NewWorkspaceMaxDeadlineSoon",
@@ -113,10 +120,11 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			deadline: buildTime,
 			// Far into the future...
 			maxDeadline: nextQuietHours.Add(24 * time.Hour),
-			newDeadline: nil,
+			newDeadline: buildTime,
 			// 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),
+			ttl:            defaultTTL, // no effect
 		},
 		{
 			name: "NewWorkspaceMaxDeadlineFuture",
@@ -125,8 +133,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			deadline: buildTime,
 			// Far into the future...
 			maxDeadline:    nextQuietHours.Add(24 * time.Hour),
-			newDeadline:    nil,
+			newDeadline:    buildTime,
 			newMaxDeadline: nextQuietHours,
+			ttl:            defaultTTL, // no effect
 		},
 		{
 			name: "DeadlineAfterNewWorkspaceMaxDeadline",
@@ -136,8 +145,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			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:    ptr.Ref(nextQuietHours),
+			newDeadline:    nextQuietHours,
 			newMaxDeadline: nextQuietHours,
+			ttl:            defaultTTL, // no effect
 		},
 		{
 			// There was a bug if a user has no quiet hours set, and autostop
@@ -151,13 +161,14 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			deadline:    buildTime.Add(time.Hour * 8),
 			maxDeadline: time.Time{}, // No max set
 			// Should be unchanged
-			newDeadline:    ptr.Ref(buildTime.Add(time.Hour * 8)),
+			newDeadline:    buildTime.Add(time.Hour * 8),
 			newMaxDeadline: time.Time{},
 			noQuietHours:   true,
 			autostopReq: &agplschedule.TemplateAutostopRequirement{
 				DaysOfWeek: 0,
 				Weeks:      0,
 			},
+			ttl: defaultTTL, // no effect
 		},
 		{
 			// A bug existed where MaxDeadline could be set, but deadline was
@@ -168,15 +179,15 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			deadline:    time.Time{},
 			maxDeadline: time.Time{}, // No max set
 			// Should be unchanged
-			newDeadline:    ptr.Ref(time.Time{}),
+			newDeadline:    time.Time{},
 			newMaxDeadline: time.Time{},
 			noQuietHours:   true,
 			autostopReq: &agplschedule.TemplateAutostopRequirement{
 				DaysOfWeek: 0,
 				Weeks:      0,
 			},
+			ttl: defaultTTL, // no effect
 		},
-
 		{
 			// Similar to 'NoDeadline' test. This has a MaxDeadline set, so
 			// the deadline of the workspace should now be set.
@@ -185,8 +196,26 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			// Start with unset times
 			deadline:       time.Time{},
 			maxDeadline:    time.Time{},
-			newDeadline:    ptr.Ref(nextQuietHours),
+			newDeadline:    nextQuietHours,
 			newMaxDeadline: nextQuietHours,
+			ttl:            defaultTTL, // no effect
+		},
+		{
+			// If the build doesn't have a max_deadline anymore, and there is no
+			// TTL anymore, then both the deadline and max_deadline should be
+			// zero.
+			name:           "NoTTLNoDeadlineNoMaxDeadline",
+			now:            buildTime,
+			deadline:       buildTime.Add(time.Hour * 8),
+			maxDeadline:    buildTime.Add(time.Hour * 8),
+			newDeadline:    time.Time{},
+			newMaxDeadline: time.Time{},
+			noQuietHours:   true,
+			autostopReq: &agplschedule.TemplateAutostopRequirement{
+				DaysOfWeek: 0,
+				Weeks:      0,
+			},
+			ttl: 0,
 		},
 	}
 
@@ -206,6 +235,7 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			t.Log("maxDeadline", c.maxDeadline)
 			t.Log("newDeadline", c.newDeadline)
 			t.Log("newMaxDeadline", c.newMaxDeadline)
+			t.Log("ttl", c.ttl)
 
 			var (
 				template = dbgen.Template(t, db, database.Template{
@@ -300,7 +330,7 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			_, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{
 				UserAutostartEnabled:     false,
 				UserAutostopEnabled:      false,
-				DefaultTTL:               0,
+				DefaultTTL:               c.ttl,
 				AutostopRequirement:      autostopReq,
 				FailureTTL:               0,
 				TimeTilDormant:           0,
@@ -312,11 +342,8 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			newBuild, err := db.GetWorkspaceBuildByID(ctx, wsBuild.ID)
 			require.NoError(t, err)
 
-			if c.newDeadline == nil {
-				c.newDeadline = &wsBuild.Deadline
-			}
-			require.WithinDuration(t, *c.newDeadline, newBuild.Deadline, time.Second)
-			require.WithinDuration(t, c.newMaxDeadline, newBuild.MaxDeadline, time.Second)
+			require.WithinDuration(t, c.newDeadline, newBuild.Deadline, time.Second, "deadline")
+			require.WithinDuration(t, c.newMaxDeadline, newBuild.MaxDeadline, time.Second, "max_deadline")
 
 			// Check that the new build has the same state as before.
 			require.Equal(t, wsBuild.ProvisionerState, newBuild.ProvisionerState, "provisioner state mismatch")

scripts/dbgen/constraint.go

@@ -0,0 +1,239 @@
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strings"
+
+	"golang.org/x/tools/imports"
+	"golang.org/x/xerrors"
+)
+
+type constraintType string
+
+const (
+	constraintTypeUnique     constraintType = "unique"
+	constraintTypeForeignKey constraintType = "foreign_key"
+	constraintTypeCheck      constraintType = "check"
+)
+
+func (c constraintType) goType() string {
+	switch c {
+	case constraintTypeUnique:
+		return "UniqueConstraint"
+	case constraintTypeForeignKey:
+		return "ForeignKeyConstraint"
+	case constraintTypeCheck:
+		return "CheckConstraint"
+	default:
+		panic(fmt.Sprintf("unknown constraint type: %s", c))
+	}
+}
+
+func (c constraintType) goTypeDescriptionPart() string {
+	switch c {
+	case constraintTypeUnique:
+		return "unique"
+	case constraintTypeForeignKey:
+		return "foreign key"
+	case constraintTypeCheck:
+		return "check"
+	default:
+		panic(fmt.Sprintf("unknown constraint type: %s", c))
+	}
+}
+
+func (c constraintType) goEnumNamePrefix() string {
+	switch c {
+	case constraintTypeUnique:
+		return "Unique"
+	case constraintTypeForeignKey:
+		return "ForeignKey"
+	case constraintTypeCheck:
+		return "Check"
+	default:
+		panic(fmt.Sprintf("unknown constraint type: %s", c))
+	}
+}
+
+type constraint struct {
+	name string
+	// comment is typically the full constraint, but for check constraints it's
+	// instead the table name.
+	comment string
+}
+
+// queryToConstraintsFn is a function that takes a query and returns zero or
+// more constraints if the query matches the wanted constraint type. If the
+// query does not match the wanted constraint type, the function should return
+// no constraints.
+type queryToConstraintsFn func(query string) ([]constraint, error)
+
+// generateConstraints does the following:
+//  1. Read the dump.sql file
+//  2. Parse the file into each query
+//  3. Pass each query to the constraintFn function
+//  4. Generate the enum from the returned constraints
+//  5. Write the generated code to the output path
+func generateConstraints(dumpPath, outputPath string, outputConstraintType constraintType, fn queryToConstraintsFn) error {
+	dump, err := os.Open(dumpPath)
+	if err != nil {
+		return err
+	}
+	defer dump.Close()
+
+	var allConstraints []constraint
+
+	dumpScanner := bufio.NewScanner(dump)
+	query := ""
+	for dumpScanner.Scan() {
+		line := strings.TrimSpace(dumpScanner.Text())
+		switch {
+		case strings.HasPrefix(line, "--"):
+		case line == "":
+		case strings.HasSuffix(line, ";"):
+			query += line
+			newConstraints, err := fn(query)
+			query = ""
+			if err != nil {
+				return xerrors.Errorf("process query %q: %w", query, err)
+			}
+			allConstraints = append(allConstraints, newConstraints...)
+		default:
+			query += line + " "
+		}
+	}
+	if err = dumpScanner.Err(); err != nil {
+		return err
+	}
+
+	s := &bytes.Buffer{}
+
+	_, _ = fmt.Fprintf(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT.
+package database
+
+// %[1]s represents a named %[2]s constraint on a table.
+type %[1]s string
+
+// %[1]s enums.
+const (
+`, outputConstraintType.goType(), outputConstraintType.goTypeDescriptionPart())
+
+	for _, c := range allConstraints {
+		constName := outputConstraintType.goEnumNamePrefix() + nameFromSnakeCase(c.name)
+		_, _ = fmt.Fprintf(s, "\t%[1]s %[2]s = %[3]q // %[4]s\n", constName, outputConstraintType.goType(), c.name, c.comment)
+	}
+	_, _ = fmt.Fprint(s, ")\n")
+
+	data, err := imports.Process(outputPath, s.Bytes(), &imports.Options{
+		Comments: true,
+	})
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(outputPath, data, 0o600)
+}
+
+// generateUniqueConstraints generates the UniqueConstraint enum.
+func generateUniqueConstraints() error {
+	localPath, err := localFilePath()
+	if err != nil {
+		return err
+	}
+	databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database")
+	dumpPath := filepath.Join(databasePath, "dump.sql")
+	outputPath := filepath.Join(databasePath, "unique_constraint.go")
+
+	fn := func(query string) ([]constraint, error) {
+		if strings.Contains(query, "UNIQUE") || strings.Contains(query, "PRIMARY KEY") {
+			name := ""
+			switch {
+			case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"):
+				name = strings.Split(query, " ")[6]
+			case strings.Contains(query, "CREATE UNIQUE INDEX"):
+				name = strings.Split(query, " ")[3]
+			default:
+				return nil, xerrors.Errorf("unknown unique constraint format: %s", query)
+			}
+			return []constraint{
+				{
+					name:    name,
+					comment: query,
+				},
+			}, nil
+		}
+		return nil, nil
+	}
+	return generateConstraints(dumpPath, outputPath, constraintTypeUnique, fn)
+}
+
+// generateForeignKeyConstraints generates the ForeignKeyConstraint enum.
+func generateForeignKeyConstraints() error {
+	localPath, err := localFilePath()
+	if err != nil {
+		return err
+	}
+	databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database")
+	dumpPath := filepath.Join(databasePath, "dump.sql")
+	outputPath := filepath.Join(databasePath, "foreign_key_constraint.go")
+
+	fn := func(query string) ([]constraint, error) {
+		if strings.Contains(query, "FOREIGN KEY") {
+			name := ""
+			switch {
+			case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"):
+				name = strings.Split(query, " ")[6]
+			default:
+				return nil, xerrors.Errorf("unknown foreign key constraint format: %s", query)
+			}
+			return []constraint{
+				{
+					name:    name,
+					comment: query,
+				},
+			}, nil
+		}
+		return []constraint{}, nil
+	}
+	return generateConstraints(dumpPath, outputPath, constraintTypeForeignKey, fn)
+}
+
+// generateCheckConstraints generates the CheckConstraint enum.
+func generateCheckConstraints() error {
+	localPath, err := localFilePath()
+	if err != nil {
+		return err
+	}
+	databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database")
+	dumpPath := filepath.Join(databasePath, "dump.sql")
+	outputPath := filepath.Join(databasePath, "check_constraint.go")
+
+	var (
+		tableRegex = regexp.MustCompile(`CREATE TABLE\s+([^\s]+)`)
+		checkRegex = regexp.MustCompile(`CONSTRAINT\s+([^\s]+)\s+CHECK`)
+	)
+	fn := func(query string) ([]constraint, error) {
+		constraints := []constraint{}
+
+		tableMatches := tableRegex.FindStringSubmatch(query)
+		if len(tableMatches) > 0 {
+			table := tableMatches[1]
+
+			// Find every CONSTRAINT xxx CHECK occurrence.
+			matches := checkRegex.FindAllStringSubmatch(query, -1)
+			for _, match := range matches {
+				constraints = append(constraints, constraint{
+					name:    match[1],
+					comment: table,
+				})
+			}
+		}
+		return constraints, nil
+	}
+
+	return generateConstraints(dumpPath, outputPath, constraintTypeCheck, fn)
+}

scripts/dbgen/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"bufio"
 	"bytes"
 	"fmt"
 	"go/format"
@@ -80,154 +79,14 @@ return %s
 		return xerrors.Errorf("generate foreign key constraints: %w", err)
 	}
 
+	err = generateCheckConstraints()
+	if err != nil {
+		return xerrors.Errorf("generate check constraints: %w", err)
+	}
+
 	return nil
 }
 
-// generateUniqueConstraints generates the UniqueConstraint enum.
-func generateUniqueConstraints() error {
-	localPath, err := localFilePath()
-	if err != nil {
-		return err
-	}
-	databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database")
-
-	dump, err := os.Open(filepath.Join(databasePath, "dump.sql"))
-	if err != nil {
-		return err
-	}
-	defer dump.Close()
-
-	var uniqueConstraints []string
-	dumpScanner := bufio.NewScanner(dump)
-	query := ""
-	for dumpScanner.Scan() {
-		line := strings.TrimSpace(dumpScanner.Text())
-		switch {
-		case strings.HasPrefix(line, "--"):
-		case line == "":
-		case strings.HasSuffix(line, ";"):
-			query += line
-			if strings.Contains(query, "UNIQUE") || strings.Contains(query, "PRIMARY KEY") {
-				uniqueConstraints = append(uniqueConstraints, query)
-			}
-			query = ""
-		default:
-			query += line + " "
-		}
-	}
-	if err = dumpScanner.Err(); err != nil {
-		return err
-	}
-
-	s := &bytes.Buffer{}
-
-	_, _ = fmt.Fprint(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT.
-package database
-`)
-	_, _ = fmt.Fprint(s, `
-// UniqueConstraint represents a named unique constraint on a table.
-type UniqueConstraint string
-
-// UniqueConstraint enums.
-const (
-`)
-	for _, query := range uniqueConstraints {
-		name := ""
-		switch {
-		case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"):
-			name = strings.Split(query, " ")[6]
-		case strings.Contains(query, "CREATE UNIQUE INDEX"):
-			name = strings.Split(query, " ")[3]
-		default:
-			return xerrors.Errorf("unknown unique constraint format: %s", query)
-		}
-		_, _ = fmt.Fprintf(s, "\tUnique%s UniqueConstraint = %q // %s\n", nameFromSnakeCase(name), name, query)
-	}
-	_, _ = fmt.Fprint(s, ")\n")
-
-	outputPath := filepath.Join(databasePath, "unique_constraint.go")
-
-	data, err := imports.Process(outputPath, s.Bytes(), &imports.Options{
-		Comments: true,
-	})
-	if err != nil {
-		return err
-	}
-	return os.WriteFile(outputPath, data, 0o600)
-}
-
-// generateForeignKeyConstraints generates the ForeignKeyConstraint enum.
-func generateForeignKeyConstraints() error {
-	localPath, err := localFilePath()
-	if err != nil {
-		return err
-	}
-	databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database")
-
-	dump, err := os.Open(filepath.Join(databasePath, "dump.sql"))
-	if err != nil {
-		return err
-	}
-	defer dump.Close()
-
-	var foreignKeyConstraints []string
-	dumpScanner := bufio.NewScanner(dump)
-	query := ""
-	for dumpScanner.Scan() {
-		line := strings.TrimSpace(dumpScanner.Text())
-		switch {
-		case strings.HasPrefix(line, "--"):
-		case line == "":
-		case strings.HasSuffix(line, ";"):
-			query += line
-			if strings.Contains(query, "FOREIGN KEY") {
-				foreignKeyConstraints = append(foreignKeyConstraints, query)
-			}
-			query = ""
-		default:
-			query += line + " "
-		}
-	}
-
-	if err := dumpScanner.Err(); err != nil {
-		return err
-	}
-
-	s := &bytes.Buffer{}
-
-	_, _ = fmt.Fprint(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT.
-package database
-`)
-	_, _ = fmt.Fprint(s, `
-// ForeignKeyConstraint represents a named foreign key constraint on a table.
-type ForeignKeyConstraint string
-
-// ForeignKeyConstraint enums.
-const (
-`)
-	for _, query := range foreignKeyConstraints {
-		name := ""
-		switch {
-		case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"):
-			name = strings.Split(query, " ")[6]
-		default:
-			return xerrors.Errorf("unknown foreign key constraint format: %s", query)
-		}
-		_, _ = fmt.Fprintf(s, "\tForeignKey%s ForeignKeyConstraint = %q // %s\n", nameFromSnakeCase(name), name, query)
-	}
-	_, _ = fmt.Fprint(s, ")\n")
-
-	outputPath := filepath.Join(databasePath, "foreign_key_constraint.go")
-
-	data, err := imports.Process(outputPath, s.Bytes(), &imports.Options{
-		Comments: true,
-	})
-	if err != nil {
-		return err
-	}
-	return os.WriteFile(outputPath, data, 0o600)
-}
-
 type stubParams struct {
 	FuncName   string
 	Parameters string