Add unit tests, prevent deadline decrease

coder · Emyrk · Nov 7, 2023 · Nov 7, 2023 · Nov 8, 2023 · Nov 8, 2023
commit 56f0d90aeccff8b33e102d3ea78aa6db11c079cd
diff --git a/coderd/activitybump_internal_test.go b/coderd/activitybump_internal_test.go
@@ -39,7 +39,9 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 		buildDeadlineOffset           *time.Duration
 		maxDeadlineOffset             *time.Duration
 		workspaceTTL                  time.Duration
+		workspaceBumpTTL              time.Duration
 		templateTTL                   time.Duration
+		templateBumpTTL               time.Duration
 		templateDisallowsUserAutostop bool
 		expectedBump                  time.Duration
 	}{
@@ -83,6 +85,16 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 			workspaceTTL:        8 * time.Hour,
 			expectedBump:        1 * time.Hour,
 		},
+		{
+			name:                "MaxDeadlineCustomBump",
+			transition:          database.WorkspaceTransitionStart,
+			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
+			buildDeadlineOffset: ptr.Ref(time.Minute), // last chance to bump!
+			maxDeadlineOffset:   ptr.Ref(4 * time.Hour),
+			workspaceTTL:        8 * time.Hour,
+			workspaceBumpTTL:    5 * time.Hour,
+			expectedBump:        5 * time.Hour,
+		},
 		{
 			// A workspace that is still running, has passed its deadline, but has not
 			// yet been auto-stopped should still bump the deadline.
@@ -101,6 +113,17 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 			buildDeadlineOffset: ptr.Ref(-time.Minute),
 			workspaceTTL:        8 * time.Hour,
 		},
+		{
+			// Deadline bump TTL can be less than the time left in the
+			// deadline. It should be a no-op.
+			name:                "DeadlineNotDecreased",
+			transition:          database.WorkspaceTransitionStart,
+			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Hour)},
+			buildDeadlineOffset: ptr.Ref(time.Hour * 5),
+			workspaceTTL:        8 * time.Hour,
+			workspaceBumpTTL:    time.Hour,
+			expectedBump:        0,
+		},
 		{
 			// A workspace built from a template that disallows user autostop should bump
 			// by the template TTL instead.
@@ -113,6 +136,18 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 			templateDisallowsUserAutostop: true,
 			expectedBump:                  8 * time.Hour,
 		},
+		{
+			name:                          "TemplateDisallowsUserAutostopCustomBump",
+			transition:                    database.WorkspaceTransitionStart,
+			templateDisallowsUserAutostop: true,
+			jobCompletedAt:                sql.NullTime{Valid: true, Time: dbtime.Now().Add(-7 * time.Hour)},
+			buildDeadlineOffset:           ptr.Ref(time.Minute),
+			templateTTL:                   time.Hour * 8,
+			templateBumpTTL:               time.Hour,
+			workspaceTTL:                  24 * time.Hour,
+			workspaceBumpTTL:              10 * time.Hour,
+			expectedBump:                  time.Hour, // Expect smaller bump
+		},
 	} {
 		tt := tt
 		for _, tz := range timezones {
@@ -147,6 +182,7 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 						OrganizationID: org.ID,
 						TemplateID:     template.ID,
 						Ttl:            sql.NullInt64{Valid: true, Int64: int64(tt.workspaceTTL)},
+						TtlBump:        int64(tt.workspaceBumpTTL),
 					})
 					job = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
 						OrganizationID: org.ID,
@@ -163,6 +199,7 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 					UpdatedAt:         dbtime.Now(),
 					AllowUserAutostop: !tt.templateDisallowsUserAutostop,
 					DefaultTTL:        int64(tt.templateTTL),
+					DefaultTTLBump:    int64(tt.templateBumpTTL),
 				}), "unexpected error updating template schedule")
 
 				var buildNumber int32 = 1
@@ -207,6 +244,7 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 				require.Equal(t, buildDeadline.UTC(), bld.Deadline.UTC(), "unexpected build deadline before bump")
 				require.Equal(t, maxDeadline.UTC(), bld.MaxDeadline.UTC(), "unexpected max deadline before bump")
 				require.Equal(t, tt.workspaceTTL, time.Duration(ws.Ttl.Int64), "unexpected workspace TTL before bump")
+				require.Equal(t, tt.workspaceBumpTTL, time.Duration(ws.TtlBump), "unexpected workspace TTL bump")
 
 				// Wait a bit before bumping as dbtime is rounded to the nearest millisecond.
 				// This should also hopefully be enough for Windows time resolution to register
