Mar 3, 2025 · Feb 19, 2025 · Feb 20, 2025 · Feb 25, 2025 · Feb 25, 2025 · Feb 26, 2025
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
 return ""
 }

 func (q *FakeQuerier) getProvisionerJobsByIDsWithQueuePositionLocked(_ context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
 // provisionerTagsetContains checks if daemonTags contain all key-value pairs from jobTags
 func provisionerTagsetContains(daemonTags, jobTags map[string]string) bool {
 for jobKey, jobValue := range jobTags {
 if daemonValue, exists := daemonTags[jobKey]; !exists || daemonValue != jobValue {
 return false
 }
 }
 return true
 }

 // GetProvisionerJobsByIDsWithQueuePosition mimics the SQL logic in pure Go
 func (q *FakeQuerier) getProvisionerJobsByIDsWithQueuePositionLockedTagBasedQueue(_ context.Context, jobIDs []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
 // Step 1: Filter provisionerJobs based on jobIDs
 filteredJobs := make(map[uuid.UUID]database.ProvisionerJob)
 for _, job := range q.provisionerJobs {
 for _, id := range jobIDs {
 if job.ID == id {
 filteredJobs[job.ID] = job
 }
 }
 }

 // Step 2: Identify pending jobs
 pendingJobs := make(map[uuid.UUID]database.ProvisionerJob)
 for _, job := range q.provisionerJobs {
 if job.JobStatus == "pending" {
 pendingJobs[job.ID] = job
 }
 }

 // Step 3: Identify pending jobs that have a matching provisioner
 matchedJobs := make(map[uuid.UUID]struct{})
 for _, job := range pendingJobs {
 for _, daemon := range q.provisionerDaemons {
 if provisionerTagsetContains(daemon.Tags, job.Tags) {
 matchedJobs[job.ID] = struct{}{}
 break
 }
 }
 }

 // Step 4: Rank pending jobs per provisioner
 jobRanks := make(map[uuid.UUID][]database.ProvisionerJob)
 for _, job := range pendingJobs {
 for _, daemon := range q.provisionerDaemons {
 if provisionerTagsetContains(daemon.Tags, job.Tags) {
 jobRanks[daemon.ID] = append(jobRanks[daemon.ID], job)
 }
 }
 }

 // Sort jobs per provisioner by CreatedAt
 for daemonID := range jobRanks {
 sort.Slice(jobRanks[daemonID], func(i, j int) bool {
 return jobRanks[daemonID][i].CreatedAt.Before(jobRanks[daemonID][j].CreatedAt)
 })
 }

 // Step 5: Compute queue position & max queue size across all provisioners
 jobQueueStats := make(map[uuid.UUID]database.GetProvisionerJobsByIDsWithQueuePositionRow)
 for _, jobs := range jobRanks {
 queueSize := int64(len(jobs)) // Queue size per provisioner
 for i, job := range jobs {
 queuePosition := int64(i + 1)

 // If the job already exists, update only if this queuePosition is better
 if existing, exists := jobQueueStats[job.ID]; exists {
 jobQueueStats[job.ID] = database.GetProvisionerJobsByIDsWithQueuePositionRow{
 ID:             job.ID,
 CreatedAt:      job.CreatedAt,
 ProvisionerJob: job,
 QueuePosition:  min(existing.QueuePosition, queuePosition),
 QueueSize:      max(existing.QueueSize, queueSize), // Take the maximum queue size across provisioners
 }
 } else {
 jobQueueStats[job.ID] = database.GetProvisionerJobsByIDsWithQueuePositionRow{
 ID:             job.ID,
 CreatedAt:      job.CreatedAt,
 ProvisionerJob: job,
 QueuePosition:  queuePosition,
 QueueSize:      queueSize,
 }
 }
 }
 }

 // Step 6: Compute the final results with minimal checks
 var results []database.GetProvisionerJobsByIDsWithQueuePositionRow
 for _, job := range filteredJobs {
 // If the job has a computed rank, use it
 if rank, found := jobQueueStats[job.ID]; found {
 results = append(results, rank)
 } else {
 // Otherwise, return (0,0) for non-pending jobs and unranked pending jobs
 results = append(results, database.GetProvisionerJobsByIDsWithQueuePositionRow{
 ID:             job.ID,
 CreatedAt:      job.CreatedAt,
 ProvisionerJob: job,
 QueuePosition:  0,
 QueueSize:      0,
 })
 }
 }

 // Step 7: Sort results by CreatedAt
 sort.Slice(results, func(i, j int) bool {
 return results[i].CreatedAt.Before(results[j].CreatedAt)
 })

 return results, nil
 }

 func (q *FakeQuerier) getProvisionerJobsByIDsWithQueuePositionLockedGlobalQueue(_ context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
 //WITH pending_jobs AS (
 //SELECT
 //id, created_at
 if ids == nil {
 ids = []uuid.UUID{}
 }
 return q.getProvisionerJobsByIDsWithQueuePositionLocked(ctx, ids)
 return q.getProvisionerJobsByIDsWithQueuePositionLockedTagBasedQueue(ctx, ids)
 }

 func (q *FakeQuerier) GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisioner(ctx context.Context, arg database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerParams) ([]database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow, error) {
 LIMIT
 sqlc.narg('limit')::int;
 */
 rowsWithQueuePosition, err := q.getProvisionerJobsByIDsWithQueuePositionLocked(ctx, nil)
 rowsWithQueuePosition, err := q.getProvisionerJobsByIDsWithQueuePositionLockedGlobalQueue(ctx, nil)
 if err != nil {
 return nil, err
 }
diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000298_provisioner_jobs_status_idx.down.sql b/coderd/database/migrations/000298_provisioner_jobs_status_idx.down.sql
 DROP INDEX idx_provisioner_jobs_status;
diff --git a/coderd/database/migrations/000298_provisioner_jobs_status_idx.up.sql b/coderd/database/migrations/000298_provisioner_jobs_status_idx.up.sql
 CREATE INDEX idx_provisioner_jobs_status ON provisioner_jobs USING btree (job_status);
Original file line number	Diff line number	Diff line change
Expand Up		@@ -1149,7 +1149,119 @@ func getOwnerFromTags(tags map[string]string) string {
		return ""
		}

		func (q *FakeQuerier) getProvisionerJobsByIDsWithQueuePositionLocked(_ context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
		// provisionerTagsetContains checks if daemonTags contain all key-value pairs from jobTags
		func provisionerTagsetContains(daemonTags, jobTags map[string]string) bool {
		for jobKey, jobValue := range jobTags {
		if daemonValue, exists := daemonTags[jobKey]; !exists \|\| daemonValue != jobValue {
		return false
		}
		}
		return true
		}

		// GetProvisionerJobsByIDsWithQueuePosition mimics the SQL logic in pure Go
		func (q *FakeQuerier) getProvisionerJobsByIDsWithQueuePositionLockedTagBasedQueue(_ context.Context, jobIDs []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
		// Step 1: Filter provisionerJobs based on jobIDs
		filteredJobs := make(map[uuid.UUID]database.ProvisionerJob)
		for _, job := range q.provisionerJobs {
		for _, id := range jobIDs {
		if job.ID == id {
		filteredJobs[job.ID] = job
		}
		}
		}

		// Step 2: Identify pending jobs
		pendingJobs := make(map[uuid.UUID]database.ProvisionerJob)
		for _, job := range q.provisionerJobs {
		if job.JobStatus == "pending" {
		pendingJobs[job.ID] = job
		}
		}

		// Step 3: Identify pending jobs that have a matching provisioner
		matchedJobs := make(map[uuid.UUID]struct{})
		for _, job := range pendingJobs {
		for _, daemon := range q.provisionerDaemons {
		if provisionerTagsetContains(daemon.Tags, job.Tags) {
		matchedJobs[job.ID] = struct{}{}
		break
		}
		}
		}

		// Step 4: Rank pending jobs per provisioner
		jobRanks := make(map[uuid.UUID][]database.ProvisionerJob)
		for _, job := range pendingJobs {
		for _, daemon := range q.provisionerDaemons {
		if provisionerTagsetContains(daemon.Tags, job.Tags) {
		jobRanks[daemon.ID] = append(jobRanks[daemon.ID], job)
		}
		}
		}

		// Sort jobs per provisioner by CreatedAt
		for daemonID := range jobRanks {
		sort.Slice(jobRanks[daemonID], func(i, j int) bool {
		return jobRanks[daemonID][i].CreatedAt.Before(jobRanks[daemonID][j].CreatedAt)
		})
		}

		// Step 5: Compute queue position & max queue size across all provisioners
		jobQueueStats := make(map[uuid.UUID]database.GetProvisionerJobsByIDsWithQueuePositionRow)
		for _, jobs := range jobRanks {
		queueSize := int64(len(jobs)) // Queue size per provisioner
		for i, job := range jobs {
		queuePosition := int64(i + 1)

		// If the job already exists, update only if this queuePosition is better
		if existing, exists := jobQueueStats[job.ID]; exists {
		jobQueueStats[job.ID] = database.GetProvisionerJobsByIDsWithQueuePositionRow{
		ID: job.ID,
		CreatedAt: job.CreatedAt,
		ProvisionerJob: job,
		QueuePosition: min(existing.QueuePosition, queuePosition),
		QueueSize: max(existing.QueueSize, queueSize), // Take the maximum queue size across provisioners
		}
		} else {
		jobQueueStats[job.ID] = database.GetProvisionerJobsByIDsWithQueuePositionRow{
		ID: job.ID,
		CreatedAt: job.CreatedAt,
		ProvisionerJob: job,
		QueuePosition: queuePosition,
		QueueSize: queueSize,
		}
		}
		}
		}

		// Step 6: Compute the final results with minimal checks
		var results []database.GetProvisionerJobsByIDsWithQueuePositionRow
		for _, job := range filteredJobs {
		// If the job has a computed rank, use it
		if rank, found := jobQueueStats[job.ID]; found {
		results = append(results, rank)
		} else {
		// Otherwise, return (0,0) for non-pending jobs and unranked pending jobs
		results = append(results, database.GetProvisionerJobsByIDsWithQueuePositionRow{
		ID: job.ID,
		CreatedAt: job.CreatedAt,
		ProvisionerJob: job,
		QueuePosition: 0,
		QueueSize: 0,
		})
		}
		}

		// Step 7: Sort results by CreatedAt
		sort.Slice(results, func(i, j int) bool {
		return results[i].CreatedAt.Before(results[j].CreatedAt)
		})

		return results, nil
		}

		func (q *FakeQuerier) getProvisionerJobsByIDsWithQueuePositionLockedGlobalQueue(_ context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
		//WITH pending_jobs AS (
		//SELECT
		//id, created_at
Expand DownExpand Up		@@ -4237,7 +4349,7 @@ func (q *FakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Conte
		if ids == nil {
		ids = []uuid.UUID{}
		}
		return q.getProvisionerJobsByIDsWithQueuePositionLocked(ctx, ids)
		return q.getProvisionerJobsByIDsWithQueuePositionLockedTagBasedQueue(ctx, ids)
		}

		func (q *FakeQuerier) GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisioner(ctx context.Context, arg database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerParams) ([]database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow, error) {
Expand DownExpand Up		@@ -4306,7 +4418,7 @@ func (q *FakeQuerier) GetProvisionerJobsByOrganizationAndStatusWithQueuePosition
		LIMIT
		sqlc.narg('limit')::int;
		*/
		rowsWithQueuePosition, err := q.getProvisionerJobsByIDsWithQueuePositionLocked(ctx, nil)
		rowsWithQueuePosition, err := q.getProvisionerJobsByIDsWithQueuePositionLockedGlobalQueue(ctx, nil)
		if err != nil {
		return nil, err
		}
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		CREATE INDEX idx_provisioner_jobs_status ON provisioner_jobs USING btree (job_status);
Copy link Member johnstcnMar 3, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. Review: it took me a minute to figure out if this is supported, but apparently it is! Thedocs on generated columns specify that only stored generated columns are supported, and these are always written to disk. TIL! Copy link ContributorAuthor evgeniy-scherbinaMar 3, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. @johnstcn yes, that's what I figured as well Also I thought about using`Hash` Index instead of`BTree`, because so far we're only using`=` or`!=` operations, but wasn't sure, maybe we'll need other operations in the future. Copy link Member johnstcnMar 3, 2025• edited Loading Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. TBH this is probably fine, almost all of our other indexes are using BTree. If you were curious, I suppose you could validate with some EXPLAINs on a sample dataset. evgeniy-scherbina reacted with thumbs up emoji