Skip to content

Commit 110102a

Browse files
evgeniy-scherbinaevgeniy-scherbina
authored
fix: optimize queue position sql query (coder#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.",

‎coderd/workspacebuilds.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -797,7 +797,10 @@ func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []datab
797797
for _, build := range workspaceBuilds {
798798
jobIDs = append(jobIDs, build.JobID)
799799
}
800-
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
800+
jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
801+
IDs: jobIDs,
802+
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
803+
})
801804
if err != nil && !errors.Is(err, sql.ErrNoRows) {
802805
return workspaceBuildsData{}, xerrors.Errorf("get provisioner jobs: %w", err)
803806
}

0 commit comments

Comments
 (0)