synced/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2025-07-15 22:20:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: correctly reject quota-violating builds (#9233)

Browse Source 
Due to a logical error in CommitQuota, all workspace Stop->Start operations
were being accepted, regardless of the Quota limit. This issue only
appeared after #9201, so this was a minor regression in main for about
3 days. This PR adds a test to make sure this kind of bug doesn't recur.

To make the new test possible, we give the echo provisioner the ability
to simulate responses to specific transitions.
This commit is contained in:
Ammar Bandukwala
2023-08-21 21:55:39 -05:00
committed by GitHub
parent 69ec8d774b
commit 545a256b57
7 changed files with 293 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
enterprise/coderd/coderd.go
View File

@ -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 {

33
enterprise/coderd/workspacequota.go
View File

@ -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 {

110
enterprise/coderd/workspacequota_test.go
View File

@ -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)
})
}