WIP

coder · ibetitsmike · May 20, 2025 · Apr 1, 2025 · May 9, 2025 · May 9, 2025
commit 0fe1404bb3c5777fa9983fe7bf268b05873e25f8
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2300,6 +2300,13 @@ func (q *querier) GetProvisionerJobByID(ctx context.Context, id uuid.UUID) (data
 	return job, nil
 }
 
+func (q *querier) GetProvisionerJobByIDForUpdate(ctx context.Context, id uuid.UUID) (database.ProvisionerJob, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceProvisionerJobs); err != nil {
+		return database.ProvisionerJob{}, err
+	}
+	return q.db.GetProvisionerJobByIDForUpdate(ctx, id)
+}
+
 func (q *querier) GetProvisionerJobTimingsByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ProvisionerJobTiming, error) {
 	_, err := q.GetProvisionerJobByID(ctx, jobID)
 	if err != nil {

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -4444,6 +4444,9 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 			VapidPrivateKey: "test",
 		}).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
 	}))
+	s.Run("GetProvisionerJobByIDForUpdate", s.Subtest(func(db database.Store, check *expects) {
+		check.Args(uuid.New()).Asserts(rbac.ResourceProvisionerJobs, policy.ActionRead).Errors(sql.ErrNoRows)
+	}))
 }
 
 func (s *MethodTestSuite) TestNotifications() {

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -4625,6 +4625,13 @@ func (q *FakeQuerier) GetProvisionerJobByID(ctx context.Context, id uuid.UUID) (
 	return q.getProvisionerJobByIDNoLock(ctx, id)
 }
 
+func (q *FakeQuerier) GetProvisionerJobByIDForUpdate(ctx context.Context, id uuid.UUID) (database.ProvisionerJob, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	return q.getProvisionerJobByIDNoLock(ctx, id)
+}
+
 func (q *FakeQuerier) GetProvisionerJobTimingsByJobID(_ context.Context, jobID uuid.UUID) ([]database.ProvisionerJobTiming, 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/provisionerjobs.sql b/coderd/database/queries/provisionerjobs.sql
@@ -41,6 +41,18 @@ FROM
 WHERE
 	id = $1;
 
+-- name: GetProvisionerJobByIDForUpdate :one
+-- Gets a single provisioner job by ID for update.
+-- This is used to securely reap jobs that have been hung/pending for a long time.
+SELECT
+	*
+FROM
+	provisioner_jobs
+WHERE
+	id = $1
+FOR UPDATE
+SKIP LOCKED;
+
 -- name: GetProvisionerJobsByIDs :many
 SELECT
 	*
@@ -288,8 +300,7 @@ WHERE
 		AND completed_at IS NULL
 	)
 ORDER BY random()
-LIMIT @max_jobs
-FOR UPDATE SKIP LOCKED;
+LIMIT @max_jobs;
 
 -- name: InsertProvisionerJobTimings :many
 INSERT INTO provisioner_job_timings (job_id, started_at, ended_at, stage, source, action, resource)

diff --git a/coderd/jobreaper/detector.go b/coderd/jobreaper/detector.go
@@ -239,18 +239,12 @@ func reapJob(ctx context.Context, log slog.Logger, db database.Store, pub pubsub
 	var lowestLogID int64
 
 	err := db.InTx(func(db database.Store) error {
-		locked, err := db.TryAcquireLock(ctx, database.GenLockID(fmt.Sprintf("reaper:%s", jobToReap.ID)))
-		if err != nil {
-			return xerrors.Errorf("acquire lock: %w", err)
-		}
-		if !locked {
-			// This error is ignored.
-			return acquireLockError{}
-		}
-
 		// Refetch the job while we hold the lock.
-		job, err := db.GetProvisionerJobByID(ctx, jobToReap.ID)
+		job, err := db.GetProvisionerJobByIDForUpdate(ctx, jobToReap.ID)
 		if err != nil {
+			if xerrors.Is(err, sql.ErrNoRows) {
+				return acquireLockError{}
+			}
 			return xerrors.Errorf("get provisioner job: %w", err)
 		}