Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit9a6dd73

Browse files
authored
feat: add managed agent license limit checks (#18937)
- Adds a query for counting managed agent workspace builds between twotimestamps- The "Actual" field in the feature entitlement for managed agents isnow populated with the value read from the database- The wsbuilder package now validates AI agent usage against the limitwhen a license is installedClosescoder/internal#777
1 parentaa1a985 commit9a6dd73

24 files changed

+586
-74
lines changed

‎cli/server.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1101,7 +1101,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
11011101
autobuildTicker:=time.NewTicker(vals.AutobuildPollInterval.Value())
11021102
deferautobuildTicker.Stop()
11031103
autobuildExecutor:=autobuild.NewExecutor(
1104-
ctx,options.Database,options.Pubsub,coderAPI.FileCache,options.PrometheusRegistry,coderAPI.TemplateScheduleStore,&coderAPI.Auditor,coderAPI.AccessControlStore,logger,autobuildTicker.C,options.NotificationsEnqueuer,coderAPI.Experiments)
1104+
ctx,options.Database,options.Pubsub,coderAPI.FileCache,options.PrometheusRegistry,coderAPI.TemplateScheduleStore,&coderAPI.Auditor,coderAPI.AccessControlStore,coderAPI.BuildUsageChecker,logger,autobuildTicker.C,options.NotificationsEnqueuer,coderAPI.Experiments)
11051105
autobuildExecutor.Run()
11061106

11071107
jobReaperTicker:=time.NewTicker(vals.JobReaperDetectorInterval.Value())

