mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
chore: add usage tracking package (#19095)
Not used in coderd yet, see stack. Adds two new packages: - `coderd/usage`: provides an interface for the "Collector" as well as a stub implementation for AGPL - `enterprise/coderd/usage`: provides an interface for the "Publisher" as well as a Tallyman implementation Relates to https://github.com/coder/internal/issues/814
This commit is contained in:
@@ -161,12 +161,13 @@ func NewWithAPI(t *testing.T, options *Options) (
|
||||
// LicenseOptions is used to generate a license for testing.
|
||||
// It supports the builder pattern for easy customization.
|
||||
type LicenseOptions struct {
|
||||
AccountType string
|
||||
AccountID string
|
||||
DeploymentIDs []string
|
||||
Trial bool
|
||||
FeatureSet codersdk.FeatureSet
|
||||
AllFeatures bool
|
||||
AccountType string
|
||||
AccountID string
|
||||
DeploymentIDs []string
|
||||
Trial bool
|
||||
FeatureSet codersdk.FeatureSet
|
||||
AllFeatures bool
|
||||
PublishUsageData bool
|
||||
// GraceAt is the time at which the license will enter the grace period.
|
||||
GraceAt time.Time
|
||||
// ExpiresAt is the time at which the license will hard expire.
|
||||
@@ -271,6 +272,13 @@ func GenerateLicense(t *testing.T, options LicenseOptions) string {
|
||||
issuedAt = time.Now().Add(-time.Minute)
|
||||
}
|
||||
|
||||
if options.AccountType == "" {
|
||||
options.AccountType = license.AccountTypeSalesforce
|
||||
}
|
||||
if options.AccountID == "" {
|
||||
options.AccountID = "test-account-id"
|
||||
}
|
||||
|
||||
c := &license.Claims{
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ID: uuid.NewString(),
|
||||
@@ -279,15 +287,16 @@ func GenerateLicense(t *testing.T, options LicenseOptions) string {
|
||||
NotBefore: jwt.NewNumericDate(options.NotBefore),
|
||||
IssuedAt: jwt.NewNumericDate(issuedAt),
|
||||
},
|
||||
LicenseExpires: jwt.NewNumericDate(options.GraceAt),
|
||||
AccountType: options.AccountType,
|
||||
AccountID: options.AccountID,
|
||||
DeploymentIDs: options.DeploymentIDs,
|
||||
Trial: options.Trial,
|
||||
Version: license.CurrentVersion,
|
||||
AllFeatures: options.AllFeatures,
|
||||
FeatureSet: options.FeatureSet,
|
||||
Features: options.Features,
|
||||
LicenseExpires: jwt.NewNumericDate(options.GraceAt),
|
||||
AccountType: options.AccountType,
|
||||
AccountID: options.AccountID,
|
||||
DeploymentIDs: options.DeploymentIDs,
|
||||
Trial: options.Trial,
|
||||
Version: license.CurrentVersion,
|
||||
AllFeatures: options.AllFeatures,
|
||||
FeatureSet: options.FeatureSet,
|
||||
Features: options.Features,
|
||||
PublishUsageData: options.PublishUsageData,
|
||||
}
|
||||
return GenerateLicenseRaw(t, c)
|
||||
}
|
||||
|
||||
@@ -584,6 +584,7 @@ type Claims struct {
|
||||
Version uint64 `json:"version"`
|
||||
Features Features `json:"features"`
|
||||
RequireTelemetry bool `json:"require_telemetry,omitempty"`
|
||||
PublishUsageData bool `json:"publish_usage_data,omitempty"`
|
||||
}
|
||||
|
||||
var _ jwt.Claims = &Claims{}
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtime"
|
||||
agplusage "github.com/coder/coder/v2/coderd/usage"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// Inserter accepts usage events and stores them in the database for publishing.
|
||||
type Inserter struct {
|
||||
clock quartz.Clock
|
||||
}
|
||||
|
||||
var _ agplusage.Inserter = &Inserter{}
|
||||
|
||||
// NewInserter creates a new database-backed usage event inserter.
|
||||
func NewInserter(opts ...InserterOptions) *Inserter {
|
||||
c := &Inserter{
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(c)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
type InserterOptions func(*Inserter)
|
||||
|
||||
// InserterWithClock sets the quartz clock to use for the inserter.
|
||||
func InserterWithClock(clock quartz.Clock) InserterOptions {
|
||||
return func(c *Inserter) {
|
||||
c.clock = clock
|
||||
}
|
||||
}
|
||||
|
||||
// InsertDiscreteUsageEvent implements agplusage.Inserter.
|
||||
func (c *Inserter) 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())
|
||||
}
|
||||
if err := event.Valid(); err != nil {
|
||||
return xerrors.Errorf("invalid %q event: %w", event.EventType(), err)
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(event.Fields())
|
||||
if err != nil {
|
||||
return xerrors.Errorf("marshal event as JSON: %w", err)
|
||||
}
|
||||
|
||||
// Duplicate events are ignored by the query, so we don't need to check the
|
||||
// error.
|
||||
return tx.InsertUsageEvent(ctx, database.InsertUsageEventParams{
|
||||
// Always generate a new UUID for discrete events.
|
||||
ID: uuid.New().String(),
|
||||
EventType: string(event.EventType()),
|
||||
EventData: jsonData,
|
||||
CreatedAt: dbtime.Time(c.clock.Now()),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package usage_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtime"
|
||||
agplusage "github.com/coder/coder/v2/coderd/usage"
|
||||
"github.com/coder/coder/v2/enterprise/coderd/usage"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
func TestInserter(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("OK", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
clock := quartz.NewMock(t)
|
||||
inserter := usage.NewInserter(usage.InserterWithClock(clock))
|
||||
|
||||
now := dbtime.Now()
|
||||
events := []struct {
|
||||
time time.Time
|
||||
event agplusage.DiscreteEvent
|
||||
}{
|
||||
{
|
||||
time: now,
|
||||
event: agplusage.DCManagedAgentsV1{
|
||||
Count: 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
time: now.Add(1 * time.Minute),
|
||||
event: agplusage.DCManagedAgentsV1{
|
||||
Count: 2,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, event := range events {
|
||||
eventJSON := jsoninate(t, event.event)
|
||||
db.EXPECT().InsertUsageEvent(ctx, gomock.Any()).DoAndReturn(
|
||||
func(ctx interface{}, params database.InsertUsageEventParams) error {
|
||||
_, err := uuid.Parse(params.ID)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, string(event.event.EventType()), params.EventType)
|
||||
assert.JSONEq(t, eventJSON, string(params.EventData))
|
||||
assert.Equal(t, event.time, params.CreatedAt)
|
||||
return nil
|
||||
},
|
||||
).Times(1)
|
||||
|
||||
clock.Set(event.time)
|
||||
err := inserter.InsertDiscreteUsageEvent(ctx, db, event.event)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("InvalidEvent", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
|
||||
// We should get an error if the event is invalid.
|
||||
inserter := usage.NewInserter()
|
||||
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
|
||||
Count: 0, // invalid
|
||||
})
|
||||
assert.ErrorContains(t, err, `invalid "dc_managed_agents_v1" event: count must be greater than 0`)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,463 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"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"
|
||||
)
|
||||
|
||||
const (
|
||||
CoderLicenseJWTHeader = "Coder-License-JWT"
|
||||
|
||||
tallymanURL = "https://tallyman-prod.coder.com"
|
||||
tallymanIngestURLV1 = tallymanURL + "/api/v1/events/ingest"
|
||||
|
||||
tallymanPublishInitialMinimumDelay = 5 * time.Minute
|
||||
// Chosen to be a prime number and not a multiple of 5 like many other
|
||||
// recurring tasks.
|
||||
tallymanPublishInterval = 17 * time.Minute
|
||||
tallymanPublishTimeout = 30 * time.Second
|
||||
tallymanPublishBatchSize = 100
|
||||
)
|
||||
|
||||
var errUsagePublishingDisabled = xerrors.New("usage publishing is not enabled by any license")
|
||||
|
||||
// Publisher publishes usage events ***somewhere***.
|
||||
type Publisher interface {
|
||||
// Close closes the publisher and waits for it to finish.
|
||||
io.Closer
|
||||
// Start starts the publisher. It must only be called once.
|
||||
Start() error
|
||||
}
|
||||
|
||||
type tallymanPublisher struct {
|
||||
ctx context.Context
|
||||
ctxCancel context.CancelFunc
|
||||
log slog.Logger
|
||||
db database.Store
|
||||
done chan struct{}
|
||||
|
||||
// Configured with options:
|
||||
ingestURL string
|
||||
httpClient *http.Client
|
||||
clock quartz.Clock
|
||||
licenseKeys map[string]ed25519.PublicKey
|
||||
initialDelay time.Duration
|
||||
}
|
||||
|
||||
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 {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
publisher := &tallymanPublisher{
|
||||
ctx: ctx,
|
||||
ctxCancel: cancel,
|
||||
log: log,
|
||||
db: db,
|
||||
done: make(chan struct{}),
|
||||
|
||||
ingestURL: tallymanIngestURLV1,
|
||||
httpClient: http.DefaultClient,
|
||||
clock: quartz.NewReal(),
|
||||
licenseKeys: coderd.Keys,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(publisher)
|
||||
}
|
||||
return publisher
|
||||
}
|
||||
|
||||
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) {
|
||||
p.httpClient = httpClient
|
||||
}
|
||||
}
|
||||
|
||||
// PublisherWithClock sets the clock to use for publishing usage events.
|
||||
func PublisherWithClock(clock quartz.Clock) TallymanPublisherOption {
|
||||
return func(p *tallymanPublisher) {
|
||||
p.clock = clock
|
||||
}
|
||||
}
|
||||
|
||||
// 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 {
|
||||
return func(p *tallymanPublisher) {
|
||||
p.ingestURL = ingestURL
|
||||
}
|
||||
}
|
||||
|
||||
// PublisherWithInitialDelay sets the initial delay for the publisher.
|
||||
func PublisherWithInitialDelay(initialDelay time.Duration) TallymanPublisherOption {
|
||||
return func(p *tallymanPublisher) {
|
||||
p.initialDelay = initialDelay
|
||||
}
|
||||
}
|
||||
|
||||
// Start implements Publisher.
|
||||
func (p *tallymanPublisher) Start() error {
|
||||
ctx := p.ctx
|
||||
deploymentID, err := p.db.GetDeploymentID(ctx)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get deployment ID: %w", err)
|
||||
}
|
||||
deploymentUUID, err := uuid.Parse(deploymentID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("parse deployment ID %q: %w", deploymentID, err)
|
||||
}
|
||||
|
||||
if p.initialDelay <= 0 {
|
||||
// Pick a random time between tallymanPublishInitialMinimumDelay and
|
||||
// tallymanPublishInterval.
|
||||
maxPlusDelay := int(tallymanPublishInterval - tallymanPublishInitialMinimumDelay)
|
||||
plusDelay, err := cryptorand.Intn(maxPlusDelay)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("could not generate random start delay: %w", err)
|
||||
}
|
||||
p.initialDelay = tallymanPublishInitialMinimumDelay + time.Duration(plusDelay)
|
||||
}
|
||||
|
||||
pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceTallymanPublisher), func(ctx context.Context) {
|
||||
p.publishLoop(ctx, deploymentUUID)
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *tallymanPublisher) publishLoop(ctx context.Context, deploymentID uuid.UUID) {
|
||||
defer close(p.done)
|
||||
|
||||
// Start the ticker with the initial delay. We will reset it to the interval
|
||||
// after the first tick.
|
||||
ticker := p.clock.NewTicker(p.initialDelay)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
|
||||
err := p.publish(ctx, deploymentID)
|
||||
if err != nil {
|
||||
p.log.Warn(ctx, "publish usage events to tallyman", slog.Error(err))
|
||||
}
|
||||
ticker.Reset(tallymanPublishInterval)
|
||||
}
|
||||
}
|
||||
|
||||
// publish publishes usage events to Tallyman in a loop until there is an error
|
||||
// (or any rejection) or there are no more events to publish.
|
||||
func (p *tallymanPublisher) publish(ctx context.Context, deploymentID uuid.UUID) error {
|
||||
for {
|
||||
publishCtx, publishCtxCancel := context.WithTimeout(ctx, tallymanPublishTimeout)
|
||||
accepted, err := p.publishOnce(publishCtx, deploymentID)
|
||||
publishCtxCancel()
|
||||
if err != nil {
|
||||
return xerrors.Errorf("publish usage events to tallyman: %w", err)
|
||||
}
|
||||
if accepted < tallymanPublishBatchSize {
|
||||
// We published less than the batch size, so we're done.
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// publishOnce publishes up to tallymanPublishBatchSize usage events to
|
||||
// tallyman. It returns the number of successfully published events.
|
||||
func (p *tallymanPublisher) publishOnce(ctx context.Context, deploymentID uuid.UUID) (int, error) {
|
||||
licenseJwt, err := p.getBestLicenseJWT(ctx)
|
||||
if xerrors.Is(err, errUsagePublishingDisabled) {
|
||||
return 0, nil
|
||||
} else if err != nil {
|
||||
return 0, xerrors.Errorf("find usage publishing license: %w", err)
|
||||
}
|
||||
|
||||
events, err := p.db.SelectUsageEventsForPublishing(ctx, dbtime.Time(p.clock.Now()))
|
||||
if err != nil {
|
||||
return 0, xerrors.Errorf("select usage events for publishing: %w", err)
|
||||
}
|
||||
if len(events) == 0 {
|
||||
// No events to publish.
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
var (
|
||||
eventIDs = make(map[string]struct{})
|
||||
tallymanReq = TallymanIngestRequestV1{
|
||||
DeploymentID: deploymentID,
|
||||
Events: make([]TallymanIngestEventV1, 0, len(events)),
|
||||
}
|
||||
)
|
||||
for _, event := range events {
|
||||
eventIDs[event.ID] = struct{}{}
|
||||
eventType := agplusage.EventType(event.EventType)
|
||||
if !eventType.Valid() {
|
||||
// This should never happen due to the check constraint in the
|
||||
// database.
|
||||
return 0, xerrors.Errorf("event %q has an invalid event type %q", event.ID, event.EventType)
|
||||
}
|
||||
tallymanReq.Events = append(tallymanReq.Events, TallymanIngestEventV1{
|
||||
ID: event.ID,
|
||||
EventType: eventType,
|
||||
EventData: event.EventData,
|
||||
CreatedAt: event.CreatedAt,
|
||||
})
|
||||
}
|
||||
if len(eventIDs) != len(events) {
|
||||
// This should never happen due to the unique constraint in the
|
||||
// database.
|
||||
return 0, xerrors.Errorf("duplicate event IDs found in events for publishing")
|
||||
}
|
||||
|
||||
resp, err := p.sendPublishRequest(ctx, licenseJwt, tallymanReq)
|
||||
allFailed := err != nil
|
||||
if err != nil {
|
||||
p.log.Warn(ctx, "failed to send publish request to tallyman", slog.F("count", len(events)), slog.Error(err))
|
||||
// Fake a response with all events temporarily rejected.
|
||||
resp = TallymanIngestResponseV1{
|
||||
AcceptedEvents: []TallymanIngestAcceptedEventV1{},
|
||||
RejectedEvents: make([]TallymanIngestRejectedEventV1, len(events)),
|
||||
}
|
||||
for i, event := range events {
|
||||
resp.RejectedEvents[i] = TallymanIngestRejectedEventV1{
|
||||
ID: event.ID,
|
||||
Message: fmt.Sprintf("failed to publish to tallyman: %v", err),
|
||||
Permanent: false,
|
||||
}
|
||||
}
|
||||
} else {
|
||||
p.log.Debug(ctx, "published usage events to tallyman", slog.F("accepted", len(resp.AcceptedEvents)), slog.F("rejected", len(resp.RejectedEvents)))
|
||||
}
|
||||
|
||||
if len(resp.AcceptedEvents)+len(resp.RejectedEvents) != len(events) {
|
||||
p.log.Warn(ctx, "tallyman returned a different number of events than we sent", slog.F("sent", len(events)), slog.F("accepted", len(resp.AcceptedEvents)), slog.F("rejected", len(resp.RejectedEvents)))
|
||||
}
|
||||
|
||||
acceptedEvents := make(map[string]*TallymanIngestAcceptedEventV1)
|
||||
rejectedEvents := make(map[string]*TallymanIngestRejectedEventV1)
|
||||
for _, event := range resp.AcceptedEvents {
|
||||
acceptedEvents[event.ID] = &event
|
||||
}
|
||||
for _, event := range resp.RejectedEvents {
|
||||
rejectedEvents[event.ID] = &event
|
||||
}
|
||||
|
||||
dbUpdate := database.UpdateUsageEventsPostPublishParams{
|
||||
Now: dbtime.Time(p.clock.Now()),
|
||||
IDs: make([]string, len(events)),
|
||||
FailureMessages: make([]string, len(events)),
|
||||
SetPublishedAts: make([]bool, len(events)),
|
||||
}
|
||||
for i, event := range events {
|
||||
dbUpdate.IDs[i] = event.ID
|
||||
if _, ok := acceptedEvents[event.ID]; ok {
|
||||
dbUpdate.FailureMessages[i] = ""
|
||||
dbUpdate.SetPublishedAts[i] = true
|
||||
continue
|
||||
}
|
||||
if rejectedEvent, ok := rejectedEvents[event.ID]; ok {
|
||||
dbUpdate.FailureMessages[i] = rejectedEvent.Message
|
||||
dbUpdate.SetPublishedAts[i] = rejectedEvent.Permanent
|
||||
continue
|
||||
}
|
||||
// It's not good if this path gets hit, but we'll handle it as if it
|
||||
// was a temporary rejection.
|
||||
dbUpdate.FailureMessages[i] = "tallyman did not include the event in the response"
|
||||
dbUpdate.SetPublishedAts[i] = false
|
||||
}
|
||||
|
||||
// Collate rejected events into a single map of ID to failure message for
|
||||
// logging. We only want to log once.
|
||||
// If all events failed, we'll log the overall error above.
|
||||
if !allFailed {
|
||||
rejectionReasonsForLog := make(map[string]string)
|
||||
for i, id := range dbUpdate.IDs {
|
||||
failureMessage := dbUpdate.FailureMessages[i]
|
||||
if failureMessage == "" {
|
||||
continue
|
||||
}
|
||||
setPublishedAt := dbUpdate.SetPublishedAts[i]
|
||||
if setPublishedAt {
|
||||
failureMessage = "permanently rejected: " + failureMessage
|
||||
}
|
||||
rejectionReasonsForLog[id] = failureMessage
|
||||
}
|
||||
if len(rejectionReasonsForLog) > 0 {
|
||||
p.log.Warn(ctx, "tallyman rejected usage events", slog.F("count", len(rejectionReasonsForLog)), slog.F("event_failure_reasons", rejectionReasonsForLog))
|
||||
}
|
||||
}
|
||||
|
||||
err = p.db.UpdateUsageEventsPostPublish(ctx, dbUpdate)
|
||||
if err != nil {
|
||||
return 0, xerrors.Errorf("update usage events post publish: %w", err)
|
||||
}
|
||||
|
||||
var returnErr error
|
||||
if len(resp.RejectedEvents) > 0 {
|
||||
returnErr = xerrors.New("some events were rejected by tallyman")
|
||||
}
|
||||
return len(resp.AcceptedEvents), returnErr
|
||||
}
|
||||
|
||||
// getBestLicenseJWT returns the best license JWT to use for the request. The
|
||||
// criteria is as follows:
|
||||
// - The license must be valid and active (after nbf, before exp)
|
||||
// - The license must have usage publishing enabled
|
||||
// The most recently issued (iat) license is chosen.
|
||||
//
|
||||
// If no licenses are found or none have usage publishing enabled,
|
||||
// errUsagePublishingDisabled is returned.
|
||||
func (p *tallymanPublisher) getBestLicenseJWT(ctx context.Context) (string, error) {
|
||||
licenses, err := p.db.GetUnexpiredLicenses(ctx)
|
||||
if err != nil {
|
||||
return "", xerrors.Errorf("get unexpired licenses: %w", err)
|
||||
}
|
||||
if len(licenses) == 0 {
|
||||
return "", errUsagePublishingDisabled
|
||||
}
|
||||
|
||||
type licenseJWTWithClaims struct {
|
||||
Claims *license.Claims
|
||||
Raw string
|
||||
}
|
||||
|
||||
var bestLicense licenseJWTWithClaims
|
||||
for _, dbLicense := range licenses {
|
||||
claims, err := license.ParseClaims(dbLicense.JWT, p.licenseKeys)
|
||||
if err != nil {
|
||||
p.log.Warn(ctx, "failed to parse license claims", slog.F("license_id", dbLicense.ID), slog.Error(err))
|
||||
continue
|
||||
}
|
||||
if claims.AccountType != license.AccountTypeSalesforce {
|
||||
// Non-Salesforce accounts cannot be tracked as they do not have a
|
||||
// trusted Salesforce opportunity ID encoded in the license.
|
||||
continue
|
||||
}
|
||||
if !claims.PublishUsageData {
|
||||
// Publishing is disabled.
|
||||
continue
|
||||
}
|
||||
|
||||
// Otherwise, if it's issued more recently, it's the best license.
|
||||
// IssuedAt is verified to be non-nil in license.ParseClaims.
|
||||
if bestLicense.Claims == nil || claims.IssuedAt.Time.After(bestLicense.Claims.IssuedAt.Time) {
|
||||
bestLicense = licenseJWTWithClaims{
|
||||
Claims: claims,
|
||||
Raw: dbLicense.JWT,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if bestLicense.Raw == "" {
|
||||
return "", errUsagePublishingDisabled
|
||||
}
|
||||
|
||||
return bestLicense.Raw, nil
|
||||
}
|
||||
|
||||
func (p *tallymanPublisher) sendPublishRequest(ctx context.Context, licenseJwt string, req TallymanIngestRequestV1) (TallymanIngestResponseV1, error) {
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return TallymanIngestResponseV1{}, err
|
||||
}
|
||||
|
||||
r, err := http.NewRequestWithContext(ctx, http.MethodPost, p.ingestURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return TallymanIngestResponseV1{}, err
|
||||
}
|
||||
r.Header.Set(CoderLicenseJWTHeader, licenseJwt)
|
||||
|
||||
resp, err := p.httpClient.Do(r)
|
||||
if err != nil {
|
||||
return TallymanIngestResponseV1{}, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
var errBody TallymanErrorV1
|
||||
if err := json.NewDecoder(resp.Body).Decode(&errBody); err != nil {
|
||||
errBody = TallymanErrorV1{
|
||||
Message: fmt.Sprintf("could not decode error response body: %v", err),
|
||||
}
|
||||
}
|
||||
return TallymanIngestResponseV1{}, xerrors.Errorf("unexpected status code %v, error: %s", resp.StatusCode, errBody.Message)
|
||||
}
|
||||
|
||||
var respBody TallymanIngestResponseV1
|
||||
if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
|
||||
return TallymanIngestResponseV1{}, xerrors.Errorf("decode response body: %w", err)
|
||||
}
|
||||
|
||||
return respBody, nil
|
||||
}
|
||||
|
||||
// Close implements Publisher.
|
||||
func (p *tallymanPublisher) Close() error {
|
||||
p.ctxCancel()
|
||||
<-p.done
|
||||
return nil
|
||||
}
|
||||
|
||||
type TallymanErrorV1 struct {
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type TallymanIngestRequestV1 struct {
|
||||
DeploymentID uuid.UUID `json:"deployment_id"`
|
||||
Events []TallymanIngestEventV1 `json:"events"`
|
||||
}
|
||||
|
||||
type TallymanIngestEventV1 struct {
|
||||
ID string `json:"id"`
|
||||
EventType agplusage.EventType `json:"event_type"`
|
||||
EventData json.RawMessage `json:"event_data"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type TallymanIngestResponseV1 struct {
|
||||
AcceptedEvents []TallymanIngestAcceptedEventV1 `json:"accepted_events"`
|
||||
RejectedEvents []TallymanIngestRejectedEventV1 `json:"rejected_events"`
|
||||
}
|
||||
|
||||
type TallymanIngestAcceptedEventV1 struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
|
||||
type TallymanIngestRejectedEventV1 struct {
|
||||
ID string `json:"id"`
|
||||
Message string `json:"message"`
|
||||
Permanent bool `json:"permanent"`
|
||||
}
|
||||
@@ -0,0 +1,729 @@
|
||||
package usage_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"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/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtime"
|
||||
agplusage "github.com/coder/coder/v2/coderd/usage"
|
||||
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
|
||||
"github.com/coder/coder/v2/enterprise/coderd/usage"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
goleak.VerifyTestMain(m, testutil.GoleakOptions...)
|
||||
}
|
||||
|
||||
// TestIntegration tests the inserter and publisher by running them with a real
|
||||
// database.
|
||||
func TestIntegration(t *testing.T) {
|
||||
t.Parallel()
|
||||
const eventCount = 3
|
||||
|
||||
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()
|
||||
|
||||
var (
|
||||
calls int
|
||||
handler func(req usage.TallymanIngestRequestV1) any
|
||||
)
|
||||
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
|
||||
calls++
|
||||
t.Logf("tallyman backend received call %d", calls)
|
||||
assert.Equal(t, deploymentID, req.DeploymentID)
|
||||
|
||||
if handler == nil {
|
||||
t.Errorf("handler is nil")
|
||||
return usage.TallymanIngestResponseV1{}
|
||||
}
|
||||
return handler(req)
|
||||
}))
|
||||
|
||||
inserter := usage.NewInserter(
|
||||
usage.InserterWithClock(clock),
|
||||
)
|
||||
// Insert an old event that should never be published.
|
||||
clock.Set(now.Add(-31 * 24 * time.Hour))
|
||||
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
|
||||
Count: 31,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Insert the events we expect to be published.
|
||||
clock.Set(now.Add(1 * time.Second))
|
||||
for i := 0; i < eventCount; i++ {
|
||||
clock.Advance(time.Second)
|
||||
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
|
||||
Count: uint64(i + 1), // nolint:gosec // these numbers are tiny and will not overflow
|
||||
})
|
||||
require.NoErrorf(t, err, "collecting event %d", i)
|
||||
}
|
||||
|
||||
publisher := usage.NewTallymanPublisher(ctx, log, db,
|
||||
usage.PublisherWithClock(clock),
|
||||
usage.PublisherWithIngestURL(ingestURL),
|
||||
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
|
||||
)
|
||||
defer publisher.Close()
|
||||
|
||||
// Start the publisher with a trap.
|
||||
tickerTrap := clock.Trap().NewTicker()
|
||||
defer tickerTrap.Close()
|
||||
startErr := make(chan error)
|
||||
go func() {
|
||||
err := publisher.Start()
|
||||
testutil.AssertSend(ctx, t, startErr, err)
|
||||
}()
|
||||
tickerCall := tickerTrap.MustWait(ctx)
|
||||
tickerCall.MustRelease(ctx)
|
||||
// The initial duration will always be some time between 5m and 17m.
|
||||
require.GreaterOrEqual(t, tickerCall.Duration, 5*time.Minute)
|
||||
require.LessOrEqual(t, tickerCall.Duration, 17*time.Minute)
|
||||
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
||||
|
||||
// Set up a trap for the ticker.Reset call.
|
||||
tickerResetTrap := clock.Trap().TickerReset()
|
||||
defer tickerResetTrap.Close()
|
||||
|
||||
// Configure the handler for the first publish. This handler will accept the
|
||||
// first event, temporarily reject the second, and permanently reject the
|
||||
// third.
|
||||
var temporarilyRejectedEventID string
|
||||
handler = func(req usage.TallymanIngestRequestV1) any {
|
||||
// On the first call, accept the first event, temporarily reject the
|
||||
// second, and permanently reject the third.
|
||||
acceptedEvents := make([]usage.TallymanIngestAcceptedEventV1, 1)
|
||||
rejectedEvents := make([]usage.TallymanIngestRejectedEventV1, 2)
|
||||
if assert.Len(t, req.Events, eventCount) {
|
||||
assert.JSONEqf(t, jsoninate(t, agplusage.DCManagedAgentsV1{
|
||||
Count: 1,
|
||||
}), string(req.Events[0].EventData), "event data did not match for event %d", 0)
|
||||
acceptedEvents[0].ID = req.Events[0].ID
|
||||
|
||||
temporarilyRejectedEventID = req.Events[1].ID
|
||||
assert.JSONEqf(t, jsoninate(t, agplusage.DCManagedAgentsV1{
|
||||
Count: 2,
|
||||
}), string(req.Events[1].EventData), "event data did not match for event %d", 1)
|
||||
rejectedEvents[0].ID = req.Events[1].ID
|
||||
rejectedEvents[0].Message = "temporarily rejected"
|
||||
rejectedEvents[0].Permanent = false
|
||||
|
||||
assert.JSONEqf(t, jsoninate(t, agplusage.DCManagedAgentsV1{
|
||||
Count: 3,
|
||||
}), string(req.Events[2].EventData), "event data did not match for event %d", 2)
|
||||
rejectedEvents[1].ID = req.Events[2].ID
|
||||
rejectedEvents[1].Message = "permanently rejected"
|
||||
rejectedEvents[1].Permanent = true
|
||||
}
|
||||
return usage.TallymanIngestResponseV1{
|
||||
AcceptedEvents: acceptedEvents,
|
||||
RejectedEvents: rejectedEvents,
|
||||
}
|
||||
}
|
||||
|
||||
// Advance the clock to the initial tick, which should trigger the first
|
||||
// publish, then wait for the reset call. The duration will always be 17m
|
||||
// for resets (only the initial tick is variable).
|
||||
clock.Advance(tickerCall.Duration)
|
||||
tickerResetCall := tickerResetTrap.MustWait(ctx)
|
||||
require.Equal(t, 17*time.Minute, tickerResetCall.Duration)
|
||||
tickerResetCall.MustRelease(ctx)
|
||||
|
||||
// The publisher should have published the events once.
|
||||
require.Equal(t, 1, calls)
|
||||
|
||||
// Set the handler for the next publish call. This call should only include
|
||||
// the temporarily rejected event from earlier. This time we'll accept it.
|
||||
handler = func(req usage.TallymanIngestRequestV1) any {
|
||||
assert.Len(t, req.Events, 1)
|
||||
acceptedEvents := make([]usage.TallymanIngestAcceptedEventV1, len(req.Events))
|
||||
for i, event := range req.Events {
|
||||
assert.Equal(t, temporarilyRejectedEventID, event.ID)
|
||||
acceptedEvents[i].ID = event.ID
|
||||
}
|
||||
return usage.TallymanIngestResponseV1{
|
||||
AcceptedEvents: acceptedEvents,
|
||||
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
|
||||
}
|
||||
}
|
||||
|
||||
// Advance the clock to the next tick and wait for the reset call.
|
||||
clock.Advance(tickerResetCall.Duration)
|
||||
tickerResetCall = tickerResetTrap.MustWait(ctx)
|
||||
tickerResetCall.MustRelease(ctx)
|
||||
|
||||
// The publisher should have published the events again.
|
||||
require.Equal(t, 2, calls)
|
||||
|
||||
// There should be no more publish calls after this, so set the handler to
|
||||
// nil.
|
||||
handler = nil
|
||||
|
||||
// Advance the clock to the next tick.
|
||||
clock.Advance(tickerResetCall.Duration)
|
||||
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
// No publish should have taken place since there are no more events to
|
||||
// publish.
|
||||
require.Equal(t, 2, calls)
|
||||
|
||||
require.NoError(t, publisher.Close())
|
||||
}
|
||||
|
||||
func TestPublisherNoEligibleLicenses(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
log := slogtest.Make(t, nil)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
clock := quartz.NewMock(t)
|
||||
|
||||
// Configure the deployment manually.
|
||||
deploymentID := uuid.New()
|
||||
db.EXPECT().GetDeploymentID(gomock.Any()).Return(deploymentID.String(), nil).Times(1)
|
||||
|
||||
var calls int
|
||||
ingestURL := fakeServer(t, tallymanHandler(t, "", func(req usage.TallymanIngestRequestV1) any {
|
||||
calls++
|
||||
return usage.TallymanIngestResponseV1{
|
||||
AcceptedEvents: []usage.TallymanIngestAcceptedEventV1{},
|
||||
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
|
||||
}
|
||||
}))
|
||||
|
||||
publisher := usage.NewTallymanPublisher(ctx, log, db,
|
||||
usage.PublisherWithClock(clock),
|
||||
usage.PublisherWithIngestURL(ingestURL),
|
||||
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
|
||||
)
|
||||
defer publisher.Close()
|
||||
|
||||
// Start the publisher with a trap.
|
||||
tickerTrap := clock.Trap().NewTicker()
|
||||
defer tickerTrap.Close()
|
||||
startErr := make(chan error)
|
||||
go func() {
|
||||
err := publisher.Start()
|
||||
testutil.RequireSend(ctx, t, startErr, err)
|
||||
}()
|
||||
tickerCall := tickerTrap.MustWait(ctx)
|
||||
tickerCall.MustRelease(ctx)
|
||||
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
||||
|
||||
// Mock zero licenses.
|
||||
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{}, nil).Times(1)
|
||||
|
||||
// Tick and wait for the reset call.
|
||||
tickerResetTrap := clock.Trap().TickerReset()
|
||||
defer tickerResetTrap.Close()
|
||||
clock.Advance(tickerCall.Duration)
|
||||
tickerResetCall := tickerResetTrap.MustWait(ctx)
|
||||
tickerResetCall.MustRelease(ctx)
|
||||
|
||||
// The publisher should not have published the events.
|
||||
require.Equal(t, 0, calls)
|
||||
|
||||
// Mock a single license with usage publishing disabled.
|
||||
licenseJWT := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: false,
|
||||
})
|
||||
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{
|
||||
{
|
||||
ID: 1,
|
||||
JWT: licenseJWT,
|
||||
UploadedAt: dbtime.Now(),
|
||||
Exp: dbtime.Now().Add(48 * time.Hour), // fake
|
||||
UUID: uuid.New(),
|
||||
},
|
||||
}, nil).Times(1)
|
||||
|
||||
// Tick and wait for the reset call.
|
||||
clock.Advance(tickerResetCall.Duration)
|
||||
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
// The publisher should still not have published the events.
|
||||
require.Equal(t, 0, calls)
|
||||
}
|
||||
|
||||
// TestPublisherClaimExpiry tests the claim query to ensure that events are not
|
||||
// claimed if they've recently been claimed by another publisher.
|
||||
func TestPublisherClaimExpiry(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
log := slogtest.Make(t, nil)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
clock := quartz.NewMock(t)
|
||||
_, licenseJWT := configureDeployment(ctx, t, db)
|
||||
now := time.Now()
|
||||
|
||||
var calls int
|
||||
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
|
||||
calls++
|
||||
return tallymanAcceptAllHandler(req)
|
||||
}))
|
||||
|
||||
inserter := usage.NewInserter(
|
||||
usage.InserterWithClock(clock),
|
||||
)
|
||||
|
||||
publisher := usage.NewTallymanPublisher(ctx, log, db,
|
||||
usage.PublisherWithClock(clock),
|
||||
usage.PublisherWithIngestURL(ingestURL),
|
||||
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
|
||||
usage.PublisherWithInitialDelay(17*time.Minute),
|
||||
)
|
||||
defer publisher.Close()
|
||||
|
||||
// Create an event that was claimed 1h-18m ago. The ticker has a forced
|
||||
// delay of 17m in this test.
|
||||
clock.Set(now)
|
||||
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
|
||||
Count: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
// Claim the event in the past. Claiming it this way via the database
|
||||
// directly means it won't be marked as published or unclaimed.
|
||||
events, err := db.SelectUsageEventsForPublishing(ctx, now.Add(-42*time.Minute))
|
||||
require.NoError(t, err)
|
||||
require.Len(t, events, 1)
|
||||
|
||||
// Start the publisher with a trap.
|
||||
tickerTrap := clock.Trap().NewTicker()
|
||||
defer tickerTrap.Close()
|
||||
startErr := make(chan error)
|
||||
go func() {
|
||||
err := publisher.Start()
|
||||
testutil.RequireSend(ctx, t, startErr, err)
|
||||
}()
|
||||
tickerCall := tickerTrap.MustWait(ctx)
|
||||
require.Equal(t, 17*time.Minute, tickerCall.Duration)
|
||||
tickerCall.MustRelease(ctx)
|
||||
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
||||
|
||||
// Set up a trap for the ticker.Reset call.
|
||||
tickerResetTrap := clock.Trap().TickerReset()
|
||||
defer tickerResetTrap.Close()
|
||||
|
||||
// Advance the clock to the initial tick, which should trigger the first
|
||||
// publish, then wait for the reset call. The duration will always be 17m
|
||||
// for resets (only the initial tick is variable).
|
||||
clock.Advance(tickerCall.Duration)
|
||||
tickerResetCall := tickerResetTrap.MustWait(ctx)
|
||||
require.Equal(t, 17*time.Minute, tickerResetCall.Duration)
|
||||
tickerResetCall.MustRelease(ctx)
|
||||
|
||||
// No events should have been published since none are eligible.
|
||||
require.Equal(t, 0, calls)
|
||||
|
||||
// Advance the clock to the next tick and wait for the reset call.
|
||||
clock.Advance(tickerResetCall.Duration)
|
||||
tickerResetCall = tickerResetTrap.MustWait(ctx)
|
||||
tickerResetCall.MustRelease(ctx)
|
||||
|
||||
// The publisher should have published the event, as it's now eligible.
|
||||
require.Equal(t, 1, calls)
|
||||
}
|
||||
|
||||
// TestPublisherMissingEvents tests that the publisher notices events that are
|
||||
// not returned by the Tallyman server and marks them as temporarily rejected.
|
||||
func TestPublisherMissingEvents(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
log := slogtest.Make(t, nil)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
_, licenseJWT := configureMockDeployment(t, db)
|
||||
clock := quartz.NewMock(t)
|
||||
now := time.Now()
|
||||
clock.Set(now)
|
||||
|
||||
var calls int
|
||||
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
|
||||
calls++
|
||||
return usage.TallymanIngestResponseV1{
|
||||
AcceptedEvents: []usage.TallymanIngestAcceptedEventV1{},
|
||||
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
|
||||
}
|
||||
}))
|
||||
|
||||
publisher := usage.NewTallymanPublisher(ctx, log, db,
|
||||
usage.PublisherWithClock(clock),
|
||||
usage.PublisherWithIngestURL(ingestURL),
|
||||
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
|
||||
)
|
||||
|
||||
// Expect the publisher to call SelectUsageEventsForPublishing, followed by
|
||||
// UpdateUsageEventsPostPublish.
|
||||
events := []database.UsageEvent{
|
||||
{
|
||||
ID: uuid.New().String(),
|
||||
EventType: string(agplusage.UsageEventTypeDCManagedAgentsV1),
|
||||
EventData: []byte(jsoninate(t, agplusage.DCManagedAgentsV1{
|
||||
Count: 1,
|
||||
})),
|
||||
CreatedAt: now,
|
||||
PublishedAt: sql.NullTime{},
|
||||
PublishStartedAt: sql.NullTime{},
|
||||
FailureMessage: sql.NullString{},
|
||||
},
|
||||
}
|
||||
db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), gomock.Any()).Return(events, nil).Times(1)
|
||||
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(ctx context.Context, params database.UpdateUsageEventsPostPublishParams) error {
|
||||
assert.Equal(t, []string{events[0].ID}, params.IDs)
|
||||
assert.Equal(t, []string{"tallyman did not include the event in the response"}, params.FailureMessages)
|
||||
assert.Equal(t, []bool{false}, params.SetPublishedAts)
|
||||
return nil
|
||||
},
|
||||
).Times(1)
|
||||
|
||||
// Start the publisher with a trap.
|
||||
tickerTrap := clock.Trap().NewTicker()
|
||||
defer tickerTrap.Close()
|
||||
startErr := make(chan error)
|
||||
go func() {
|
||||
err := publisher.Start()
|
||||
testutil.RequireSend(ctx, t, startErr, err)
|
||||
}()
|
||||
tickerCall := tickerTrap.MustWait(ctx)
|
||||
tickerCall.MustRelease(ctx)
|
||||
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
||||
|
||||
// Tick and wait for the reset call.
|
||||
tickerResetTrap := clock.Trap().TickerReset()
|
||||
defer tickerResetTrap.Close()
|
||||
clock.Advance(tickerCall.Duration)
|
||||
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
// The publisher should have published the events once.
|
||||
require.Equal(t, 1, calls)
|
||||
|
||||
require.NoError(t, publisher.Close())
|
||||
}
|
||||
|
||||
func TestPublisherLicenseSelection(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
log := slogtest.Make(t, nil)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
clock := quartz.NewMock(t)
|
||||
now := time.Now()
|
||||
|
||||
// Configure the deployment manually.
|
||||
deploymentID := uuid.New()
|
||||
db.EXPECT().GetDeploymentID(gomock.Any()).Return(deploymentID.String(), nil).Times(1)
|
||||
|
||||
// Insert multiple licenses:
|
||||
// 1. PublishUsageData false, type=salesforce, iat 30m ago (ineligible, publish not enabled)
|
||||
// 2. PublishUsageData true, type=trial, iat 1h ago (ineligible, not salesforce)
|
||||
// 3. PublishUsageData true, type=salesforce, iat 30m ago, exp 10m ago (ineligible, expired)
|
||||
// 4. PublishUsageData true, type=salesforce, iat 1h ago (eligible)
|
||||
// 5. PublishUsageData true, type=salesforce, iat 30m ago (eligible, and newer!)
|
||||
badLicense1 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: false,
|
||||
IssuedAt: now.Add(-30 * time.Minute),
|
||||
})
|
||||
badLicense2 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: true,
|
||||
IssuedAt: now.Add(-1 * time.Hour),
|
||||
AccountType: "trial",
|
||||
})
|
||||
badLicense3 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: true,
|
||||
IssuedAt: now.Add(-30 * time.Minute),
|
||||
ExpiresAt: now.Add(-10 * time.Minute),
|
||||
})
|
||||
badLicense4 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: true,
|
||||
IssuedAt: now.Add(-1 * time.Hour),
|
||||
})
|
||||
expectedLicense := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: true,
|
||||
IssuedAt: now.Add(-30 * time.Minute),
|
||||
})
|
||||
// GetUnexpiredLicenses is not supposed to return expired licenses, but for
|
||||
// the purposes of this test we're going to do it anyway.
|
||||
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{
|
||||
{
|
||||
ID: 1,
|
||||
JWT: badLicense1,
|
||||
Exp: now.Add(48 * time.Hour), // fake times, the JWT should be checked
|
||||
UUID: uuid.New(),
|
||||
UploadedAt: now,
|
||||
},
|
||||
{
|
||||
ID: 2,
|
||||
JWT: badLicense2,
|
||||
Exp: now.Add(48 * time.Hour),
|
||||
UUID: uuid.New(),
|
||||
UploadedAt: now,
|
||||
},
|
||||
{
|
||||
ID: 3,
|
||||
JWT: badLicense3,
|
||||
Exp: now.Add(48 * time.Hour),
|
||||
UUID: uuid.New(),
|
||||
UploadedAt: now,
|
||||
},
|
||||
{
|
||||
ID: 4,
|
||||
JWT: badLicense4,
|
||||
Exp: now.Add(48 * time.Hour),
|
||||
UUID: uuid.New(),
|
||||
UploadedAt: now,
|
||||
},
|
||||
{
|
||||
ID: 5,
|
||||
JWT: expectedLicense,
|
||||
Exp: now.Add(48 * time.Hour),
|
||||
UUID: uuid.New(),
|
||||
UploadedAt: now,
|
||||
},
|
||||
}, nil)
|
||||
|
||||
called := false
|
||||
ingestURL := fakeServer(t, tallymanHandler(t, expectedLicense, func(req usage.TallymanIngestRequestV1) any {
|
||||
called = true
|
||||
assert.Equal(t, deploymentID, req.DeploymentID)
|
||||
return tallymanAcceptAllHandler(req)
|
||||
}))
|
||||
|
||||
publisher := usage.NewTallymanPublisher(ctx, log, db,
|
||||
usage.PublisherWithClock(clock),
|
||||
usage.PublisherWithIngestURL(ingestURL),
|
||||
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
|
||||
)
|
||||
defer publisher.Close()
|
||||
|
||||
// Start the publisher with a trap.
|
||||
tickerTrap := clock.Trap().NewTicker()
|
||||
defer tickerTrap.Close()
|
||||
startErr := make(chan error)
|
||||
go func() {
|
||||
err := publisher.Start()
|
||||
testutil.RequireSend(ctx, t, startErr, err)
|
||||
}()
|
||||
tickerCall := tickerTrap.MustWait(ctx)
|
||||
tickerCall.MustRelease(ctx)
|
||||
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
||||
|
||||
// Mock events to be published.
|
||||
events := []database.UsageEvent{
|
||||
{
|
||||
ID: uuid.New().String(),
|
||||
EventType: string(agplusage.UsageEventTypeDCManagedAgentsV1),
|
||||
EventData: []byte(jsoninate(t, agplusage.DCManagedAgentsV1{
|
||||
Count: 1,
|
||||
})),
|
||||
},
|
||||
}
|
||||
db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), gomock.Any()).Return(events, nil).Times(1)
|
||||
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(ctx context.Context, params database.UpdateUsageEventsPostPublishParams) error {
|
||||
assert.Equal(t, []string{events[0].ID}, params.IDs)
|
||||
assert.Equal(t, []string{""}, params.FailureMessages)
|
||||
assert.Equal(t, []bool{true}, params.SetPublishedAts)
|
||||
return nil
|
||||
},
|
||||
).Times(1)
|
||||
|
||||
// Tick and wait for the reset call.
|
||||
tickerResetTrap := clock.Trap().TickerReset()
|
||||
defer tickerResetTrap.Close()
|
||||
clock.Advance(tickerCall.Duration)
|
||||
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
// The publisher should have published the events once.
|
||||
require.True(t, called)
|
||||
}
|
||||
|
||||
func TestPublisherTallymanError(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
log := slogtest.Make(t, nil)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
clock := quartz.NewMock(t)
|
||||
now := time.Now()
|
||||
clock.Set(now)
|
||||
|
||||
_, licenseJWT := configureMockDeployment(t, db)
|
||||
const errorMessage = "tallyman error"
|
||||
var calls int
|
||||
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
|
||||
calls++
|
||||
return usage.TallymanErrorV1{
|
||||
Message: errorMessage,
|
||||
}
|
||||
}))
|
||||
|
||||
publisher := usage.NewTallymanPublisher(ctx, log, db,
|
||||
usage.PublisherWithClock(clock),
|
||||
usage.PublisherWithIngestURL(ingestURL),
|
||||
usage.PublisherWithLicenseKeys(coderdenttest.Keys),
|
||||
)
|
||||
defer publisher.Close()
|
||||
|
||||
// Start the publisher with a trap.
|
||||
tickerTrap := clock.Trap().NewTicker()
|
||||
defer tickerTrap.Close()
|
||||
startErr := make(chan error)
|
||||
go func() {
|
||||
err := publisher.Start()
|
||||
testutil.RequireSend(ctx, t, startErr, err)
|
||||
}()
|
||||
tickerCall := tickerTrap.MustWait(ctx)
|
||||
tickerCall.MustRelease(ctx)
|
||||
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
||||
|
||||
// Mock events to be published.
|
||||
events := []database.UsageEvent{
|
||||
{
|
||||
ID: uuid.New().String(),
|
||||
EventType: string(agplusage.UsageEventTypeDCManagedAgentsV1),
|
||||
EventData: []byte(jsoninate(t, agplusage.DCManagedAgentsV1{
|
||||
Count: 1,
|
||||
})),
|
||||
},
|
||||
}
|
||||
db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), gomock.Any()).Return(events, nil).Times(1)
|
||||
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(ctx context.Context, params database.UpdateUsageEventsPostPublishParams) error {
|
||||
assert.Equal(t, []string{events[0].ID}, params.IDs)
|
||||
assert.Contains(t, params.FailureMessages[0], errorMessage)
|
||||
assert.Equal(t, []bool{false}, params.SetPublishedAts)
|
||||
return nil
|
||||
},
|
||||
).Times(1)
|
||||
|
||||
// Tick and wait for the reset call.
|
||||
tickerResetTrap := clock.Trap().TickerReset()
|
||||
defer tickerResetTrap.Close()
|
||||
clock.Advance(tickerCall.Duration)
|
||||
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
// The publisher should have published the events once.
|
||||
require.Equal(t, 1, calls)
|
||||
}
|
||||
|
||||
func jsoninate(t *testing.T, v any) string {
|
||||
t.Helper()
|
||||
if e, ok := v.(agplusage.Event); ok {
|
||||
v = e.Fields()
|
||||
}
|
||||
buf, err := json.Marshal(v)
|
||||
require.NoError(t, err)
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
func configureDeployment(ctx context.Context, t *testing.T, db database.Store) (uuid.UUID, string) {
|
||||
t.Helper()
|
||||
deploymentID := uuid.New()
|
||||
err := db.InsertDeploymentID(ctx, deploymentID.String())
|
||||
require.NoError(t, err)
|
||||
|
||||
licenseRaw := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: true,
|
||||
})
|
||||
_, err = db.InsertLicense(ctx, database.InsertLicenseParams{
|
||||
UploadedAt: dbtime.Now(),
|
||||
JWT: licenseRaw,
|
||||
Exp: dbtime.Now().Add(48 * time.Hour),
|
||||
UUID: uuid.New(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
return deploymentID, licenseRaw
|
||||
}
|
||||
|
||||
func configureMockDeployment(t *testing.T, db *dbmock.MockStore) (uuid.UUID, string) {
|
||||
t.Helper()
|
||||
deploymentID := uuid.New()
|
||||
db.EXPECT().GetDeploymentID(gomock.Any()).Return(deploymentID.String(), nil).Times(1)
|
||||
|
||||
licenseRaw := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
||||
PublishUsageData: true,
|
||||
})
|
||||
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{
|
||||
{
|
||||
ID: 1,
|
||||
UploadedAt: dbtime.Now(),
|
||||
JWT: licenseRaw,
|
||||
Exp: dbtime.Now().Add(48 * time.Hour),
|
||||
UUID: uuid.New(),
|
||||
},
|
||||
}, nil)
|
||||
|
||||
return deploymentID, licenseRaw
|
||||
}
|
||||
|
||||
func fakeServer(t *testing.T, handler http.Handler) string {
|
||||
t.Helper()
|
||||
server := httptest.NewServer(handler)
|
||||
t.Cleanup(server.Close)
|
||||
return server.URL
|
||||
}
|
||||
|
||||
func tallymanHandler(t *testing.T, expectLicenseJWT string, handler func(req usage.TallymanIngestRequestV1) any) http.Handler {
|
||||
t.Helper()
|
||||
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
|
||||
t.Helper()
|
||||
licenseJWT := r.Header.Get(usage.CoderLicenseJWTHeader)
|
||||
if expectLicenseJWT != "" && !assert.Equal(t, expectLicenseJWT, licenseJWT, "license JWT in request did not match") {
|
||||
rw.WriteHeader(http.StatusUnauthorized)
|
||||
err := json.NewEncoder(rw).Encode(usage.TallymanErrorV1{
|
||||
Message: "license JWT in request did not match",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return
|
||||
}
|
||||
|
||||
var req usage.TallymanIngestRequestV1
|
||||
err := json.NewDecoder(r.Body).Decode(&req)
|
||||
require.NoError(t, err)
|
||||
|
||||
resp := handler(req)
|
||||
switch resp.(type) {
|
||||
case usage.TallymanErrorV1:
|
||||
rw.WriteHeader(http.StatusInternalServerError)
|
||||
default:
|
||||
rw.WriteHeader(http.StatusOK)
|
||||
}
|
||||
err = json.NewEncoder(rw).Encode(resp)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func tallymanAcceptAllHandler(req usage.TallymanIngestRequestV1) usage.TallymanIngestResponseV1 {
|
||||
acceptedEvents := make([]usage.TallymanIngestAcceptedEventV1, len(req.Events))
|
||||
for i, event := range req.Events {
|
||||
acceptedEvents[i].ID = event.ID
|
||||
}
|
||||
|
||||
return usage.TallymanIngestResponseV1{
|
||||
AcceptedEvents: acceptedEvents,
|
||||
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user