fix(coderd/database): avoid clobbering workspace build state (#9826)

Fixes #9823.

- Decomposes UpdateWorkspaceBuildByID into UpdateWorkspaceBuildProvisionerStateByID and UpdateWorkspaceBuildDeadlineByID.
- Replaces existing invocations of UpdateWorkspaceBuildByID with the newer queries where applicable.
- Modifies GetActiveWorkspaceBuildsByTemplateID to not return incomplete workspace builds.

coderd/activitybump_test.go

@@ -77,12 +77,11 @@ func TestWorkspaceActivityBump(t *testing.T) {
 			dbBuild, err := db.GetWorkspaceBuildByID(ctx, workspace.LatestBuild.ID)
 			require.NoError(t, err)
 
-			err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
-				ID:               workspace.LatestBuild.ID,
-				UpdatedAt:        dbtime.Now(),
-				ProvisionerState: dbBuild.ProvisionerState,
-				Deadline:         dbBuild.Deadline,
-				MaxDeadline:      dbtime.Now().Add(maxTTL),
+			err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+				ID:          workspace.LatestBuild.ID,
+				UpdatedAt:   dbtime.Now(),
+				Deadline:    dbBuild.Deadline,
+				MaxDeadline: dbtime.Now().Add(maxTTL),
 			})
 			require.NoError(t, err)
 		}

coderd/database/dbauthz/dbauthz.go

@@ -2675,7 +2675,15 @@ func (q *querier) UpdateWorkspaceAutostart(ctx context.Context, arg database.Upd
 	return update(q.log, q.auth, fetch, q.db.UpdateWorkspaceAutostart)(ctx, arg)
 }
 
-func (q *querier) UpdateWorkspaceBuildByID(ctx context.Context, arg database.UpdateWorkspaceBuildByIDParams) error {
+// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
+func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
+	if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceSystem); err != nil {
+		return err
+	}
+	return q.db.UpdateWorkspaceBuildCostByID(ctx, arg)
+}
+
+func (q *querier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg database.UpdateWorkspaceBuildDeadlineByIDParams) error {
 	build, err := q.db.GetWorkspaceBuildByID(ctx, arg.ID)
 	if err != nil {
 		return err
@@ -2685,20 +2693,19 @@ func (q *querier) UpdateWorkspaceBuildByID(ctx context.Context, arg database.Upd
 	if err != nil {
 		return err
 	}
+
 	err = q.authorizeContext(ctx, rbac.ActionUpdate, workspace.RBACObject())
 	if err != nil {
 		return err
 	}
-
-	return q.db.UpdateWorkspaceBuildByID(ctx, arg)
+	return q.db.UpdateWorkspaceBuildDeadlineByID(ctx, arg)
 }
 
-// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
-func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
+func (q *querier) UpdateWorkspaceBuildProvisionerStateByID(ctx context.Context, arg database.UpdateWorkspaceBuildProvisionerStateByIDParams) error {
 	if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceSystem); err != nil {
 		return err
 	}
-	return q.db.UpdateWorkspaceBuildCostByID(ctx, arg)
+	return q.db.UpdateWorkspaceBuildProvisionerStateByID(ctx, arg)
 }
 
 // Deprecated: Use SoftDeleteWorkspaceByID

coderd/database/dbauthz/dbauthz_test.go

@@ -1232,14 +1232,13 @@ func (s *MethodTestSuite) TestWorkspace() {
 			ID: ws.ID,
 		}).Asserts(ws, rbac.ActionUpdate).Returns()
 	}))
