Skip to content

Commit 9a6dd73

Browse files
deansheatherdeansheather
authored
feat: add managed agent license limit checks (#18937)
- Adds a query for counting managed agent workspace builds between two timestamps - The "Actual" field in the feature entitlement for managed agents is now populated with the value read from the database - The wsbuilder package now validates AI agent usage against the limit when a license is installed Closes coder/internal#777
1 parent aa1a985 commit 9a6dd73

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
defer autobuildTicker.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
statsCh chan<- Stats
@@ -65,7 +66,7 @@ type Stats struct {
6566
}
6667

6768
// New returns a new wsactions executor.
68-
func NewExecutor(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+
func NewExecutor(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
if nextTransition != "" {
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+
var buildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
567+
var noopUsageChecker 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+
var buildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
369+
var noopUsageChecker 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
return q.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+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace); err != nil {
2199+
return 0, xerrors.Errorf("authorize read all workspaces: %w", err)
2200+
}
2201+
return q.db.GetManagedAgentCount(ctx, arg)
2202+
}
2203+
21962204
func (q *querier) GetNotificationMessagesByStatus(ctx context.Context, arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage, error) {
21972205
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceNotificationMessage); err != nil {
21982206
return nil, 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 about customizing 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 about customizing 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 about customizing 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 about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy