coder · deansheather · Oct 4, 2023 · johnstcn · Oct 4, 2023 · johnstcn
diff --git a/cli/schedule_test.go b/cli/schedule_test.go
@@ -243,9 +243,8 @@ func TestScheduleOverride(t *testing.T) {
 		require.NoError(t, err)
 		expectedDeadline := time.Now().Add(10 * time.Hour)
 
-		// Assert test invariant: workspace build has a deadline set equal to now plus ttl
-		initDeadline := time.Now().Add(time.Duration(*workspace.TTLMillis) * time.Millisecond)
-		require.WithinDuration(t, initDeadline, workspace.LatestBuild.Deadline.Time, time.Minute)
+		// Assert test invariant: workspace build has a deadline set
+		require.False(t, workspace.LatestBuild.Deadline.IsZero())
 
 		inv, root := clitest.New(t, cmdArgs...)
 		clitest.SetupConfig(t, client, root)
@@ -283,10 +282,6 @@ func TestScheduleOverride(t *testing.T) {
 		workspace, err = client.Workspace(ctx, workspace.ID)
 		require.NoError(t, err)
 
-		// Assert test invariant: workspace build has a deadline set equal to now plus ttl
-		initDeadline := time.Now().Add(time.Duration(*workspace.TTLMillis) * time.Millisecond)
-		require.WithinDuration(t, initDeadline, workspace.LatestBuild.Deadline.Time, time.Minute)
-
 		inv, root := clitest.New(t, cmdArgs...)
 		clitest.SetupConfig(t, client, root)
 		inv.Stdout = stdoutBuf

diff --git a/coderd/activitybump_internal_test.go b/coderd/activitybump_internal_test.go
@@ -32,13 +32,15 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 	}
 
 	for _, tt := range []struct {
-		name                string
-		transition          database.WorkspaceTransition
-		jobCompletedAt      sql.NullTime
-		buildDeadlineOffset *time.Duration
-		maxDeadlineOffset   *time.Duration
-		workspaceTTL        time.Duration
-		expectedBump        time.Duration
+		name                         string
+		transition                   database.WorkspaceTransition
+		jobCompletedAt               sql.NullTime
+		buildDeadlineOffset          *time.Duration
+		maxDeadlineOffset            *time.Duration
+		workspaceTTL                 time.Duration
+		templateDisallowUserAutostop bool // inverted
+		templateTTL                  time.Duration
+		expectedBump                 time.Duration
 	}{
 		{
 			name:                "NotFinishedYet",
@@ -69,6 +71,7 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
 			buildDeadlineOffset: ptr.Ref(8*time.Hour - 24*time.Minute),
 			workspaceTTL:        8 * time.Hour,
+			templateTTL:         6 * time.Hour, // unused
 			expectedBump:        8 * time.Hour,
 		},
 		{
@@ -99,6 +102,18 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 			workspaceTTL:        8 * time.Hour,
 			expectedBump:        0,
 		},
+		{
+			// If the template disallows user autostop, then the TTL of the
+			// template should be used.
+			name:                         "TemplateDisallowsUserAutostop",
+			transition:                   database.WorkspaceTransitionStart,
+			jobCompletedAt:               sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
+			buildDeadlineOffset:          ptr.Ref(8*time.Hour - 24*time.Minute),
+			workspaceTTL:                 6 * time.Hour,
+			templateDisallowUserAutostop: true,
+			templateTTL:                  8 * time.Hour,
+			expectedBump:                 8 * time.Hour,
+		},
 	} {
 		tt := tt
 		for _, tz := range timezones {
@@ -144,6 +159,15 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 					buildID = uuid.New()
 				)
 
+				err := db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{
+					ID:                template.ID,
+					UpdatedAt:         dbtime.Now(),
+					AllowUserAutostop: !tt.templateDisallowUserAutostop,
+					DefaultTTL:        int64(tt.templateTTL),
+					// The other fields don't matter.
+				})
+				require.NoError(t, err)
+
 				var buildNumber int32 = 1
 				// Insert a number of previous workspace builds.
 				for i := 0; i < 5; i++ {
@@ -162,7 +186,7 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 				if tt.maxDeadlineOffset != nil {
 					maxDeadline = now.Add(*tt.maxDeadlineOffset)
 				}
-				err := db.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
+				err = db.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
 					ID:                buildID,
 					CreatedAt:         dbtime.Now(),
 					UpdatedAt:         dbtime.Now(),

diff --git a/coderd/autobuild/lifecycle_executor.go b/coderd/autobuild/lifecycle_executor.go
@@ -19,7 +19,6 @@ import (
 	"github.com/coder/coder/v2/coderd/database/provisionerjobs"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/schedule"
-	"github.com/coder/coder/v2/coderd/schedule/cron"
 	"github.com/coder/coder/v2/coderd/wsbuilder"
 	"github.com/coder/coder/v2/codersdk"
 )
@@ -116,7 +115,7 @@ func (e *Executor) runOnce(t time.Time) Stats {
 	// NOTE: If a workspace build is created with a given TTL and then the user either
 	//       changes or unsets the TTL, the deadline for the workspace build will not
 	//       have changed. This behavior is as expected per #2229.
-	workspaces, err := e.db.GetWorkspacesEligibleForTransition(e.ctx, t)
+	workspaces, err := e.db.GetWorkspacesEligibleForTransition(e.ctx)
 	if err != nil {
 		e.log.Error(e.ctx, "get workspaces for autostart or autostop", slog.Error(err))
 		return stats
@@ -183,7 +182,26 @@ func (e *Executor) runOnce(t time.Time) Stats {
 					}
 				}
 
-				//  Transition the workspace to dormant if it has breached the template's
+				if reason == database.BuildReasonBump {
+					newDeadline := shouldBump(ws, latestBuild, latestJob, templateSchedule, currentTick)
+					if !newDeadline.IsZero() {
+						err := tx.UpdateWorkspaceBuildDeadlineByID(e.ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+							ID:          latestBuild.ID,
+							UpdatedAt:   dbtime.Now(),
+							Deadline:    newDeadline,
+							MaxDeadline: latestBuild.MaxDeadline,
+						})
+						if err != nil {
+							log.Error(e.ctx, "unable to bump workspace deadline",
+								slog.F("new_deadline", newDeadline),
+								slog.Error(err),
+							)
+							return nil
+						}
+					}
+				}
+
+				// Transition the workspace to dormant if it has breached the template's
 				// threshold for inactivity.
 				if reason == database.BuildReasonAutolock {
 					ws, err = tx.UpdateWorkspaceDormantDeletingAt(e.ctx, database.UpdateWorkspaceDormantDeletingAtParams{
@@ -295,6 +313,8 @@ func getNextTransition(
 
 	case isEligibleForDelete(ws, templateSchedule, currentTick):
 		return database.WorkspaceTransitionDelete, database.BuildReasonAutodelete, nil
+	case !shouldBump(ws, latestBuild, latestJob, templateSchedule, currentTick).IsZero():
+		return "", database.BuildReasonBump, nil
 	default:
 		return "", "", xerrors.Errorf("last transition not valid for autostart or autostop")
 	}
@@ -320,14 +340,11 @@ func isEligibleForAutostart(ws database.Workspace, build database.WorkspaceBuild
 
 	// If autostart isn't enabled, or the schedule isn't valid/populated we can't
 	// autostart the workspace.
-	if !templateSchedule.UserAutostartEnabled || !ws.AutostartSchedule.Valid || ws.AutostartSchedule.String == "" {
+	sched, err := schedule.GetWorkspaceAutostartSchedule(templateSchedule, ws)
+	if err != nil || sched == nil {
 		return false
 	}
 
-	sched, err := cron.Weekly(ws.AutostartSchedule.String)
-	if err != nil {
-		return false
-	}
 	// Round down to the nearest minute, as this is the finest granularity cron supports.
 	// Truncate is probably not necessary here, but doing it anyway to be sure.
 	nextTransition := sched.Next(build.CreatedAt).Truncate(time.Minute)
@@ -385,3 +402,49 @@ func isEligibleForFailedStop(build database.WorkspaceBuild, job database.Provisi
 		job.CompletedAt.Valid &&
 		currentTick.Sub(job.CompletedAt.Time) > templateSchedule.FailureTTL
 }
+
+// shouldBump returns a non-zero time if the workspace deadline should
+// should be bumped.
+func shouldBump(ws database.Workspace, build database.WorkspaceBuild, job database.ProvisionerJob, templateSchedule schedule.TemplateScheduleOptions, currentTick time.Time) time.Time {
+	if build.Transition != database.WorkspaceTransitionStart {
+		return time.Time{}
+	}
+	if db2sdk.ProvisionerJobStatus(job) != codersdk.ProvisionerJobSucceeded {
+		return time.Time{}
+	}
+	if build.Deadline.IsZero() {
+		return time.Time{}
+	}
+	if currentTick.After(build.Deadline.Add(time.Hour)) {
+		// If the workspace has already breached its deadline + 1h, we should
+		// not bump it, because we don't want to race autostop code.
+		return time.Time{}
+	}
+
+	ttl := schedule.WorkspaceTTL(templateSchedule, ws)
+	if ttl <= 0 {
+		return time.Time{}
+	}
+
+	// If autostart isn't enabled, or the schedule isn't valid/populated we
+	// don't care about bumping it.
+	sched, err := schedule.GetWorkspaceAutostartSchedule(templateSchedule, ws)
+	if err != nil || sched == nil {
+		return time.Time{}
+	}
+
+	newDeadline := schedule.MaybeBumpDeadline(sched, build.Deadline, ttl)
+	if newDeadline.IsZero() {
+		return time.Time{}
+	}
+	if newDeadline.Before(build.Deadline) {
+		// Don't reduce the deadline.
+		return time.Time{}
+	}
+	if !build.MaxDeadline.IsZero() && newDeadline.After(build.MaxDeadline) {
+		// Don't exceed the max deadline.
+		return time.Time{}
+	}
+
+	return newDeadline
+}
diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go
@@ -16,12 +16,15 @@ import (
 	"github.com/coder/coder/v2/coderd/autobuild"
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/schedule"
 	"github.com/coder/coder/v2/coderd/schedule/cron"
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/provisioner/echo"
 	"github.com/coder/coder/v2/provisionersdk/proto"
+	"github.com/coder/coder/v2/testutil"
 )
 
 func TestExecutorAutostartOK(t *testing.T) {
@@ -701,6 +704,131 @@ func TestExecutorFailedWorkspace(t *testing.T) {
 	})
 }
 
+func TestExecutorBumpWorkspace(t *testing.T) {
+	t.Parallel()
+
+	midnight := time.Date(2023, 10, 4, 0, 0, 0, 0, time.UTC)
+
+	cases := []struct {
+		name              string
+		autostartSchedule string
+		now               time.Time // defaults to 1h before the deadline
+		deadline          time.Time
+		maxDeadline       time.Time
+		ttl               time.Duration
+		expected          time.Time
+	}{
+		{
+			name:              "not eligible",
+			autostartSchedule: "CRON_TZ=UTC 0 0 * * *",
+			deadline:          midnight.Add(time.Hour * -2),
+			ttl:               time.Hour * 10,
+			// not eligible because the deadline is over an hour before the next
+			// autostart time
+			expected: time.Time{},
+		},
+		{
+			name:              "autostart before deadline by 1h",
+			autostartSchedule: "CRON_TZ=UTC 0 0 * * *",
+			deadline:          midnight.Add(time.Hour),
+			ttl:               time.Hour * 10,
+			expected:          midnight.Add(time.Hour * 10),
+		},
+		{
+			name:              "autostart before deadline by 9h",
+			autostartSchedule: "CRON_TZ=UTC 0 0 * * *",
+			deadline:          midnight.Add(time.Hour * 9),
+			ttl:               time.Hour * 10,
+			// should still be bumped
+			expected: midnight.Add(time.Hour * 10),
+		},
+		{
+			name:              "eligible but exceeds next next autostart",
+			autostartSchedule: "CRON_TZ=UTC 0 0 * * *",
+			deadline:          midnight.Add(time.Hour * 1),
+			// ttl causes next autostart + 25h to exceed the next next autostart
+			ttl: time.Hour * 25,
+			// should not be bumped to avoid infinite bumping every day
+			expected: time.Time{},
+		},
+		{
+			name:              "deadline is 1h before autostart",
+			autostartSchedule: "CRON_TZ=UTC 0 0 * * *",
+			deadline:          midnight.Add(time.Hour * -1).Add(time.Minute),
+			ttl:               time.Hour * 10,
+			// should still be bumped
+			expected: midnight.Add(time.Hour * 10),
+		},
+	}
+
+	for _, c := range cases {
+		c := c
+		t.Run(c.name, func(t *testing.T) {
+			t.Parallel()
+
+			var (
+				tickCh     = make(chan time.Time)
+				statsCh    = make(chan autobuild.Stats)
+				db, pubsub = dbtestutil.NewDB(t)
+
+				client = coderdtest.New(t, &coderdtest.Options{
+					AutobuildTicker:          tickCh,
+					AutobuildStats:           statsCh,
+					Database:                 db,
+					Pubsub:                   pubsub,
+					IncludeProvisionerDaemon: true,
+					TemplateScheduleStore: schedule.MockTemplateScheduleStore{
+						GetFn: func(_ context.Context, _ database.Store, _ uuid.UUID) (schedule.TemplateScheduleOptions, error) {
+							return schedule.TemplateScheduleOptions{
+								UserAutostartEnabled: true,
+								UserAutostopEnabled:  true,
+								DefaultTTL:           0,
+								AutostopRequirement:  schedule.TemplateAutostopRequirement{},
+							}, nil
+						},
+					},
+				})
+				workspace = mustProvisionWorkspace(t, client, func(cwr *codersdk.CreateWorkspaceRequest) {
+					cwr.AutostartSchedule = ptr.Ref(c.autostartSchedule)
+					cwr.TTLMillis = ptr.Ref(c.ttl.Milliseconds())
+				})
+			)
+
+			// Set the deadline and max deadline to what the test expects.
+			ctx := testutil.Context(t, testutil.WaitLong)
+			err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+				ID:          workspace.LatestBuild.ID,
+				UpdatedAt:   dbtime.Now(),
+				Deadline:    c.deadline,
+				MaxDeadline: c.maxDeadline,
+			})
+			require.NoError(t, err)
+
+			if c.now.IsZero() {
+				c.now = c.deadline.Add(-time.Hour)
+			}
+
+			go func() {
+				tickCh <- c.now
+				close(tickCh)
+			}()
+			stats := <-statsCh
+			require.NoError(t, stats.Error)
+			require.Len(t, stats.Transitions, 0)
+
+			// Get the latest workspace build.
+			build, err := client.WorkspaceBuild(ctx, workspace.LatestBuild.ID)
+			require.NoError(t, err)
+
+			// Should be bumped to the expected time.
+			if c.expected.IsZero() {
+				c.expected = c.deadline
+			}
+			require.WithinDuration(t, c.expected, build.Deadline.Time, time.Minute)
+		})
+	}
+}
+
 // TestExecutorInactiveWorkspace test AGPL functionality which mainly
 // ensures that autostop actions as a result of an inactive workspace
 // do not trigger.

diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -1915,8 +1915,8 @@ func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesP
 	return q.db.GetAuthorizedWorkspaces(ctx, arg, prep)
 }
 
-func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.Workspace, error) {
-	return q.db.GetWorkspacesEligibleForTransition(ctx, now)
+func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context) ([]database.Workspace, error) {
+	return q.db.GetWorkspacesEligibleForTransition(ctx)
 }
 
 func (q *querier) InsertAPIKey(ctx context.Context, arg database.InsertAPIKeyParams) (database.APIKey, error) {