-	s.Run("UpdateWorkspaceBuildByID", s.Subtest(func(db database.Store, check *expects) {
+	s.Run("UpdateWorkspaceBuildDeadlineByID", s.Subtest(func(db database.Store, check *expects) {
 		ws := dbgen.Workspace(s.T(), db, database.Workspace{})
 		build := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: uuid.New()})
-		check.Args(database.UpdateWorkspaceBuildByIDParams{
-			ID:               build.ID,
-			UpdatedAt:        build.UpdatedAt,
-			Deadline:         build.Deadline,
-			ProvisionerState: []byte{},
+		check.Args(database.UpdateWorkspaceBuildDeadlineByIDParams{
+			ID:        build.ID,
+			UpdatedAt: build.UpdatedAt,
+			Deadline:  build.Deadline,
 		}).Asserts(ws, rbac.ActionUpdate)
 	}))
 	s.Run("SoftDeleteWorkspaceByID", s.Subtest(func(db database.Store, check *expects) {
@@ -1378,6 +1377,14 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 			DailyCost: 10,
 		}).Asserts(rbac.ResourceSystem, rbac.ActionUpdate)
 	}))
+	s.Run("UpdateWorkspaceBuildProvisionerStateByID", s.Subtest(func(db database.Store, check *expects) {
+		ws := dbgen.Workspace(s.T(), db, database.Workspace{})
+		build := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: uuid.New()})
+		check.Args(database.UpdateWorkspaceBuildProvisionerStateByIDParams{
+			ID:               build.ID,
+			ProvisionerState: []byte("testing"),
+		}).Asserts(rbac.ResourceSystem, rbac.ActionUpdate)
+	}))
 	s.Run("UpsertLastUpdateCheck", s.Subtest(func(db database.Store, check *expects) {
 		check.Args("value").Asserts(rbac.ResourceSystem, rbac.ActionUpdate)
 	}))

coderd/database/dbfake/dbfake.go

@@ -5854,28 +5854,6 @@ func (q *FakeQuerier) UpdateWorkspaceAutostart(_ context.Context, arg database.U
 	return sql.ErrNoRows
 }
 
