Skip to content

Commit 110102a

Browse files
fix: optimize queue position sql query (#17974)
Use only `online provisioner daemons` for `GetProvisionerJobsByIDsWithQueuePosition` query. It should improve performance of the query.
1 parent 2bcbd9b commit 110102a

File tree

11 files changed

+88
-36
lines changed

11 files changed

+88
-36
lines changed

coderd/database/dbauthz/dbauthz.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2341,7 +2341,7 @@ func (q *querier) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID)
23412341
return provisionerJobs, nil
23422342
}
23432343

2344-
func (q *querier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
2344+
func (q *querier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids database.GetProvisionerJobsByIDsWithQueuePositionParams) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
23452345
// TODO: Remove this once we have a proper rbac check for provisioner jobs.
23462346
// Details in https://github.com/coder/coder/issues/16160
23472347
return q.db.GetProvisionerJobsByIDsWithQueuePosition(ctx, ids)

coderd/database/dbauthz/dbauthz_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4345,7 +4345,7 @@ func (s *MethodTestSuite) TestSystemFunctions() {
43454345
check.Args([]uuid.UUID{uuid.New()}).Asserts(rbac.ResourceSystem, policy.ActionRead)
43464346
}))
43474347
s.Run("GetProvisionerJobsByIDsWithQueuePosition", s.Subtest(func(db database.Store, check *expects) {
4348-
check.Args([]uuid.UUID{}).Asserts()
4348+
check.Args(database.GetProvisionerJobsByIDsWithQueuePositionParams{}).Asserts()
43494349
}))
43504350
s.Run("GetReplicaByID", s.Subtest(func(db database.Store, check *expects) {
43514351
check.Args(uuid.New()).Asserts(rbac.ResourceSystem, policy.ActionRead).Errors(sql.ErrNoRows)

coderd/database/dbmem/dbmem.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4684,14 +4684,14 @@ func (q *FakeQuerier) GetProvisionerJobsByIDs(_ context.Context, ids []uuid.UUID
46844684
return jobs, nil
46854685
}
46864686

4687-
func (q *FakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
4687+
func (q *FakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, arg database.GetProvisionerJobsByIDsWithQueuePositionParams) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
46884688
q.mutex.RLock()
46894689
defer q.mutex.RUnlock()
46904690

4691-
if ids == nil {
4692-
ids = []uuid.UUID{}
4691+
if arg.IDs == nil {
4692+
arg.IDs = []uuid.UUID{}
46934693
}
4694-
return q.getProvisionerJobsByIDsWithQueuePositionLockedTagBasedQueue(ctx, ids)
4694+
return q.getProvisionerJobsByIDsWithQueuePositionLockedTagBasedQueue(ctx, arg.IDs)
46954695
}
46964696

46974697
func (q *FakeQuerier) GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisioner(ctx context.Context, arg database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerParams) ([]database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow, error) {

coderd/database/dbmetrics/querymetrics.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

coderd/database/dbmock/dbmock.go

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

coderd/database/querier.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

coderd/database/querier_test.go

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ import (
1515
"github.com/stretchr/testify/require"
1616

1717
"cdr.dev/slog/sloggers/slogtest"
18-
1918
"github.com/coder/coder/v2/coderd/coderdtest"
2019
"github.com/coder/coder/v2/coderd/database"
2120
"github.com/coder/coder/v2/coderd/database/db2sdk"
@@ -27,6 +26,7 @@ import (
2726
"github.com/coder/coder/v2/coderd/database/migrations"
2827
"github.com/coder/coder/v2/coderd/httpmw"
2928
"github.com/coder/coder/v2/coderd/prebuilds"
29+
"github.com/coder/coder/v2/coderd/provisionerdserver"
3030
"github.com/coder/coder/v2/coderd/rbac"
3131
"github.com/coder/coder/v2/coderd/rbac/policy"
3232
"github.com/coder/coder/v2/provisionersdk"
@@ -1268,7 +1268,10 @@ func TestQueuePosition(t *testing.T) {
12681268
Tags: database.StringMap{},
12691269
})
12701270

1271-
queued, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
1271+
queued, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
1272+
IDs: jobIDs,
1273+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
1274+
})
12721275
require.NoError(t, err)
12731276
require.Len(t, queued, jobCount)
12741277
sort.Slice(queued, func(i, j int) bool {
@@ -1296,7 +1299,10 @@ func TestQueuePosition(t *testing.T) {
12961299
require.NoError(t, err)
12971300
require.Equal(t, jobs[0].ID, job.ID)
12981301

1299-
queued, err = db.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
1302+
queued, err = db.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
1303+
IDs: jobIDs,
1304+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
1305+
})
13001306
require.NoError(t, err)
13011307
require.Len(t, queued, jobCount)
13021308
sort.Slice(queued, func(i, j int) bool {
@@ -2550,7 +2556,10 @@ func TestGetProvisionerJobsByIDsWithQueuePosition(t *testing.T) {
25502556
}
25512557

25522558
// When: we fetch the jobs by their IDs
2553-
actualJobs, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, filteredJobIDs)
2559+
actualJobs, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
2560+
IDs: filteredJobIDs,
2561+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
2562+
})
25542563
require.NoError(t, err)
25552564
require.Len(t, actualJobs, len(filteredJobs), "should return all unskipped jobs")
25562565

@@ -2693,7 +2702,10 @@ func TestGetProvisionerJobsByIDsWithQueuePosition_MixedStatuses(t *testing.T) {
26932702
}
26942703

26952704
// When: we fetch the jobs by their IDs
2696-
actualJobs, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
2705+
actualJobs, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
2706+
IDs: jobIDs,
2707+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
2708+
})
26972709
require.NoError(t, err)
26982710
require.Len(t, actualJobs, len(allJobs), "should return all jobs")
26992711

@@ -2788,7 +2800,10 @@ func TestGetProvisionerJobsByIDsWithQueuePosition_OrderValidation(t *testing.T)
27882800
}
27892801

27902802
// When: we fetch the jobs by their IDs
2791-
actualJobs, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
2803+
actualJobs, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
2804+
IDs: jobIDs,
2805+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
2806+
})
27922807
require.NoError(t, err)
27932808
require.Len(t, actualJobs, len(allJobs), "should return all jobs")
27942809

