Skip to content

Commit 97dbd4d

Browse files
ammarioammario
authored
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
1 parent 3fb7892 commit 97dbd4d

File tree

101 files changed

+1569
-839
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

101 files changed

+1569
-839
lines changed

‎cli/deployment/config.go

-6
Original file line numberDiff line numberDiff line change
@@ -362,12 +362,6 @@ func newConfig() *codersdk.DeploymentConfig {
362362
Enterprise: true,
363363
Secret: true,
364364
},
365-
UserWorkspaceQuota: &codersdk.DeploymentConfigField[int]{
366-
Name: "User Workspace Quota",
367-
Usage: "Enables and sets a limit on how many workspaces each user can create.",
368-
Flag: "user-workspace-quota",
369-
Enterprise: true,
370-
},
371365
Provisioner: &codersdk.ProvisionerConfig{
372366
Daemons: &codersdk.DeploymentConfigField[int]{
373367
Name: "Provisioner Daemons",

‎cli/deployment/config_test.go

+3-5
Original file line numberDiff line numberDiff line change
@@ -79,16 +79,14 @@ func TestConfig(t *testing.T) {
7979
}, {
8080
Name: "Enterprise",
8181
Env: map[string]string{
82-
"CODER_AUDIT_LOGGING": "false",
83-
"CODER_BROWSER_ONLY": "true",
84-
"CODER_SCIM_API_KEY": "some-key",
85-
"CODER_USER_WORKSPACE_QUOTA": "10",
82+
"CODER_AUDIT_LOGGING": "false",
83+
"CODER_BROWSER_ONLY": "true",
84+
"CODER_SCIM_API_KEY": "some-key",
8685
},
8786
Valid: func(config *codersdk.DeploymentConfig) {
8887
require.Equal(t, config.AuditLogging.Value, false)
8988
require.Equal(t, config.BrowserOnly.Value, true)
9089
require.Equal(t, config.SCIMAPIKey.Value, "some-key")
91-
require.Equal(t, config.UserWorkspaceQuota.Value, 10)
9290
},
9391
}, {
9492
Name: "TLS",

‎coderd/activitybump.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ func activityBumpWorkspace(log slog.Logger, db database.Store, workspace databas
6363
return xerrors.Errorf("update workspace build: %w", err)
6464
}
6565
return nil
66-
})
66+
}, nil)
6767
if err != nil {
6868
log.Error(
6969
ctx, "bump failed",

‎coderd/autobuild/executor/lifecycle_executor.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ func (e *Executor) runOnce(t time.Time) Stats {
177177
}
178178

179179
return nil
180-
})
180+
}, nil)
181181
if err != nil {
182182
log.Error(e.ctx, "workspace scheduling failed", slog.Error(err))
183183
}

‎coderd/coderd.go