-func (q *FakeQuerier) UpdateWorkspaceBuildByID(_ context.Context, arg database.UpdateWorkspaceBuildByIDParams) error {
-	if err := validateDatabaseType(arg); err != nil {
-		return err
-	}
-
-	q.mutex.Lock()
-	defer q.mutex.Unlock()
-
-	for index, workspaceBuild := range q.workspaceBuilds {
-		if workspaceBuild.ID != arg.ID {
-			continue
-		}
-		workspaceBuild.UpdatedAt = arg.UpdatedAt
-		workspaceBuild.ProvisionerState = arg.ProvisionerState
-		workspaceBuild.Deadline = arg.Deadline
-		workspaceBuild.MaxDeadline = arg.MaxDeadline
-		q.workspaceBuilds[index] = workspaceBuild
-		return nil
-	}
-	return sql.ErrNoRows
-}
-
 func (q *FakeQuerier) UpdateWorkspaceBuildCostByID(_ context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
 	if err := validateDatabaseType(arg); err != nil {
 		return err
@@ -5895,6 +5873,51 @@ func (q *FakeQuerier) UpdateWorkspaceBuildCostByID(_ context.Context, arg databa
 	return sql.ErrNoRows
 }
 
+func (q *FakeQuerier) UpdateWorkspaceBuildDeadlineByID(_ context.Context, arg database.UpdateWorkspaceBuildDeadlineByIDParams) error {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	for idx, build := range q.workspaceBuilds {
+		if build.ID != arg.ID {
+			continue
+		}
+		build.Deadline = arg.Deadline
+		build.MaxDeadline = arg.MaxDeadline
+		build.UpdatedAt = arg.UpdatedAt
+		q.workspaceBuilds[idx] = build
+		return nil
+	}
+
+	return sql.ErrNoRows
+}
+
+func (q *FakeQuerier) UpdateWorkspaceBuildProvisionerStateByID(_ context.Context, arg database.UpdateWorkspaceBuildProvisionerStateByIDParams) error {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	for idx, build := range q.workspaceBuilds {
+		if build.ID != arg.ID {
+			continue
+		}
+		build.ProvisionerState = arg.ProvisionerState
+		build.UpdatedAt = arg.UpdatedAt
+		q.workspaceBuilds[idx] = build
+		return nil
+	}
+
+	return sql.ErrNoRows
+}
+
 func (q *FakeQuerier) UpdateWorkspaceDeletedByID(_ context.Context, arg database.UpdateWorkspaceDeletedByIDParams) error {
 	if err := validateDatabaseType(arg); err != nil {
 		return err

coderd/database/dbmetrics/dbmetrics.go

@@ -1649,13 +1649,6 @@ func (m metricsStore) UpdateWorkspaceAutostart(ctx context.Context, arg database
 	return err
 }
 
-func (m metricsStore) UpdateWorkspaceBuildByID(ctx context.Context, arg database.UpdateWorkspaceBuildByIDParams) error {
-	start := time.Now()
-	err := m.s.UpdateWorkspaceBuildByID(ctx, arg)
-	m.queryLatencies.WithLabelValues("UpdateWorkspaceBuildByID").Observe(time.Since(start).Seconds())
-	return err
-}
-
 func (m metricsStore) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
 	start := time.Now()
 	err := m.s.UpdateWorkspaceBuildCostByID(ctx, arg)
@@ -1663,6 +1656,20 @@ func (m metricsStore) UpdateWorkspaceBuildCostByID(ctx context.Context, arg data
 	return err
 }
 
+func (m metricsStore) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg database.UpdateWorkspaceBuildDeadlineByIDParams) error {
+	start := time.Now()
+	r0 := m.s.UpdateWorkspaceBuildDeadlineByID(ctx, arg)
+	m.queryLatencies.WithLabelValues("UpdateWorkspaceBuildDeadlineByID").Observe(time.Since(start).Seconds())
+	return r0
+}
+
+func (m metricsStore) UpdateWorkspaceBuildProvisionerStateByID(ctx context.Context, arg database.UpdateWorkspaceBuildProvisionerStateByIDParams) error {
+	start := time.Now()
+	r0 := m.s.UpdateWorkspaceBuildProvisionerStateByID(ctx, arg)
+	m.queryLatencies.WithLabelValues("UpdateWorkspaceBuildProvisionerStateByID").Observe(time.Since(start).Seconds())
+	return r0
+}
+
 func (m metricsStore) UpdateWorkspaceDeletedByID(ctx context.Context, arg database.UpdateWorkspaceDeletedByIDParams) error {
 	start := time.Now()
 	err := m.s.UpdateWorkspaceDeletedByID(ctx, arg)

coderd/database/dbmock/dbmock.go

@@ -3467,20 +3467,6 @@ func (mr *MockStoreMockRecorder) UpdateWorkspaceAutostart(arg0, arg1 interface{}
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspaceAutostart", reflect.TypeOf((*MockStore)(nil).UpdateWorkspaceAutostart), arg0, arg1)
 }
 
-// UpdateWorkspaceBuildByID mocks base method.
-func (m *MockStore) UpdateWorkspaceBuildByID(arg0 context.Context, arg1 database.UpdateWorkspaceBuildByIDParams) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "UpdateWorkspaceBuildByID", arg0, arg1)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// UpdateWorkspaceBuildByID indicates an expected call of UpdateWorkspaceBuildByID.
-func (mr *MockStoreMockRecorder) UpdateWorkspaceBuildByID(arg0, arg1 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspaceBuildByID", reflect.TypeOf((*MockStore)(nil).UpdateWorkspaceBuildByID), arg0, arg1)
-}
-
 // UpdateWorkspaceBuildCostByID mocks base method.
 func (m *MockStore) UpdateWorkspaceBuildCostByID(arg0 context.Context, arg1 database.UpdateWorkspaceBuildCostByIDParams) error {
 	m.ctrl.T.Helper()
@@ -3495,6 +3481,34 @@ func (mr *MockStoreMockRecorder) UpdateWorkspaceBuildCostByID(arg0, arg1 interfa
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspaceBuildCostByID", reflect.TypeOf((*MockStore)(nil).UpdateWorkspaceBuildCostByID), arg0, arg1)
 }
 
+// UpdateWorkspaceBuildDeadlineByID mocks base method.
+func (m *MockStore) UpdateWorkspaceBuildDeadlineByID(arg0 context.Context, arg1 database.UpdateWorkspaceBuildDeadlineByIDParams) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateWorkspaceBuildDeadlineByID", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateWorkspaceBuildDeadlineByID indicates an expected call of UpdateWorkspaceBuildDeadlineByID.
+func (mr *MockStoreMockRecorder) UpdateWorkspaceBuildDeadlineByID(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspaceBuildDeadlineByID", reflect.TypeOf((*MockStore)(nil).UpdateWorkspaceBuildDeadlineByID), arg0, arg1)
+}
+
+// UpdateWorkspaceBuildProvisionerStateByID mocks base method.
+func (m *MockStore) UpdateWorkspaceBuildProvisionerStateByID(arg0 context.Context, arg1 database.UpdateWorkspaceBuildProvisionerStateByIDParams) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateWorkspaceBuildProvisionerStateByID", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateWorkspaceBuildProvisionerStateByID indicates an expected call of UpdateWorkspaceBuildProvisionerStateByID.
+func (mr *MockStoreMockRecorder) UpdateWorkspaceBuildProvisionerStateByID(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspaceBuildProvisionerStateByID", reflect.TypeOf((*MockStore)(nil).UpdateWorkspaceBuildProvisionerStateByID), arg0, arg1)
+}
+
 // UpdateWorkspaceDeletedByID mocks base method.
 func (m *MockStore) UpdateWorkspaceDeletedByID(arg0 context.Context, arg1 database.UpdateWorkspaceDeletedByIDParams) error {
 	m.ctrl.T.Helper()

coderd/database/querier.go

@@ -305,8 +305,9 @@ type sqlcQuerier interface {
 	UpdateWorkspaceAgentStartupByID(ctx context.Context, arg UpdateWorkspaceAgentStartupByIDParams) error
 	UpdateWorkspaceAppHealthByID(ctx context.Context, arg UpdateWorkspaceAppHealthByIDParams) error
 	UpdateWorkspaceAutostart(ctx context.Context, arg UpdateWorkspaceAutostartParams) error
-	UpdateWorkspaceBuildByID(ctx context.Context, arg UpdateWorkspaceBuildByIDParams) error
 	UpdateWorkspaceBuildCostByID(ctx context.Context, arg UpdateWorkspaceBuildCostByIDParams) error
+	UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg UpdateWorkspaceBuildDeadlineByIDParams) error
+	UpdateWorkspaceBuildProvisionerStateByID(ctx context.Context, arg UpdateWorkspaceBuildProvisionerStateByIDParams) error
 	UpdateWorkspaceDeletedByID(ctx context.Context, arg UpdateWorkspaceDeletedByIDParams) error
 	UpdateWorkspaceDormantDeletingAt(ctx context.Context, arg UpdateWorkspaceDormantDeletingAtParams) (Workspace, error)
 	UpdateWorkspaceLastUsedAt(ctx context.Context, arg UpdateWorkspaceLastUsedAtParams) error

coderd/database/queries.sql.go

@@ -8489,8 +8489,13 @@ FROM (
 JOIN
 	workspace_build_with_user AS wb
 	ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number
+JOIN
+	provisioner_jobs AS pj
+	ON wb.job_id = pj.id
 WHERE
 	wb.transition = 'start'::workspace_transition
+AND
+	pj.completed_at IS NOT NULL
 `
 
 func (q *sqlQuerier) GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceBuild, error) {
@@ -8979,37 +8984,6 @@ func (q *sqlQuerier) InsertWorkspaceBuild(ctx context.Context, arg InsertWorkspa
 	return err
 }
 
-const updateWorkspaceBuildByID = `-- name: UpdateWorkspaceBuildByID :exec
-UPDATE
-	workspace_builds
-SET
-	updated_at = $2,
-	provisioner_state = $3,
-	deadline = $4,
-	max_deadline = $5
-WHERE
-	id = $1
-`
-
-type UpdateWorkspaceBuildByIDParams struct {
-	ID               uuid.UUID `db:"id" json:"id"`
-	UpdatedAt        time.Time `db:"updated_at" json:"updated_at"`
-	ProvisionerState []byte    `db:"provisioner_state" json:"provisioner_state"`
-	Deadline         time.Time `db:"deadline" json:"deadline"`
-	MaxDeadline      time.Time `db:"max_deadline" json:"max_deadline"`
-}
-
-func (q *sqlQuerier) UpdateWorkspaceBuildByID(ctx context.Context, arg UpdateWorkspaceBuildByIDParams) error {
-	_, err := q.db.ExecContext(ctx, updateWorkspaceBuildByID,
-		arg.ID,
-		arg.UpdatedAt,
-		arg.ProvisionerState,
-		arg.Deadline,
-		arg.MaxDeadline,
-	)
-	return err
-}
-
 const updateWorkspaceBuildCostByID = `-- name: UpdateWorkspaceBuildCostByID :exec
 UPDATE
 	workspace_builds
@@ -9029,6 +9003,53 @@ func (q *sqlQuerier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg Updat
 	return err
 }
 
+const updateWorkspaceBuildDeadlineByID = `-- name: UpdateWorkspaceBuildDeadlineByID :exec
+UPDATE
+	workspace_builds
+SET
+	deadline = $1::timestamptz,
+	max_deadline = $2::timestamptz,
+	updated_at = $3::timestamptz
+WHERE id = $4::uuid
+`
+
+type UpdateWorkspaceBuildDeadlineByIDParams struct {
+	Deadline    time.Time `db:"deadline" json:"deadline"`
+	MaxDeadline time.Time `db:"max_deadline" json:"max_deadline"`
+	UpdatedAt   time.Time `db:"updated_at" json:"updated_at"`
+	ID          uuid.UUID `db:"id" json:"id"`
+}
+
+func (q *sqlQuerier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg UpdateWorkspaceBuildDeadlineByIDParams) error {
+	_, err := q.db.ExecContext(ctx, updateWorkspaceBuildDeadlineByID,
+		arg.Deadline,
+		arg.MaxDeadline,
+		arg.UpdatedAt,
+		arg.ID,
+	)
+	return err
+}
+
+const updateWorkspaceBuildProvisionerStateByID = `-- name: UpdateWorkspaceBuildProvisionerStateByID :exec
+UPDATE
+	workspace_builds
+SET
+	provisioner_state = $1::bytea,
+	updated_at = $2::timestamptz
+WHERE id = $3::uuid
+`
+
+type UpdateWorkspaceBuildProvisionerStateByIDParams struct {
+	ProvisionerState []byte    `db:"provisioner_state" json:"provisioner_state"`
+	UpdatedAt        time.Time `db:"updated_at" json:"updated_at"`
+	ID               uuid.UUID `db:"id" json:"id"`
+}
+
+func (q *sqlQuerier) UpdateWorkspaceBuildProvisionerStateByID(ctx context.Context, arg UpdateWorkspaceBuildProvisionerStateByIDParams) error {
+	_, err := q.db.ExecContext(ctx, updateWorkspaceBuildProvisionerStateByID, arg.ProvisionerState, arg.UpdatedAt, arg.ID)
+	return err
+}
+
 const getWorkspaceResourceByID = `-- name: GetWorkspaceResourceByID :one
 SELECT
 	id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost

coderd/database/queries/workspacebuilds.sql

@@ -125,17 +125,6 @@ INSERT INTO
 VALUES
 	($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13);
 
--- name: UpdateWorkspaceBuildByID :exec
-UPDATE
-	workspace_builds
-SET
-	updated_at = $2,
-	provisioner_state = $3,
-	deadline = $4,
-	max_deadline = $5
-WHERE
-	id = $1;
-
 -- name: UpdateWorkspaceBuildCostByID :exec
 UPDATE
 	workspace_builds
@@ -144,6 +133,23 @@ SET
 WHERE
 	id = $1;
 
+-- name: UpdateWorkspaceBuildDeadlineByID :exec
+UPDATE
+	workspace_builds
+SET
+	deadline = @deadline::timestamptz,
+	max_deadline = @max_deadline::timestamptz,
+	updated_at = @updated_at::timestamptz
+WHERE id = @id::uuid;
+
+-- name: UpdateWorkspaceBuildProvisionerStateByID :exec
+UPDATE
+	workspace_builds
+SET
+	provisioner_state = @provisioner_state::bytea,
+	updated_at = @updated_at::timestamptz
+WHERE id = @id::uuid;
+
 -- name: GetActiveWorkspaceBuildsByTemplateID :many
 SELECT wb.*
 FROM (
@@ -166,5 +172,10 @@ FROM (
 JOIN
 	workspace_build_with_user AS wb
 	ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number
+JOIN
+	provisioner_jobs AS pj
+	ON wb.job_id = pj.id
 WHERE
-	wb.transition = 'start'::workspace_transition;
+	wb.transition = 'start'::workspace_transition
+AND
+	pj.completed_at IS NOT NULL;

coderd/provisionerdserver/provisionerdserver.go

@@ -832,16 +832,23 @@ func (s *server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto.
 			}
 
 			if jobType.WorkspaceBuild.State != nil {
-				err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+				err = db.UpdateWorkspaceBuildProvisionerStateByID(ctx, database.UpdateWorkspaceBuildProvisionerStateByIDParams{
 					ID:               input.WorkspaceBuildID,
 					UpdatedAt:        dbtime.Now(),
 					ProvisionerState: jobType.WorkspaceBuild.State,
-					Deadline:         build.Deadline,
-					MaxDeadline:      build.MaxDeadline,
 				})
 				if err != nil {
 					return xerrors.Errorf("update workspace build state: %w", err)
 				}
+				err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+					ID:          input.WorkspaceBuildID,
+					UpdatedAt:   dbtime.Now(),
+					Deadline:    build.Deadline,
+					MaxDeadline: build.MaxDeadline,
+				})
+				if err != nil {
+					return xerrors.Errorf("update workspace build deadline: %w", err)
+				}
 			}
 
 			return nil
@@ -1114,15 +1121,22 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
 			if err != nil {
 				return xerrors.Errorf("update provisioner job: %w", err)
 			}
-			err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+			err = db.UpdateWorkspaceBuildProvisionerStateByID(ctx, database.UpdateWorkspaceBuildProvisionerStateByIDParams{
 				ID:               workspaceBuild.ID,
-				Deadline:         autoStop.Deadline,
-				MaxDeadline:      autoStop.MaxDeadline,
 				ProvisionerState: jobType.WorkspaceBuild.State,
 				UpdatedAt:        now,
 			})
 			if err != nil {
-				return xerrors.Errorf("update workspace build: %w", err)
+				return xerrors.Errorf("update workspace build provisioner state: %w", err)
+			}
+			err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+				ID:          workspaceBuild.ID,
+				Deadline:    autoStop.Deadline,
+				MaxDeadline: autoStop.MaxDeadline,
+				UpdatedAt:   now,
+			})
+			if err != nil {
+				return xerrors.Errorf("update workspace build deadline: %w", err)
 			}
 
 			agentTimeouts := make(map[time.Duration]bool) // A set of agent timeouts.

coderd/unhanger/detector.go

@@ -333,12 +333,10 @@ func unhangJob(ctx context.Context, log slog.Logger, db database.Store, pub pubs
 					return xerrors.Errorf("get previous workspace build: %w", err)
 				}
 				if err == nil {
-					err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+					err = db.UpdateWorkspaceBuildProvisionerStateByID(ctx, database.UpdateWorkspaceBuildProvisionerStateByIDParams{
 						ID:               build.ID,
 						UpdatedAt:        dbtime.Now(),
 						ProvisionerState: prevBuild.ProvisionerState,
-						Deadline:         time.Time{},
-						MaxDeadline:      time.Time{},
 					})
 					if err != nil {
 						return xerrors.Errorf("update workspace build by id: %w", err)

coderd/workspaces.go

@@ -941,12 +941,11 @@ func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {
 			return xerrors.New("Cannot extend workspace: deadline is beyond max deadline imposed by template")
 		}
 
-		if err := s.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
-			ID:               build.ID,
-			UpdatedAt:        build.UpdatedAt,
-			ProvisionerState: build.ProvisionerState,
-			Deadline:         newDeadline,
-			MaxDeadline:      build.MaxDeadline,
+		if err := s.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+			ID:          build.ID,
+			UpdatedAt:   dbtime.Now(),
+			Deadline:    newDeadline,
+			MaxDeadline: build.MaxDeadline,
 		}); err != nil {
 			code = http.StatusInternalServerError
 			resp.Message = "Failed to extend workspace deadline."

enterprise/coderd/schedule/template.go

@@ -278,8 +278,8 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Conte
 		autostop.Deadline = autostop.MaxDeadline
 	}
 
-	// Update the workspace build.
-	err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
+	// Update the workspace build deadline.
+	err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
 		ID:          build.ID,
 		UpdatedAt:   now,
 		Deadline:    autostop.Deadline,

enterprise/coderd/schedule/template_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database/dbgen"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	agplschedule "github.com/coder/coder/v2/coderd/schedule"
+	"github.com/coder/coder/v2/cryptorand"
 	"github.com/coder/coder/v2/enterprise/coderd/schedule"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -164,9 +165,13 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 					JobID:             job.ID,
 					InitiatorID:       user.ID,
 					TemplateVersionID: templateVersion.ID,
+					ProvisionerState:  []byte(must(cryptorand.String(64))),
 				})
 			)
 
+			// Assert test invariant: workspace build state must not be empty
+			require.NotEmpty(t, wsBuild.ProvisionerState, "provisioner state must not be empty")
+
 			acquiredJob, err := db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
 				StartedAt: sql.NullTime{
 					Time:  buildTime,
@@ -191,12 +196,11 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			})
 			require.NoError(t, err)
 
-			err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
-				ID:               wsBuild.ID,
-				UpdatedAt:        buildTime,
-				ProvisionerState: []byte{},
-				Deadline:         c.deadline,
-				MaxDeadline:      c.maxDeadline,
+			err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+				ID:          wsBuild.ID,
+				UpdatedAt:   buildTime,
+				Deadline:    c.deadline,
+				MaxDeadline: c.maxDeadline,
 			})
 			require.NoError(t, err)
 
@@ -240,6 +244,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 			}
 			require.WithinDuration(t, c.newDeadline, newBuild.Deadline, time.Second)
 			require.WithinDuration(t, c.newMaxDeadline, newBuild.MaxDeadline, time.Second)
+
+			// Check that the new build has the same state as before.
+			require.Equal(t, wsBuild.ProvisionerState, newBuild.ProvisionerState, "provisioner state mismatch")
 		})
 	}
 }
@@ -420,20 +427,26 @@ func TestTemplateUpdateBuildDeadlinesSkip(t *testing.T) {
 			JobID:             job.ID,
 			InitiatorID:       user.ID,
 			TemplateVersionID: templateVersion.ID,
+			ProvisionerState:  []byte(must(cryptorand.String(64))),
 		})
 
-		err := db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
-			ID:               wsBuild.ID,
-			UpdatedAt:        buildTime,
-			ProvisionerState: []byte{},
-			Deadline:         originalMaxDeadline,
-			MaxDeadline:      originalMaxDeadline,
+		// Assert test invariant: workspace build state must not be empty
+		require.NotEmpty(t, wsBuild.ProvisionerState, "provisioner state must not be empty")
+
+		err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
+			ID:          wsBuild.ID,
+			UpdatedAt:   buildTime,
+			Deadline:    originalMaxDeadline,
+			MaxDeadline: originalMaxDeadline,
 		})
 		require.NoError(t, err)
 
 		wsBuild, err = db.GetWorkspaceBuildByID(ctx, wsBuild.ID)
 		require.NoError(t, err)
 
+		// Assert test invariant: workspace build state must not be empty
+		require.NotEmpty(t, wsBuild.ProvisionerState, "provisioner state must not be empty")
+
 		builds[i].wsBuild = wsBuild
 
 		if !b.buildStarted {
@@ -519,5 +532,14 @@ func TestTemplateUpdateBuildDeadlinesSkip(t *testing.T) {
 			assert.WithinDuration(t, originalMaxDeadline, newBuild.Deadline, time.Second, msg)
 			assert.WithinDuration(t, originalMaxDeadline, newBuild.MaxDeadline, time.Second, msg)
 		}
+
+		assert.Equal(t, builds[i].wsBuild.ProvisionerState, newBuild.ProvisionerState, "provisioner state mismatch")
 	}
 }
+
+func must[V any](v V, err error) V {
+	if err != nil {
+		panic(err)
+	}
+	return v
+}