opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: wire up usage tracking for managed agents (#19096)

Browse Source 
Wires up the usage collector and publisher to coderd.

Relates to coder/internal#814
This commit is contained in:
Dean Sheather
2025-08-20 23:38:09 +10:00
committed by GitHub
parent dd867bd743
commit 6eb02d1c2a
43 changed files with 540 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/usage/inserter.go
+12 -11
View File
@@ -13,16 +13,17 @@ import (
"github.com/coder/quartz"
)
// Inserter accepts usage events and stores them in the database for publishing.
type Inserter struct {
// dbCollector collects usage events and stores them in the database for
// publishing.
type dbCollector struct {
clock quartz.Clock
}
var _ agplusage.Inserter = &Inserter{}
var _ agplusage.Inserter = &dbCollector{}
// NewInserter creates a new database-backed usage event inserter.
func NewInserter(opts ...InserterOptions) *Inserter {
c := &Inserter{
// NewDBInserter creates a new database-backed usage event inserter.
func NewDBInserter(opts ...InserterOption) agplusage.Inserter {
c := &dbCollector{
clock: quartz.NewReal(),
}
for _, opt := range opts {
@@ -31,17 +32,17 @@ func NewInserter(opts ...InserterOptions) *Inserter {
return c
}
type InserterOptions func(*Inserter)
type InserterOption func(*dbCollector)
// InserterWithClock sets the quartz clock to use for the inserter.
func InserterWithClock(clock quartz.Clock) InserterOptions {
return func(c *Inserter) {
func InserterWithClock(clock quartz.Clock) InserterOption {
return func(c *dbCollector) {
c.clock = clock
}
}
// InsertDiscreteUsageEvent implements agplusage.Inserter.
func (c *Inserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event agplusage.DiscreteEvent) error {
func (i *dbCollector) InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event agplusage.DiscreteEvent) error {
if !event.EventType().IsDiscrete() {
return xerrors.Errorf("event type %q is not a discrete event", event.EventType())
}
@@ -61,6 +62,6 @@ func (c *Inserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.Sto
ID: uuid.New().String(),
EventType: string(event.EventType()),
EventData: jsonData,
CreatedAt: dbtime.Time(c.clock.Now()),
CreatedAt: dbtime.Time(i.clock.Now()),
})
}
enterprise/coderd/usage/inserter_test.go
+4 -4
View File
@@ -28,7 +28,7 @@ func TestInserter(t *testing.T) {
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
clock := quartz.NewMock(t)
inserter := usage.NewInserter(usage.InserterWithClock(clock))
inserter := usage.NewDBInserter(usage.InserterWithClock(clock))
now := dbtime.Now()
events := []struct {
@@ -51,8 +51,8 @@ func TestInserter(t *testing.T) {
for _, event := range events {
eventJSON := jsoninate(t, event.event)
db.EXPECT().InsertUsageEvent(ctx, gomock.Any()).DoAndReturn(
func(ctx interface{}, params database.InsertUsageEventParams) error {
db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn(
func(ctx any, params database.InsertUsageEventParams) error {
_, err := uuid.Parse(params.ID)
assert.NoError(t, err)
assert.Equal(t, string(event.event.EventType()), params.EventType)
@@ -76,7 +76,7 @@ func TestInserter(t *testing.T) {
db := dbmock.NewMockStore(ctrl)
// We should get an error if the event is invalid.
inserter := usage.NewInserter()
inserter := usage.NewDBInserter()
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
Count: 0, // invalid
})
enterprise/coderd/usage/publisher.go
+27 -26
View File
@@ -15,11 +15,11 @@ import (
"cdr.dev/slog"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/pproflabel"
agplusage "github.com/coder/coder/v2/coderd/usage"
"github.com/coder/coder/v2/cryptorand"
"github.com/coder/coder/v2/enterprise/coderd"
"github.com/coder/coder/v2/enterprise/coderd/license"
"github.com/coder/quartz"
)
@@ -49,17 +49,17 @@ type Publisher interface {
}
type tallymanPublisher struct {
ctx context.Context
ctxCancel context.CancelFunc
log slog.Logger
db database.Store
done chan struct{}
ctx context.Context
ctxCancel context.CancelFunc
log slog.Logger
db database.Store
licenseKeys map[string]ed25519.PublicKey
done chan struct{}
// Configured with options:
ingestURL string
httpClient *http.Client
clock quartz.Clock
licenseKeys map[string]ed25519.PublicKey
initialDelay time.Duration
}
@@ -67,19 +67,21 @@ var _ Publisher = &tallymanPublisher{}
// NewTallymanPublisher creates a Publisher that publishes usage events to
// Coder's Tallyman service.
func NewTallymanPublisher(ctx context.Context, log slog.Logger, db database.Store, opts ...TallymanPublisherOption) Publisher {
func NewTallymanPublisher(ctx context.Context, log slog.Logger, db database.Store, keys map[string]ed25519.PublicKey, opts ...TallymanPublisherOption) Publisher {
ctx, cancel := context.WithCancel(ctx)
publisher := &tallymanPublisher{
ctx: ctx,
ctxCancel: cancel,
log: log,
db: db,
done: make(chan struct{}),
ctx = dbauthz.AsUsagePublisher(ctx) //nolint:gocritic // we intentionally want to be able to process usage events
ingestURL: tallymanIngestURLV1,
httpClient: http.DefaultClient,
clock: quartz.NewReal(),
licenseKeys: coderd.Keys,
publisher := &tallymanPublisher{
ctx: ctx,
ctxCancel: cancel,
log: log,
db: db,
licenseKeys: keys,
done: make(chan struct{}),
ingestURL: tallymanIngestURLV1,
httpClient: http.DefaultClient,
clock: quartz.NewReal(),
}
for _, opt := range opts {
opt(publisher)
@@ -92,6 +94,9 @@ type TallymanPublisherOption func(*tallymanPublisher)
// PublisherWithHTTPClient sets the HTTP client to use for publishing usage events.
func PublisherWithHTTPClient(httpClient *http.Client) TallymanPublisherOption {
return func(p *tallymanPublisher) {
if httpClient == nil {
httpClient = http.DefaultClient
}
p.httpClient = httpClient
}
}
@@ -103,14 +108,6 @@ func PublisherWithClock(clock quartz.Clock) TallymanPublisherOption {
}
}
// PublisherWithLicenseKeys sets the license public keys to use for license
// validation.
func PublisherWithLicenseKeys(keys map[string]ed25519.PublicKey) TallymanPublisherOption {
return func(p *tallymanPublisher) {
p.licenseKeys = keys
}
}
// PublisherWithIngestURL sets the ingest URL to use for publishing usage
// events.
func PublisherWithIngestURL(ingestURL string) TallymanPublisherOption {
@@ -149,6 +146,10 @@ func (p *tallymanPublisher) Start() error {
p.initialDelay = tallymanPublishInitialMinimumDelay + time.Duration(plusDelay)
}
if len(p.licenseKeys) == 0 {
return xerrors.New("no license keys provided")
}
pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceTallymanPublisher), func(ctx context.Context) {
p.publishLoop(ctx, deploymentUUID)
})
enterprise/coderd/usage/publisher_test.go
+16 -14
View File
@@ -10,16 +10,20 @@ import (
"time"
"github.com/google/uuid"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/goleak"
"go.uber.org/mock/gomock"
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbmock"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/rbac"
agplusage "github.com/coder/coder/v2/coderd/usage"
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
"github.com/coder/coder/v2/enterprise/coderd/usage"
@@ -40,6 +44,7 @@ func TestIntegration(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitLong)
log := slogtest.Make(t, nil)
db, _ := dbtestutil.NewDB(t)
clock := quartz.NewMock(t)
deploymentID, licenseJWT := configureDeployment(ctx, t, db)
now := time.Now()
@@ -60,7 +65,7 @@ func TestIntegration(t *testing.T) {
return handler(req)
}))
inserter := usage.NewInserter(
inserter := usage.NewDBInserter(
usage.InserterWithClock(clock),
)
// Insert an old event that should never be published.
@@ -80,10 +85,12 @@ func TestIntegration(t *testing.T) {
require.NoErrorf(t, err, "collecting event %d", i)
}
publisher := usage.NewTallymanPublisher(ctx, log, db,
// Wrap the publisher's DB in a dbauthz to ensure that the publisher has
// enough permissions.
authzDB := dbauthz.New(db, rbac.NewAuthorizer(prometheus.NewRegistry()), log, coderdtest.AccessControlStorePointer())
publisher := usage.NewTallymanPublisher(ctx, log, authzDB, coderdenttest.Keys,
usage.PublisherWithClock(clock),
usage.PublisherWithIngestURL(ingestURL),
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
)
defer publisher.Close()
@@ -212,10 +219,9 @@ func TestPublisherNoEligibleLicenses(t *testing.T) {
}
}))
publisher := usage.NewTallymanPublisher(ctx, log, db,
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
usage.PublisherWithClock(clock),
usage.PublisherWithIngestURL(ingestURL),
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
)
defer publisher.Close()
@@ -283,14 +289,13 @@ func TestPublisherClaimExpiry(t *testing.T) {
return tallymanAcceptAllHandler(req)
}))
inserter := usage.NewInserter(
inserter := usage.NewDBInserter(
usage.InserterWithClock(clock),
)
publisher := usage.NewTallymanPublisher(ctx, log, db,
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
usage.PublisherWithClock(clock),
usage.PublisherWithIngestURL(ingestURL),
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
usage.PublisherWithInitialDelay(17*time.Minute),
)
defer publisher.Close()
@@ -367,10 +372,9 @@ func TestPublisherMissingEvents(t *testing.T) {
}
}))
publisher := usage.NewTallymanPublisher(ctx, log, db,
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
usage.PublisherWithClock(clock),
usage.PublisherWithIngestURL(ingestURL),
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
)
// Expect the publisher to call SelectUsageEventsForPublishing, followed by
@@ -510,10 +514,9 @@ func TestPublisherLicenseSelection(t *testing.T) {
return tallymanAcceptAllHandler(req)
}))
publisher := usage.NewTallymanPublisher(ctx, log, db,
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
usage.PublisherWithClock(clock),
usage.PublisherWithIngestURL(ingestURL),
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
)
defer publisher.Close()
@@ -579,10 +582,9 @@ func TestPublisherTallymanError(t *testing.T) {
}
}))
publisher := usage.NewTallymanPublisher(ctx, log, db,
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
usage.PublisherWithClock(clock),
usage.PublisherWithIngestURL(ingestURL),
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
)
defer publisher.Close()