Implement Quotas v3 (#5012)

* provisioner/terraform: add cost to resource_metadata

* provisionerd/runner: use Options struct

* Complete provisionerd implementation

* Add quota_allowance to groups

* Combine Quota and RBAC licenses

* Add Opts to InTx
This commit is contained in:
Ammar Bandukwala
2022-11-14 11:57:33 -06:00
committed by GitHub
parent 3fb7892c07
commit 97dbd4dc5d
101 changed files with 1577 additions and 847 deletions

View File

@ -63,7 +63,7 @@ func activityBumpWorkspace(log slog.Logger, db database.Store, workspace databas
return xerrors.Errorf("update workspace build: %w", err)
}
return nil
})
}, nil)
if err != nil {
log.Error(
ctx, "bump failed",

View File

@ -177,7 +177,7 @@ func (e *Executor) runOnce(t time.Time) Stats {
}
return nil
})
}, nil)
if err != nil {
log.Error(e.ctx, "workspace scheduling failed", slog.Error(err))
}

View File

@ -40,9 +40,9 @@ import (
"github.com/coder/coder/coderd/rbac"
"github.com/coder/coder/coderd/telemetry"
"github.com/coder/coder/coderd/tracing"
"github.com/coder/coder/coderd/workspacequota"
"github.com/coder/coder/coderd/wsconncache"
"github.com/coder/coder/codersdk"
"github.com/coder/coder/provisionerd/proto"
"github.com/coder/coder/site"
"github.com/coder/coder/tailnet"
)
@ -66,7 +66,6 @@ type Options struct {
CacheDir string
Auditor audit.Auditor
WorkspaceQuotaEnforcer workspacequota.Enforcer
AgentConnectionUpdateFrequency time.Duration
AgentInactiveDisconnectTimeout time.Duration
// APIRateLimit is the minutely throughput rate limit per user or ip.
@ -145,9 +144,6 @@ func New(options *Options) *API {
if options.Auditor == nil {
options.Auditor = audit.NewNop()
}
if options.WorkspaceQuotaEnforcer == nil {
options.WorkspaceQuotaEnforcer = workspacequota.NewNop()
}
siteCacheDir := options.CacheDir
if siteCacheDir != "" {
@ -174,12 +170,10 @@ func New(options *Options) *API {
Authorizer: options.Authorizer,
Logger: options.Logger,
},
metricsCache: metricsCache,
Auditor: atomic.Pointer[audit.Auditor]{},
WorkspaceQuotaEnforcer: atomic.Pointer[workspacequota.Enforcer]{},
metricsCache: metricsCache,
Auditor: atomic.Pointer[audit.Auditor]{},
}
api.Auditor.Store(&options.Auditor)
api.WorkspaceQuotaEnforcer.Store(&options.WorkspaceQuotaEnforcer)
api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgentTailnet, 0)
api.TailnetCoordinator.Store(&options.TailnetCoordinator)
oauthConfigs := &httpmw.OAuth2Configs{
@ -590,8 +584,8 @@ type API struct {
ID uuid.UUID
Auditor atomic.Pointer[audit.Auditor]
WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool]
WorkspaceQuotaEnforcer atomic.Pointer[workspacequota.Enforcer]
TailnetCoordinator atomic.Pointer[tailnet.Coordinator]
QuotaCommitter atomic.Pointer[proto.QuotaCommitter]
HTTPAuth *HTTPAuthorizer
// APIHandler serves "/api/v2"

View File

@ -528,7 +528,8 @@ func AwaitWorkspaceBuildJob(t *testing.T, client *codersdk.Client, build uuid.UU
t.Logf("waiting for workspace build job %s", build)
var workspaceBuild codersdk.WorkspaceBuild
require.Eventually(t, func() bool {
workspaceBuild, err := client.WorkspaceBuild(context.Background(), build)
var err error
workspaceBuild, err = client.WorkspaceBuild(context.Background(), build)
return assert.NoError(t, err) && workspaceBuild.Job.CompletedAt != nil
}, testutil.WaitShort, testutil.IntervalFast)
return workspaceBuild

View File

@ -121,7 +121,7 @@ func (*fakeQuerier) Ping(_ context.Context) (time.Duration, error) {
}
// InTx doesn't rollback data properly for in-memory yet.
func (q *fakeQuerier) InTx(fn func(database.Store) error) error {
func (q *fakeQuerier) InTx(fn func(database.Store) error, _ *sql.TxOptions) error {
q.mutex.Lock()
defer q.mutex.Unlock()
return fn(&fakeQuerier{mutex: inTxMutex{}, data: q.data})
@ -2246,6 +2246,7 @@ func (q *fakeQuerier) InsertWorkspaceResource(_ context.Context, arg database.In
Name: arg.Name,
Hide: arg.Hide,
Icon: arg.Icon,
DailyCost: arg.DailyCost,
}
q.provisionerJobResources = append(q.provisionerJobResources, resource)
return resource, nil
@ -2757,6 +2758,20 @@ func (q *fakeQuerier) UpdateWorkspaceBuildByID(_ context.Context, arg database.U
}
return database.WorkspaceBuild{}, sql.ErrNoRows
}
func (q *fakeQuerier) UpdateWorkspaceBuildCostByID(_ context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) (database.WorkspaceBuild, error) {
q.mutex.Lock()
defer q.mutex.Unlock()
for index, workspaceBuild := range q.workspaceBuilds {
if workspaceBuild.ID != arg.ID {
continue
}
workspaceBuild.DailyCost = arg.DailyCost
q.workspaceBuilds[index] = workspaceBuild
return workspaceBuild, nil
}
return database.WorkspaceBuild{}, sql.ErrNoRows
}
func (q *fakeQuerier) UpdateWorkspaceDeletedByID(_ context.Context, arg database.UpdateWorkspaceDeletedByIDParams) error {
q.mutex.Lock()
@ -2858,6 +2873,7 @@ func (q *fakeQuerier) UpdateGroupByID(_ context.Context, arg database.UpdateGrou
if group.ID == arg.ID {
group.Name = arg.Name
group.AvatarURL = arg.AvatarURL
group.QuotaAllowance = arg.QuotaAllowance
q.groups[i] = group
return group, nil
}
@ -3230,6 +3246,7 @@ func (q *fakeQuerier) InsertGroup(_ context.Context, arg database.InsertGroupPar
Name: arg.Name,
OrganizationID: arg.OrganizationID,
AvatarURL: arg.AvatarURL,
QuotaAllowance: arg.QuotaAllowance,
}
q.groups = append(q.groups, group)
@ -3430,3 +3447,46 @@ func (q *fakeQuerier) UpdateGitAuthLink(_ context.Context, arg database.UpdateGi
}
return nil
}
func (q *fakeQuerier) GetQuotaAllowanceForUser(_ context.Context, userID uuid.UUID) (int64, error) {
q.mutex.Lock()
defer q.mutex.Unlock()
var sum int64
for _, member := range q.groupMembers {
if member.UserID != userID {
continue
}
for _, group := range q.groups {
if group.ID == member.GroupID {
sum += int64(group.QuotaAllowance)
}
}
}
return sum, nil
}
func (q *fakeQuerier) GetQuotaConsumedForUser(_ context.Context, userID uuid.UUID) (int64, error) {
q.mutex.Lock()
defer q.mutex.Unlock()
var sum int64
for _, workspace := range q.workspaces {
if workspace.OwnerID != userID {
continue
}
if workspace.Deleted {
continue
}
var lastBuild database.WorkspaceBuild
for _, build := range q.workspaceBuilds {
if build.WorkspaceID != workspace.ID {
continue
}
if build.CreatedAt.After(lastBuild.CreatedAt) {
lastBuild = build
}
}
sum += int64(lastBuild.DailyCost)
}
return sum, nil
}

View File

@ -38,7 +38,7 @@ func TestInTx(t *testing.T) {
})
assert.NoError(t, err)
return nil
})
}, nil)
assert.NoError(t, err)
}()
var nums []int

