coder · DanielleMaywood · Dec 6, 2024 · Dec 4, 2024 · Dec 5, 2024 · Dec 5, 2024
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -4138,6 +4138,17 @@ func (q *querier) UpdateWorkspacesDormantDeletingAtByTemplateID(ctx context.Cont
 	return q.db.UpdateWorkspacesDormantDeletingAtByTemplateID(ctx, arg)
 }
 
+func (q *querier) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg database.UpdateWorkspacesTTLByTemplateIDParams) error {
+	template, err := q.db.GetTemplateByID(ctx, arg.TemplateID)
+	if err != nil {
+		return xerrors.Errorf("get template by id: %w", err)
+	}
+	if err := q.authorizeContext(ctx, policy.ActionUpdate, template); err != nil {
+		return err
+	}
+	return q.db.UpdateWorkspacesTTLByTemplateID(ctx, arg)
+}
+
 func (q *querier) UpsertAnnouncementBanners(ctx context.Context, value string) error {
 	if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
 		return err

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1017,6 +1017,12 @@ func (s *MethodTestSuite) TestTemplate() {
 			TemplateID: t1.ID,
 		}).Asserts(t1, policy.ActionUpdate)
 	}))
+	s.Run("UpdateWorkspacesTTLByTemplateID", s.Subtest(func(db database.Store, check *expects) {
+		t1 := dbgen.Template(s.T(), db, database.Template{})
+		check.Args(database.UpdateWorkspacesTTLByTemplateIDParams{
+			TemplateID: t1.ID,
+		}).Asserts(t1, policy.ActionUpdate)
+	}))
 	s.Run("UpdateTemplateActiveVersionByID", s.Subtest(func(db database.Store, check *expects) {
 		t1 := dbgen.Template(s.T(), db, database.Template{
 			ActiveVersionID: uuid.New(),

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -10192,6 +10192,26 @@ func (q *FakeQuerier) UpdateWorkspacesDormantDeletingAtByTemplateID(_ context.Co
 	return affectedRows, nil
 }
 
+func (q *FakeQuerier) UpdateWorkspacesTTLByTemplateID(_ context.Context, arg database.UpdateWorkspacesTTLByTemplateIDParams) error {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	for i, ws := range q.workspaces {
+		if ws.TemplateID != arg.TemplateID {
+			continue
+		}
+
+		q.workspaces[i].Ttl = arg.Ttl
+	}
+
+	return nil
+}
+
 func (q *FakeQuerier) UpsertAnnouncementBanners(_ context.Context, data string) error {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -501,6 +501,14 @@ SET
 WHERE
 	id = $1;
 
+-- name: UpdateWorkspacesTTLByTemplateID :exec
+UPDATE
+		workspaces
+SET
+		ttl = $2
+WHERE
+		template_id = $1;
+
 -- name: UpdateWorkspaceLastUsedAt :exec
 UPDATE
 	workspaces

diff --git a/coderd/schedule/template.go b/coderd/schedule/template.go
@@ -2,6 +2,7 @@ package schedule
 
 import (
 	"context"
+	"database/sql"
 	"time"
 
 	"github.com/google/uuid"
@@ -228,6 +229,23 @@ func (*agplTemplateScheduleStore) Set(ctx context.Context, db database.Store, tp
 			return xerrors.Errorf("update template schedule: %w", err)
 		}
 
+		// Users running the AGPL version are unable to customize their workspaces
+		// autostop, so we want to keep their workspaces in track with any template
+		// TTL changes.
+		if tpl.DefaultTTL != int64(opts.DefaultTTL) {
+			var ttl sql.NullInt64
+			if opts.DefaultTTL != 0 {
+				ttl = sql.NullInt64{Valid: true, Int64: int64(opts.DefaultTTL)}
+			}
+
+			if err = db.UpdateWorkspacesTTLByTemplateID(ctx, database.UpdateWorkspacesTTLByTemplateIDParams{
+				TemplateID: tpl.ID,
+				Ttl:        ttl,
+			}); err != nil {
+				return xerrors.Errorf("update workspace ttl by template id %q: %w", tpl.ID, err)
+			}
+		}
+
 		template, err = db.GetTemplateByID(ctx, tpl.ID)
 		if err != nil {
 			return xerrors.Errorf("fetch updated template: %w", err)

diff --git a/coderd/schedule/template_test.go b/coderd/schedule/template_test.go
@@ -0,0 +1,150 @@
+package schedule_test
+
+import (
+	"database/sql"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"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/testutil"
+)
+
+func TestTemplateTTL(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		fromTTL  time.Duration
+		toTTL    time.Duration
+		expected sql.NullInt64
+	}{
+		{
+			name:     "ModifyTTLDurationDown",
+			fromTTL:  24 * time.Hour,
+			toTTL:    1 * time.Hour,
+			expected: sql.NullInt64{Valid: true, Int64: int64(1 * time.Hour)},
+		},
+		{
+			name:     "ModifyTTLDurationUp",
+			fromTTL:  24 * time.Hour,
+			toTTL:    36 * time.Hour,
+			expected: sql.NullInt64{Valid: true, Int64: int64(36 * time.Hour)},
+		},
+		{
+			name:     "ModifyTTLDurationSame",
+			fromTTL:  24 * time.Hour,
+			toTTL:    24 * time.Hour,
+			expected: sql.NullInt64{Valid: true, Int64: int64(24 * time.Hour)},
+		},
+		{
+			name:     "DisableTTL",
+			fromTTL:  24 * time.Hour,
+			toTTL:    0,
+			expected: sql.NullInt64{},
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			var (
+				db, _ = dbtestutil.NewDB(t)
+				ctx   = testutil.Context(t, testutil.WaitLong)
+				user  = dbgen.User(t, db, database.User{})
+				file  = dbgen.File(t, db, database.File{CreatedBy: user.ID})
+				// Create first template
+				templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+					FileID:      file.ID,
+					InitiatorID: user.ID,
+					Tags:        database.StringMap{"foo": "bar"},
+				})
+				templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
+					CreatedBy:      user.ID,
+					JobID:          templateJob.ID,
+					OrganizationID: templateJob.OrganizationID,
+				})
+				template = dbgen.Template(t, db, database.Template{
+					ActiveVersionID: templateVersion.ID,
+					CreatedBy:       user.ID,
+					OrganizationID:  templateJob.OrganizationID,
+				})
+				// Create second template
+				otherTTL         = tt.fromTTL + 6*time.Hour
+				otherTemplateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+					FileID:      file.ID,
+					InitiatorID: user.ID,
+					Tags:        database.StringMap{"foo": "bar"},
+				})
+				otherTemplateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
+					CreatedBy:      user.ID,
+					JobID:          otherTemplateJob.ID,
+					OrganizationID: otherTemplateJob.OrganizationID,
+				})
+				otherTemplate = dbgen.Template(t, db, database.Template{
+					ActiveVersionID: otherTemplateVersion.ID,
+					CreatedBy:       user.ID,
+					OrganizationID:  otherTemplateJob.OrganizationID,
+				})
+			)
+
+			templateScheduleStore := schedule.NewAGPLTemplateScheduleStore()
+
+			// Set both template's default TTL
+			template, err := templateScheduleStore.Set(ctx, db, template, schedule.TemplateScheduleOptions{
+				DefaultTTL: tt.fromTTL,
+			})
+			require.NoError(t, err)
+			otherTemplate, err = templateScheduleStore.Set(ctx, db, otherTemplate, schedule.TemplateScheduleOptions{
+				DefaultTTL: otherTTL,
+			})
+			require.NoError(t, err)
+
+			// We create two workspaces here, one with the template we're modifying, the
+			// other with a different template. We want to ensure we only modify one
+			// of the workspaces.
+			workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+				OwnerID:        user.ID,
+				TemplateID:     template.ID,
+				OrganizationID: templateJob.OrganizationID,
+				LastUsedAt:     dbtime.Now(),
+				Ttl:            sql.NullInt64{Valid: true, Int64: int64(tt.fromTTL)},
+			})
+			otherWorkspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+				OwnerID:        user.ID,
+				TemplateID:     otherTemplate.ID,
+				OrganizationID: otherTemplateJob.OrganizationID,
+				LastUsedAt:     dbtime.Now(),
+				Ttl:            sql.NullInt64{Valid: true, Int64: int64(otherTTL)},
+			})
+
+			// Ensure the workspace's start with the correct TTLs
+			require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(tt.fromTTL)}, workspace.Ttl)
+			require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(otherTTL)}, otherWorkspace.Ttl)
+
+			// Update _only_ the primary template's TTL
+			_, err = templateScheduleStore.Set(ctx, db, template, schedule.TemplateScheduleOptions{
+				DefaultTTL: tt.toTTL,
+			})
+			require.NoError(t, err)
+
+			// Verify the primary workspace's TTL has been updated.
+			ws, err := db.GetWorkspaceByID(ctx, workspace.ID)
+			require.NoError(t, err)
+			require.Equal(t, tt.expected, ws.Ttl)
+
+			// Verify that the other workspace's TTL has not been touched.
+			ws, err = db.GetWorkspaceByID(ctx, otherWorkspace.ID)
+			require.NoError(t, err)
+			require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(otherTTL)}, ws.Ttl)
+		})
+	}
+}
diff --git a/enterprise/coderd/schedule/template.go b/enterprise/coderd/schedule/template.go
@@ -195,6 +195,23 @@ func (s *EnterpriseTemplateScheduleStore) Set(ctx context.Context, db database.S
 			return xerrors.Errorf("get updated template schedule: %w", err)
 		}
 
+		// Update all workspace's TTL using this template if either of the following:
+		//   - The template's AllowUserAutostop has just been disabled
+		//   - The template's TTL has been modified and AllowUserAutostop is disabled
+		if !opts.UserAutostopEnabled && (tpl.AllowUserAutostop || int64(opts.DefaultTTL) != tpl.DefaultTTL) {
+			var ttl sql.NullInt64
+			if opts.DefaultTTL != 0 {
+				ttl = sql.NullInt64{Valid: true, Int64: int64(opts.DefaultTTL)}
+			}
+
+			if err = tx.UpdateWorkspacesTTLByTemplateID(ctx, database.UpdateWorkspacesTTLByTemplateIDParams{
+				TemplateID: template.ID,
+				Ttl:        ttl,
+			}); err != nil {
+				return xerrors.Errorf("update workspaces ttl by template id %q: %w", template.ID, err)
+			}
+		}
+
 		// Recalculate max_deadline and deadline for all running workspace
 		// builds on this template.
 		err = s.updateWorkspaceBuilds(ctx, tx, template)