opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add queue_position and queue_size to provisioner jobs (#8074)

Browse Source
This commit is contained in:
Kyle Carberry
2023-06-20 15:07:18 -05:00
committed by GitHub
parent bbb0fab1de
commit 69f911dfd5
26 changed files with 417 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+5
View File
@@ -1087,6 +1087,11 @@ func (q *querier) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID)
return q.db.GetProvisionerJobsByIDs(ctx, ids)
}
// TODO: we need to add a provisioner job resource
func (q *querier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
return q.db.GetProvisionerJobsByIDsWithQueuePosition(ctx, ids)
}
// TODO: We need to create a ProvisionerJob resource type
func (q *querier) GetProvisionerJobsCreatedAfter(ctx context.Context, createdAt time.Time) ([]database.ProvisionerJob, error) {
// if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil {
coderd/database/dbfake/dbfake.go
+32
View File
@@ -2051,6 +2051,38 @@ func (q *fakeQuerier) GetProvisionerJobsByIDs(_ context.Context, ids []uuid.UUID
return jobs, nil
}
func (q *fakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(_ context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
jobs := make([]database.GetProvisionerJobsByIDsWithQueuePositionRow, 0)
queuePosition := int64(1)
for _, job := range q.provisionerJobs {
for _, id := range ids {
if id == job.ID {
job := database.GetProvisionerJobsByIDsWithQueuePositionRow{
ProvisionerJob: job,
}
if !job.ProvisionerJob.StartedAt.Valid {
job.QueuePosition = queuePosition
}
jobs = append(jobs, job)
break
}
}
if !job.StartedAt.Valid {
queuePosition++
}
}
for _, job := range jobs {
if !job.ProvisionerJob.StartedAt.Valid {
// Set it to the max position!
job.QueueSize = queuePosition
}
}
return jobs, nil
}
func (q *fakeQuerier) GetProvisionerJobsCreatedAfter(_ context.Context, after time.Time) ([]database.ProvisionerJob, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
coderd/database/dbmetrics/dbmetrics.go
+7
View File
@@ -507,6 +507,13 @@ func (m metricsStore) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UU
return jobs, err
}
func (m metricsStore) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
start := time.Now()
r0, r1 := m.s.GetProvisionerJobsByIDsWithQueuePosition(ctx, ids)
m.queryLatencies.WithLabelValues("GetProvisionerJobsByIDsWithQueuePosition").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m metricsStore) GetProvisionerJobsCreatedAfter(ctx context.Context, createdAt time.Time) ([]database.ProvisionerJob, error) {
start := time.Now()
jobs, err := m.s.GetProvisionerJobsCreatedAfter(ctx, createdAt)
coderd/database/dbmock/dbmock.go
+15
View File
@@ -928,6 +928,21 @@ func (mr *MockStoreMockRecorder) GetProvisionerJobsByIDs(arg0, arg1 interface{})
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetProvisionerJobsByIDs", reflect.TypeOf((*MockStore)(nil).GetProvisionerJobsByIDs), arg0, arg1)
}
// GetProvisionerJobsByIDsWithQueuePosition mocks base method.
func (m *MockStore) GetProvisionerJobsByIDsWithQueuePosition(arg0 context.Context, arg1 []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetProvisionerJobsByIDsWithQueuePosition", arg0, arg1)
ret0, _ := ret[0].([]database.GetProvisionerJobsByIDsWithQueuePositionRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetProvisionerJobsByIDsWithQueuePosition indicates an expected call of GetProvisionerJobsByIDsWithQueuePosition.
func (mr *MockStoreMockRecorder) GetProvisionerJobsByIDsWithQueuePosition(arg0, arg1 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetProvisionerJobsByIDsWithQueuePosition", reflect.TypeOf((*MockStore)(nil).GetProvisionerJobsByIDsWithQueuePosition), arg0, arg1)
}
// GetProvisionerJobsCreatedAfter mocks base method.
func (m *MockStore) GetProvisionerJobsCreatedAfter(arg0 context.Context, arg1 time.Time) ([]database.ProvisionerJob, error) {
m.ctrl.T.Helper()
coderd/database/models.go
+1 -1
View File
@@ -1,6 +1,6 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.17.2
// sqlc v1.18.0
package database
coderd/database/querier.go
+2 -1
View File
@@ -1,6 +1,6 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.17.2
// sqlc v1.18.0
package database
@@ -90,6 +90,7 @@ type sqlcQuerier interface {
GetProvisionerDaemons(ctx context.Context) ([]ProvisionerDaemon, error)
GetProvisionerJobByID(ctx context.Context, id uuid.UUID) (ProvisionerJob, error)
GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]ProvisionerJob, error)
GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids []uuid.UUID) ([]GetProvisionerJobsByIDsWithQueuePositionRow, error)
GetProvisionerJobsCreatedAfter(ctx context.Context, createdAt time.Time) ([]ProvisionerJob, error)
GetProvisionerLogsAfterID(ctx context.Context, arg GetProvisionerLogsAfterIDParams) ([]ProvisionerJobLog, error)
GetQuotaAllowanceForUser(ctx context.Context, userID uuid.UUID) (int64, error)
coderd/database/querier_test.go
+75
View File
@@ -5,6 +5,8 @@ package database_test
import (
"context"
"database/sql"
"encoding/json"
"sort"
"testing"
"time"
@@ -312,3 +314,76 @@ func TestDefaultProxy(t *testing.T) {
require.NoError(t, err, "get deployment id")
require.Equal(t, depID, found)
}
func TestQueuePosition(t *testing.T) {
t.Parallel()
if testing.Short() {
t.SkipNow()
}
sqlDB := testSQLDB(t)
err := migrations.Up(sqlDB)
require.NoError(t, err)
db := database.New(sqlDB)
ctx := testutil.Context(t, testutil.WaitLong)
org := dbgen.Organization(t, db, database.Organization{})
jobCount := 10
jobs := []database.ProvisionerJob{}
jobIDs := []uuid.UUID{}
for i := 0; i < jobCount; i++ {
job := dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
OrganizationID: org.ID,
Tags: database.StringMap{},
})
jobs = append(jobs, job)
jobIDs = append(jobIDs, job.ID)
// We need a slight amount of time between each insertion to ensure that
// the queue position is correct... it's sorted by `created_at`.
time.Sleep(time.Millisecond)
}
queued, err := db.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
require.NoError(t, err)
require.Len(t, queued, jobCount)
sort.Slice(queued, func(i, j int) bool {
return queued[i].QueuePosition < queued[j].QueuePosition
})
// Ensure that the queue positions are correct based on insertion ID!
for index, job := range queued {
require.Equal(t, job.QueuePosition, int64(index+1))
require.Equal(t, job.ProvisionerJob.ID, jobs[index].ID)
}
job, err := db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
StartedAt: sql.NullTime{
Time: database.Now(),
Valid: true,
},
Types: database.AllProvisionerTypeValues(),
WorkerID: uuid.NullUUID{
UUID: uuid.New(),
Valid: true,
},
Tags: json.RawMessage("{}"),
})
require.NoError(t, err)
require.Equal(t, jobs[0].ID, job.ID)
queued, err = db.GetProvisionerJobsByIDsWithQueuePosition(ctx, jobIDs)
require.NoError(t, err)
require.Len(t, queued, jobCount)
sort.Slice(queued, func(i, j int) bool {
return queued[i].QueuePosition < queued[j].QueuePosition
})
// Ensure that queue positions are updated now that the first job has been acquired!
for index, job := range queued {
if index == 0 {
require.Equal(t, job.QueuePosition, int64(0))
continue
}
require.Equal(t, job.QueuePosition, int64(index))
require.Equal(t, job.ProvisionerJob.ID, jobs[index].ID)
}
}
coderd/database/queries.sql.go
+84 -1
View File
@@ -1,6 +1,6 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.17.2
// sqlc v1.18.0
package database
@@ -2287,6 +2287,89 @@ func (q *sqlQuerier) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUI
return items, nil
}
const getProvisionerJobsByIDsWithQueuePosition = `-- name: GetProvisionerJobsByIDsWithQueuePosition :many
WITH unstarted_jobs AS (
SELECT
id, created_at
FROM
provisioner_jobs
WHERE
started_at IS NULL
),
queue_position AS (
SELECT
id,
ROW_NUMBER() OVER (ORDER BY created_at ASC) AS queue_position
FROM
unstarted_jobs
),
queue_size AS (
SELECT COUNT(*) as count FROM unstarted_jobs
)
SELECT
pj.id, pj.created_at, pj.updated_at, pj.started_at, pj.canceled_at, pj.completed_at, pj.error, pj.organization_id, pj.initiator_id, pj.provisioner, pj.storage_method, pj.type, pj.input, pj.worker_id, pj.file_id, pj.tags, pj.error_code, pj.trace_metadata,
COALESCE(qp.queue_position, 0) AS queue_position,
COALESCE(qs.count, 0) AS queue_size
FROM
provisioner_jobs pj
LEFT JOIN
queue_position qp ON qp.id = pj.id
LEFT JOIN
queue_size qs ON TRUE
WHERE
pj.id = ANY($1 :: uuid [ ])
`
type GetProvisionerJobsByIDsWithQueuePositionRow struct {
ProvisionerJob ProvisionerJob `db:"provisionerjob" json:"provisionerjob"`
QueuePosition int64 `db:"queue_position" json:"queue_position"`
QueueSize int64 `db:"queue_size" json:"queue_size"`
}
func (q *sqlQuerier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids []uuid.UUID) ([]GetProvisionerJobsByIDsWithQueuePositionRow, error) {
rows, err := q.db.QueryContext(ctx, getProvisionerJobsByIDsWithQueuePosition, pq.Array(ids))
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetProvisionerJobsByIDsWithQueuePositionRow
for rows.Next() {
var i GetProvisionerJobsByIDsWithQueuePositionRow
if err := rows.Scan(
&i.ProvisionerJob.ID,
&i.ProvisionerJob.CreatedAt,
&i.ProvisionerJob.UpdatedAt,
&i.ProvisionerJob.StartedAt,
&i.ProvisionerJob.CanceledAt,
&i.ProvisionerJob.CompletedAt,
&i.ProvisionerJob.Error,
&i.ProvisionerJob.OrganizationID,
&i.ProvisionerJob.InitiatorID,
&i.ProvisionerJob.Provisioner,
&i.ProvisionerJob.StorageMethod,
&i.ProvisionerJob.Type,
&i.ProvisionerJob.Input,
&i.ProvisionerJob.WorkerID,
&i.ProvisionerJob.FileID,
&i.ProvisionerJob.Tags,
&i.ProvisionerJob.ErrorCode,
&i.ProvisionerJob.TraceMetadata,
&i.QueuePosition,
&i.QueueSize,
); 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 getProvisionerJobsCreatedAfter = `-- name: GetProvisionerJobsCreatedAfter :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 FROM provisioner_jobs WHERE created_at > $1
`
coderd/database/queries/provisionerjobs.sql
+32
View File
@@ -47,6 +47,38 @@ FROM
WHERE
id = ANY(@ids :: uuid [ ]);
-- name: GetProvisionerJobsByIDsWithQueuePosition :many
WITH unstarted_jobs AS (
SELECT
id, created_at
FROM
provisioner_jobs
WHERE
started_at IS NULL
),
queue_position AS (
SELECT
id,
ROW_NUMBER() OVER (ORDER BY created_at ASC) AS queue_position
FROM
unstarted_jobs
),
queue_size AS (
SELECT COUNT(*) as count FROM unstarted_jobs
)
SELECT
sqlc.embed(pj),
COALESCE(qp.queue_position, 0) AS queue_position,
COALESCE(qs.count, 0) AS queue_size
FROM
provisioner_jobs pj
LEFT JOIN
queue_position qp ON qp.id = pj.id
LEFT JOIN
queue_size qs ON TRUE
WHERE
pj.id = ANY(@ids :: uuid [ ]);
-- name: GetProvisionerJobsCreatedAfter :many
SELECT * FROM provisioner_jobs WHERE created_at > $1;