coder · evgeniy-scherbina · Mar 3, 2025 · Feb 19, 2025 · Feb 20, 2025 · Feb 25, 2025
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -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
@@ -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) {
@@ -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
 	}

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
@@ -0,0 +1 @@
+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
@@ -0,0 +1 @@
+CREATE INDEX idx_provisioner_jobs_status ON provisioner_jobs USING btree (job_status);
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 johnstcn Mar 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! The docs on generated columns specify that only stored generated columns are supported, and these are always written to disk. TIL! Copy link Contributor Author evgeniy-scherbina Mar 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 johnstcn Mar 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