feat: add migrations and queries to support prebuilds (#16891)

Depends on https://github.com/coder/coder/pull/16916 _(change base to
`main` once it is merged)_

Closes https://github.com/coder/internal/issues/514

_This is one of several PRs to decompose the `dk/prebuilds` feature
branch into separate PRs to merge into `main`._

---------

Signed-off-by: Danny Kopping <dannykopping@gmail.com>
Co-authored-by: Danny Kopping <dannykopping@gmail.com>
Co-authored-by: evgeniy-scherbina <evgeniy.shcherbina.es@gmail.com>
This commit is contained in:
Sas Swart
2025-04-03 10:58:30 +02:00
committed by GitHub
parent 4aa45a5c43
commit 99c6f235eb
22 changed files with 2110 additions and 59 deletions

View File

@ -5961,9 +5961,413 @@ func (q *sqlQuerier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.
return items, nil
}
const claimPrebuiltWorkspace = `-- name: ClaimPrebuiltWorkspace :one
UPDATE workspaces w
SET owner_id = $1::uuid,
name = $2::text,
updated_at = NOW()
WHERE w.id IN (
SELECT p.id
FROM workspace_prebuilds p
INNER JOIN workspace_latest_builds b ON b.workspace_id = p.id
INNER JOIN templates t ON p.template_id = t.id
WHERE (b.transition = 'start'::workspace_transition
AND b.job_status IN ('succeeded'::provisioner_job_status))
-- The prebuilds system should never try to claim a prebuild for an inactive template version.
-- Nevertheless, this filter is here as a defensive measure:
AND b.template_version_id = t.active_version_id
AND p.current_preset_id = $3::uuid
AND p.ready
LIMIT 1 FOR UPDATE OF p SKIP LOCKED -- Ensure that a concurrent request will not select the same prebuild.
)
RETURNING w.id, w.name
`
type ClaimPrebuiltWorkspaceParams struct {
NewUserID uuid.UUID `db:"new_user_id" json:"new_user_id"`
NewName string `db:"new_name" json:"new_name"`
PresetID uuid.UUID `db:"preset_id" json:"preset_id"`
}
type ClaimPrebuiltWorkspaceRow struct {
ID uuid.UUID `db:"id" json:"id"`
Name string `db:"name" json:"name"`
}
func (q *sqlQuerier) ClaimPrebuiltWorkspace(ctx context.Context, arg ClaimPrebuiltWorkspaceParams) (ClaimPrebuiltWorkspaceRow, error) {
row := q.db.QueryRowContext(ctx, claimPrebuiltWorkspace, arg.NewUserID, arg.NewName, arg.PresetID)
var i ClaimPrebuiltWorkspaceRow
err := row.Scan(&i.ID, &i.Name)
return i, err
}
const countInProgressPrebuilds = `-- name: CountInProgressPrebuilds :many
SELECT t.id AS template_id, wpb.template_version_id, wpb.transition, COUNT(wpb.transition)::int AS count
FROM workspace_latest_builds wlb
INNER JOIN workspace_prebuild_builds wpb ON wpb.id = wlb.id
-- We only need these counts for active template versions.
-- It doesn't influence whether we create or delete prebuilds
-- for inactive template versions. This is because we never create
-- prebuilds for inactive template versions, we always delete
-- running prebuilds for inactive template versions, and we ignore
-- prebuilds that are still building.
INNER JOIN templates t ON t.active_version_id = wlb.template_version_id
WHERE wlb.job_status IN ('pending'::provisioner_job_status, 'running'::provisioner_job_status)
GROUP BY t.id, wpb.template_version_id, wpb.transition
`
type CountInProgressPrebuildsRow struct {
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
Transition WorkspaceTransition `db:"transition" json:"transition"`
Count int32 `db:"count" json:"count"`
}
// CountInProgressPrebuilds returns the number of in-progress prebuilds, grouped by template version ID and transition.
// Prebuild considered in-progress if it's in the "starting", "stopping", or "deleting" state.
func (q *sqlQuerier) CountInProgressPrebuilds(ctx context.Context) ([]CountInProgressPrebuildsRow, error) {
rows, err := q.db.QueryContext(ctx, countInProgressPrebuilds)
if err != nil {
return nil, err
}
defer rows.Close()
var items []CountInProgressPrebuildsRow
for rows.Next() {
var i CountInProgressPrebuildsRow
if err := rows.Scan(
&i.TemplateID,
&i.TemplateVersionID,
&i.Transition,
&i.Count,
); 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 getPrebuildMetrics = `-- name: GetPrebuildMetrics :many
SELECT
t.name as template_name,
tvp.name as preset_name,
o.name as organization_name,
COUNT(*) as created_count,
COUNT(*) FILTER (WHERE pj.job_status = 'failed'::provisioner_job_status) as failed_count,
COUNT(*) FILTER (
WHERE w.owner_id != 'c42fdf75-3097-471c-8c33-fb52454d81c0'::uuid -- The system user responsible for prebuilds.
) as claimed_count
FROM workspaces w
INNER JOIN workspace_prebuild_builds wpb ON wpb.workspace_id = w.id
INNER JOIN templates t ON t.id = w.template_id
INNER JOIN template_version_presets tvp ON tvp.id = wpb.template_version_preset_id
INNER JOIN provisioner_jobs pj ON pj.id = wpb.job_id
INNER JOIN organizations o ON o.id = w.organization_id
WHERE NOT t.deleted AND wpb.build_number = 1
GROUP BY t.name, tvp.name, o.name
ORDER BY t.name, tvp.name, o.name
`
type GetPrebuildMetricsRow struct {
TemplateName string `db:"template_name" json:"template_name"`
PresetName string `db:"preset_name" json:"preset_name"`
OrganizationName string `db:"organization_name" json:"organization_name"`
CreatedCount int64 `db:"created_count" json:"created_count"`
FailedCount int64 `db:"failed_count" json:"failed_count"`
ClaimedCount int64 `db:"claimed_count" json:"claimed_count"`
}
func (q *sqlQuerier) GetPrebuildMetrics(ctx context.Context) ([]GetPrebuildMetricsRow, error) {
rows, err := q.db.QueryContext(ctx, getPrebuildMetrics)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetPrebuildMetricsRow
for rows.Next() {
var i GetPrebuildMetricsRow
if err := rows.Scan(
&i.TemplateName,
&i.PresetName,
&i.OrganizationName,
&i.CreatedCount,
&i.FailedCount,
&i.ClaimedCount,
); 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 getPresetsBackoff = `-- name: GetPresetsBackoff :many
WITH filtered_builds AS (
-- Only select builds which are for prebuild creations
SELECT wlb.template_version_id, wlb.created_at, tvp.id AS preset_id, wlb.job_status, tvp.desired_instances
FROM template_version_presets tvp
INNER JOIN workspace_latest_builds wlb ON wlb.template_version_preset_id = tvp.id
INNER JOIN workspaces w ON wlb.workspace_id = w.id
INNER JOIN template_versions tv ON wlb.template_version_id = tv.id
INNER JOIN templates t ON tv.template_id = t.id AND t.active_version_id = tv.id
WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a prebuild configuration.
AND wlb.transition = 'start'::workspace_transition
AND w.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'
),
time_sorted_builds AS (
-- Group builds by preset, then sort each group by created_at.
SELECT fb.template_version_id, fb.created_at, fb.preset_id, fb.job_status, fb.desired_instances,
ROW_NUMBER() OVER (PARTITION BY fb.preset_id ORDER BY fb.created_at DESC) as rn
FROM filtered_builds fb
),
failed_count AS (
-- Count failed builds per preset in the given period
SELECT preset_id, COUNT(*) AS num_failed
FROM filtered_builds
WHERE job_status = 'failed'::provisioner_job_status
AND created_at >= $1::timestamptz
GROUP BY preset_id
)
SELECT
tsb.template_version_id,
tsb.preset_id,
COALESCE(fc.num_failed, 0)::int AS num_failed,
MAX(tsb.created_at)::timestamptz AS last_build_at
FROM time_sorted_builds tsb
LEFT JOIN failed_count fc ON fc.preset_id = tsb.preset_id
WHERE tsb.rn <= tsb.desired_instances -- Fetch the last N builds, where N is the number of desired instances; if any fail, we backoff
AND tsb.job_status = 'failed'::provisioner_job_status
AND created_at >= $1::timestamptz
GROUP BY tsb.template_version_id, tsb.preset_id, fc.num_failed
`
type GetPresetsBackoffRow struct {
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
PresetID uuid.UUID `db:"preset_id" json:"preset_id"`
NumFailed int32 `db:"num_failed" json:"num_failed"`
LastBuildAt time.Time `db:"last_build_at" json:"last_build_at"`
}
// GetPresetsBackoff groups workspace builds by preset ID.
// Each preset is associated with exactly one template version ID.
// For each group, the query checks up to N of the most recent jobs that occurred within the
// lookback period, where N equals the number of desired instances for the corresponding preset.
// If at least one of the job within a group has failed, we should backoff on the corresponding preset ID.
// Query returns a list of preset IDs for which we should backoff.
// Only active template versions with configured presets are considered.
// We also return the number of failed workspace builds that occurred during the lookback period.
//
// NOTE:
// - To **decide whether to back off**, we look at up to the N most recent builds (within the defined lookback period).
// - To **calculate the number of failed builds**, we consider all builds within the defined lookback period.
//
// The number of failed builds is used downstream to determine the backoff duration.
func (q *sqlQuerier) GetPresetsBackoff(ctx context.Context, lookback time.Time) ([]GetPresetsBackoffRow, error) {
rows, err := q.db.QueryContext(ctx, getPresetsBackoff, lookback)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetPresetsBackoffRow
for rows.Next() {
var i GetPresetsBackoffRow
if err := rows.Scan(
&i.TemplateVersionID,
&i.PresetID,
&i.NumFailed,
&i.LastBuildAt,
); 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 getRunningPrebuiltWorkspaces = `-- name: GetRunningPrebuiltWorkspaces :many
SELECT
p.id,
p.name,
p.template_id,
b.template_version_id,
p.current_preset_id AS current_preset_id,
p.ready,
p.created_at
FROM workspace_prebuilds p
INNER JOIN workspace_latest_builds b ON b.workspace_id = p.id
WHERE (b.transition = 'start'::workspace_transition
AND b.job_status = 'succeeded'::provisioner_job_status)
`
type GetRunningPrebuiltWorkspacesRow struct {
ID uuid.UUID `db:"id" json:"id"`
Name string `db:"name" json:"name"`
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
CurrentPresetID uuid.NullUUID `db:"current_preset_id" json:"current_preset_id"`
Ready bool `db:"ready" json:"ready"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
}
func (q *sqlQuerier) GetRunningPrebuiltWorkspaces(ctx context.Context) ([]GetRunningPrebuiltWorkspacesRow, error) {
rows, err := q.db.QueryContext(ctx, getRunningPrebuiltWorkspaces)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetRunningPrebuiltWorkspacesRow
for rows.Next() {
var i GetRunningPrebuiltWorkspacesRow
if err := rows.Scan(
&i.ID,
&i.Name,
&i.TemplateID,
&i.TemplateVersionID,
&i.CurrentPresetID,
&i.Ready,
&i.CreatedAt,
); 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 getTemplatePresetsWithPrebuilds = `-- name: GetTemplatePresetsWithPrebuilds :many
SELECT
t.id AS template_id,
t.name AS template_name,
o.name AS organization_name,
tv.id AS template_version_id,
tv.name AS template_version_name,
tv.id = t.active_version_id AS using_active_version,
tvp.id,
tvp.name,
tvp.desired_instances AS desired_instances,
t.deleted,
t.deprecated != '' AS deprecated
FROM templates t
INNER JOIN template_versions tv ON tv.template_id = t.id
INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id
INNER JOIN organizations o ON o.id = t.organization_id
WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a prebuild configuration.
AND (t.id = $1::uuid OR $1 IS NULL)
`
type GetTemplatePresetsWithPrebuildsRow struct {
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
TemplateName string `db:"template_name" json:"template_name"`
OrganizationName string `db:"organization_name" json:"organization_name"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
TemplateVersionName string `db:"template_version_name" json:"template_version_name"`
UsingActiveVersion bool `db:"using_active_version" json:"using_active_version"`
ID uuid.UUID `db:"id" json:"id"`
Name string `db:"name" json:"name"`
DesiredInstances sql.NullInt32 `db:"desired_instances" json:"desired_instances"`
Deleted bool `db:"deleted" json:"deleted"`
Deprecated bool `db:"deprecated" json:"deprecated"`
}
// GetTemplatePresetsWithPrebuilds retrieves template versions with configured presets and prebuilds.
// It also returns the number of desired instances for each preset.
// If template_id is specified, only template versions associated with that template will be returned.
func (q *sqlQuerier) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]GetTemplatePresetsWithPrebuildsRow, error) {
rows, err := q.db.QueryContext(ctx, getTemplatePresetsWithPrebuilds, templateID)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetTemplatePresetsWithPrebuildsRow
for rows.Next() {
var i GetTemplatePresetsWithPrebuildsRow
if err := rows.Scan(
&i.TemplateID,
&i.TemplateName,
&i.OrganizationName,
&i.TemplateVersionID,
&i.TemplateVersionName,
&i.UsingActiveVersion,
&i.ID,
&i.Name,
&i.DesiredInstances,
&i.Deleted,
&i.Deprecated,
); 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 getPresetByID = `-- name: GetPresetByID :one
SELECT tvp.id, tvp.template_version_id, tvp.name, tvp.created_at, tvp.desired_instances, tvp.invalidate_after_secs, tv.template_id, tv.organization_id FROM
template_version_presets tvp
INNER JOIN template_versions tv ON tvp.template_version_id = tv.id
WHERE tvp.id = $1
`
type GetPresetByIDRow struct {
ID uuid.UUID `db:"id" json:"id"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
Name string `db:"name" json:"name"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
DesiredInstances sql.NullInt32 `db:"desired_instances" json:"desired_instances"`
InvalidateAfterSecs sql.NullInt32 `db:"invalidate_after_secs" json:"invalidate_after_secs"`
TemplateID uuid.NullUUID `db:"template_id" json:"template_id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
}
func (q *sqlQuerier) GetPresetByID(ctx context.Context, presetID uuid.UUID) (GetPresetByIDRow, error) {
row := q.db.QueryRowContext(ctx, getPresetByID, presetID)
var i GetPresetByIDRow
err := row.Scan(
&i.ID,
&i.TemplateVersionID,
&i.Name,
&i.CreatedAt,
&i.DesiredInstances,
&i.InvalidateAfterSecs,
&i.TemplateID,
&i.OrganizationID,
)
return i, err
}
const getPresetByWorkspaceBuildID = `-- name: GetPresetByWorkspaceBuildID :one
SELECT
template_version_presets.id, template_version_presets.template_version_id, template_version_presets.name, template_version_presets.created_at
template_version_presets.id, template_version_presets.template_version_id, template_version_presets.name, template_version_presets.created_at, template_version_presets.desired_instances, template_version_presets.invalidate_after_secs
FROM
template_version_presets
INNER JOIN workspace_builds ON workspace_builds.template_version_preset_id = template_version_presets.id
@ -5979,6 +6383,8 @@ func (q *sqlQuerier) GetPresetByWorkspaceBuildID(ctx context.Context, workspaceB
&i.TemplateVersionID,
&i.Name,
&i.CreatedAt,
&i.DesiredInstances,
&i.InvalidateAfterSecs,
)
return i, err
}
@ -6023,7 +6429,7 @@ func (q *sqlQuerier) GetPresetParametersByTemplateVersionID(ctx context.Context,
const getPresetsByTemplateVersionID = `-- name: GetPresetsByTemplateVersionID :many
SELECT
id, template_version_id, name, created_at
id, template_version_id, name, created_at, desired_instances, invalidate_after_secs
FROM
template_version_presets
WHERE
@ -6044,6 +6450,8 @@ func (q *sqlQuerier) GetPresetsByTemplateVersionID(ctx context.Context, template
&i.TemplateVersionID,
&i.Name,
&i.CreatedAt,
&i.DesiredInstances,
&i.InvalidateAfterSecs,
); err != nil {
return nil, err
}
@ -6059,26 +6467,46 @@ func (q *sqlQuerier) GetPresetsByTemplateVersionID(ctx context.Context, template
}
const insertPreset = `-- name: InsertPreset :one
INSERT INTO
template_version_presets (template_version_id, name, created_at)
VALUES
($1, $2, $3) RETURNING id, template_version_id, name, created_at
INSERT INTO template_version_presets (
template_version_id,
name,
created_at,
desired_instances,
invalidate_after_secs
)
VALUES (
$1,
$2,
$3,
$4,
$5
) RETURNING id, template_version_id, name, created_at, desired_instances, invalidate_after_secs
`
type InsertPresetParams struct {
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
Name string `db:"name" json:"name"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
Name string `db:"name" json:"name"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
DesiredInstances sql.NullInt32 `db:"desired_instances" json:"desired_instances"`
InvalidateAfterSecs sql.NullInt32 `db:"invalidate_after_secs" json:"invalidate_after_secs"`
}
func (q *sqlQuerier) InsertPreset(ctx context.Context, arg InsertPresetParams) (TemplateVersionPreset, error) {
row := q.db.QueryRowContext(ctx, insertPreset, arg.TemplateVersionID, arg.Name, arg.CreatedAt)
row := q.db.QueryRowContext(ctx, insertPreset,
arg.TemplateVersionID,
arg.Name,
arg.CreatedAt,
arg.DesiredInstances,
arg.InvalidateAfterSecs,
)
var i TemplateVersionPreset
err := row.Scan(
&i.ID,
&i.TemplateVersionID,
&i.Name,
&i.CreatedAt,
&i.DesiredInstances,
&i.InvalidateAfterSecs,
)
return i, err
}