+4-10
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,9 @@ import (
4040
"github.com/coder/coder/coderd/rbac"
4141
"github.com/coder/coder/coderd/telemetry"
4242
"github.com/coder/coder/coderd/tracing"
43-
"github.com/coder/coder/coderd/workspacequota"
4443
"github.com/coder/coder/coderd/wsconncache"
4544
"github.com/coder/coder/codersdk"
45+
"github.com/coder/coder/provisionerd/proto"
4646
"github.com/coder/coder/site"
4747
"github.com/coder/coder/tailnet"
4848
)
@@ -66,7 +66,6 @@ type Options struct {
6666
CacheDir string
6767

6868
Auditor audit.Auditor
69-
WorkspaceQuotaEnforcer workspacequota.Enforcer
7069
AgentConnectionUpdateFrequency time.Duration
7170
AgentInactiveDisconnectTimeout time.Duration
7271
// APIRateLimit is the minutely throughput rate limit per user or ip.
@@ -145,9 +144,6 @@ func New(options *Options) *API {
145144
if options.Auditor == nil {
146145
options.Auditor = audit.NewNop()
147146
}
148-
if options.WorkspaceQuotaEnforcer == nil {
149-
options.WorkspaceQuotaEnforcer = workspacequota.NewNop()
150-
}
151147

152148
siteCacheDir := options.CacheDir
153149
if siteCacheDir != "" {
@@ -174,12 +170,10 @@ func New(options *Options) *API {
174170
Authorizer: options.Authorizer,
175171
Logger: options.Logger,
176172
},
177-
metricsCache: metricsCache,
178-
Auditor: atomic.Pointer[audit.Auditor]{},
179-
WorkspaceQuotaEnforcer: atomic.Pointer[workspacequota.Enforcer]{},
173+
metricsCache: metricsCache,
174+
Auditor: atomic.Pointer[audit.Auditor]{},
180175
}
181176
api.Auditor.Store(&options.Auditor)
182-
api.WorkspaceQuotaEnforcer.Store(&options.WorkspaceQuotaEnforcer)
183177
api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgentTailnet, 0)
184178
api.TailnetCoordinator.Store(&options.TailnetCoordinator)
185179
oauthConfigs := &httpmw.OAuth2Configs{
@@ -590,8 +584,8 @@ type API struct {
590584
ID uuid.UUID
591585
Auditor atomic.Pointer[audit.Auditor]
592586
WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool]
593-
WorkspaceQuotaEnforcer atomic.Pointer[workspacequota.Enforcer]
594587
TailnetCoordinator atomic.Pointer[tailnet.Coordinator]
588+
QuotaCommitter atomic.Pointer[proto.QuotaCommitter]
595589
HTTPAuth *HTTPAuthorizer
596590

597591
// APIHandler serves "/api/v2"

‎coderd/coderdtest/coderdtest.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -528,7 +528,8 @@ func AwaitWorkspaceBuildJob(t *testing.T, client *codersdk.Client, build uuid.UU
528528
t.Logf("waiting for workspace build job %s", build)
529529
var workspaceBuild codersdk.WorkspaceBuild
530530
require.Eventually(t, func() bool {
531-
workspaceBuild, err := client.WorkspaceBuild(context.Background(), build)
531+
var err error
532+
workspaceBuild, err = client.WorkspaceBuild(context.Background(), build)
532533
return assert.NoError(t, err) && workspaceBuild.Job.CompletedAt != nil
533534
}, testutil.WaitShort, testutil.IntervalFast)
534535
return workspaceBuild

‎coderd/database/databasefake/databasefake.go

+61-1
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ func (*fakeQuerier) Ping(_ context.Context) (time.Duration, error) {
121121
}
122122

123123
// InTx doesn't rollback data properly for in-memory yet.
124-
func (q *fakeQuerier) InTx(fn func(database.Store) error) error {
124+
func (q *fakeQuerier) InTx(fn func(database.Store) error, _ *sql.TxOptions) error {
125125
q.mutex.Lock()
126126
defer q.mutex.Unlock()
127127
return fn(&fakeQuerier{mutex: inTxMutex{}, data: q.data})
@@ -2246,6 +2246,7 @@ func (q *fakeQuerier) InsertWorkspaceResource(_ context.Context, arg database.In
22462246
Name: arg.Name,
22472247
Hide: arg.Hide,
22482248
Icon: arg.Icon,
2249+
DailyCost: arg.DailyCost,
22492250
}
22502251
q.provisionerJobResources = append(q.provisionerJobResources, resource)
22512252
return resource, nil
@@ -2757,6 +2758,20 @@ func (q *fakeQuerier) UpdateWorkspaceBuildByID(_ context.Context, arg database.U
27572758
}
27582759
return database.WorkspaceBuild{}, sql.ErrNoRows
27592760
}
2761+
func (q *fakeQuerier) UpdateWorkspaceBuildCostByID(_ context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) (database.WorkspaceBuild, error) {
2762+
q.mutex.Lock()
2763+
defer q.mutex.Unlock()
2764+
2765+
for index, workspaceBuild := range q.workspaceBuilds {
2766+
if workspaceBuild.ID != arg.ID {
2767+
continue
2768+
}
2769+
workspaceBuild.DailyCost = arg.DailyCost
2770+
q.workspaceBuilds[index] = workspaceBuild
2771+
return workspaceBuild, nil
2772+
}
2773+
return database.WorkspaceBuild{}, sql.ErrNoRows
2774+
}
27602775

27612776
func (q *fakeQuerier) UpdateWorkspaceDeletedByID(_ context.Context, arg database.UpdateWorkspaceDeletedByIDParams) error {
27622777
q.mutex.Lock()
@@ -2858,6 +2873,7 @@ func (q *fakeQuerier) UpdateGroupByID(_ context.Context, arg database.UpdateGrou
28582873
if group.ID == arg.ID {
28592874
group.Name = arg.Name
28602875
group.AvatarURL = arg.AvatarURL
2876+
group.QuotaAllowance = arg.QuotaAllowance
28612877
q.groups[i] = group
28622878
return group, nil
28632879
}
@@ -3230,6 +3246,7 @@ func (q *fakeQuerier) InsertGroup(_ context.Context, arg database.InsertGroupPar
32303246
Name: arg.Name,
32313247
OrganizationID: arg.OrganizationID,
32323248
AvatarURL: arg.AvatarURL,
3249+
QuotaAllowance: arg.QuotaAllowance,
32333250
}
32343251

32353252
q.groups = append(q.groups, group)
@@ -3430,3 +3447,46 @@ func (q *fakeQuerier) UpdateGitAuthLink(_ context.Context, arg database.UpdateGi
34303447
}
34313448
return nil
34323449
}
3450+
3451+
func (q *fakeQuerier) GetQuotaAllowanceForUser(_ context.Context, userID uuid.UUID) (int64, error) {
3452+
q.mutex.Lock()
3453+
defer q.mutex.Unlock()
3454+
var sum int64
3455+
for _, member := range q.groupMembers {
3456+
if member.UserID != userID {
3457+
continue
3458+
}
3459+
for _, group := range q.groups {
3460+
if group.ID == member.GroupID {
3461+
sum += int64(group.QuotaAllowance)
3462+
}
3463+
}
3464+
}
3465+
return sum, nil
3466+
}
3467+
3468+
func (q *fakeQuerier) GetQuotaConsumedForUser(_ context.Context, userID uuid.UUID) (int64, error) {
3469+
q.mutex.Lock()
3470+
defer q.mutex.Unlock()
3471+
var sum int64
3472+
for _, workspace := range q.workspaces {
3473+
if workspace.OwnerID != userID {
3474+
continue
3475+
}
3476+
if workspace.Deleted {
3477+
continue
3478+
}
3479+
3480+
var lastBuild database.WorkspaceBuild
3481+
for _, build := range q.workspaceBuilds {
3482+
if build.WorkspaceID != workspace.ID {
3483+
continue
3484+
}
3485+
if build.CreatedAt.After(lastBuild.CreatedAt) {
3486+
lastBuild = build
3487+
}
3488+
}
3489+
sum += int64(lastBuild.DailyCost)
3490+
}
3491+
return sum, nil
3492+
}

‎coderd/database/databasefake/databasefake_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ func TestInTx(t *testing.T) {
3838
})
3939
assert.NoError(t, err)
4040
return nil
41-
})
41+
}, nil)
4242
assert.NoError(t, err)
4343
}()
4444
var nums []int

