fix(coderd/database): consider tag sets when calculating queue position (#16685)

Relates to https://github.com/coder/coder/issues/15843

## PR Contents

- Reimplementation of the `GetProvisionerJobsByIDsWithQueuePosition` SQL
query to **take into account** provisioner job tags and provisioner
daemon tags.
- Unit tests covering different **tag sets**, **job statuses**, and
**job ordering** scenarios.

## Notes

- The original row order is preserved by introducing the `ordinality`
field.
- Unnecessary rows are filtered as early as possible to ensure that
expensive joins operate on a smaller dataset.
- A "fake" join with `provisioner_jobs` is added at the end to ensure
`sqlc.embed` compiles successfully.
- **Backward compatibility is preserved**—only the SQL query has been
updated, while the Go code remains unchanged.
This commit is contained in:
Yevhenii Shcherbina
2025-03-03 10:02:18 -05:00
committed by GitHub
parent 7637d39528
commit b85ba586ee
7 changed files with 658 additions and 67 deletions

View File

@ -6627,45 +6627,69 @@ func (q *sqlQuerier) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUI
}
const getProvisionerJobsByIDsWithQueuePosition = `-- name: GetProvisionerJobsByIDsWithQueuePosition :many
WITH pending_jobs AS (
SELECT
id, created_at
FROM
provisioner_jobs
WHERE
started_at IS NULL
AND
canceled_at IS NULL
AND
completed_at IS NULL
AND
error IS NULL
WITH filtered_provisioner_jobs AS (
-- Step 1: Filter provisioner_jobs
SELECT
id, created_at
FROM
provisioner_jobs
WHERE
id = ANY($1 :: uuid [ ]) -- Apply filter early to reduce dataset size before expensive JOIN
),
queue_position AS (
SELECT
id,
ROW_NUMBER() OVER (ORDER BY created_at ASC) AS queue_position
FROM
pending_jobs
pending_jobs AS (
-- Step 2: Extract only pending jobs
SELECT
id, created_at, tags
FROM
provisioner_jobs
WHERE
job_status = 'pending'
),
queue_size AS (
SELECT COUNT(*) AS count FROM pending_jobs
ranked_jobs AS (
-- Step 3: Rank only pending jobs based on provisioner availability
SELECT
pj.id,
pj.created_at,
ROW_NUMBER() OVER (PARTITION BY pd.id ORDER BY pj.created_at ASC) AS queue_position,
COUNT(*) OVER (PARTITION BY pd.id) AS queue_size
FROM
pending_jobs pj
INNER JOIN provisioner_daemons pd
ON provisioner_tagset_contains(pd.tags, pj.tags) -- Join only on the small pending set
),
final_jobs AS (
-- Step 4: Compute best queue position and max queue size per job
SELECT
fpj.id,
fpj.created_at,
COALESCE(MIN(rj.queue_position), 0) :: BIGINT AS queue_position, -- Best queue position across provisioners
COALESCE(MAX(rj.queue_size), 0) :: BIGINT AS queue_size -- Max queue size across provisioners
FROM
filtered_provisioner_jobs fpj -- Use the pre-filtered dataset instead of full provisioner_jobs
LEFT JOIN ranked_jobs rj
ON fpj.id = rj.id -- Join with the ranking jobs CTE to assign a rank to each specified provisioner job.
GROUP BY
fpj.id, fpj.created_at
)
SELECT
-- Step 5: Final SELECT with INNER JOIN provisioner_jobs
fj.id,
fj.created_at,
pj.id, pj.created_at, pj.updated_at, pj.started_at, pj.canceled_at, pj.completed_at, pj.error, pj.organization_id, pj.initiator_id, pj.provisioner, pj.storage_method, pj.type, pj.input, pj.worker_id, pj.file_id, pj.tags, pj.error_code, pj.trace_metadata, pj.job_status,
COALESCE(qp.queue_position, 0) AS queue_position,
COALESCE(qs.count, 0) AS queue_size
fj.queue_position,
fj.queue_size
FROM
provisioner_jobs pj
LEFT JOIN
queue_position qp ON qp.id = pj.id
LEFT JOIN
queue_size qs ON TRUE
WHERE
pj.id = ANY($1 :: uuid [ ])
final_jobs fj
INNER JOIN provisioner_jobs pj
ON fj.id = pj.id -- Ensure we retrieve full details from ` + "`" + `provisioner_jobs` + "`" + `.
-- JOIN with pj is required for sqlc.embed(pj) to compile successfully.
ORDER BY
fj.created_at
`
type GetProvisionerJobsByIDsWithQueuePositionRow struct {
ID uuid.UUID `db:"id" json:"id"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
ProvisionerJob ProvisionerJob `db:"provisioner_job" json:"provisioner_job"`
QueuePosition int64 `db:"queue_position" json:"queue_position"`
QueueSize int64 `db:"queue_size" json:"queue_size"`
@ -6681,6 +6705,8 @@ func (q *sqlQuerier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Contex
for rows.Next() {
var i GetProvisionerJobsByIDsWithQueuePositionRow
if err := rows.Scan(
&i.ID,
&i.CreatedAt,
&i.ProvisionerJob.ID,
&i.ProvisionerJob.CreatedAt,
&i.ProvisionerJob.UpdatedAt,