‎coderd/autobuild/lifecycle_executor.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ type Executor struct {
4242
templateScheduleStore*atomic.Pointer[schedule.TemplateScheduleStore]
4343
accessControlStore*atomic.Pointer[dbauthz.AccessControlStore]
4444
auditor*atomic.Pointer[audit.Auditor]
45+
buildUsageChecker*atomic.Pointer[wsbuilder.UsageChecker]
4546
log slog.Logger
4647
tick<-chan time.Time
4748
statsChchan<-Stats
@@ -65,7 +66,7 @@ type Stats struct {
6566
}
6667

6768
// New returns a new wsactions executor.
68-
funcNewExecutor(ctx context.Context,db database.Store,ps pubsub.Pubsub,fc*files.Cache,reg prometheus.Registerer,tss*atomic.Pointer[schedule.TemplateScheduleStore],auditor*atomic.Pointer[audit.Auditor],acs*atomic.Pointer[dbauthz.AccessControlStore],log slog.Logger,tick<-chan time.Time,enqueuer notifications.Enqueuer,exp codersdk.Experiments)*Executor {
69+
funcNewExecutor(ctx context.Context,db database.Store,ps pubsub.Pubsub,fc*files.Cache,reg prometheus.Registerer,tss*atomic.Pointer[schedule.TemplateScheduleStore],auditor*atomic.Pointer[audit.Auditor],acs*atomic.Pointer[dbauthz.AccessControlStore],buildUsageChecker*atomic.Pointer[wsbuilder.UsageChecker],log slog.Logger,tick<-chan time.Time,enqueuer notifications.Enqueuer,exp codersdk.Experiments)*Executor {
6970
factory:=promauto.With(reg)
7071
le:=&Executor{
7172
//nolint:gocritic // Autostart has a limited set of permissions.
@@ -78,6 +79,7 @@ func NewExecutor(ctx context.Context, db database.Store, ps pubsub.Pubsub, fc *f
7879
log:log.Named("autobuild"),
7980
auditor:auditor,
8081
accessControlStore:acs,
82+
buildUsageChecker:buildUsageChecker,
8183
notificationsEnqueuer:enqueuer,
8284
reg:reg,
8385
experiments:exp,
@@ -279,7 +281,7 @@ func (e *Executor) runOnce(t time.Time) Stats {
279281
}
280282

281283
ifnextTransition!="" {
282-
builder:=wsbuilder.New(ws,nextTransition).
284+
builder:=wsbuilder.New(ws,nextTransition,*e.buildUsageChecker.Load()).
283285
SetLastWorkspaceBuildInTx(&latestBuild).
284286
SetLastWorkspaceBuildJobInTx(&latestJob).
285287
Experiments(e.experiments).

‎coderd/coderd.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121

2222
"github.com/coder/coder/v2/coderd/oauth2provider"
2323
"github.com/coder/coder/v2/coderd/prebuilds"
24+
"github.com/coder/coder/v2/coderd/wsbuilder"
2425

2526
"github.com/andybalholm/brotli"
2627
"github.com/go-chi/chi/v5"
@@ -559,6 +560,13 @@ func New(options *Options) *API {
559560
// bugs that may only occur when a key isn't precached in tests and the latency cost is minimal.
560561
cryptokeys.StartRotator(ctx,options.Logger,options.Database)
561562

563+
// AGPL uses a no-op build usage checker as there are no license
564+
// entitlements to enforce. This is swapped out in
565+
// enterprise/coderd/coderd.go.
566+
varbuildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
567+
varnoopUsageChecker wsbuilder.UsageChecker= wsbuilder.NoopUsageChecker{}
568+
buildUsageChecker.Store(&noopUsageChecker)
569+
562570
api:=&API{
563571
ctx:ctx,
564572
cancel:cancel,
@@ -579,6 +587,7 @@ func New(options *Options) *API {
579587
TemplateScheduleStore:options.TemplateScheduleStore,
580588
UserQuietHoursScheduleStore:options.UserQuietHoursScheduleStore,
581589
AccessControlStore:options.AccessControlStore,
590+
BuildUsageChecker:&buildUsageChecker,
582591
FileCache:files.New(options.PrometheusRegistry,options.Authorizer),
583592
Experiments:experiments,
584593
WebpushDispatcher:options.WebPushDispatcher,
@@ -1650,6 +1659,9 @@ type API struct {
16501659
FileCache*files.Cache
16511660
PrebuildsClaimer atomic.Pointer[prebuilds.Claimer]
16521661
PrebuildsReconciler atomic.Pointer[prebuilds.ReconciliationOrchestrator]
1662+
// BuildUsageChecker is a pointer as it's passed around to multiple
1663+
// components.
1664+
BuildUsageChecker*atomic.Pointer[wsbuilder.UsageChecker]
16531665

16541666
UpdatesProvider tailnet.WorkspaceUpdatesProvider
16551667

‎coderd/coderdtest/coderdtest.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ import (
5555
"cdr.dev/slog/sloggers/slogtest"
5656
"github.com/coder/coder/v2/archive"
5757
"github.com/coder/coder/v2/coderd/files"
58+
"github.com/coder/coder/v2/coderd/wsbuilder"
5859
"github.com/coder/quartz"
5960

6061
"github.com/coder/coder/v2/coderd"
@@ -364,6 +365,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
364365
}
365366
connectionLogger.Store(&options.ConnectionLogger)
366367

368+
varbuildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
369+
varnoopUsageChecker wsbuilder.UsageChecker= wsbuilder.NoopUsageChecker{}
370+
buildUsageChecker.Store(&noopUsageChecker)
371+
367372
ctx,cancelFunc:=context.WithCancel(context.Background())
368373
experiments:=coderd.ReadExperiments(*options.Logger,options.DeploymentValues.Experiments)
369374
lifecycleExecutor:=autobuild.NewExecutor(
@@ -375,6 +380,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
375380
&templateScheduleStore,
376381
&auditor,
377382
accessControlStore,
383+
&buildUsageChecker,
378384
*options.Logger,
379385
options.AutobuildTicker,
380386
options.NotificationsEnqueuer,

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2193,6 +2193,14 @@ func (q *querier) GetLogoURL(ctx context.Context) (string, error) {
21932193
returnq.db.GetLogoURL(ctx)
21942194
}
21952195

2196+
func (q*querier)GetManagedAgentCount(ctx context.Context,arg database.GetManagedAgentCountParams) (int64,error) {
2197+
// Must be able to read all workspaces to check usage.
2198+
iferr:=q.authorizeContext(ctx,policy.ActionRead,rbac.ResourceWorkspace);err!=nil {
2199+
return0,xerrors.Errorf("authorize read all workspaces: %w",err)
2200+
}
2201+
returnq.db.GetManagedAgentCount(ctx,arg)
2202+
}
2203+
21962204
func (q*querier)GetNotificationMessagesByStatus(ctx context.Context,arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage,error) {
21972205
iferr:=q.authorizeContext(ctx,policy.ActionRead,rbac.ResourceNotificationMessage);err!=nil {
21982206
returnnil,err

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,20 +17,18 @@ import (
1717
"golang.org/x/xerrors"
1818

1919
"cdr.dev/slog"
20-
21-
"github.com/coder/coder/v2/coderd/database/db2sdk"
22-
"github.com/coder/coder/v2/coderd/notifications"
23-
"github.com/coder/coder/v2/coderd/rbac/policy"
24-
"github.com/coder/coder/v2/codersdk"
25-
2620
"github.com/coder/coder/v2/coderd/coderdtest"
2721
"github.com/coder/coder/v2/coderd/database"
22+
"github.com/coder/coder/v2/coderd/database/db2sdk"
2823
"github.com/coder/coder/v2/coderd/database/dbauthz"
2924
"github.com/coder/coder/v2/coderd/database/dbgen"
3025
"github.com/coder/coder/v2/coderd/database/dbtestutil"
3126
"github.com/coder/coder/v2/coderd/database/dbtime"
27+
"github.com/coder/coder/v2/coderd/notifications"
3228
"github.com/coder/coder/v2/coderd/rbac"
29+
"github.com/coder/coder/v2/coderd/rbac/policy"
3330
"github.com/coder/coder/v2/coderd/util/slice"
31+
"github.com/coder/coder/v2/codersdk"
3432
"github.com/coder/coder/v2/provisionersdk"
3533
"github.com/coder/coder/v2/testutil"
3634
)
@@ -903,6 +901,14 @@ func (s *MethodTestSuite) TestLicense() {
903901
require.NoError(s.T(),err)
904902
check.Args().Asserts().Returns("value")
905903
}))
904+
s.Run("GetManagedAgentCount",s.Subtest(func(db database.Store,check*expects) {
905+
start:=dbtime.Now()
906+
end:=start.Add(time.Hour)
907+
check.Args(database.GetManagedAgentCountParams{
908+
StartTime:start,
909+
EndTime:end,
910+
}).Asserts(rbac.ResourceWorkspace,policy.ActionRead).Returns(int64(0))
911+
}))
906912
}
907913

908914
func (s*MethodTestSuite)TestOrganization() {

‎coderd/database/dbmetrics/querymetrics.go‎

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

‎coderd/database/dbmock/dbmock.go‎

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

‎coderd/database/querier.go‎

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

‎coderd/database/queries.sql.go‎

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

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp