fix(coderd): make activitybump aware of default template ttl (coder#10253)

johnstcn · deansheather · web-flow · commit d56f49f619b7 · 2023-10-13T13:53:02.000+01:00
The refactored ActivityBump query did not take into account the
template-level TTL, resulting in potentially incorrect bump
amounts for workspaces that have both a user-defined and template-
defined TTL that differ.

This change is ported over from PR#10035 to reduce the overall
size of that PR.

Also includes a drive-by unit test in autobuild for checking template autostop/TTL.

Co-authored-by: Dean Sheather &lt;dean@deansheather.com&gt;
diff --git a/coderd/activitybump_internal_test.go b/coderd/activitybump_internal_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/testutil"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -32,13 +33,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
+		templateTTL                   time.Duration
+		templateDisallowsUserAutostop bool
+		expectedBump                  time.Duration
 	}{
 		{
 			name:                "NotFinishedYet",
@@ -97,7 +100,18 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-time.Minute)},
 			buildDeadlineOffset: ptr.Ref(-time.Minute),
 			workspaceTTL:        8 * time.Hour,
-			expectedBump:        0,
+		},
+		{
+			// A workspace built from a template that disallows user autostop should bump
+			// by the template TTL instead.
+			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,
+			templateTTL:                   8 * time.Hour,
+			templateDisallowsUserAutostop: true,
+			expectedBump:                  8 * time.Hour,
 		},
 	} {
 		tt := tt
@@ -144,6 +158,13 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 					buildID = uuid.New()
 				)
 