@@ -236,8 +274,8 @@ func Test_ActivityBumpWorkspace(t *testing.T) {
 				// Assert that the bump occurred between start and end.
 				expectedDeadlineStart := start.Add(tt.expectedBump)
 				expectedDeadlineEnd := end.Add(tt.expectedBump)
-				require.GreaterOrEqual(t, updatedBuild.Deadline, expectedDeadlineStart, "new deadline should be greater than or equal to start")
-				require.LessOrEqual(t, updatedBuild.Deadline, expectedDeadlineEnd, "new deadline should be lesser than or equal to end")
+				require.GreaterOrEqualf(t, updatedBuild.Deadline, expectedDeadlineStart, "new deadline should be greater than or equal to start, pre-bump: %s, diff_from_last_deadline: %s", buildDeadline.String(), updatedBuild.Deadline.Sub(buildDeadline))
+				require.LessOrEqualf(t, updatedBuild.Deadline, expectedDeadlineEnd, "new deadline should be lesser than or equal to end, pre-bump: %s, diff_from_last_deadline: %s", buildDeadline.String(), updatedBuild.Deadline.Sub(buildDeadline))
 			})
 		}
 	}

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -178,6 +178,7 @@ func Workspace(t testing.TB, db database.Store, orig database.Workspace) databas
 		Name:              takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
 		AutostartSchedule: orig.AutostartSchedule,
 		Ttl:               orig.Ttl,
