mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +08:00
fix: do not enforce managed agent limit for non-task workspaces (#21689)
Only task workspaces have the checks in wsbuilder for violating the managed agent caps in the license. Stopped tasks that are resumed with a regular workspace start **still count as usage**.
This commit is contained in:
@@ -83,6 +83,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/schedule"
|
||||
"github.com/coder/coder/v2/coderd/telemetry"
|
||||
"github.com/coder/coder/v2/coderd/updatecheck"
|
||||
"github.com/coder/coder/v2/coderd/usage"
|
||||
"github.com/coder/coder/v2/coderd/util/namesgenerator"
|
||||
"github.com/coder/coder/v2/coderd/util/ptr"
|
||||
"github.com/coder/coder/v2/coderd/webpush"
|
||||
@@ -190,6 +191,7 @@ type Options struct {
|
||||
TelemetryReporter telemetry.Reporter
|
||||
|
||||
ProvisionerdServerMetrics *provisionerdserver.Metrics
|
||||
UsageInserter usage.Inserter
|
||||
}
|
||||
|
||||
// New constructs a codersdk client connected to an in-memory API instance.
|
||||
@@ -270,6 +272,11 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
|
||||
}
|
||||
}
|
||||
|
||||
var usageInserter *atomic.Pointer[usage.Inserter]
|
||||
if options.UsageInserter != nil {
|
||||
usageInserter = &atomic.Pointer[usage.Inserter]{}
|
||||
usageInserter.Store(&options.UsageInserter)
|
||||
}
|
||||
if options.Database == nil {
|
||||
options.Database, options.Pubsub = dbtestutil.NewDB(t)
|
||||
}
|
||||
@@ -563,6 +570,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
|
||||
Database: options.Database,
|
||||
Pubsub: options.Pubsub,
|
||||
ExternalAuthConfigs: options.ExternalAuthConfigs,
|
||||
UsageInserter: usageInserter,
|
||||
|
||||
Auditor: options.Auditor,
|
||||
ConnectionLogger: options.ConnectionLogger,
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
package coderdtest
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/usage"
|
||||
"github.com/coder/coder/v2/coderd/usage/usagetypes"
|
||||
)
|
||||
|
||||
var _ usage.Inserter = (*UsageInserter)(nil)
|
||||
|
||||
type UsageInserter struct {
|
||||
sync.Mutex
|
||||
events []usagetypes.DiscreteEvent
|
||||
}
|
||||
|
||||
func NewUsageInserter() *UsageInserter {
|
||||
return &UsageInserter{
|
||||
events: []usagetypes.DiscreteEvent{},
|
||||
}
|
||||
}
|
||||
|
||||
func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
u.events = append(u.events, event)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *UsageInserter) GetEvents() []usagetypes.DiscreteEvent {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
eventsCopy := make([]usagetypes.DiscreteEvent, len(u.events))
|
||||
copy(eventsCopy, u.events)
|
||||
return eventsCopy
|
||||
}
|
||||
|
||||
func (u *UsageInserter) Reset() {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
u.events = []usagetypes.DiscreteEvent{}
|
||||
}
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
@@ -85,13 +84,15 @@ type Builder struct {
|
||||
templateVersionPresetParameterValues *[]database.TemplateVersionPresetParameter
|
||||
parameterRender dynamicparameters.Renderer
|
||||
workspaceTags *map[string]string
|
||||
task *database.Task
|
||||
hasTask *bool // A workspace without a task will have a nil `task` and false `hasTask`.
|
||||
|
||||
prebuiltWorkspaceBuildStage sdkproto.PrebuiltWorkspaceBuildStage
|
||||
verifyNoLegacyParametersOnce bool
|
||||
}
|
||||
|
||||
type UsageChecker interface {
|
||||
CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (UsageCheckResponse, error)
|
||||
CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, transition database.WorkspaceTransition) (UsageCheckResponse, error)
|
||||
}
|
||||
|
||||
type UsageCheckResponse struct {
|
||||
@@ -103,7 +104,7 @@ type NoopUsageChecker struct{}
|
||||
|
||||
var _ UsageChecker = NoopUsageChecker{}
|
||||
|
||||
func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ database.WorkspaceTransition) (UsageCheckResponse, error) {
|
||||
func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ *database.Task, _ database.WorkspaceTransition) (UsageCheckResponse, error) {
|
||||
return UsageCheckResponse{
|
||||
Permitted: true,
|
||||
}, nil
|
||||
@@ -487,8 +488,12 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
|
||||
return BuildError{code, "insert workspace build", err}
|
||||
}
|
||||
|
||||
task, err := b.getWorkspaceTask()
|
||||
if err != nil {
|
||||
return BuildError{http.StatusInternalServerError, "get task by workspace id", err}
|
||||
}
|
||||
// If this is a task workspace, link it to the latest workspace build.
|
||||
if task, err := store.GetTaskByWorkspaceID(b.ctx, b.workspace.ID); err == nil {
|
||||
if task != nil {
|
||||
_, err = store.UpsertTaskWorkspaceApp(b.ctx, database.UpsertTaskWorkspaceAppParams{
|
||||
TaskID: task.ID,
|
||||
WorkspaceBuildNumber: buildNum,
|
||||
@@ -498,8 +503,6 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
|
||||
if err != nil {
|
||||
return BuildError{http.StatusInternalServerError, "upsert task workspace app", err}
|
||||
}
|
||||
} else if !errors.Is(err, sql.ErrNoRows) {
|
||||
return BuildError{http.StatusInternalServerError, "get task by workspace id", err}
|
||||
}
|
||||
|
||||
err = store.InsertWorkspaceBuildParameters(b.ctx, database.InsertWorkspaceBuildParametersParams{
|
||||
@@ -632,6 +635,27 @@ func (b *Builder) getTemplateVersionID() (uuid.UUID, error) {
|
||||
return bld.TemplateVersionID, nil
|
||||
}
|
||||
|
||||
// getWorkspaceTask returns the task associated with the workspace, if any.
|
||||
// If no task exists, it returns (nil, nil).
|
||||
func (b *Builder) getWorkspaceTask() (*database.Task, error) {
|
||||
if b.hasTask != nil {
|
||||
return b.task, nil
|
||||
}
|
||||
t, err := b.store.GetTaskByWorkspaceID(b.ctx, b.workspace.ID)
|
||||
if err != nil {
|
||||
if xerrors.Is(err, sql.ErrNoRows) {
|
||||
b.hasTask = ptr.Ref(false)
|
||||
//nolint:nilnil // No task exists.
|
||||
return nil, nil
|
||||
}
|
||||
return nil, xerrors.Errorf("get task: %w", err)
|
||||
}
|
||||
|
||||
b.task = &t
|
||||
b.hasTask = ptr.Ref(true)
|
||||
return b.task, nil
|
||||
}
|
||||
|
||||
func (b *Builder) getTemplateTerraformValues() (*database.TemplateVersionTerraformValue, error) {
|
||||
if b.terraformValues != nil {
|
||||
return b.terraformValues, nil
|
||||
@@ -1313,7 +1337,12 @@ func (b *Builder) checkUsage() error {
|
||||
return BuildError{http.StatusInternalServerError, "Failed to fetch template version", err}
|
||||
}
|
||||
|
||||
resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion, b.trans)
|
||||
task, err := b.getWorkspaceTask()
|
||||
if err != nil {
|
||||
return BuildError{http.StatusInternalServerError, "Failed to fetch workspace task", err}
|
||||
}
|
||||
|
||||
resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion, task, b.trans)
|
||||
if err != nil {
|
||||
return BuildError{http.StatusInternalServerError, "Failed to check build usage", err}
|
||||
}
|
||||
|
||||
@@ -570,6 +570,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
|
||||
mDB := expectDB(t,
|
||||
// Inputs
|
||||
withTemplate,
|
||||
withNoTask,
|
||||
withInactiveVersionNoParams(),
|
||||
withLastBuildFound,
|
||||
withTemplateVersionVariables(inactiveVersionID, nil),
|
||||
@@ -605,6 +606,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
|
||||
withTemplate,
|
||||
withInactiveVersion(richParameters),
|
||||
withLastBuildFound,
|
||||
withNoTask,
|
||||
withTemplateVersionVariables(inactiveVersionID, nil),
|
||||
withRichParameters(initialBuildParameters),
|
||||
withParameterSchemas(inactiveJobID, nil),
|
||||
@@ -1049,7 +1051,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
|
||||
|
||||
var calls int64
|
||||
fakeUsageChecker := &fakeUsageChecker{
|
||||
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
|
||||
checkBuildUsageFunc: func(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ *database.Task, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
|
||||
atomic.AddInt64(&calls, 1)
|
||||
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
|
||||
},
|
||||
@@ -1126,7 +1128,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
|
||||
|
||||
var calls int64
|
||||
fakeUsageChecker := &fakeUsageChecker{
|
||||
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
|
||||
checkBuildUsageFunc: func(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ *database.Task, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
|
||||
atomic.AddInt64(&calls, 1)
|
||||
return c.response, c.responseErr
|
||||
},
|
||||
@@ -1134,6 +1136,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
|
||||
|
||||
mDB := expectDB(t,
|
||||
withTemplate,
|
||||
withNoTask,
|
||||
withInactiveVersionNoParams(),
|
||||
)
|
||||
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
|
||||
@@ -1577,11 +1580,11 @@ func expectFindMatchingPresetID(id uuid.UUID, err error) func(mTx *dbmock.MockSt
|
||||
}
|
||||
|
||||
type fakeUsageChecker struct {
|
||||
checkBuildUsageFunc func(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error)
|
||||
checkBuildUsageFunc func(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error)
|
||||
}
|
||||
|
||||
func (f *fakeUsageChecker) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
|
||||
return f.checkBuildUsageFunc(ctx, store, templateVersion, transition)
|
||||
func (f *fakeUsageChecker) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
|
||||
return f.checkBuildUsageFunc(ctx, store, templateVersion, task, transition)
|
||||
}
|
||||
|
||||
func withNoTask(mTx *dbmock.MockStore) {
|
||||
|
||||
Reference in New Issue
Block a user