+				require.NoError(t, db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{
+					ID:                template.ID,
+					UpdatedAt:         dbtime.Now(),
+					AllowUserAutostop: !tt.templateDisallowsUserAutostop,
+					DefaultTTL:        int64(tt.templateTTL),
+				}), "unexpected error updating template schedule")
+
 				var buildNumber int32 = 1
 				// Insert a number of previous workspace builds.
 				for i := 0; i < 5; i++ {
@@ -202,13 +223,13 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 				require.NoError(t, err, "unexpected error getting latest workspace build")
 				require.Equal(t, bld.MaxDeadline.UTC(), updatedBuild.MaxDeadline.UTC(), "max_deadline should not have changed")
 				if tt.expectedBump == 0 {
-					require.Equal(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should not have bumped updated_at")
-					require.Equal(t, bld.Deadline.UTC(), updatedBuild.Deadline.UTC(), "should not have bumped deadline")
+					assert.Equal(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should not have bumped updated_at")
+					assert.Equal(t, bld.Deadline.UTC(), updatedBuild.Deadline.UTC(), "should not have bumped deadline")
 					return
 				}
-				require.NotEqual(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should have bumped updated_at")
+				assert.NotEqual(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should have bumped updated_at")
 				if tt.maxDeadlineOffset != nil {
-					require.Equal(t, bld.MaxDeadline.UTC(), updatedBuild.MaxDeadline.UTC(), "new deadline must equal original max deadline")
+					assert.Equal(t, bld.MaxDeadline.UTC(), updatedBuild.MaxDeadline.UTC(), "new deadline must equal original max deadline")
 					return
 				}
 
diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go
@@ -731,6 +731,51 @@ func TestExecutorAutostartTemplateDisabled(t *testing.T) {
 	assert.Len(t, stats.Transitions, 0)
 }
 
+func TestExecutorAutostopTemplateDisabled(t *testing.T) {
+	t.Parallel()
+
+	// Given: we have a workspace built from a template that disallows user autostop
+	var (
+		sched   = mustSchedule(t, "CRON_TZ=UTC 0 * * * *")
+		tickCh  = make(chan time.Time)
+		statsCh = make(chan autobuild.Stats)
+
+		client = coderdtest.New(t, &coderdtest.Options{
+			AutobuildTicker:          tickCh,
+			IncludeProvisionerDaemon: true,
+			AutobuildStats:           statsCh,
+			// We are using a mock store here as the AGPL store does not implement this.
+			TemplateScheduleStore: schedule.MockTemplateScheduleStore{
+				GetFn: func(_ context.Context, _ database.Store, _ uuid.UUID) (schedule.TemplateScheduleOptions, error) {
+					return schedule.TemplateScheduleOptions{
+						UserAutostopEnabled: false,
+						DefaultTTL:          time.Hour,
+					}, nil
+				},
+			},
+		})
+		// Given: we have a user with a workspace configured to autostart some time in the future
+		workspace = mustProvisionWorkspace(t, client, func(cwr *codersdk.CreateWorkspaceRequest) {
+			cwr.TTLMillis = ptr.Ref(8 * time.Hour.Milliseconds())
+		})
+	)
+
+	// When: we create the workspace
+	// Then: the deadline should be set to the template default TTL
+	assert.WithinDuration(t, workspace.LatestBuild.CreatedAt.Add(time.Hour), workspace.LatestBuild.Deadline.Time, time.Minute)
+
+	// When: the autobuild executor ticks before the next scheduled time
+	go func() {
+		tickCh <- sched.Next(workspace.LatestBuild.CreatedAt).Add(time.Minute)
+		close(tickCh)
+	}()
+
+	// Then: nothing should happen
+	stats := <-statsCh
+	assert.NoError(t, stats.Error)
+	assert.Len(t, stats.Transitions, 0)
+}
+
 // TestExecutorFailedWorkspace test AGPL functionality which mainly
 // ensures that autostop actions as a result of a failed workspace
 // build do not trigger.
diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -814,8 +814,26 @@ func (q *FakeQuerier) ActivityBumpWorkspace(ctx context.Context, workspaceID uui
 		if q.workspaceBuilds[i].Deadline.IsZero() {
 			return nil
 		}
+
+		// Check the template default TTL.
+		template, err := q.getTemplateByIDNoLock(ctx, workspace.TemplateID)
+		if err != nil {
+			return err
+		}
+
+		var ttlDur time.Duration
+		if workspace.Ttl.Valid {
+			ttlDur = time.Duration(workspace.Ttl.Int64)
+		}
+		if !template.AllowUserAutostop {
+			ttlDur = time.Duration(template.DefaultTTL)
+		}
+		if ttlDur <= 0 {
+			// There's no TTL set anymore, so we don't know the bump duration.
+			return nil
+		}
+
 		// Only bump if 5% of the deadline has passed.
-		ttlDur := time.Duration(workspace.Ttl.Int64)
 		ttlDur95 := ttlDur - (ttlDur / 20)
 		minBumpDeadline := q.workspaceBuilds[i].Deadline.Add(-ttlDur95)
 		if now.Before(minBumpDeadline) {
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/activitybump.sql b/coderd/database/queries/activitybump.sql
@@ -10,12 +10,20 @@ WITH latest AS (
 		workspace_builds.max_deadline::timestamp with time zone AS build_max_deadline,
 		workspace_builds.transition AS build_transition,
 		provisioner_jobs.completed_at::timestamp with time zone AS job_completed_at,
-		(workspaces.ttl / 1000 / 1000 / 1000 || ' seconds')::interval AS ttl_interval
+		(
+			CASE
+				WHEN templates.allow_user_autostop
+				THEN (workspaces.ttl / 1000 / 1000 / 1000 || ' seconds')::interval
+				ELSE (templates.default_ttl / 1000 / 1000 / 1000 || ' seconds')::interval
+			END
+		) AS ttl_interval
 	FROM workspace_builds
 	JOIN provisioner_jobs
 		ON provisioner_jobs.id = workspace_builds.job_id
 	JOIN workspaces
 		ON workspaces.id = workspace_builds.workspace_id
+	JOIN templates
+		ON templates.id = workspaces.template_id
 	WHERE workspace_builds.workspace_id = @workspace_id::uuid
 	ORDER BY workspace_builds.build_number DESC
 	LIMIT 1
@@ -33,6 +41,8 @@ FROM latest l
 WHERE wb.id = l.build_id
 AND l.job_completed_at IS NOT NULL
 AND l.build_transition = 'start'
+-- We only bump if the raw interval is positive and non-zero.
+AND l.ttl_interval > '0 seconds'::interval
 -- We only bump if workspace shutdown is manual.
 AND l.build_deadline != '0001-01-01 00:00:00+00'
 -- We only bump when 5% of the deadline has elapsed.