coderd/database/queries.sql.go

Lines changed: 15 additions & 6 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

coderd/database/queries/provisionerjobs.sql

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -80,17 +80,21 @@ pending_jobs AS (
8080
WHERE
8181
job_status = 'pending'
8282
),
83+
online_provisioner_daemons AS (
84+
SELECT id, tags FROM provisioner_daemons pd
85+
WHERE pd.last_seen_at IS NOT NULL AND pd.last_seen_at >= (NOW() - (@stale_interval_ms::bigint || ' ms')::interval)
86+
),
8387
ranked_jobs AS (
8488
-- Step 3: Rank only pending jobs based on provisioner availability
8589
SELECT
8690
pj.id,
8791
pj.created_at,
88-
ROW_NUMBER() OVER (PARTITION BY pd.id ORDER BY pj.created_at ASC) AS queue_position,
89-
COUNT(*) OVER (PARTITION BY pd.id) AS queue_size
92+
ROW_NUMBER() OVER (PARTITION BY opd.id ORDER BY pj.created_at ASC) AS queue_position,
93+
COUNT(*) OVER (PARTITION BY opd.id) AS queue_size
9094
FROM
9195
pending_jobs pj
92-
INNER JOIN provisioner_daemons pd
93-
ON provisioner_tagset_contains(pd.tags, pj.tags) -- Join only on the small pending set
96+
INNER JOIN online_provisioner_daemons opd
97+
ON provisioner_tagset_contains(opd.tags, pj.tags) -- Join only on the small pending set
9498
),
9599
final_jobs AS (
96100
-- Step 4: Compute best queue position and max queue size per job

coderd/templateversions.go

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,10 @@ func (api *API) templateVersion(rw http.ResponseWriter, r *http.Request) {
5353
ctx := r.Context()
5454
templateVersion := httpmw.TemplateVersionParam(r)
5555

56-
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, []uuid.UUID{templateVersion.JobID})
56+
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
57+
IDs: []uuid.UUID{templateVersion.JobID},
58+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
59+
})
5760
if err != nil || len(jobs) == 0 {
5861
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
5962
Message: "Internal error fetching provisioner job.",
@@ -182,7 +185,10 @@ func (api *API) patchTemplateVersion(rw http.ResponseWriter, r *http.Request) {
182185
return
183186
}
184187

185-
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, []uuid.UUID{templateVersion.JobID})
188+
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
189+
IDs: []uuid.UUID{templateVersion.JobID},
190+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
191+
})
186192
if err != nil || len(jobs) == 0 {
187193
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
188194
Message: "Internal error fetching provisioner job.",
@@ -733,7 +739,10 @@ func (api *API) fetchTemplateVersionDryRunJob(rw http.ResponseWriter, r *http.Re
733739
return database.GetProvisionerJobsByIDsWithQueuePositionRow{}, false
734740
}
735741

736-
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, []uuid.UUID{jobUUID})
742+
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
743+
IDs: []uuid.UUID{jobUUID},
744+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
745+
})
737746
if httpapi.Is404Error(err) {
738747
httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
739748
Message: fmt.Sprintf("Provisioner job %q not found.", jobUUID),
@@ -865,7 +874,10 @@ func (api *API) templateVersionsByTemplate(rw http.ResponseWriter, r *http.Reque
865874
for _, version := range versions {
866875
jobIDs = append(jobIDs, version.JobID)
867876
}
868-
jobs, err := store.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
877+
jobs, err := store.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
878+
IDs: jobIDs,
879+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
880+
})
869881
if err != nil {
870882
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
871883
Message: "Internal error fetching provisioner job.",
@@ -933,7 +945,10 @@ func (api *API) templateVersionByName(rw http.ResponseWriter, r *http.Request) {
933945
})
934946
return
935947
}
936-
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, []uuid.UUID{templateVersion.JobID})
948+
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
949+
IDs: []uuid.UUID{templateVersion.JobID},
950+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
951+
})
937952
if err != nil || len(jobs) == 0 {
938953
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
939954
Message: "Internal error fetching provisioner job.",
@@ -1013,7 +1028,10 @@ func (api *API) templateVersionByOrganizationTemplateAndName(rw http.ResponseWri
10131028
})
10141029
return
10151030
}
1016-
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, []uuid.UUID{templateVersion.JobID})
1031+
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
1032+
IDs: []uuid.UUID{templateVersion.JobID},
1033+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
1034+
})
10171035
if err != nil || len(jobs) == 0 {
10181036
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
10191037
Message: "Internal error fetching provisioner job.",
@@ -1115,7 +1133,10 @@ func (api *API) previousTemplateVersionByOrganizationTemplateAndName(rw http.Res
11151133
return
11161134
}
11171135

1118-
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, []uuid.UUID{previousTemplateVersion.JobID})
1136+
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
1137+
IDs: []uuid.UUID{previousTemplateVersion.JobID},
1138+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
1139+
})
11191140
if err != nil || len(jobs) == 0 {
11201141
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
11211142
Message: "Internal error fetching provisioner job.",

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy