mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
chore: add tallyman events for ai seat tracking (#22689)
AI seat tracking inserted as heartbeat into usage table.
This commit is contained in:
Steven Masley
@@ -0,0 +1,215 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math/rand"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/pproflabel"
|
||||
agplusage "github.com/coder/coder/v2/coderd/usage"
|
||||
"github.com/coder/coder/v2/coderd/usage/usagetypes"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// epoch is a fixed reference point for aligning interval boundaries.
|
||||
// All replicas use this same epoch so their buckets are identical.
|
||||
var epoch = time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
const (
|
||||
cronDateFormat = "2006-01-02_15:04:05"
|
||||
)
|
||||
|
||||
// HeartbeatFunc generates a heartbeat event and its stable ID.
|
||||
// It is called periodically by the cron. Returning an error skips
|
||||
// the insert for that tick and logs a warning.
|
||||
type HeartbeatFunc func(ctx context.Context) (event usagetypes.HeartbeatEvent, err error)
|
||||
|
||||
// CronJob defines a periodic heartbeat job.
|
||||
type CronJob struct {
|
||||
// Name is a human-readable label used in logs.
|
||||
Name string
|
||||
// Interval is the base duration between ticks.
|
||||
Interval time.Duration
|
||||
// EventType must match the events generated by the Fn.
|
||||
EventType usagetypes.UsageEventType
|
||||
// Jitter is the maximum random delay added after the boundary.
|
||||
// The actual offset is uniformly distributed in [0, Jitter).
|
||||
// This staggers replicas so one is likely to complete the work
|
||||
// before others attempt it, allowing them to skip via the
|
||||
// existence check (heartbeat inserts are idempotent).
|
||||
Jitter time.Duration
|
||||
// Fn produces the heartbeat event.
|
||||
Fn HeartbeatFunc
|
||||
}
|
||||
|
||||
// Cron runs registered CronJobs on the dbInserter's clock. Stopping
|
||||
// the context passed to Start cancels all jobs. Daemon restarts
|
||||
// naturally restart the timers since Start() creates them fresh —
|
||||
// there is no state to persist or recover.
|
||||
type Cron struct {
|
||||
clock quartz.Clock
|
||||
log slog.Logger
|
||||
db database.Store
|
||||
ins agplusage.Inserter
|
||||
jobs []CronJob
|
||||
|
||||
// cancel cancels the context on all running jobs. If the ctx passed into `Start`
|
||||
// is canceled, the jobs will also stop.
|
||||
cancel context.CancelFunc
|
||||
|
||||
// wg ensures all job goroutines have exited before Close returns.
|
||||
wg sync.WaitGroup
|
||||
|
||||
// startOnce ensures Start is idempotent.
|
||||
startOnce sync.Once
|
||||
started atomic.Bool
|
||||
}
|
||||
|
||||
// NewCron creates a Cron that periodically generates and inserts
|
||||
// heartbeat events. The clock controls all timers so that tests can
|
||||
// advance time deterministically via quartz.Mock.
|
||||
func NewCron(clock quartz.Clock, log slog.Logger, db database.Store, ins agplusage.Inserter) *Cron {
|
||||
return &Cron{
|
||||
clock: clock,
|
||||
log: log,
|
||||
db: db,
|
||||
ins: ins,
|
||||
}
|
||||
}
|
||||
|
||||
// Register adds a job. It must be called before Start; calling it
|
||||
// after Start returns an error.
|
||||
func (c *Cron) Register(job CronJob) error {
|
||||
if !job.EventType.IsHeartbeat() {
|
||||
return xerrors.New("event type must be a heartbeat type")
|
||||
}
|
||||
if c.started.Load() {
|
||||
return xerrors.New("cannot register a job after Start has been called")
|
||||
}
|
||||
c.jobs = append(c.jobs, job)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Start launches a goroutine per job. Subsequent calls are no-ops.
|
||||
// On daemon restart a new Cron should be created.
|
||||
func (c *Cron) Start(ctx context.Context) {
|
||||
c.startOnce.Do(func() {
|
||||
c.started.Store(true)
|
||||
ctx, c.cancel = context.WithCancel(ctx)
|
||||
for _, job := range c.jobs {
|
||||
c.wg.Add(1)
|
||||
pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceUsageEventCron, "job", job.Name), func(ctx context.Context) {
|
||||
c.run(ctx, job)
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Close cancels all jobs and waits for goroutines to exit.
|
||||
func (c *Cron) Close() error {
|
||||
if c.cancel != nil {
|
||||
c.cancel()
|
||||
}
|
||||
c.wg.Wait()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Cron) run(ctx context.Context, job CronJob) {
|
||||
//nolint:gocritic // We are a publisher in this function
|
||||
ctx = dbauthz.AsUsagePublisher(ctx)
|
||||
defer c.wg.Done()
|
||||
for {
|
||||
boundary, delay := nextTick(c.clock.Now(), job.Interval, job.Jitter)
|
||||
|
||||
// Use a quartz timer so the wait honors ctx cancellation and
|
||||
// tests can advance time deterministically.
|
||||
timer := c.clock.NewTimer(delay, job.Name)
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
if !timer.Stop() {
|
||||
// Drain the channel if the timer already fired.
|
||||
<-timer.C
|
||||
}
|
||||
return
|
||||
case <-timer.C:
|
||||
}
|
||||
|
||||
// Use the boundary (not wall-clock "now") for the stable ID
|
||||
// so all replicas targeting the same boundary produce the
|
||||
// same key.
|
||||
stableID := string(job.EventType) + ":" + boundary.UTC().Format(cronDateFormat)
|
||||
|
||||
// Skip if this bucket was already recorded — avoids running
|
||||
// the potentially expensive heartbeat function for a
|
||||
// duplicate.
|
||||
exists, err := c.db.UsageEventExistsByID(ctx, stableID)
|
||||
if err != nil {
|
||||
c.log.Warn(ctx, "cron heartbeat existence check failed",
|
||||
slog.F("job", job.Name),
|
||||
slog.Error(err),
|
||||
)
|
||||
continue
|
||||
}
|
||||
if exists {
|
||||
c.log.Debug(ctx, "cron heartbeat already recorded, skipping",
|
||||
slog.F("job", job.Name),
|
||||
slog.F("id", stableID),
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
event, err := job.Fn(ctx)
|
||||
if err != nil {
|
||||
c.log.Error(ctx, "cron heartbeat func failed",
|
||||
slog.F("job", job.Name),
|
||||
slog.Error(err),
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
if event.EventType() != job.EventType {
|
||||
c.log.Error(ctx, "cron heartbeat func returned wrong event type",
|
||||
slog.F("job", job.Name),
|
||||
slog.F("expected", job.EventType),
|
||||
slog.F("actual", event.EventType()),
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, event); err != nil {
|
||||
c.log.Warn(ctx, "cron heartbeat insert failed",
|
||||
slog.F("job", job.Name),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// nextTick computes the delay until the next epoch-aligned boundary
|
||||
// for the given interval, plus a random jitter in [0, jitter). It
|
||||
// returns the target boundary and the total delay from now.
|
||||
func nextTick(now time.Time, interval, jitter time.Duration) (boundary time.Time, delay time.Duration) {
|
||||
boundary = nextBoundary(now, interval)
|
||||
delay = boundary.Sub(now)
|
||||
if jitter > 0 {
|
||||
//nolint:gosec // Jitter does not need cryptographic randomness.
|
||||
delay += time.Duration(rand.Int63n(int64(jitter)))
|
||||
}
|
||||
return boundary, delay
|
||||
}
|
||||
|
||||
// nextBoundary returns the first multiple of interval (relative to
|
||||
// epoch) that is strictly after t.
|
||||
func nextBoundary(t time.Time, interval time.Duration) time.Time {
|
||||
since := t.Sub(epoch)
|
||||
n := since / interval
|
||||
return epoch.Add((n + 1) * interval)
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNextBoundary(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tcs := []struct {
|
||||
name string
|
||||
T time.Time
|
||||
interval time.Duration
|
||||
expected time.Time
|
||||
}{
|
||||
{
|
||||
name: "exactly_on_boundary",
|
||||
T: time.Date(2023, 1, 1, 8, 0, 0, 0, time.UTC),
|
||||
interval: 4 * time.Hour,
|
||||
// On a boundary → returns the next one.
|
||||
expected: time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "1ns_after_boundary",
|
||||
T: time.Date(2023, 1, 1, 8, 0, 0, 1, time.UTC),
|
||||
interval: 4 * time.Hour,
|
||||
expected: time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "1ns_before_boundary",
|
||||
T: time.Date(2023, 1, 1, 7, 59, 59, 999999999, time.UTC),
|
||||
interval: 4 * time.Hour,
|
||||
expected: time.Date(2023, 1, 1, 8, 0, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "mid_interval",
|
||||
T: time.Date(2023, 1, 1, 10, 0, 0, 0, time.UTC),
|
||||
interval: 4 * time.Hour,
|
||||
expected: time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "5min_interval",
|
||||
T: time.Date(2026, 3, 13, 14, 2, 30, 0, time.UTC),
|
||||
interval: 5 * time.Minute,
|
||||
expected: time.Date(2026, 3, 13, 14, 5, 0, 0, time.UTC),
|
||||
},
|
||||
{
|
||||
name: "1hr_interval",
|
||||
T: time.Date(2026, 6, 15, 9, 45, 0, 0, time.UTC),
|
||||
interval: 1 * time.Hour,
|
||||
expected: time.Date(2026, 6, 15, 10, 0, 0, 0, time.UTC),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := nextBoundary(tc.T, tc.interval)
|
||||
require.Equal(t, tc.expected, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNextTick(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("NoJitter", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 3, 13, 14, 2, 30, 0, time.UTC)
|
||||
interval := 4 * time.Hour
|
||||
|
||||
boundary, delay := nextTick(now, interval, 0)
|
||||
|
||||
expectedBoundary := time.Date(2026, 3, 13, 16, 0, 0, 0, time.UTC)
|
||||
require.Equal(t, expectedBoundary, boundary)
|
||||
require.Equal(t, boundary.Sub(now), delay)
|
||||
})
|
||||
|
||||
t.Run("WithJitter", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Date(2026, 3, 13, 14, 2, 30, 0, time.UTC)
|
||||
interval := 4 * time.Hour
|
||||
jitter := 10 * time.Minute
|
||||
|
||||
boundary, delay := nextTick(now, interval, jitter)
|
||||
|
||||
expectedBoundary := time.Date(2026, 3, 13, 16, 0, 0, 0, time.UTC)
|
||||
require.Equal(t, expectedBoundary, boundary)
|
||||
|
||||
base := boundary.Sub(now)
|
||||
require.GreaterOrEqual(t, delay, base,
|
||||
"delay must be at least the base distance to boundary")
|
||||
require.Less(t, delay, base+jitter,
|
||||
"delay must be less than base + jitter")
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
package usage_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/usage/usagetypes"
|
||||
"github.com/coder/coder/v2/enterprise/coderd/usage"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
func TestCron(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("BasicTick", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
clock := quartz.NewMock(t)
|
||||
|
||||
// The existence check should return false so the event gets
|
||||
// inserted.
|
||||
db.EXPECT().UsageEventExistsByID(gomock.Any(), gomock.Any()).
|
||||
Return(false, nil).AnyTimes()
|
||||
|
||||
inserted := make(chan database.InsertUsageEventParams, 1)
|
||||
db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, params database.InsertUsageEventParams) error {
|
||||
inserted <- params
|
||||
return nil
|
||||
}).AnyTimes()
|
||||
|
||||
inserter := usage.NewDBInserter(usage.InserterWithClock(clock))
|
||||
cron := usage.NewCron(clock, slogtest.Make(t, nil), db, inserter)
|
||||
require.NoError(t, cron.Register(usage.CronJob{
|
||||
Name: "test-job",
|
||||
Interval: 5 * time.Minute,
|
||||
EventType: usagetypes.UsageEventTypeHBAISeatsV1,
|
||||
Fn: func(_ context.Context) (usagetypes.HeartbeatEvent, error) {
|
||||
return usagetypes.HBAISeats{Count: 42}, nil
|
||||
},
|
||||
}))
|
||||
|
||||
timerTrap := clock.Trap().NewTimer("test-job")
|
||||
|
||||
cron.Start(ctx)
|
||||
defer cron.Close()
|
||||
defer timerTrap.Close()
|
||||
|
||||
// Wait for timer creation, then fire it. The delay is the
|
||||
// time until the next epoch-aligned boundary for the 5-minute
|
||||
// interval — we don't assert the exact value since it depends
|
||||
// on the mock clock's current time.
|
||||
timerCall := timerTrap.MustWait(ctx)
|
||||
timerCall.MustRelease(ctx)
|
||||
clock.Advance(timerCall.Duration)
|
||||
|
||||
// Verify the event was inserted with an epoch-aligned ID.
|
||||
select {
|
||||
case params := <-inserted:
|
||||
assert.Contains(t, params.ID, "hb_ai_seats_v1:")
|
||||
case <-ctx.Done():
|
||||
t.Fatal("timed out waiting for insert")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestAISeatsHeartbeat checks that AISeatsHeartbeat returns the
|
||||
// correct event type and count.
|
||||
func TestAISeatsHeartbeat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
|
||||
db.EXPECT().GetActiveAISeatCount(gomock.Any()).Return(int64(42), nil)
|
||||
|
||||
fn := usage.AISeatsHeartbeat(db)
|
||||
event, err := fn(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify the event type and count.
|
||||
hb, ok := event.(usagetypes.HBAISeats)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, int64(42), hb.Count)
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/usage/usagetypes"
|
||||
)
|
||||
|
||||
const (
|
||||
AISeatsInterval = 4 * time.Hour
|
||||
)
|
||||
|
||||
// AISeatsHeartbeat returns a HeartbeatFunc that queries the active
|
||||
// AI seat count and emits it as an HBAISeats heartbeat event.
|
||||
func AISeatsHeartbeat(db database.Store) HeartbeatFunc {
|
||||
return func(ctx context.Context) (usagetypes.HeartbeatEvent, error) {
|
||||
//nolint:gocritic // We are a publisher in this function
|
||||
ctx = dbauthz.AsUsagePublisher(ctx)
|
||||
count, err := db.GetActiveAISeatCount(ctx)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("get active AI seat count: %w", err)
|
||||
}
|
||||
|
||||
return usagetypes.HBAISeats{Count: count}, nil
|
||||
}
|
||||
}
|
||||
@@ -66,3 +66,27 @@ func (i *dbInserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.S
|
||||
CreatedAt: dbtime.Time(i.clock.Now()),
|
||||
})
|
||||
}
|
||||
|
||||
// InsertHeartbeatUsageEvent implements agplusage.Inserter.
|
||||
func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error {
|
||||
if !event.EventType().IsHeartbeat() {
|
||||
return xerrors.Errorf("event type %q is not a heartbeat 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{
|
||||
ID: id,
|
||||
EventType: string(event.EventType()),
|
||||
EventData: jsonData,
|
||||
CreatedAt: dbtime.Time(i.clock.Now()),
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user