coder · ammario · Aug 22, 2023 · Aug 22, 2023
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -498,7 +498,10 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 
 	if initial, changed, enabled := featureChanged(codersdk.FeatureTemplateRBAC); shouldUpdate(initial, changed, enabled) {
 		if enabled {
-			committer := committer{Database: api.Database}
+			committer := committer{
+				Log:      api.Logger.Named("quota_committer"),
+				Database: api.Database,
+			}
 			ptr := proto.QuotaCommitter(&committer)
 			api.AGPL.QuotaCommitter.Store(&ptr)
 		} else {

diff --git a/enterprise/coderd/workspacequota.go b/enterprise/coderd/workspacequota.go
@@ -3,10 +3,12 @@ package coderd
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"net/http"
 
 	"github.com/google/uuid"
-	"golang.org/x/xerrors"
+
+	"cdr.dev/slog"
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/httpapi"
@@ -17,6 +19,7 @@ import (
 )
 
 type committer struct {
+	Log      slog.Logger
 	Database database.Store
 }
 
@@ -28,12 +31,12 @@ func (c *committer) CommitQuota(
 		return nil, err
 	}
 
-	build, err := c.Database.GetWorkspaceBuildByJobID(ctx, jobID)
+	nextBuild, err := c.Database.GetWorkspaceBuildByJobID(ctx, jobID)
 	if err != nil {
 		return nil, err
 	}
 
-	workspace, err := c.Database.GetWorkspaceByID(ctx, build.WorkspaceID)
+	workspace, err := c.Database.GetWorkspaceByID(ctx, nextBuild.WorkspaceID)
 	if err != nil {
 		return nil, err
 	}
@@ -58,25 +61,35 @@ func (c *committer) CommitQuota(
 		// If the new build will reduce overall quota consumption, then we
 		// allow it even if the user is over quota.
 		netIncrease := true
-		previousBuild, err := s.GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx, database.GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams{
+		prevBuild, err := s.GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx, database.GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams{
 			WorkspaceID: workspace.ID,
-			BuildNumber: build.BuildNumber - 1,
+			BuildNumber: nextBuild.BuildNumber - 1,
 		})
 		if err == nil {
-			if build.DailyCost < previousBuild.DailyCost {
-				netIncrease = false
-			}
-		} else if !xerrors.Is(err, sql.ErrNoRows) {
+			netIncrease = request.DailyCost >= prevBuild.DailyCost
+			c.Log.Debug(
+				ctx, "previous build cost",
+				slog.F("prev_cost", prevBuild.DailyCost),
+				slog.F("next_cost", request.DailyCost),
+				slog.F("net_increase", netIncrease),
+			)
+		} else if !errors.Is(err, sql.ErrNoRows) {
 			return err
 		}
 
 		newConsumed := int64(request.DailyCost) + consumed
 		if newConsumed > budget && netIncrease {
+			c.Log.Debug(
+				ctx, "over quota, rejecting",
+				slog.F("prev_consumed", consumed),
+				slog.F("next_consumed", newConsumed),
+				slog.F("budget", budget),
+			)
 			return nil
 		}
 
 		err = s.UpdateWorkspaceBuildCostByID(ctx, database.UpdateWorkspaceBuildCostByIDParams{
-			ID:        build.ID,
+			ID:        nextBuild.ID,
 			DailyCost: request.DailyCost,
 		})
 		if err != nil {

diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
@@ -31,12 +32,13 @@ func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, con
 }
 
 func TestWorkspaceQuota(t *testing.T) {
-	// TODO: refactor for new impl
-
 	t.Parallel()
 
-	t.Run("BlocksBuild", func(t *testing.T) {
+	// This first test verifies the behavior of creating and deleting workspaces.
+	// It also tests multi-group quota stacking and the everyone group.
+	t.Run("CreateDelete", func(t *testing.T) {
 		t.Parallel()
+
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
 		max := 1
@@ -49,8 +51,6 @@ func TestWorkspaceQuota(t *testing.T) {
 			},
 		})
 		coderdtest.NewProvisionerDaemon(t, api.AGPL)
-		coderdtest.NewProvisionerDaemon(t, api.AGPL)
-		coderdtest.NewProvisionerDaemon(t, api.AGPL)
 
 		verifyQuota(ctx, t, client, 0, 0)
 
@@ -157,4 +157,104 @@ func TestWorkspaceQuota(t *testing.T) {
 		verifyQuota(ctx, t, client, 4, 4)
 		require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
 	})
+
+	t.Run("StartStop", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+		max := 1
+		client, _, api, user := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+			UserWorkspaceQuota: max,
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureTemplateRBAC: 1,
+				},
+			},
+		})
+		coderdtest.NewProvisionerDaemon(t, api.AGPL)
+
+		verifyQuota(ctx, t, client, 0, 0)
+
+		// Patch the 'Everyone' group to verify its quota allowance is being accounted for.
+		_, err := client.PatchGroup(ctx, user.OrganizationID, codersdk.PatchGroupRequest{
+			QuotaAllowance: ptr.Ref(4),
+		})
+		require.NoError(t, err)
+		verifyQuota(ctx, t, client, 0, 4)
+
+		stopResp := []*proto.Provision_Response{{
+			Type: &proto.Provision_Response_Complete{
+				Complete: &proto.Provision_Complete{
+					Resources: []*proto.Resource{{
+						Name:      "example",
+						Type:      "aws_instance",
+						DailyCost: 1,
+					}},
+				},
+			},
+		}}
+
+		startResp := []*proto.Provision_Response{{
+			Type: &proto.Provision_Response_Complete{
+				Complete: &proto.Provision_Complete{
+					Resources: []*proto.Resource{{
+						Name:      "example",
+						Type:      "aws_instance",
+						DailyCost: 2,
+					}},
+				},
+			},
+		}}
+
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse: echo.ParseComplete,
+			ProvisionPlanMap: map[proto.WorkspaceTransition][]*proto.Provision_Response{
+				proto.WorkspaceTransition_START: startResp,
+				proto.WorkspaceTransition_STOP:  stopResp,
+			},
+			ProvisionApplyMap: map[proto.WorkspaceTransition][]*proto.Provision_Response{
+				proto.WorkspaceTransition_START: startResp,
+				proto.WorkspaceTransition_STOP:  stopResp,
+			},
+		})
+
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+
+		// Spin up two workspaces.
+		var wg sync.WaitGroup
+		var workspaces []codersdk.Workspace
+		for i := 0; i < 2; i++ {
+			workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+			workspaces = append(workspaces, workspace)
+			build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+			assert.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
+		}
+		wg.Wait()
+		verifyQuota(ctx, t, client, 4, 4)
+
+		// Next one must fail
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		require.Contains(t, build.Job.Error, "quota")
+
+		// Consumed shouldn't bump
+		verifyQuota(ctx, t, client, 4, 4)
+		require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status)
+
+		build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStop)
+		build = coderdtest.AwaitWorkspaceBuildJob(t, client, build.ID)
+
+		// Quota goes down one
+		verifyQuota(ctx, t, client, 3, 4)
+		require.Equal(t, codersdk.WorkspaceStatusStopped, build.Status)
+
+		build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStart)
+		build = coderdtest.AwaitWorkspaceBuildJob(t, client, build.ID)
+
+		// Quota goes back up
+		verifyQuota(ctx, t, client, 4, 4)
+		require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
+	})
 }
diff --git a/enterprise/tailnet/workspaceproxy.go b/enterprise/tailnet/workspaceproxy.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"net"
 	"time"
 
@@ -26,6 +27,9 @@ func ServeWorkspaceProxy(ctx context.Context, conn net.Conn, ma agpl.MultiAgentC
 		var msg wsproxysdk.CoordinateMessage
 		err := decoder.Decode(&msg)
 		if err != nil {
+			if errors.Is(err, net.ErrClosed) {
+				return nil
+			}
 			return xerrors.Errorf("read json: %w", err)
 		}