+		TtlBump:           orig.TtlBump,
 		AutomaticUpdates:  takeFirst(orig.AutomaticUpdates, database.AutomaticUpdatesNever),
 	})
 	require.NoError(t, err, "insert workspace")

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -826,13 +826,13 @@ func (q *FakeQuerier) ActivityBumpWorkspace(ctx context.Context, workspaceID uui
 		if workspace.Ttl.Valid {
 			ttlDur = time.Duration(workspace.Ttl.Int64)
 			if workspace.TtlBump > 0 {
-				ttlDur += time.Duration(workspace.TtlBump)
+				ttlDur = time.Duration(workspace.TtlBump)
 			}
 		}
 		if !template.AllowUserAutostop {
 			ttlDur = time.Duration(template.DefaultTTL)
 			if template.DefaultTTLBump > 0 {
-				ttlDur += time.Duration(template.DefaultTTLBump)
+				ttlDur = time.Duration(template.DefaultTTLBump)
 			}
 		}
 		if ttlDur <= 0 {
@@ -850,6 +850,10 @@ func (q *FakeQuerier) ActivityBumpWorkspace(ctx context.Context, workspaceID uui
 		// Bump.
 		newDeadline := now.Add(ttlDur)
 		q.workspaceBuilds[i].UpdatedAt = now
+		if newDeadline.Before(q.workspaceBuilds[i].Deadline) {
+			// Never shorten a deadline
+			newDeadline = q.workspaceBuilds[i].Deadline
+		}
 		if !q.workspaceBuilds[i].MaxDeadline.IsZero() {
 			q.workspaceBuilds[i].Deadline = minTime(newDeadline, q.workspaceBuilds[i].MaxDeadline)
 		} else {

diff --git a/coderd/database/queries/activitybump.sql b/coderd/database/queries/activitybump.sql
@@ -42,11 +42,12 @@ UPDATE
 	workspace_builds wb
 SET
 	updated_at = NOW(),
-	deadline = CASE
+	-- Never shorten a deadline with an activity bump.
+	deadline = MAX(wb.deadline, CASE
 		WHEN l.build_max_deadline = '0001-01-01 00:00:00+00'
 		THEN NOW() + l.ttl_interval
 		ELSE LEAST(NOW() + l.ttl_interval, l.build_max_deadline)
-	END
+	END)
 FROM latest l
 WHERE wb.id = l.build_id
 AND l.job_completed_at IS NOT NULL

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -780,11 +780,6 @@ func (api *API) putWorkspaceTTL(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	if req.TTLBumpMillis == nil {
-		// Always default to the existing value if the request omits.
-		req.TTLBumpMillis = &workspace.TtlBump
-	}
-
 	var dbTTL sql.NullInt64
 	var dbTTLBump time.Duration
 
@@ -811,7 +806,9 @@ func (api *API) putWorkspaceTTL(rw http.ResponseWriter, r *http.Request) {
 			return codersdk.ValidationError{Field: "ttl_ms", Detail: validityErr.Error()}
 		}
 
-		dbTTLBump, validityErr = validWorkspaceTTLBumpMillis(req.TTLBumpMillis, templateSchedule.DefaultTTLBump)
+		// The TTL bump defaults to 0, which means the bump amount is the
+		// TTL.
+		dbTTLBump, validityErr = validWorkspaceTTLBumpMillis(req.TTLBumpMillis, 0)
 		if validityErr != nil {
 			return codersdk.ValidationError{Field: "ttl_bump_ms", Detail: validityErr.Error()}
 		}

diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go
@@ -1807,21 +1807,29 @@ func TestWorkspaceUpdateTTL(t *testing.T) {
 	testCases := []struct {
 		name           string
 		ttlMillis      *int64
+		ttlBumpMillis  *int64
 		expectedError  string
 		modifyTemplate func(*codersdk.CreateTemplateRequest)
+
+		expectedTTLMillis     *int64
+		expectedTTLBumpMillis int64
 	}{
 		{
 			name:          "disable ttl",
 			ttlMillis:     nil,
 			expectedError: "",
 			modifyTemplate: func(ctr *codersdk.CreateTemplateRequest) {
 				ctr.DefaultTTLMillis = ptr.Ref((8 * time.Hour).Milliseconds())
+				ctr.DefaultTTLBumpMillis = ptr.Ref((5 * time.Minute).Milliseconds())
 			},
+			expectedTTLMillis:     nil,
+			expectedTTLBumpMillis: 0,
 		},
 		{
-			name:          "update ttl",
-			ttlMillis:     ptr.Ref(12 * time.Hour.Milliseconds()),
-			expectedError: "",
+			name:              "update ttl",
+			ttlMillis:         ptr.Ref(12 * time.Hour.Milliseconds()),
+			expectedTTLMillis: ptr.Ref(12 * time.Hour.Milliseconds()),
+			expectedError:     "",
 			modifyTemplate: func(ctr *codersdk.CreateTemplateRequest) {
 				ctr.DefaultTTLMillis = ptr.Ref((8 * time.Hour).Milliseconds())
 			},
@@ -1832,20 +1840,34 @@ func TestWorkspaceUpdateTTL(t *testing.T) {
 			expectedError: "time until shutdown must be at least one minute",
 		},
 		{
-			name:          "minimum ttl",
-			ttlMillis:     ptr.Ref(time.Minute.Milliseconds()),
-			expectedError: "",
+			name:              "minimum ttl",
+			ttlMillis:         ptr.Ref(time.Minute.Milliseconds()),
+			expectedTTLMillis: ptr.Ref(time.Minute.Milliseconds()),
+			expectedError:     "",
 		},
 		{
-			name:          "maximum ttl",
-			ttlMillis:     ptr.Ref((24 * 30 * time.Hour).Milliseconds()),
-			expectedError: "",
+			name:              "maximum ttl",
+			ttlMillis:         ptr.Ref((24 * 30 * time.Hour).Milliseconds()),
+			expectedTTLMillis: ptr.Ref((24 * 30 * time.Hour).Milliseconds()),
+			expectedError:     "",
 		},
 		{
 			name:          "above maximum ttl",
 			ttlMillis:     ptr.Ref((24*30*time.Hour + time.Minute).Milliseconds()),
 			expectedError: "time until shutdown must be less than 30 days",
 		},
+		{
+			name:          "update ttl bump",
+			ttlMillis:     ptr.Ref(12 * time.Hour.Milliseconds()),
+			ttlBumpMillis: ptr.Ref(12 * time.Minute.Milliseconds()),
+			expectedError: "",
+			modifyTemplate: func(request *codersdk.CreateTemplateRequest) {
+				request.DefaultTTLMillis = ptr.Ref((8 * time.Hour).Milliseconds())
+				request.DefaultTTLBumpMillis = ptr.Ref((5 * time.Minute).Milliseconds())
+			},
+			expectedTTLMillis:     ptr.Ref(12 * time.Hour.Milliseconds()),
+			expectedTTLBumpMillis: 12 * time.Minute.Milliseconds(),
+		},
 	}
 
 	for _, testCase := range testCases {
@@ -1875,7 +1897,8 @@ func TestWorkspaceUpdateTTL(t *testing.T) {
 			defer cancel()
 
 			err := client.UpdateWorkspaceTTL(ctx, workspace.ID, codersdk.UpdateWorkspaceTTLRequest{
-				TTLMillis: testCase.ttlMillis,
+				TTLMillis:     testCase.ttlMillis,
+				TTLBumpMillis: testCase.ttlBumpMillis,
 			})
 
 			if testCase.expectedError != "" {
@@ -1888,7 +1911,8 @@ func TestWorkspaceUpdateTTL(t *testing.T) {
 			updated, err := client.Workspace(ctx, workspace.ID)
 			require.NoError(t, err, "fetch updated workspace")
 
-			require.Equal(t, testCase.ttlMillis, updated.TTLMillis, "expected autostop ttl to equal requested")
+			require.EqualValues(t, testCase.expectedTTLMillis, updated.TTLMillis, "expected autostop ttl to equal requested")
+			require.Equal(t, testCase.expectedTTLBumpMillis, updated.TTLBumpMillis, "expected autostop ttl bump to equal requested")
 
 			require.Eventually(t, func() bool {
 				if len(auditor.AuditLogs()) != 7 {