‎coderd/database/db.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ type Store interface {
2626
customQuerier
2727

2828
Ping(ctx context.Context) (time.Duration, error)
29-
InTx(func(Store) error) error
29+
InTx(func(Store) error, *sql.TxOptions) error
3030
}
3131

3232
// DBTX represents a database connection or transaction.
@@ -68,7 +68,7 @@ func (q *sqlQuerier) Ping(ctx context.Context) (time.Duration, error) {
6868
}
6969

7070
// InTx performs database operations inside a transaction.
71-
func (q *sqlQuerier) InTx(function func(Store) error) error {
71+
func (q *sqlQuerier) InTx(function func(Store) error, txOpts *sql.TxOptions) error {
7272
if _, ok := q.db.(*sqlx.Tx); ok {
7373
// If the current inner "db" is already a transaction, we just reuse it.
7474
// 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 {
8080
return nil
8181
}
8282

83-
transaction, err := q.sdb.BeginTxx(context.Background(), nil)
83+
transaction, err := q.sdb.BeginTxx(context.Background(), txOpts)
8484
if err != nil {
8585
return xerrors.Errorf("begin transaction: %w", err)
8686
}

‎coderd/database/db_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,8 @@ func TestNestedInTx(t *testing.T) {
4343
LoginType: database.LoginTypeGithub,
4444
})
4545
return err
46-
})
47-
})
46+
}, nil)
47+
}, nil)
4848
require.NoError(t, err, "outer tx: %w", err)
4949

5050
user, err := db.GetUserByID(context.Background(), uid)

‎coderd/database/dbtestutil/db.go

+10-3
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,16 @@ func NewDB(t *testing.T) (database.Store, database.Pubsub) {
1919
db := databasefake.New()
2020
pubsub := database.NewPubsubInMemory()
2121
if os.Getenv("DB") != "" {
22-
connectionURL, closePg, err := postgres.Open()
23-
require.NoError(t, err)
24-
t.Cleanup(closePg)
22+
connectionURL := os.Getenv("CODER_PG_CONNECTION_URL")
23+
if connectionURL == "" {
24+
var (
25+
err error
26+
closePg func()
27+
)
28+
connectionURL, closePg, err = postgres.Open()
29+
require.NoError(t, err)
30+
t.Cleanup(closePg)
31+
}
2532
sqlDB, err := sql.Open("postgres", connectionURL)
2633
require.NoError(t, err)
2734
t.Cleanup(func() {

‎coderd/database/dump.sql

+6-3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000076_cost.down.sql

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
ALTER TABLE workspace_builds DROP COLUMN daily_cost;
2+
ALTER TABLE workspace_resources DROP COLUMN daily_cost;
3+
ALTER TABLE groups DROP COLUMN quota_allowance;

‎coderd/database/migrations/000076_cost.up.sql

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
ALTER TABLE workspace_builds ADD COLUMN daily_cost int NOT NULL DEFAULT 0;
2+
ALTER TABLE workspace_resources ADD COLUMN daily_cost int NOT NULL DEFAULT 0;
3+
ALTER TABLE groups ADD COLUMN quota_allowance int NOT NULL DEFAULT 0;

‎coderd/database/models.go

+3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

+3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)