From 13c7d7505e4cc44d652ed56608a74b7e5d1623bc Mon Sep 17 00:00:00 2001 From: Spike Curtis Date: Thu, 20 Aug 2026 09:37:47 +0000 Subject: [PATCH] feat(coderd): add queue-position-free provisioner job query --- coderd/database/dbauthz/dbauthz.go | 6 +++ coderd/database/dbauthz/dbauthz_test.go | 5 ++ coderd/database/dbmetrics/querymetrics.go | 8 +++ coderd/database/dbmock/dbmock.go | 15 ++++++ coderd/database/querier.go | 5 ++ coderd/database/queries.sql.go | 60 +++++++++++++++++++++ coderd/database/queries/provisionerjobs.sql | 14 +++++ coderd/workspacebuilds.go | 35 ++++++++++-- coderd/wsrelateddata_internal_test.go | 13 ++++- 9 files changed, 156 insertions(+), 5 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 632efc4c8e1..0c0d6376d78 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4528,6 +4528,12 @@ func (q *querier) GetProvisionerJobTimingsByJobID(ctx context.Context, jobID uui return q.db.GetProvisionerJobTimingsByJobID(ctx, jobID) } +func (q *querier) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]database.ProvisionerJob, error) { + // TODO: Remove this once we have a proper rbac check for provisioner jobs. + // Details in https://github.com/coder/coder/issues/16160 + return q.db.GetProvisionerJobsByIDs(ctx, ids) +} + func (q *querier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids database.GetProvisionerJobsByIDsWithQueuePositionParams) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) { // TODO: Remove this once we have a proper rbac check for provisioner jobs. // Details in https://github.com/coder/coder/issues/16160 diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 738b4a2de99..b0d9e00905b 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -5427,6 +5427,11 @@ func (s *MethodTestSuite) TestSystemFunctions() { dbm.EXPECT().GetWorkspaceAgentLogSourcesByAgentIDs(gomock.Any(), ids).Return([]database.WorkspaceAgentLogSource{}, nil).AnyTimes() check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead) })) + s.Run("GetProvisionerJobsByIDs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + ids := []uuid.UUID{} + dbm.EXPECT().GetProvisionerJobsByIDs(gomock.Any(), ids).Return([]database.ProvisionerJob{}, nil).AnyTimes() + check.Args(ids).Asserts() + })) s.Run("GetProvisionerJobsByIDsWithQueuePosition", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { arg := database.GetProvisionerJobsByIDsWithQueuePositionParams{} dbm.EXPECT().GetProvisionerJobsByIDsWithQueuePosition(gomock.Any(), arg).Return([]database.GetProvisionerJobsByIDsWithQueuePositionRow{}, nil).AnyTimes() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 492686145ad..4680e59233a 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2769,6 +2769,14 @@ func (m queryMetricsStore) GetProvisionerJobTimingsByJobID(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]database.ProvisionerJob, error) { + start := time.Now() + r0, r1 := m.s.GetProvisionerJobsByIDs(ctx, ids) + m.queryLatencies.WithLabelValues("GetProvisionerJobsByIDs").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetProvisionerJobsByIDs").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, arg database.GetProvisionerJobsByIDsWithQueuePositionParams) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) { start := time.Now() r0, r1 := m.s.GetProvisionerJobsByIDsWithQueuePosition(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 0c08caa669e..815d88d1253 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -5174,6 +5174,21 @@ func (mr *MockStoreMockRecorder) GetProvisionerJobTimingsByJobID(ctx, jobID any) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetProvisionerJobTimingsByJobID", reflect.TypeOf((*MockStore)(nil).GetProvisionerJobTimingsByJobID), ctx, jobID) } +// GetProvisionerJobsByIDs mocks base method. +func (m *MockStore) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]database.ProvisionerJob, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetProvisionerJobsByIDs", ctx, ids) + ret0, _ := ret[0].([]database.ProvisionerJob) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetProvisionerJobsByIDs indicates an expected call of GetProvisionerJobsByIDs. +func (mr *MockStoreMockRecorder) GetProvisionerJobsByIDs(ctx, ids any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetProvisionerJobsByIDs", reflect.TypeOf((*MockStore)(nil).GetProvisionerJobsByIDs), ctx, ids) +} + // GetProvisionerJobsByIDsWithQueuePosition mocks base method. func (m *MockStore) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, arg database.GetProvisionerJobsByIDsWithQueuePositionParams) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index ac1945f20f0..02849cb6605 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -753,6 +753,11 @@ type sqlcQuerier interface { // Blocks until the row is available for update. GetProvisionerJobByIDWithLock(ctx context.Context, id uuid.UUID) (ProvisionerJob, error) GetProvisionerJobTimingsByJobID(ctx context.Context, jobID uuid.UUID) ([]ProvisionerJobTiming, error) + // Fetches provisioner jobs by their IDs without computing queue position or + // queue size. Callers that do not need the queue position should prefer this + // over GetProvisionerJobsByIDsWithQueuePosition, whose window functions over + // pending jobs and provisioner daemons are comparatively expensive. + GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]ProvisionerJob, error) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, arg GetProvisionerJobsByIDsWithQueuePositionParams) ([]GetProvisionerJobsByIDsWithQueuePositionRow, error) GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisioner(ctx context.Context, arg GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerParams) ([]GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow, error) GetProvisionerJobsCreatedAfter(ctx context.Context, createdAt time.Time) ([]ProvisionerJob, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index e364a9400ea..e001e4aa6c1 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -23219,6 +23219,66 @@ func (q *sqlQuerier) GetProvisionerJobTimingsByJobID(ctx context.Context, jobID return items, nil } +const getProvisionerJobsByIDs = `-- name: GetProvisionerJobsByIDs :many +SELECT + id, created_at, updated_at, started_at, canceled_at, completed_at, error, organization_id, initiator_id, provisioner, storage_method, type, input, worker_id, file_id, tags, error_code, trace_metadata, job_status, logs_length, logs_overflowed +FROM + provisioner_jobs +WHERE + id = ANY($1 :: uuid [ ]) +ORDER BY + created_at +` + +// Fetches provisioner jobs by their IDs without computing queue position or +// queue size. Callers that do not need the queue position should prefer this +// over GetProvisionerJobsByIDsWithQueuePosition, whose window functions over +// pending jobs and provisioner daemons are comparatively expensive. +func (q *sqlQuerier) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]ProvisionerJob, error) { + rows, err := q.db.QueryContext(ctx, getProvisionerJobsByIDs, pq.Array(ids)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ProvisionerJob + for rows.Next() { + var i ProvisionerJob + if err := rows.Scan( + &i.ID, + &i.CreatedAt, + &i.UpdatedAt, + &i.StartedAt, + &i.CanceledAt, + &i.CompletedAt, + &i.Error, + &i.OrganizationID, + &i.InitiatorID, + &i.Provisioner, + &i.StorageMethod, + &i.Type, + &i.Input, + &i.WorkerID, + &i.FileID, + &i.Tags, + &i.ErrorCode, + &i.TraceMetadata, + &i.JobStatus, + &i.LogsLength, + &i.LogsOverflowed, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getProvisionerJobsByIDsWithQueuePosition = `-- name: GetProvisionerJobsByIDsWithQueuePosition :many WITH filtered_provisioner_jobs AS ( -- Step 1: Filter provisioner_jobs diff --git a/coderd/database/queries/provisionerjobs.sql b/coderd/database/queries/provisionerjobs.sql index 1b30e1edee3..cf1aa43fc58 100644 --- a/coderd/database/queries/provisionerjobs.sql +++ b/coderd/database/queries/provisionerjobs.sql @@ -67,6 +67,20 @@ WHERE id = $1 FOR UPDATE; +-- name: GetProvisionerJobsByIDs :many +-- Fetches provisioner jobs by their IDs without computing queue position or +-- queue size. Callers that do not need the queue position should prefer this +-- over GetProvisionerJobsByIDsWithQueuePosition, whose window functions over +-- pending jobs and provisioner daemons are comparatively expensive. +SELECT + * +FROM + provisioner_jobs +WHERE + id = ANY(@ids :: uuid [ ]) +ORDER BY + created_at; + -- name: GetProvisionerJobsByIDsWithQueuePosition :many WITH filtered_provisioner_jobs AS ( -- Step 1: Filter provisioner_jobs diff --git a/coderd/workspacebuilds.go b/coderd/workspacebuilds.go index e42dc1119b5..d570fc77d67 100644 --- a/coderd/workspacebuilds.go +++ b/coderd/workspacebuilds.go @@ -1110,6 +1110,36 @@ type workspaceBuildsData struct { provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow } +// provisionerJobsByIDs fetches provisioner jobs by ID, shaped as +// GetProvisionerJobsByIDsWithQueuePositionRow so callers can treat the result +// uniformly. When the selection requests the queue position it uses +// GetProvisionerJobsByIDsWithQueuePosition, whose queue position and size are +// computed with window functions over pending jobs and provisioner daemons and +// are comparatively expensive. Otherwise, it uses the cheaper +// GetProvisionerJobsByIDs and leaves QueuePosition and QueueSize zero. +func (api *API) provisionerJobsByIDs(ctx context.Context, jobIDs []uuid.UUID, cfg jobRelated) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) { + if cfg.QueuePosition { + return api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{ + IDs: jobIDs, + StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(), + }) + } + + provisionerJobs, err := api.Database.GetProvisionerJobsByIDs(ctx, jobIDs) + if err != nil { + return nil, err + } + jobs := make([]database.GetProvisionerJobsByIDsWithQueuePositionRow, 0, len(provisionerJobs)) + for _, job := range provisionerJobs { + jobs = append(jobs, database.GetProvisionerJobsByIDsWithQueuePositionRow{ + ID: job.ID, + CreatedAt: job.CreatedAt, + ProvisionerJob: job, + }) + } + return jobs, nil +} + func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []database.WorkspaceBuild, cfg latestBuildRelated) (workspaceBuildsData, error) { jobIDs := make([]uuid.UUID, 0, len(workspaceBuilds)) for _, build := range workspaceBuilds { @@ -1122,10 +1152,7 @@ func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []datab ) if cfg.Job != nil { var err error - jobs, err = api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{ - IDs: jobIDs, - StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(), - }) + jobs, err = api.provisionerJobsByIDs(ctx, jobIDs, *cfg.Job) if err != nil && !errors.Is(err, sql.ErrNoRows) { return workspaceBuildsData{}, xerrors.Errorf("get provisioner jobs: %w", err) } diff --git a/coderd/wsrelateddata_internal_test.go b/coderd/wsrelateddata_internal_test.go index ff7c586975e..e95846df180 100644 --- a/coderd/wsrelateddata_internal_test.go +++ b/coderd/wsrelateddata_internal_test.go @@ -63,6 +63,12 @@ func TestWorkspaceBuildsDataQueryGating(t *testing.T) { app := database.WorkspaceApp{ID: uuid.New(), AgentID: agent.ID} expectJob := func(db *dbmock.MockStore) { + db.EXPECT().GetProvisionerJobsByIDs(gomock.Any(), gomock.Any()). + Return([]database.ProvisionerJob{}, nil) + db.EXPECT().GetEligibleProvisionerDaemonsByProvisionerJobIDs(gomock.Any(), gomock.Any()). + Return([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow{}, nil) + } + expectJobWithQueuePosition := func(db *dbmock.MockStore) { db.EXPECT().GetProvisionerJobsByIDsWithQueuePosition(gomock.Any(), gomock.Any()). Return([]database.GetProvisionerJobsByIDsWithQueuePositionRow{}, nil) db.EXPECT().GetEligibleProvisionerDaemonsByProvisionerJobIDs(gomock.Any(), gomock.Any()). @@ -100,6 +106,11 @@ func TestWorkspaceBuildsDataQueryGating(t *testing.T) { cfg: latestBuildRelated{Job: &jobRelated{}}, setup: expectJob, }, + { + name: "JobWithQueuePosition", + cfg: latestBuildRelated{Job: &jobRelated{QueuePosition: true}}, + setup: expectJobWithQueuePosition, + }, { name: "TemplateVersion", cfg: latestBuildRelated{TemplateVersion: true}, @@ -171,7 +182,7 @@ func TestWorkspaceBuildsDataQueryGating(t *testing.T) { name: "All", cfg: allLatestBuildRelated(), setup: func(db *dbmock.MockStore) { - expectJob(db) + expectJobWithQueuePosition(db) expectTemplateVersion(db) expectResources(db) expectAgents(db)