View File

@ -26,7 +26,7 @@ type Store interface {
customQuerier
Ping(ctx context.Context) (time.Duration, error)
InTx(func(Store) error) error
InTx(func(Store) error, *sql.TxOptions) error
}
// DBTX represents a database connection or transaction.
@ -68,7 +68,7 @@ func (q *sqlQuerier) Ping(ctx context.Context) (time.Duration, error) {
}
// InTx performs database operations inside a transaction.
func (q *sqlQuerier) InTx(function func(Store) error) error {
func (q *sqlQuerier) InTx(function func(Store) error, txOpts *sql.TxOptions) error {
if _, ok := q.db.(*sqlx.Tx); ok {
// If the current inner "db" is already a transaction, we just reuse it.
// We do not need to handle commit/rollback as the outer tx will handle
@ -80,7 +80,7 @@ func (q *sqlQuerier) InTx(function func(Store) error) error {
return nil
}
transaction, err := q.sdb.BeginTxx(context.Background(), nil)
transaction, err := q.sdb.BeginTxx(context.Background(), txOpts)
if err != nil {
return xerrors.Errorf("begin transaction: %w", err)
}

View File

@ -43,8 +43,8 @@ func TestNestedInTx(t *testing.T) {
LoginType: database.LoginTypeGithub,
})
return err
})
})
}, nil)
}, nil)
require.NoError(t, err, "outer tx: %w", err)
user, err := db.GetUserByID(context.Background(), uid)

View File

