Skip to content

Commit 69c73b2

Browse files
f0sself0ssel
authored
feat: workspace quotas (#4184)
1 parent f9b7588 commit 69c73b2

28 files changed

+712
-83
lines changed

‎coderd/coderd.go

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import (
3535
"github.com/coder/coder/coderd/rbac"
3636
"github.com/coder/coder/coderd/telemetry"
3737
"github.com/coder/coder/coderd/tracing"
38+
"github.com/coder/coder/coderd/workspacequota"
3839
"github.com/coder/coder/coderd/wsconncache"
3940
"github.com/coder/coder/codersdk"
4041
"github.com/coder/coder/site"
@@ -55,6 +56,7 @@ type Options struct {
5556
CacheDir string
5657

5758
Auditor audit.Auditor
59+
WorkspaceQuotaEnforcer workspacequota.Enforcer
5860
AgentConnectionUpdateFrequency time.Duration
5961
AgentInactiveDisconnectTimeout time.Duration
6062
// APIRateLimit is the minutely throughput rate limit per user or ip.
@@ -120,6 +122,9 @@ func New(options *Options) *API {
120122
if options.Auditor == nil {
121123
options.Auditor = audit.NewNop()
122124
}
125+
if options.WorkspaceQuotaEnforcer == nil {
126+
options.WorkspaceQuotaEnforcer = workspacequota.NewNop()
127+
}
123128

124129
siteCacheDir := options.CacheDir
125130
if siteCacheDir != "" {
@@ -145,10 +150,12 @@ func New(options *Options) *API {
145150
Authorizer: options.Authorizer,
146151
Logger: options.Logger,
147152
},
148-
metricsCache: metricsCache,
149-
Auditor: atomic.Pointer[audit.Auditor]{},
153+
metricsCache: metricsCache,
154+
Auditor: atomic.Pointer[audit.Auditor]{},
155+
WorkspaceQuotaEnforcer: atomic.Pointer[workspacequota.Enforcer]{},
150156
}
151157
api.Auditor.Store(&options.Auditor)
158+
api.WorkspaceQuotaEnforcer.Store(&options.WorkspaceQuotaEnforcer)
152159
api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgentTailnet, 0)
153160
api.derpServer = derp.NewServer(key.NewNode(), tailnet.Logger(options.Logger))
154161
oauthConfigs := &httpmw.OAuth2Configs{
@@ -516,6 +523,7 @@ type API struct {
516523
*Options
517524
Auditor atomic.Pointer[audit.Auditor]
518525
WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool]
526+
WorkspaceQuotaEnforcer atomic.Pointer[workspacequota.Enforcer]
519527
HTTPAuth *HTTPAuthorizer
520528

521529
// APIHandler serves "/api/v2"

‎coderd/database/databasefake/databasefake.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -698,6 +698,22 @@ func (q *fakeQuerier) GetWorkspaceBuildByID(_ context.Context, id uuid.UUID) (da
698698
return database.WorkspaceBuild{}, sql.ErrNoRows
699699
}
700700

701+
func (q *fakeQuerier) GetWorkspaceCountByUserID(_ context.Context, id uuid.UUID) (int64, error) {
702+
q.mutex.RLock()
703+
defer q.mutex.RUnlock()
704+
var count int64
705+
for _, workspace := range q.workspaces {
706+
if workspace.OwnerID.String() == id.String() {
707+
if workspace.Deleted {
708+
continue
709+
}
710+
711+
count++
712+
}
713+
}
714+
return count, nil
715+
}
716+
701717
func (q *fakeQuerier) GetWorkspaceBuildByJobID(_ context.Context, jobID uuid.UUID) (database.WorkspaceBuild, error) {
702718
q.mutex.RLock()
703719
defer q.mutex.RUnlock()

‎coderd/database/querier.go

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

Lines changed: 18 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaces.sql

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,16 @@ WHERE
7474
GROUP BY
7575
template_id;
7676

77+
-- name: GetWorkspaceCountByUserID :one
78+
SELECT
79+
COUNT(id)
80+
FROM
81+
workspaces
82+
WHERE
83+
owner_id = @owner_id
84+
-- Ignore deleted workspaces
85+
AND deleted != true;
86+
7787
-- name: InsertWorkspace :one
7888
INSERT INTO
7989
workspaces (

‎coderd/workspacequota/workspacequota.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package workspacequota
2+
3+
type Enforcer interface {
4+
UserWorkspaceLimit() int
5+
CanCreateWorkspace(count int) bool
6+
}
7+
8+
type nop struct{}
9+
10+
func NewNop() Enforcer {
11+
return &nop{}
12+
}
13+
14+
func (*nop) UserWorkspaceLimit() int {
15+
return 0
16+
}
17+
func (*nop) CanCreateWorkspace(_ int) bool {
18+
return true
19+
}

‎coderd/workspaces.go

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,25 @@ func (api *API) postWorkspacesByOrganization(rw http.ResponseWriter, r *http.Req
317317
return
318318
}
319319

320+
workspaceCount, err := api.Database.GetWorkspaceCountByUserID(ctx, user.ID)
321+
if err != nil {
322+
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
323+
Message: "Internal error fetching workspace count.",
324+
Detail: err.Error(),
325+
})
326+
return
327+
}
328+
329+
// make sure the user has not hit their quota limit
330+
e := *api.WorkspaceQuotaEnforcer.Load()
331+
canCreate := e.CanCreateWorkspace(int(workspaceCount))
332+
if !canCreate {
333+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
334+
Message: fmt.Sprintf("User workspace limit of %d is already reached.", e.UserWorkspaceLimit()),
335+
})
336+
return
337+
}
338+
320339
templateVersion, err := api.Database.GetTemplateVersionByID(ctx, template.ActiveVersionID)
321340
if err != nil {
322341
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
@@ -352,8 +371,10 @@ func (api *API) postWorkspacesByOrganization(rw http.ResponseWriter, r *http.Req
352371
return
353372
}
354373

355-
var provisionerJob database.ProvisionerJob
356-
var workspaceBuild database.WorkspaceBuild
374+
var (
375+
provisionerJob database.ProvisionerJob
376+
workspaceBuild database.WorkspaceBuild
377+
)
357378
err = api.Database.InTx(func(db database.Store) error {
358379
now := database.Now()
359380
workspaceBuildID := uuid.New()

‎codersdk/features.go

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,20 @@ const (
1515
)
1616

1717
const (
18-
FeatureUserLimit = "user_limit"
19-
FeatureAuditLog = "audit_log"
20-
FeatureBrowserOnly = "browser_only"
21-
FeatureSCIM = "scim"
18+
FeatureUserLimit = "user_limit"
19+
FeatureAuditLog = "audit_log"
20+
FeatureBrowserOnly = "browser_only"
21+
FeatureSCIM = "scim"
22+
FeatureWorkspaceQuota = "workspace_quota"
2223
)
2324

24-
var FeatureNames = []string{FeatureUserLimit, FeatureAuditLog, FeatureBrowserOnly, FeatureSCIM}
25+
var FeatureNames = []string{
26+
FeatureUserLimit,
27+
FeatureAuditLog,
28+
FeatureBrowserOnly,
29+
FeatureSCIM,
30+
FeatureWorkspaceQuota,
31+
}
2532

2633
type Feature struct {
2734
Entitlement Entitlement `json:"entitlement"`

‎codersdk/workspacequota.go

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
package codersdk
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"net/http"
8+
)
9+
10+
type WorkspaceQuota struct {
11+
UserWorkspaceCount int `json:"user_workspace_count"`
12+
UserWorkspaceLimit int `json:"user_workspace_limit"`
13+
}
14+
15+
func (c *Client) WorkspaceQuota(ctx context.Context, userID string) (WorkspaceQuota, error) {
16+
res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspace-quota/%s", userID), nil)
17+
if err != nil {
18+
return WorkspaceQuota{}, err
19+
}
20+
defer res.Body.Close()
21+
if res.StatusCode != http.StatusOK {
22+
return WorkspaceQuota{}, readBodyAsError(res)
23+
}
24+
var quota WorkspaceQuota
25+
return quota, json.NewDecoder(res.Body).Decode(&quota)
26+
}

‎enterprise/cli/features_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,14 +57,16 @@ func TestFeaturesList(t *testing.T) {
5757
var entitlements codersdk.Entitlements
5858
err := json.Unmarshal(buf.Bytes(), &entitlements)
5959
require.NoError(t, err, "unmarshal JSON output")
60-
assert.Len(t, entitlements.Features, 3)
60+
assert.Len(t, entitlements.Features, 4)
6161
assert.Empty(t, entitlements.Warnings)
6262
assert.Equal(t, codersdk.EntitlementNotEntitled,
6363
entitlements.Features[codersdk.FeatureUserLimit].Entitlement)
6464
assert.Equal(t, codersdk.EntitlementNotEntitled,
6565
entitlements.Features[codersdk.FeatureAuditLog].Entitlement)
6666
assert.Equal(t, codersdk.EntitlementNotEntitled,
6767
entitlements.Features[codersdk.FeatureBrowserOnly].Entitlement)
68+
assert.Equal(t, codersdk.EntitlementNotEntitled,
69+
entitlements.Features[codersdk.FeatureWorkspaceQuota].Entitlement)
6870
assert.False(t, entitlements.HasLicense)
6971
})
7072
}

0 commit comments

Comments
 (0)