coder · ammario · Nov 14, 2022 · Nov 10, 2022 · Nov 10, 2022 · Nov 10, 2022
diff --git a/cli/deployment/config.go b/cli/deployment/config.go
@@ -362,12 +362,6 @@ func newConfig() *codersdk.DeploymentConfig {
 			Enterprise: true,
 			Secret:     true,
 		},
-		UserWorkspaceQuota: &codersdk.DeploymentConfigField[int]{
-			Name:       "User Workspace Quota",
-			Usage:      "Enables and sets a limit on how many workspaces each user can create.",
-			Flag:       "user-workspace-quota",
-			Enterprise: true,
-		},
 		Provisioner: &codersdk.ProvisionerConfig{
 			Daemons: &codersdk.DeploymentConfigField[int]{
 				Name:    "Provisioner Daemons",

diff --git a/cli/deployment/config_test.go b/cli/deployment/config_test.go
@@ -79,16 +79,14 @@ func TestConfig(t *testing.T) {
 	}, {
 		Name: "Enterprise",
 		Env: map[string]string{
-			"CODER_AUDIT_LOGGING":        "false",
-			"CODER_BROWSER_ONLY":         "true",
-			"CODER_SCIM_API_KEY":         "some-key",
-			"CODER_USER_WORKSPACE_QUOTA": "10",
+			"CODER_AUDIT_LOGGING": "false",
+			"CODER_BROWSER_ONLY":  "true",
+			"CODER_SCIM_API_KEY":  "some-key",
 		},
 		Valid: func(config *codersdk.DeploymentConfig) {
 			require.Equal(t, config.AuditLogging.Value, false)
 			require.Equal(t, config.BrowserOnly.Value, true)
 			require.Equal(t, config.SCIMAPIKey.Value, "some-key")
-			require.Equal(t, config.UserWorkspaceQuota.Value, 10)
 		},
 	}, {
 		Name: "TLS",

diff --git a/coderd/activitybump.go b/coderd/activitybump.go
@@ -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",

diff --git a/coderd/autobuild/executor/lifecycle_executor.go b/coderd/autobuild/executor/lifecycle_executor.go
@@ -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))
 			}

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -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"

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -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

diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -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
+}
diff --git a/coderd/database/databasefake/databasefake_test.go b/coderd/database/databasefake/databasefake_test.go
@@ -38,7 +38,7 @@ func TestInTx(t *testing.T) {
 			})
 			assert.NoError(t, err)
 			return nil
-		})
+		}, nil)
 		assert.NoError(t, err)
 	}()
 	var nums []int

diff --git a/coderd/database/db.go b/coderd/database/db.go
@@ -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)
 	}

diff --git a/coderd/database/db_test.go b/coderd/database/db_test.go
@@ -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)

diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
@@ -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() {

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000076_cost.down.sql b/coderd/database/migrations/000076_cost.down.sql
@@ -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;
diff --git a/coderd/database/migrations/000076_cost.up.sql b/coderd/database/migrations/000076_cost.up.sql
@@ -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;
diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
-Original file line number
+Diff line change
@@ Expand Up / @@ -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))
     			}
@@ Expand Down @@