@ -19,9 +19,16 @@ func NewDB(t *testing.T) (database.Store, database.Pubsub) {
db := databasefake.New()
pubsub := database.NewPubsubInMemory()
if os.Getenv("DB") != "" {
connectionURL, closePg, err := postgres.Open()
require.NoError(t, err)
t.Cleanup(closePg)
connectionURL := os.Getenv("CODER_PG_CONNECTION_URL")
if connectionURL == "" {
var (
err error
closePg func()
)
connectionURL, closePg, err = postgres.Open()
require.NoError(t, err)
t.Cleanup(closePg)
}
sqlDB, err := sql.Open("postgres", connectionURL)
require.NoError(t, err)
t.Cleanup(func() {

View File

@ -192,7 +192,8 @@ CREATE TABLE groups (
id uuid NOT NULL,
name text NOT NULL,
organization_id uuid NOT NULL,
avatar_url text DEFAULT ''::text NOT NULL
avatar_url text DEFAULT ''::text NOT NULL,
quota_allowance integer DEFAULT 0 NOT NULL
);
CREATE TABLE licenses (
@ -444,7 +445,8 @@ CREATE TABLE workspace_builds (
provisioner_state bytea,
job_id uuid NOT NULL,
deadline timestamp with time zone DEFAULT '0001-01-01 00:00:00+00'::timestamp with time zone NOT NULL,
reason build_reason DEFAULT 'initiator'::build_reason NOT NULL
reason build_reason DEFAULT 'initiator'::build_reason NOT NULL,
daily_cost integer DEFAULT 0 NOT NULL
);
CREATE TABLE workspace_resource_metadata (
@ -463,7 +465,8 @@ CREATE TABLE workspace_resources (
name character varying(64) NOT NULL,
hide boolean DEFAULT false NOT NULL,
icon character varying(256) DEFAULT ''::character varying NOT NULL,
instance_type character varying(256)
instance_type character varying(256),
daily_cost integer DEFAULT 0 NOT NULL
);
CREATE TABLE workspaces (

View File

@ -0,0 +1,3 @@
ALTER TABLE workspace_builds DROP COLUMN daily_cost;
ALTER TABLE workspace_resources DROP COLUMN daily_cost;
ALTER TABLE groups DROP COLUMN quota_allowance;

View File

@ -0,0 +1,3 @@
ALTER TABLE workspace_builds ADD COLUMN daily_cost int NOT NULL DEFAULT 0;
ALTER TABLE workspace_resources ADD COLUMN daily_cost int NOT NULL DEFAULT 0;
ALTER TABLE groups ADD COLUMN quota_allowance int NOT NULL DEFAULT 0;

View File

@ -454,6 +454,7 @@ type Group struct {
Name string `db:"name" json:"name"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
QuotaAllowance int32 `db:"quota_allowance" json:"quota_allowance"`
}
type GroupMember struct {
@ -700,6 +701,7 @@ type WorkspaceBuild struct {
JobID uuid.UUID `db:"job_id" json:"job_id"`
Deadline time.Time `db:"deadline" json:"deadline"`
Reason BuildReason `db:"reason" json:"reason"`
DailyCost int32 `db:"daily_cost" json:"daily_cost"`
}
type WorkspaceResource struct {
@ -712,6 +714,7 @@ type WorkspaceResource struct {
Hide bool `db:"hide" json:"hide"`
Icon string `db:"icon" json:"icon"`
InstanceType sql.NullString `db:"instance_type" json:"instance_type"`
DailyCost int32 `db:"daily_cost" json:"daily_cost"`
}
type WorkspaceResourceMetadatum struct {

View File

@ -72,6 +72,8 @@ type sqlcQuerier interface {
GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]ProvisionerJob, error)
GetProvisionerJobsCreatedAfter(ctx context.Context, createdAt time.Time) ([]ProvisionerJob, error)
GetProvisionerLogsByIDBetween(ctx context.Context, arg GetProvisionerLogsByIDBetweenParams) ([]ProvisionerJobLog, error)
GetQuotaAllowanceForUser(ctx context.Context, userID uuid.UUID) (int64, error)
GetQuotaConsumedForUser(ctx context.Context, ownerID uuid.UUID) (int64, error)
GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]Replica, error)
GetTemplateAverageBuildTime(ctx context.Context, arg GetTemplateAverageBuildTimeParams) (GetTemplateAverageBuildTimeRow, error)
GetTemplateByID(ctx context.Context, id uuid.UUID) (Template, error)
@ -187,6 +189,7 @@ type sqlcQuerier interface {
UpdateWorkspaceAppHealthByID(ctx context.Context, arg UpdateWorkspaceAppHealthByIDParams) error
UpdateWorkspaceAutostart(ctx context.Context, arg UpdateWorkspaceAutostartParams) error
UpdateWorkspaceBuildByID(ctx context.Context, arg UpdateWorkspaceBuildByIDParams) (WorkspaceBuild, error)
UpdateWorkspaceBuildCostByID(ctx context.Context, arg UpdateWorkspaceBuildCostByIDParams) (WorkspaceBuild, error)
UpdateWorkspaceDeletedByID(ctx context.Context, arg UpdateWorkspaceDeletedByIDParams) error
UpdateWorkspaceLastUsedAt(ctx context.Context, arg UpdateWorkspaceLastUsedAtParams) error
UpdateWorkspaceTTL(ctx context.Context, arg UpdateWorkspaceTTLParams) error

View File

@ -1079,7 +1079,7 @@ func (q *sqlQuerier) GetAllOrganizationMembers(ctx context.Context, organization
const getGroupByID = `-- name: GetGroupByID :one
SELECT
id, name, organization_id, avatar_url
id, name, organization_id, avatar_url, quota_allowance
FROM
groups
WHERE
@ -1096,13 +1096,14 @@ func (q *sqlQuerier) GetGroupByID(ctx context.Context, id uuid.UUID) (Group, err
&i.Name,
&i.OrganizationID,
&i.AvatarURL,
&i.QuotaAllowance,
)
return i, err
}
const getGroupByOrgAndName = `-- name: GetGroupByOrgAndName :one
SELECT
id, name, organization_id, avatar_url
id, name, organization_id, avatar_url, quota_allowance
FROM
groups
WHERE
@ -1126,6 +1127,7 @@ func (q *sqlQuerier) GetGroupByOrgAndName(ctx context.Context, arg GetGroupByOrg
&i.Name,
&i.OrganizationID,
&i.AvatarURL,
&i.QuotaAllowance,
)
return i, err
}
@ -1185,7 +1187,7 @@ func (q *sqlQuerier) GetGroupMembers(ctx context.Context, groupID uuid.UUID) ([]
const getGroupsByOrganizationID = `-- name: GetGroupsByOrganizationID :many
SELECT
id, name, organization_id, avatar_url
id, name, organization_id, avatar_url, quota_allowance
FROM
groups
WHERE
@ -1208,6 +1210,7 @@ func (q *sqlQuerier) GetGroupsByOrganizationID(ctx context.Context, organization
&i.Name,
&i.OrganizationID,
&i.AvatarURL,
&i.QuotaAllowance,
); err != nil {
return nil, err
}
@ -1224,7 +1227,7 @@ func (q *sqlQuerier) GetGroupsByOrganizationID(ctx context.Context, organization
const getUserGroups = `-- name: GetUserGroups :many
SELECT
groups.id, groups.name, groups.organization_id, groups.avatar_url
groups.id, groups.name, groups.organization_id, groups.avatar_url, groups.quota_allowance
FROM
groups
JOIN
@ -1249,6 +1252,7 @@ func (q *sqlQuerier) GetUserGroups(ctx context.Context, userID uuid.UUID) ([]Gro
&i.Name,
&i.OrganizationID,
&i.AvatarURL,
&i.QuotaAllowance,
); err != nil {
return nil, err
}
@ -1270,7 +1274,7 @@ INSERT INTO groups (
organization_id
)
VALUES
( $1, 'Everyone', $1) RETURNING id, name, organization_id, avatar_url
( $1, 'Everyone', $1) RETURNING id, name, organization_id, avatar_url, quota_allowance
`
// We use the organization_id as the id
@ -1284,6 +1288,7 @@ func (q *sqlQuerier) InsertAllUsersGroup(ctx context.Context, organizationID uui
&i.Name,
&i.OrganizationID,
&i.AvatarURL,
&i.QuotaAllowance,
)
return i, err
}
@ -1293,10 +1298,11 @@ INSERT INTO groups (
id,
name,
organization_id,
avatar_url
avatar_url,
quota_allowance
)
VALUES
( $1, $2, $3, $4) RETURNING id, name, organization_id, avatar_url
( $1, $2, $3, $4, $5) RETURNING id, name, organization_id, avatar_url, quota_allowance
`
type InsertGroupParams struct {
@ -1304,6 +1310,7 @@ type InsertGroupParams struct {
Name string `db:"name" json:"name"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
QuotaAllowance int32 `db:"quota_allowance" json:"quota_allowance"`
}
func (q *sqlQuerier) InsertGroup(ctx context.Context, arg InsertGroupParams) (Group, error) {
@ -1312,6 +1319,7 @@ func (q *sqlQuerier) InsertGroup(ctx context.Context, arg InsertGroupParams) (Gr
arg.Name,
arg.OrganizationID,
arg.AvatarURL,
arg.QuotaAllowance,
)
var i Group
err := row.Scan(
@ -1319,6 +1327,7 @@ func (q *sqlQuerier) InsertGroup(ctx context.Context, arg InsertGroupParams) (Gr
&i.Name,
&i.OrganizationID,
&i.AvatarURL,
&i.QuotaAllowance,
)
return i, err
}
@ -1346,26 +1355,34 @@ UPDATE
groups
SET
name = $1,
avatar_url = $2
avatar_url = $2,
quota_allowance = $3
WHERE
id = $3
RETURNING id, name, organization_id, avatar_url
id = $4
RETURNING id, name, organization_id, avatar_url, quota_allowance
`
type UpdateGroupByIDParams struct {
Name string `db:"name" json:"name"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
ID uuid.UUID `db:"id" json:"id"`
Name string `db:"name" json:"name"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
QuotaAllowance int32 `db:"quota_allowance" json:"quota_allowance"`
ID uuid.UUID `db:"id" json:"id"`
}
func (q *sqlQuerier) UpdateGroupByID(ctx context.Context, arg UpdateGroupByIDParams) (Group, error) {
row := q.db.QueryRowContext(ctx, updateGroupByID, arg.Name, arg.AvatarURL, arg.ID)
row := q.db.QueryRowContext(ctx, updateGroupByID,
arg.Name,
arg.AvatarURL,
arg.QuotaAllowance,
arg.ID,
)
var i Group
err := row.Scan(
&i.ID,
&i.Name,
&i.OrganizationID,
&i.AvatarURL,
&i.QuotaAllowance,
)
return i, err
}
@ -2770,6 +2787,53 @@ func (q *sqlQuerier) UpdateProvisionerJobWithCompleteByID(ctx context.Context, a
return err
}
const getQuotaAllowanceForUser = `-- name: GetQuotaAllowanceForUser :one
SELECT
coalesce(SUM(quota_allowance), 0)::BIGINT
FROM
group_members gm
JOIN groups g ON
g.id = gm.group_id
WHERE
user_id = $1
`
func (q *sqlQuerier) GetQuotaAllowanceForUser(ctx context.Context, userID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, getQuotaAllowanceForUser, userID)
var column_1 int64
err := row.Scan(&column_1)
return column_1, err
}
const getQuotaConsumedForUser = `-- name: GetQuotaConsumedForUser :one
WITH latest_builds AS (
SELECT
DISTINCT ON
(workspace_id) id,
workspace_id,
daily_cost
FROM
workspace_builds wb
ORDER BY
workspace_id,
created_at DESC
)
SELECT
coalesce(SUM(daily_cost), 0)::BIGINT
FROM
workspaces
JOIN latest_builds ON
latest_builds.workspace_id = workspaces.id
WHERE NOT deleted AND workspaces.owner_id = $1
`
func (q *sqlQuerier) GetQuotaConsumedForUser(ctx context.Context, ownerID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, getQuotaConsumedForUser, ownerID)
var column_1 int64
err := row.Scan(&column_1)
return column_1, err
}
const deleteReplicasUpdatedBefore = `-- name: DeleteReplicasUpdatedBefore :exec
DELETE FROM replicas WHERE updated_at < $1
`
@ -5135,7 +5199,7 @@ func (q *sqlQuerier) UpdateWorkspaceAppHealthByID(ctx context.Context, arg Updat
const getLatestWorkspaceBuildByWorkspaceID = `-- name: GetLatestWorkspaceBuildByWorkspaceID :one
SELECT
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
FROM
workspace_builds
WHERE
@ -5162,12 +5226,13 @@ func (q *sqlQuerier) GetLatestWorkspaceBuildByWorkspaceID(ctx context.Context, w
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
)
return i, err
}
const getLatestWorkspaceBuilds = `-- name: GetLatestWorkspaceBuilds :many
SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason
SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason, wb.daily_cost
FROM (
SELECT
workspace_id, MAX(build_number) as max_build_number
@ -5203,6 +5268,7 @@ func (q *sqlQuerier) GetLatestWorkspaceBuilds(ctx context.Context) ([]WorkspaceB
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
); err != nil {
return nil, err
}
@ -5218,7 +5284,7 @@ func (q *sqlQuerier) GetLatestWorkspaceBuilds(ctx context.Context) ([]WorkspaceB
}
const getLatestWorkspaceBuildsByWorkspaceIDs = `-- name: GetLatestWorkspaceBuildsByWorkspaceIDs :many
SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason
SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason, wb.daily_cost
FROM (
SELECT
workspace_id, MAX(build_number) as max_build_number
@ -5256,6 +5322,7 @@ func (q *sqlQuerier) GetLatestWorkspaceBuildsByWorkspaceIDs(ctx context.Context,
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
); err != nil {
return nil, err
}
@ -5272,7 +5339,7 @@ func (q *sqlQuerier) GetLatestWorkspaceBuildsByWorkspaceIDs(ctx context.Context,
const getWorkspaceBuildByID = `-- name: GetWorkspaceBuildByID :one
SELECT
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
FROM
workspace_builds
WHERE
@ -5297,13 +5364,14 @@ func (q *sqlQuerier) GetWorkspaceBuildByID(ctx context.Context, id uuid.UUID) (W
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
)
return i, err
}
const getWorkspaceBuildByJobID = `-- name: GetWorkspaceBuildByJobID :one
SELECT
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
FROM
workspace_builds
WHERE
@ -5328,13 +5396,14 @@ func (q *sqlQuerier) GetWorkspaceBuildByJobID(ctx context.Context, jobID uuid.UU
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
)
return i, err
}
const getWorkspaceBuildByWorkspaceIDAndBuildNumber = `-- name: GetWorkspaceBuildByWorkspaceIDAndBuildNumber :one
SELECT
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
FROM
workspace_builds
WHERE
@ -5363,13 +5432,14 @@ func (q *sqlQuerier) GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx context.Co
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
)
return i, err
}
const getWorkspaceBuildsByWorkspaceID = `-- name: GetWorkspaceBuildsByWorkspaceID :many
SELECT
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason
id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
FROM
workspace_builds
WHERE
@ -5437,6 +5507,7 @@ func (q *sqlQuerier) GetWorkspaceBuildsByWorkspaceID(ctx context.Context, arg Ge
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
); err != nil {
return nil, err
}
@ -5452,7 +5523,7 @@ func (q *sqlQuerier) GetWorkspaceBuildsByWorkspaceID(ctx context.Context, arg Ge
}
const getWorkspaceBuildsCreatedAfter = `-- name: GetWorkspaceBuildsCreatedAfter :many
SELECT id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason FROM workspace_builds WHERE created_at > $1
SELECT id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost FROM workspace_builds WHERE created_at > $1
`
func (q *sqlQuerier) GetWorkspaceBuildsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceBuild, error) {
@ -5477,6 +5548,7 @@ func (q *sqlQuerier) GetWorkspaceBuildsCreatedAfter(ctx context.Context, created
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
); err != nil {
return nil, err
}
@ -5508,7 +5580,7 @@ INSERT INTO
reason
)
VALUES
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
`
type InsertWorkspaceBuildParams struct {
@ -5555,6 +5627,7 @@ func (q *sqlQuerier) InsertWorkspaceBuild(ctx context.Context, arg InsertWorkspa
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
)
return i, err
}
@ -5567,7 +5640,7 @@ SET
provisioner_state = $3,
deadline = $4
WHERE
id = $1 RETURNING id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason
id = $1 RETURNING id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
`
type UpdateWorkspaceBuildByIDParams struct {
@ -5598,13 +5671,49 @@ func (q *sqlQuerier) UpdateWorkspaceBuildByID(ctx context.Context, arg UpdateWor
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
)
return i, err
}
const updateWorkspaceBuildCostByID = `-- name: UpdateWorkspaceBuildCostByID :one
UPDATE
workspace_builds
SET
daily_cost = $2
WHERE
id = $1 RETURNING id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost
`
type UpdateWorkspaceBuildCostByIDParams struct {
ID uuid.UUID `db:"id" json:"id"`
DailyCost int32 `db:"daily_cost" json:"daily_cost"`
}
func (q *sqlQuerier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg UpdateWorkspaceBuildCostByIDParams) (WorkspaceBuild, error) {
row := q.db.QueryRowContext(ctx, updateWorkspaceBuildCostByID, arg.ID, arg.DailyCost)
var i WorkspaceBuild
err := row.Scan(
&i.ID,
&i.CreatedAt,
&i.UpdatedAt,
&i.WorkspaceID,
&i.TemplateVersionID,
&i.BuildNumber,
&i.Transition,
&i.InitiatorID,
&i.ProvisionerState,
&i.JobID,
&i.Deadline,
&i.Reason,
&i.DailyCost,
)
return i, err
}
const getWorkspaceResourceByID = `-- name: GetWorkspaceResourceByID :one
SELECT
id, created_at, job_id, transition, type, name, hide, icon, instance_type
id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost
FROM
workspace_resources
WHERE
@ -5624,6 +5733,7 @@ func (q *sqlQuerier) GetWorkspaceResourceByID(ctx context.Context, id uuid.UUID)
&i.Hide,
&i.Icon,
&i.InstanceType,
&i.DailyCost,
)
return i, err
}
@ -5738,7 +5848,7 @@ func (q *sqlQuerier) GetWorkspaceResourceMetadataCreatedAfter(ctx context.Contex
const getWorkspaceResourcesByJobID = `-- name: GetWorkspaceResourcesByJobID :many
SELECT
id, created_at, job_id, transition, type, name, hide, icon, instance_type
id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost
FROM
workspace_resources
WHERE
@ -5764,6 +5874,7 @@ func (q *sqlQuerier) GetWorkspaceResourcesByJobID(ctx context.Context, jobID uui
&i.Hide,
&i.Icon,
&i.InstanceType,
&i.DailyCost,
); err != nil {
return nil, err
}
@ -5780,7 +5891,7 @@ func (q *sqlQuerier) GetWorkspaceResourcesByJobID(ctx context.Context, jobID uui
const getWorkspaceResourcesByJobIDs = `-- name: GetWorkspaceResourcesByJobIDs :many
SELECT
id, created_at, job_id, transition, type, name, hide, icon, instance_type
id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost
FROM
workspace_resources
WHERE
@ -5806,6 +5917,7 @@ func (q *sqlQuerier) GetWorkspaceResourcesByJobIDs(ctx context.Context, ids []uu
&i.Hide,
&i.Icon,
&i.InstanceType,
&i.DailyCost,
); err != nil {
return nil, err
}
@ -5821,7 +5933,7 @@ func (q *sqlQuerier) GetWorkspaceResourcesByJobIDs(ctx context.Context, ids []uu
}
const getWorkspaceResourcesCreatedAfter = `-- name: GetWorkspaceResourcesCreatedAfter :many
SELECT id, created_at, job_id, transition, type, name, hide, icon, instance_type FROM workspace_resources WHERE created_at > $1
SELECT id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost FROM workspace_resources WHERE created_at > $1
`
func (q *sqlQuerier) GetWorkspaceResourcesCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceResource, error) {
@ -5843,6 +5955,7 @@ func (q *sqlQuerier) GetWorkspaceResourcesCreatedAfter(ctx context.Context, crea
&i.Hide,
&i.Icon,
&i.InstanceType,
&i.DailyCost,
); err != nil {
return nil, err
}
@ -5859,9 +5972,9 @@ func (q *sqlQuerier) GetWorkspaceResourcesCreatedAfter(ctx context.Context, crea
const insertWorkspaceResource = `-- name: InsertWorkspaceResource :one
INSERT INTO
workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, instance_type)
workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost)
VALUES
($1, $2, $3, $4, $5, $6, $7, $8, $9) RETURNING id, created_at, job_id, transition, type, name, hide, icon, instance_type
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost
`
type InsertWorkspaceResourceParams struct {
@ -5874,6 +5987,7 @@ type InsertWorkspaceResourceParams struct {
Hide bool `db:"hide" json:"hide"`
Icon string `db:"icon" json:"icon"`
InstanceType sql.NullString `db:"instance_type" json:"instance_type"`
DailyCost int32 `db:"daily_cost" json:"daily_cost"`
}
func (q *sqlQuerier) InsertWorkspaceResource(ctx context.Context, arg InsertWorkspaceResourceParams) (WorkspaceResource, error) {
@ -5887,6 +6001,7 @@ func (q *sqlQuerier) InsertWorkspaceResource(ctx context.Context, arg InsertWork
arg.Hide,
arg.Icon,
arg.InstanceType,
arg.DailyCost,
)
var i WorkspaceResource
err := row.Scan(
@ -5899,6 +6014,7 @@ func (q *sqlQuerier) InsertWorkspaceResource(ctx context.Context, arg InsertWork
&i.Hide,
&i.Icon,
&i.InstanceType,
&i.DailyCost,
)
return i, err
}

View File

@ -75,10 +75,11 @@ INSERT INTO groups (
id,
name,
organization_id,
avatar_url
avatar_url,
quota_allowance
)
VALUES
( $1, $2, $3, $4) RETURNING *;
( $1, $2, $3, $4, $5) RETURNING *;
-- We use the organization_id as the id
-- for simplicity since all users is
@ -97,9 +98,10 @@ UPDATE
groups
SET
name = $1,
avatar_url = $2
avatar_url = $2,
quota_allowance = $3
WHERE
id = $3
id = $4
RETURNING *;
-- name: InsertGroupMember :exec

View File

@ -0,0 +1,30 @@
-- name: GetQuotaAllowanceForUser :one
SELECT
coalesce(SUM(quota_allowance), 0)::BIGINT
FROM
group_members gm
JOIN groups g ON
g.id = gm.group_id
WHERE
user_id = $1;
-- name: GetQuotaConsumedForUser :one
WITH latest_builds AS (
SELECT
DISTINCT ON
(workspace_id) id,
workspace_id,
daily_cost
FROM
workspace_builds wb
ORDER BY
workspace_id,
created_at DESC
)
SELECT
coalesce(SUM(daily_cost), 0)::BIGINT
FROM
workspaces
JOIN latest_builds ON
latest_builds.workspace_id = workspaces.id
WHERE NOT deleted AND workspaces.owner_id = $1;

View File

@ -133,3 +133,12 @@ SET
deadline = $4
WHERE
id = $1 RETURNING *;
-- name: UpdateWorkspaceBuildCostByID :one
UPDATE
workspace_builds
SET
daily_cost = $2
WHERE
id = $1 RETURNING *;

View File

@ -27,9 +27,9 @@ SELECT * FROM workspace_resources WHERE created_at > $1;
-- name: InsertWorkspaceResource :one
INSERT INTO
workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, instance_type)
workspace_resources (id, created_at, job_id, transition, type, name, hide, icon, instance_type, daily_cost)
VALUES
($1, $2, $3, $4, $5, $6, $7, $8, $9) RETURNING *;
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING *;
-- name: GetWorkspaceResourceMetadataByResourceID :many
SELECT

View File

@ -88,7 +88,7 @@ func (api *API) postOrganizations(rw http.ResponseWriter, r *http.Request) {
return xerrors.Errorf("create %q group: %w", database.AllUsersGroup, err)
}
return nil
})
}, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error inserting organization member.",

View File

@ -86,6 +86,7 @@ func (api *API) ListenProvisionerDaemon(ctx context.Context, acquireJobDebounce
Telemetry: api.Telemetry,
Logger: api.Logger.Named(fmt.Sprintf("provisionerd-%s", daemon.Name)),
AcquireJobDebounce: acquireJobDebounce,
QuotaCommitter: &api.QuotaCommitter,
})
if err != nil {
return nil, err

View File

@ -9,6 +9,7 @@ import (
"net/url"
"reflect"
"sync"
"sync/atomic"
"time"
"github.com/google/uuid"
@ -34,13 +35,14 @@ var (
)
type Server struct {
AccessURL *url.URL
ID uuid.UUID
Logger slog.Logger
Provisioners []database.ProvisionerType
Database database.Store
Pubsub database.Pubsub
Telemetry telemetry.Reporter
AccessURL *url.URL
ID uuid.UUID
Logger slog.Logger
Provisioners []database.ProvisionerType
Database database.Store
Pubsub database.Pubsub
Telemetry telemetry.Reporter
QuotaCommitter *atomic.Pointer[proto.QuotaCommitter]
AcquireJobDebounce time.Duration
}
@ -252,6 +254,35 @@ func (server *Server) AcquireJob(ctx context.Context, _ *proto.Empty) (*proto.Ac
return protoJob, err
}
func (server *Server) CommitQuota(ctx context.Context, request *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error) {
jobID, err := uuid.Parse(request.JobId)
if err != nil {
return nil, xerrors.Errorf("parse job id: %w", err)
}
job, err := server.Database.GetProvisionerJobByID(ctx, jobID)
if err != nil {
return nil, xerrors.Errorf("get job: %w", err)
}
if !job.WorkerID.Valid {
return nil, xerrors.New("job isn't running yet")
}
if job.WorkerID.UUID.String() != server.ID.String() {
return nil, xerrors.New("you don't own this job")
}
q := server.QuotaCommitter.Load()
if q == nil {
// We're probably in community edition or a test.
return &proto.CommitQuotaResponse{
Budget: -1,
Ok: true,
}, nil
}
return (*q).CommitQuota(ctx, request)
}
func (server *Server) UpdateJob(ctx context.Context, request *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
parsedID, err := uuid.Parse(request.JobId)
if err != nil {
@ -620,7 +651,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
}
return nil
})
}, nil)
if err != nil {
return nil, xerrors.Errorf("complete job: %w", err)
}
@ -690,6 +721,7 @@ func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.
Name: protoResource.Name,
Hide: protoResource.Hide,
Icon: protoResource.Icon,
DailyCost: protoResource.DailyCost,
InstanceType: sql.NullString{
String: protoResource.InstanceType,
Valid: protoResource.InstanceType != "",

View File

@ -745,8 +745,9 @@ func TestInsertWorkspaceResource(t *testing.T) {
db := databasefake.New()
job := uuid.New()
err := insert(db, job, &sdkproto.Resource{
Name: "something",
Type: "aws_instance",
Name: "something",
Type: "aws_instance",
DailyCost: 10,
Agents: []*sdkproto.Agent{{
Name: "dev",
Env: map[string]string{
@ -767,6 +768,7 @@ func TestInsertWorkspaceResource(t *testing.T) {
resources, err := db.GetWorkspaceResourcesByJobID(ctx, job)
require.NoError(t, err)
require.Len(t, resources, 1)
require.EqualValues(t, 10, resources[0].DailyCost)
agents, err := db.GetWorkspaceAgentsByResourceIDs(ctx, []uuid.UUID{resources[0].ID})
require.NoError(t, err)
require.Len(t, agents, 1)

View File

@ -290,7 +290,7 @@ func (api *API) postTemplateByOrganization(rw http.ResponseWriter, r *http.Reque
template = api.convertTemplate(dbTemplate, 0, createdByNameMap[dbTemplate.ID.String()])
return nil
})
}, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error inserting template.",
@ -511,7 +511,7 @@ func (api *API) patchTemplateMeta(rw http.ResponseWriter, r *http.Request) {
}
return nil
})
}, nil)
if err != nil {
httpapi.InternalServerError(rw, err)
return
@ -690,7 +690,7 @@ func (api *API) autoImportTemplate(ctx context.Context, opts autoImportTemplateO
}
return nil
})
}, nil)
return template, err
}

View File

@ -538,7 +538,7 @@ func (api *API) templateVersionsByTemplate(rw http.ResponseWriter, r *http.Reque
}
return nil
})
}, nil)
if err != nil {
return
}
@ -651,7 +651,7 @@ func (api *API) patchActiveTemplateVersion(rw http.ResponseWriter, r *http.Reque
return xerrors.Errorf("update active version: %w", err)
}
return nil
})
}, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error updating active template version.",
@ -852,7 +852,7 @@ func (api *API) postTemplateVersionsByOrganization(rw http.ResponseWriter, r *ht
return xerrors.Errorf("insert template version: %w", err)
}
return nil
})
}, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: err.Error(),

View File

@ -520,7 +520,7 @@ func (api *API) oauthLogin(r *http.Request, params oauthLoginParams) (*http.Cook
}
return nil
})
}, nil)
if err != nil {
return nil, xerrors.Errorf("in tx: %w", err)
}

View File

@ -700,7 +700,7 @@ func (api *API) putUserPassword(rw http.ResponseWriter, r *http.Request) {
}
return nil
})
}, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error updating user's password.",
@ -1147,7 +1147,7 @@ func (api *API) CreateUser(ctx context.Context, store database.Store, req Create
return xerrors.Errorf("create organization member: %w", err)
}
return nil
})
}, nil)
}
func convertUser(user database.User, organizationIDs []uuid.UUID) codersdk.User {

View File

@ -136,7 +136,7 @@ func (api *API) workspaceBuilds(rw http.ResponseWriter, r *http.Request) {
}
return nil
})
}, nil)
if err != nil {
return
}
@ -536,7 +536,7 @@ func (api *API) postWorkspaceBuilds(rw http.ResponseWriter, r *http.Request) {
}
return nil
})
}, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error inserting workspace build.",
@ -931,6 +931,7 @@ func (api *API) convertWorkspaceBuild(
Reason: codersdk.BuildReason(build.Reason),
Resources: apiResources,
Status: convertWorkspaceStatus(apiJob.Status, transition),
DailyCost: build.DailyCost,
}, nil
}
@ -974,6 +975,7 @@ func convertWorkspaceResource(resource database.WorkspaceResource, agents []code
Icon: resource.Icon,
Agents: agents,
Metadata: convertedMetadata,
DailyCost: resource.DailyCost,
}
}

View File

@ -1,19 +0,0 @@
package workspacequota
type Enforcer interface {
UserWorkspaceLimit() int
CanCreateWorkspace(count int) bool
}
type nop struct{}
func NewNop() Enforcer {
return &nop{}
}
func (*nop) UserWorkspaceLimit() int {
return 0
}
func (*nop) CanCreateWorkspace(_ int) bool {
return true
}

View File

@ -342,25 +342,6 @@ func (api *API) postWorkspacesByOrganization(rw http.ResponseWriter, r *http.Req
return
}
workspaceCount, err := api.Database.GetWorkspaceCountByUserID(ctx, user.ID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace count.",
Detail: err.Error(),
})
return
}
// make sure the user has not hit their quota limit
e := *api.WorkspaceQuotaEnforcer.Load()
canCreate := e.CanCreateWorkspace(int(workspaceCount))
if !canCreate {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: fmt.Sprintf("User workspace limit of %d is already reached.", e.UserWorkspaceLimit()),
})
return
}
templateVersion, err := api.Database.GetTemplateVersionByID(ctx, template.ActiveVersionID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
@ -479,7 +460,7 @@ func (api *API) postWorkspacesByOrganization(rw http.ResponseWriter, r *http.Req
return xerrors.Errorf("insert workspace build: %w", err)
}
return nil
})
}, nil)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error creating workspace.",
@ -710,7 +691,7 @@ func (api *API) putWorkspaceTTL(rw http.ResponseWriter, r *http.Request) {
}
return nil
})
}, nil)
if err != nil {
resp := codersdk.Response{
Message: "Error updating workspace time until shutdown.",
@ -807,7 +788,7 @@ func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {
resp.Message = "Deadline updated to " + newDeadline.Format(time.RFC3339) + "."
return nil
})
}, nil)
if err != nil {
api.Logger.Info(ctx, "extending workspace", slog.Error(err))
}