feat: add hourly hb_agent_runtime_v1 usage events for Coder Agent runtime (#27312)

closes CODAGT-839
closes CODAGT-843
closes CODAGT-773

## Summary

Adds a new heartbeat usage event type, `hb_agent_runtime_v1`, measuring
the total agent-loop runtime of Coder Agents (chats) per UTC hour, plus
a reconciler that generates one event per hour with self-healing
backfill over a trailing 7-day window. Events flow to Tallyman through
the existing publisher unchanged. This measures the new Coder Agents
(the `chats` tables), not the deprecated Tasks counted by
`dc_managed_agents_v1`.

Independent of #27508, which fixes the dead ai-seats cron registration.
Both PRs carry the identical `usage_event` create permission hunk for
the usage-publisher subject (this feature's generator and the ai-seats
cron each need it for heartbeat inserts), so they can land in either
order and the overlap merges cleanly.

> [!WARNING]
> **Do not include this in a release until Tallyman accepts
`hb_agent_runtime_v1`.** The publisher marks permanently rejected events
as done-forever, and the generator then sees those buckets as complete
locally, so their usage would be silently and permanently lost.

## Details

Each event's payload is `{"runtime_ms": N}`: the sum of
`chat_messages.runtime_ms` for messages created in the hour bucket `[H,
H+1)`, across all chats (sub-agents, API-created, archived, and
soft-deleted messages included). Events use deterministic IDs
(`hb_agent_runtime_v1:<bucket start>`) with `created_at` set to the
bucket start, so concurrent replicas race safely via `ON CONFLICT (id)
DO NOTHING` without locking, and daily rollups attribute backfilled
hours to the correct day. Idle hours produce zero-valued events. A
bucket becomes eligible 5 minutes after it closes; hours missing for
longer than the 7-day window are forfeited, which can only undercount.

Note that this makes `usage_events.created_at` explicitly the *event
occurrence time* rather than the row insertion time; the two only
diverge for backfilled events. It already behaved as the occurrence
timestamp (it drives the daily rollup day and is shipped to
Tallyman/Metronome as the event timestamp), and the migration now
documents this with a `COMMENT ON COLUMN`, which also surfaces as a Go
doc comment on `UsageEvent.CreatedAt`.

The new `usage.Generator` runs unconditionally in enterprise builds; the
`publish_usage_data` license flag continues to gate egress only, so
air-gapped deployments still fill their local ledger. The
`aggregate_usage_event()` trigger sums `runtime_ms` per day into
`usage_events_daily` (unlike `hb_ai_seats_v1`, which takes the daily
max).

`InsertHeartbeatUsageEvent` now takes an explicit `createdAt` so
generators can backfill historical buckets; the cron passes
`clock.Now()` to preserve its existing behavior.

## Tallyman follow-up

<details>
<summary>Prompt for the Tallyman-repo agent</summary>

> **Task**: Add support for the new Coder usage event type
`hb_agent_runtime_v1` so Tallyman accepts, validates, and forwards it to
Metronome.
>
> **Background**: coder/coder PR (this PR) adds hourly heartbeat events
measuring Coder Agent runtime. Events arrive via the existing
`/api/v1/events/ingest` endpoint with: `event_type:
"hb_agent_runtime_v1"`, `event_data: {"runtime_ms": <int64 >= 0>}`,
deterministic `id` of the form `hb_agent_runtime_v1:2026-07-15_14:00:00`
(UTC hour bucket start), and `created_at` set to the bucket start (may
be up to ~8 days in the past due to backfill; within Metronome's 34-day
dedup window). Zero-value events are normal (idle hours).
>
> **Work**:
> 1. Update Tallyman's vendored/imported `coderd/usage/usagetypes` (or
equivalent) to the coder/coder commit that adds
`UsageEventTypeHBAgentRuntimeV1` and `HBAgentRuntime`.
> 2. Ensure ingestion validation accepts the type (`Valid()` switches)
and rejects negative `runtime_ms`.
> 3. Ensure Metronome forwarding maps the event with transaction ID
derived from the event `id` as for existing types, passing `runtime_ms`
through as the property for a SUM-aggregated billable metric ("Coder
Agent Hours" = `SUM(runtime_ms) / 3,600,000`).
> 4. Do NOT permanently reject unknown-but-well-formed future `hb_*`
types if avoidable; at minimum confirm current behavior for unknown
types (temporary vs permanent rejection) and report it.
> 5. Tests: ingest accept/validate, dedup by ID, Metronome payload
mapping.
>
> **Constraint**: this must be deployed to tallyman-prod **before** any
coder/coder release containing the event generator; coderd treats
permanent rejections as terminal per event.

</details>
This commit is contained in:
Jaayden Halko
2026-07-30 15:37:45 +08:00
committed by GitHub
parent 2b28515d9b
commit 54d5eb7ec2
26 changed files with 1402 additions and 20 deletions
+18 -7
View File
@@ -3,6 +3,7 @@ package coderdtest
import (
"context"
"sync"
"time"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/usage"
@@ -11,10 +12,16 @@ import (
var _ usage.Inserter = (*UsageInserter)(nil)
type HeartbeatEvent struct {
ID string
CreatedAt time.Time
Event usagetypes.HeartbeatEvent
}
type UsageInserter struct {
sync.Mutex
discreteEvents []usagetypes.DiscreteEvent
heartbeatEvents []usagetypes.HeartbeatEvent
heartbeatEvents []HeartbeatEvent
seenHeartbeats map[string]struct{}
}
@@ -22,7 +29,7 @@ func NewUsageInserter() *UsageInserter {
return &UsageInserter{
discreteEvents: []usagetypes.DiscreteEvent{},
seenHeartbeats: map[string]struct{}{},
heartbeatEvents: []usagetypes.HeartbeatEvent{},
heartbeatEvents: []HeartbeatEvent{},
}
}
@@ -33,7 +40,7 @@ func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.S
return nil
}
func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, event usagetypes.HeartbeatEvent) error {
func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error {
u.Lock()
defer u.Unlock()
if _, seen := u.seenHeartbeats[id]; seen {
@@ -41,14 +48,18 @@ func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.
}
u.seenHeartbeats[id] = struct{}{}
u.heartbeatEvents = append(u.heartbeatEvents, event)
u.heartbeatEvents = append(u.heartbeatEvents, HeartbeatEvent{
ID: id,
CreatedAt: createdAt,
Event: event,
})
return nil
}
func (u *UsageInserter) GetHeartbeatEvents() []usagetypes.HeartbeatEvent {
func (u *UsageInserter) GetHeartbeatEvents() []HeartbeatEvent {
u.Lock()
defer u.Unlock()
eventsCopy := make([]usagetypes.HeartbeatEvent, len(u.heartbeatEvents))
eventsCopy := make([]HeartbeatEvent, len(u.heartbeatEvents))
copy(eventsCopy, u.heartbeatEvents)
return eventsCopy
}
@@ -72,5 +83,5 @@ func (u *UsageInserter) Reset() {
defer u.Unlock()
u.seenHeartbeats = map[string]struct{}{}
u.discreteEvents = []usagetypes.DiscreteEvent{}
u.heartbeatEvents = []usagetypes.HeartbeatEvent{}
u.heartbeatEvents = []HeartbeatEvent{}
}
+20 -3
View File
@@ -671,9 +671,9 @@ var (
Site: rbac.Permissions(map[string][]policy.Action{
rbac.ResourceLicense.Type: {policy.ActionRead},
rbac.ResourceAiSeat.Type: {policy.ActionRead}, // Required for GetActiveAISeatCount.
// The usage publisher doesn't create events, just
// reads/processes them.
rbac.ResourceUsageEvent.Type: {policy.ActionRead, policy.ActionUpdate},
// Create is required to insert heartbeat usage events
// under this subject.
rbac.ResourceUsageEvent.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
}),
User: []rbac.Permission{},
ByOrgID: map[string]rbac.OrgPermissions{},
@@ -4992,6 +4992,16 @@ func (q *querier) GetTemplatesWithFilter(ctx context.Context, arg database.GetTe
return q.db.GetAuthorizedTemplates(ctx, arg, prep)
}
func (q *querier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) {
// This query exists solely to compute hb_agent_runtime_v1 usage event
// payloads and returns a bare sum with no chat content, so it is gated
// on usage event creation rather than on reading chats.
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceUsageEvent); err != nil {
return 0, err
}
return q.db.GetTotalChatMessageRuntimeMsInRange(ctx, arg)
}
func (q *querier) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil {
return 0, err
@@ -6901,6 +6911,13 @@ func (q *querier) ListTasks(ctx context.Context, arg database.ListTasksParams) (
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ListTasks)(ctx, arg)
}
func (q *querier) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil {
return nil, err
}
return q.db.ListUsageEventCreatedAtsByTypeSince(ctx, arg)
}
func (q *querier) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) {
u, err := q.db.GetUserByID(ctx, userID)
if err != nil {
+56
View File
@@ -6717,6 +6717,27 @@ func (s *MethodTestSuite) TestUsageEvents() {
EndDate: time.Time{},
}).Asserts(rbac.ResourceUsageEvent, policy.ActionRead)
}))
s.Run("ListUsageEventCreatedAtsByTypeSince", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
params := database.ListUsageEventCreatedAtsByTypeSinceParams{
EventType: "hb_agent_runtime_v1",
Since: dbtime.Now(),
}
db.EXPECT().ListUsageEventCreatedAtsByTypeSince(gomock.Any(), params).Return([]time.Time{}, nil)
check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionRead)
}))
// GetTotalChatMessageRuntimeMsInRange exists solely to compute usage
// event payloads, so it asserts usage event creation rather than chat
// read permissions.
s.Run("GetTotalChatMessageRuntimeMsInRange", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
params := database.GetTotalChatMessageRuntimeMsInRangeParams{
StartTime: time.Time{},
EndTime: time.Time{},
}
db.EXPECT().GetTotalChatMessageRuntimeMsInRange(gomock.Any(), params).Return(int64(0), nil)
check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionCreate)
}))
}
// Ensures that the prebuilds actor may never insert an api key.
@@ -6735,6 +6756,41 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) {
require.True(t, dbauthz.IsNotAuthorizedError(err))
}
// TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied mechanically
// checks the invariant the query's authz gate relies on: it exposes a
// deployment-wide aggregate behind usage_event create at site scope, which no
// human-assignable role holds. Owner is excluded from usage_event via
// allPermsExcept in roles.go; org roles such as org-admin do carry
// usage_event permissions, but only at org scope, which cannot satisfy a
// site-scoped check. If either of those ever changes, this test fails.
func TestGetTotalChatMessageRuntimeMsInRange_HumanRolesDenied(t *testing.T) {
t.Parallel()
orgID := uuid.New()
var roles []rbac.RoleIdentifier
for _, role := range rbac.SiteBuiltInRoles() {
roles = append(roles, role.Identifier)
}
for _, role := range rbac.OrganizationRoles(orgID) {
roles = append(roles, role.Identifier)
}
require.NotEmpty(t, roles)
for _, role := range roles {
subj := rbac.Subject{
ID: uuid.NewString(),
Roles: rbac.RoleIdentifiers{role},
Scope: rbac.ScopeAll,
}
ctx := dbauthz.As(testutil.Context(t, testutil.WaitShort), subj)
mDB := dbmock.NewMockStore(gomock.NewController(t))
mDB.EXPECT().Wrappers().Times(1).Return([]string{})
dbz := dbauthz.New(mDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), slogtest.Make(t, nil), coderdtest.AccessControlStorePointer())
_, err := dbz.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{})
require.True(t, dbauthz.IsNotAuthorizedError(err), "role %s must be denied", role)
}
}
func (s *MethodTestSuite) TestAIBridge() {
s.Run("InsertAIBridgeInterception", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
initID := uuid.UUID{3}
+16
View File
@@ -3177,6 +3177,14 @@ func (m queryMetricsStore) GetTemplatesWithFilter(ctx context.Context, arg datab
return r0, r1
}
func (m queryMetricsStore) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) {
start := time.Now()
r0, r1 := m.s.GetTotalChatMessageRuntimeMsInRange(ctx, arg)
m.queryLatencies.WithLabelValues("GetTotalChatMessageRuntimeMsInRange").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTotalChatMessageRuntimeMsInRange").Inc()
return r0, r1
}
func (m queryMetricsStore) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) {
start := time.Now()
r0, r1 := m.s.GetTotalUsageDCManagedAgentsV1(ctx, arg)
@@ -4857,6 +4865,14 @@ func (m queryMetricsStore) ListTasks(ctx context.Context, arg database.ListTasks
return r0, r1
}
func (m queryMetricsStore) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) {
start := time.Now()
r0, r1 := m.s.ListUsageEventCreatedAtsByTypeSince(ctx, arg)
m.queryLatencies.WithLabelValues("ListUsageEventCreatedAtsByTypeSince").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListUsageEventCreatedAtsByTypeSince").Inc()
return r0, r1
}
func (m queryMetricsStore) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) {
start := time.Now()
r0, r1 := m.s.ListUserChatCompactionThresholds(ctx, userID)
+30
View File
@@ -5938,6 +5938,21 @@ func (mr *MockStoreMockRecorder) GetTemplatesWithFilter(ctx, arg any) *gomock.Ca
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplatesWithFilter", reflect.TypeOf((*MockStore)(nil).GetTemplatesWithFilter), ctx, arg)
}
// GetTotalChatMessageRuntimeMsInRange mocks base method.
func (m *MockStore) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg database.GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetTotalChatMessageRuntimeMsInRange", ctx, arg)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetTotalChatMessageRuntimeMsInRange indicates an expected call of GetTotalChatMessageRuntimeMsInRange.
func (mr *MockStoreMockRecorder) GetTotalChatMessageRuntimeMsInRange(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTotalChatMessageRuntimeMsInRange", reflect.TypeOf((*MockStore)(nil).GetTotalChatMessageRuntimeMsInRange), ctx, arg)
}
// GetTotalUsageDCManagedAgentsV1 mocks base method.
func (m *MockStore) GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg database.GetTotalUsageDCManagedAgentsV1Params) (int64, error) {
m.ctrl.T.Helper()
@@ -9146,6 +9161,21 @@ func (mr *MockStoreMockRecorder) ListTasks(ctx, arg any) *gomock.Call {
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListTasks", reflect.TypeOf((*MockStore)(nil).ListTasks), ctx, arg)
}
// ListUsageEventCreatedAtsByTypeSince mocks base method.
func (m *MockStore) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg database.ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ListUsageEventCreatedAtsByTypeSince", ctx, arg)
ret0, _ := ret[0].([]time.Time)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ListUsageEventCreatedAtsByTypeSince indicates an expected call of ListUsageEventCreatedAtsByTypeSince.
func (mr *MockStoreMockRecorder) ListUsageEventCreatedAtsByTypeSince(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListUsageEventCreatedAtsByTypeSince", reflect.TypeOf((*MockStore)(nil).ListUsageEventCreatedAtsByTypeSince), ctx, arg)
}
// ListUserChatCompactionThresholds mocks base method.
func (m *MockStore) ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]database.UserConfig, error) {
m.ctrl.T.Helper()
+13 -2
View File
@@ -728,7 +728,7 @@ CREATE FUNCTION aggregate_usage_event() RETURNS trigger
AS $$
BEGIN
-- Check for supported event types and throw error for unknown types.
IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1') THEN
IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1') THEN
RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type;
END IF;
@@ -756,6 +756,13 @@ BEGIN
COALESCE((NEW.event_data->>'count')::bigint, 0)
)
)
-- Hourly runtime heartbeats: sum the runtime per day.
WHEN NEW.event_type IN ('hb_agent_runtime_v1') THEN
jsonb_build_object(
'runtime_ms',
COALESCE((usage_events_daily.usage_data->>'runtime_ms')::bigint, 0) +
COALESCE((NEW.event_data->>'runtime_ms')::bigint, 0)
)
END;
RETURN NEW;
@@ -3529,7 +3536,7 @@ CREATE TABLE usage_events (
publish_started_at timestamp with time zone,
published_at timestamp with time zone,
failure_message text,
CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_v1'::text])))
CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_v1'::text, 'hb_agent_runtime_v1'::text])))
);
COMMENT ON TABLE usage_events IS 'usage_events contains usage data that is collected from the product and potentially shipped to the usage collector service.';
@@ -3540,6 +3547,8 @@ COMMENT ON COLUMN usage_events.event_type IS 'The usage event type with version.
COMMENT ON COLUMN usage_events.event_data IS 'Event payload. Determined by the matching usage struct for this event type.';
COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.';
COMMENT ON COLUMN usage_events.publish_started_at IS 'Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired.';
COMMENT ON COLUMN usage_events.published_at IS 'Set to a timestamp when the event is successfully (or permanently unsuccessfully) published to the usage collector service. If set, the event should never be attempted to be published again.';
@@ -4870,6 +4879,8 @@ CREATE INDEX idx_template_versions_has_ai_task ON template_versions USING btree
CREATE UNIQUE INDEX idx_unique_preset_name ON template_version_presets USING btree (name, template_version_id);
CREATE INDEX idx_usage_events_agent_runtime ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_agent_runtime_v1'::text);
CREATE INDEX idx_usage_events_ai_seats ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_ai_seats_v1'::text);
CREATE INDEX idx_usage_events_select_for_publishing ON usage_events USING btree (published_at, publish_started_at, created_at);
@@ -0,0 +1,50 @@
COMMENT ON COLUMN usage_events.created_at IS NULL;
DROP INDEX IF EXISTS idx_usage_events_agent_runtime;
-- Remove hb_agent_runtime_v1 rows so the previous constraint can be restored.
DELETE FROM usage_events WHERE event_type = 'hb_agent_runtime_v1';
DELETE FROM usage_events_daily WHERE event_type = 'hb_agent_runtime_v1';
ALTER TABLE usage_events
DROP CONSTRAINT usage_event_type_check,
ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1'));
-- Restores the 000444 version of the function.
CREATE OR REPLACE FUNCTION aggregate_usage_event()
RETURNS TRIGGER AS $$
BEGIN
-- Check for supported event types and throw error for unknown types.
IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1') THEN
RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type;
END IF;
INSERT INTO usage_events_daily (day, event_type, usage_data)
VALUES (
date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date,
NEW.event_type,
NEW.event_data
)
ON CONFLICT (day, event_type) DO UPDATE SET
usage_data = CASE
-- Handle simple counter events by summing the count.
WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN
jsonb_build_object(
'count',
COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) +
COALESCE((NEW.event_data->>'count')::bigint, 0)
)
-- Heartbeat events: keep the max value seen that day
WHEN NEW.event_type IN ('hb_ai_seats_v1') THEN
jsonb_build_object(
'count',
GREATEST(
COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0),
COALESCE((NEW.event_data->>'count')::bigint, 0)
)
)
END;
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
@@ -0,0 +1,55 @@
ALTER TABLE usage_events
DROP CONSTRAINT usage_event_type_check,
ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1'));
COMMENT ON COLUMN usage_events.created_at IS 'The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.';
-- Serves the usage generator's scan for missing hourly buckets.
CREATE INDEX idx_usage_events_agent_runtime
ON usage_events (event_type, created_at)
WHERE event_type = 'hb_agent_runtime_v1';
CREATE OR REPLACE FUNCTION aggregate_usage_event()
RETURNS TRIGGER AS $$
BEGIN
-- Check for supported event types and throw error for unknown types.
IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1', 'hb_agent_runtime_v1') THEN
RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type;
END IF;
INSERT INTO usage_events_daily (day, event_type, usage_data)
VALUES (
date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date,
NEW.event_type,
NEW.event_data
)
ON CONFLICT (day, event_type) DO UPDATE SET
usage_data = CASE
-- Handle simple counter events by summing the count.
WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN
jsonb_build_object(
'count',
COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) +
COALESCE((NEW.event_data->>'count')::bigint, 0)
)
-- Heartbeat events: keep the max value seen that day
WHEN NEW.event_type IN ('hb_ai_seats_v1') THEN
jsonb_build_object(
'count',
GREATEST(
COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0),
COALESCE((NEW.event_data->>'count')::bigint, 0)
)
)
-- Hourly runtime heartbeats: sum the runtime per day.
WHEN NEW.event_type IN ('hb_agent_runtime_v1') THEN
jsonb_build_object(
'runtime_ms',
COALESCE((usage_events_daily.usage_data->>'runtime_ms')::bigint, 0) +
COALESCE((NEW.event_data->>'runtime_ms')::bigint, 0)
)
END;
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
@@ -0,0 +1,20 @@
INSERT INTO usage_events (
id,
event_type,
event_data,
created_at,
publish_started_at,
published_at,
failure_message
)
VALUES
-- Unpublished hb_agent_runtime_v1 event.
(
'hb_agent_runtime_v1:2023-06-01_00:00:00',
'hb_agent_runtime_v1',
'{"runtime_ms":3600000}',
'2023-06-01 00:00:00+00',
NULL,
NULL,
NULL
);
+2 -1
View File
@@ -6142,7 +6142,8 @@ type UsageEvent struct {
EventType string `db:"event_type" json:"event_type"`
// Event payload. Determined by the matching usage struct for this event type.
EventData json.RawMessage `db:"event_data" json:"event_data"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
// The time the usage occurred, which is not necessarily the time the row was inserted. Events that measure a time bucket (e.g. hb_agent_runtime_v1) always set this to the bucket start, regardless of when the row was inserted. This timestamp determines the day used by the daily rollup trigger and is sent to the usage collector service as the event timestamp.
CreatedAt time.Time `db:"created_at" json:"created_at"`
// Set to a timestamp while the event is being published by a Coder replica to the usage collector service. Used to avoid duplicate publishes by multiple replicas. Timestamps older than 1 hour are considered expired.
PublishStartedAt sql.NullTime `db:"publish_started_at" json:"publish_started_at"`
// Set to a timestamp when the event is successfully (or permanently unsuccessfully) published to the usage collector service. If set, the event should never be attempted to be published again.
+5
View File
@@ -867,6 +867,9 @@ type sqlcQuerier interface {
GetTemplateVersionsCreatedAfter(ctx context.Context, createdAt time.Time) ([]TemplateVersion, error)
GetTemplates(ctx context.Context) ([]Template, error)
GetTemplatesWithFilter(ctx context.Context, arg GetTemplatesWithFilterParams) ([]Template, error)
// Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes
// soft-deleted messages and messages from all chats.
GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error)
// Gets the total number of managed agents created between two dates. Uses the
// aggregate table to avoid large scans or a complex index on the usage_events
// table.
@@ -1257,6 +1260,8 @@ type sqlcQuerier interface {
ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error)
ListProvisionerKeysByOrganizationExcludeReserved(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error)
ListTasks(ctx context.Context, arg ListTasksParams) ([]Task, error)
// Used by the usage generator to find missing heartbeat buckets.
ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error)
ListUserChatCompactionThresholds(ctx context.Context, userID uuid.UUID) ([]UserConfig, error)
ListUserChatPersonalModelOverrides(ctx context.Context, userID uuid.UUID) ([]ListUserChatPersonalModelOverridesRow, error)
// Returns metadata only (no value or value_key_id) for the
+167
View File
@@ -10717,6 +10717,62 @@ func TestUsageEventsTrigger(t *testing.T) {
require.Len(t, rows, 3)
})
t.Run("HeartbeatAgentRuntime", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
insert := func(id, eventType, eventData string, createdAt time.Time) {
t.Helper()
err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
ID: id,
EventType: eventType,
EventData: []byte(eventData),
CreatedAt: createdAt,
})
require.NoError(t, err)
}
requireDaily := func(wantUsageData ...string) {
t.Helper()
rows := getDailyRows(ctx, sqlDB)
require.Len(t, rows, len(wantUsageData))
for i, want := range wantUsageData {
require.JSONEq(t, want, string(rows[i].UsageData))
}
}
day1 := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
day2 := time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC)
insert("hb_agent_runtime_v1:2025-01-01_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 1000}`, day1)
requireDaily(`{"runtime_ms": 1000}`)
// Unlike hb_ai_seats_v1, hourly runtime events are summed per day.
insert("hb_agent_runtime_v1:2025-01-01_12:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 500}`, day1.Add(12*time.Hour))
requireDaily(`{"runtime_ms": 1500}`)
// Zero-valued events (idle hours) do not change the sum.
insert("hb_agent_runtime_v1:2025-01-01_18:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 0}`, day1.Add(18*time.Hour))
requireDaily(`{"runtime_ms": 1500}`)
insert("hb_agent_runtime_v1:2025-01-02_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 250}`, day2)
requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`)
// Re-inserting a bucket must not double-count it. The daily rollup
// sums runtime_ms, so idempotency rests on the aggregate trigger
// being AFTER INSERT: Postgres does not fire it for rows suppressed
// by ON CONFLICT (id) DO NOTHING. Concurrent replicas and backfill
// re-runs both take this path.
insert("hb_agent_runtime_v1:2025-01-01_00:00:00", "hb_agent_runtime_v1", `{"runtime_ms": 1000}`, day1)
requireDaily(`{"runtime_ms": 1500}`, `{"runtime_ms": 250}`)
// A different event type on the same day gets its own daily row.
insert("hb-seats-1", "hb_ai_seats_v1", `{"count": 3}`, day2)
rows := getDailyRows(ctx, sqlDB)
require.Len(t, rows, 3)
})
t.Run("UnknownEventType", func(t *testing.T) {
t.Parallel()
@@ -10750,6 +10806,117 @@ func TestUsageEventsTrigger(t *testing.T) {
})
}
func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
rangeStart := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC)
rangeEnd := rangeStart.Add(time.Hour)
total, err := db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{
StartTime: rangeStart,
EndTime: rangeEnd,
})
require.NoError(t, err)
require.EqualValues(t, 0, total)
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
_ = dbgen.ChatProvider(t, db, database.ChatProvider{
Provider: "openai",
DisplayName: "OpenAI",
})
mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
Model: "test-model",
ContextLimit: 8192,
})
chat1 := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: mc.ID,
})
chat2 := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: mc.ID,
})
insertMessage := func(chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) {
t.Helper()
msg := dbgen.ChatMessage(t, db, database.ChatMessage{
ChatID: chatID,
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
ModelConfigID: uuid.NullUUID{UUID: mc.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
RuntimeMs: sql.NullInt64{Int64: runtimeMs, Valid: true},
})
_, err := sqlDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1, deleted = $2 WHERE id = $3", createdAt, deleted, msg.ID)
require.NoError(t, err)
}
// Counted: on the inclusive start boundary, in the middle (across two
// chats), soft-deleted, and just before the exclusive end boundary.
insertMessage(chat1.ID, 1, rangeStart, false)
insertMessage(chat2.ID, 2, rangeStart.Add(30*time.Minute), false)
insertMessage(chat1.ID, 4, rangeStart.Add(45*time.Minute), true)
insertMessage(chat1.ID, 8, rangeEnd.Add(-time.Second), false)
// Not counted: before the range, on the exclusive end boundary, and a
// NULL runtime (runtime 0 is stored as NULL).
insertMessage(chat1.ID, 16, rangeStart.Add(-time.Second), false)
insertMessage(chat1.ID, 32, rangeEnd, false)
insertMessage(chat1.ID, 0, rangeStart.Add(10*time.Minute), false)
total, err = db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{
StartTime: rangeStart,
EndTime: rangeEnd,
})
require.NoError(t, err)
require.EqualValues(t, 15, total)
}
func TestListUsageEventCreatedAtsByTypeSince(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
since := time.Date(2025, 3, 10, 0, 0, 0, 0, time.UTC)
insertEvent := func(id, eventType string, eventData string, createdAt time.Time) {
t.Helper()
err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
ID: id,
EventType: eventType,
EventData: []byte(eventData),
CreatedAt: createdAt,
})
require.NoError(t, err)
}
// Matching type: one before since (excluded), one exactly at since
// (included), one after (included).
insertEvent("rt-old", "hb_agent_runtime_v1", `{"runtime_ms": 1}`, since.Add(-time.Hour))
insertEvent("rt-at", "hb_agent_runtime_v1", `{"runtime_ms": 2}`, since)
insertEvent("rt-new", "hb_agent_runtime_v1", `{"runtime_ms": 3}`, since.Add(time.Hour))
// Different type after since: excluded.
insertEvent("seats-new", "hb_ai_seats_v1", `{"count": 1}`, since.Add(time.Hour))
createdAts, err := db.ListUsageEventCreatedAtsByTypeSince(ctx, database.ListUsageEventCreatedAtsByTypeSinceParams{
EventType: "hb_agent_runtime_v1",
Since: since,
})
require.NoError(t, err)
require.Len(t, createdAts, 2)
normalized := make([]time.Time, len(createdAts))
for i, ts := range createdAts {
normalized[i] = ts.UTC()
}
require.ElementsMatch(t, []time.Time{since, since.Add(time.Hour)}, normalized)
}
func TestListTasks(t *testing.T) {
t.Parallel()
+58
View File
@@ -10090,6 +10090,28 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time
return items, nil
}
const getTotalChatMessageRuntimeMsInRange = `-- name: GetTotalChatMessageRuntimeMsInRange :one
SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms
FROM chat_messages cm
WHERE cm.created_at >= $1::timestamptz
AND cm.created_at < $2::timestamptz
AND cm.runtime_ms IS NOT NULL
`
type GetTotalChatMessageRuntimeMsInRangeParams struct {
StartTime time.Time `db:"start_time" json:"start_time"`
EndTime time.Time `db:"end_time" json:"end_time"`
}
// Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes
// soft-deleted messages and messages from all chats.
func (q *sqlQuerier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, arg GetTotalChatMessageRuntimeMsInRangeParams) (int64, error) {
row := q.db.QueryRowContext(ctx, getTotalChatMessageRuntimeMsInRange, arg.StartTime, arg.EndTime)
var total_runtime_ms int64
err := row.Scan(&total_runtime_ms)
return total_runtime_ms, err
}
const getUserChatSpendInPeriod = `-- name: GetUserChatSpendInPeriod :one
SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros
FROM chat_messages cm
@@ -29088,6 +29110,42 @@ func (q *sqlQuerier) InsertUsageEvent(ctx context.Context, arg InsertUsageEventP
return err
}
const listUsageEventCreatedAtsByTypeSince = `-- name: ListUsageEventCreatedAtsByTypeSince :many
SELECT created_at
FROM usage_events
WHERE event_type = $1
AND created_at >= $2::timestamptz
`
type ListUsageEventCreatedAtsByTypeSinceParams struct {
EventType string `db:"event_type" json:"event_type"`
Since time.Time `db:"since" json:"since"`
}
// Used by the usage generator to find missing heartbeat buckets.
func (q *sqlQuerier) ListUsageEventCreatedAtsByTypeSince(ctx context.Context, arg ListUsageEventCreatedAtsByTypeSinceParams) ([]time.Time, error) {
rows, err := q.db.QueryContext(ctx, listUsageEventCreatedAtsByTypeSince, arg.EventType, arg.Since)
if err != nil {
return nil, err
}
defer rows.Close()
var items []time.Time
for rows.Next() {
var created_at time.Time
if err := rows.Scan(&created_at); err != nil {
return nil, err
}
items = append(items, created_at)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const selectUsageEventsForPublishing = `-- name: SelectUsageEventsForPublishing :many
WITH usage_events AS (
UPDATE
+9
View File
@@ -2483,6 +2483,15 @@ WHERE c.owner_id = @user_id::uuid
AND cm.created_at < @end_time::timestamptz
AND cm.total_cost_micros IS NOT NULL;
-- name: GetTotalChatMessageRuntimeMsInRange :one
-- Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes
-- soft-deleted messages and messages from all chats.
SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms
FROM chat_messages cm
WHERE cm.created_at >= @start_time::timestamptz
AND cm.created_at < @end_time::timestamptz
AND cm.runtime_ms IS NOT NULL;
-- name: CountEnabledModelsWithoutPricing :one
-- Counts enabled, non-deleted model configs that lack both input and
-- output pricing in their JSONB options.cost configuration.
+7
View File
@@ -20,6 +20,13 @@ SELECT EXISTS(
SELECT 1 FROM usage_events WHERE id = @id
)::bool;
-- name: ListUsageEventCreatedAtsByTypeSince :many
-- Used by the usage generator to find missing heartbeat buckets.
SELECT created_at
FROM usage_events
WHERE event_type = @event_type
AND created_at >= @since::timestamptz;
-- name: SelectUsageEventsForPublishing :many
WITH usage_events AS (
UPDATE
+1
View File
@@ -38,6 +38,7 @@ const (
// ServiceWorkspaceBuildOrchestrator fulfills workspace build
// orchestrations once their parent build reaches a terminal state.
ServiceWorkspaceBuildOrchestrator = "workspace-build-orchestrator"
ServiceUsageEventGenerator = "usage-event-generator"
RequestTypeTag = "coder_request_type"
)
+7 -2
View File
@@ -2,6 +2,7 @@ package usage
import (
"context"
"time"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/usage/usagetypes"
@@ -28,7 +29,11 @@ type Inserter interface {
//
// Inserts with the same `id` must be idempotent. The database enforces this by
// ignoring duplicate records.
InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error
//
// Generators that backfill historical buckets pass the bucket start as
// `createdAt` rather than the insertion time. `createdAt` must be
// non-zero; implementations reject the zero time rather than storing it.
InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error
}
// AGPLInserter is a no-op implementation of Inserter.
@@ -48,6 +53,6 @@ func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store
// InsertHeartbeatUsageEvent is a no-op implementation of
// InsertHeartbeatUsageEvent.
func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ usagetypes.HeartbeatEvent) error {
func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ time.Time, _ usagetypes.HeartbeatEvent) error {
return nil
}
+44
View File
@@ -30,6 +30,7 @@ type UsageEventType string
const (
UsageEventTypeDCManagedAgentsV1 UsageEventType = "dc_managed_agents_v1"
UsageEventTypeHBAISeatsV1 UsageEventType = "hb_ai_seats_v1"
UsageEventTypeHBAgentRuntimeV1 UsageEventType = "hb_agent_runtime_v1"
)
func (e UsageEventType) Valid() bool {
@@ -38,6 +39,8 @@ func (e UsageEventType) Valid() bool {
return true
case UsageEventTypeHBAISeatsV1:
return true
case UsageEventTypeHBAgentRuntimeV1:
return true
default:
return false
}
@@ -105,6 +108,12 @@ func ParseEventWithType(eventType UsageEventType, data json.RawMessage) (Event,
return nil, err
}
return event, nil
case UsageEventTypeHBAgentRuntimeV1:
var event HBAgentRuntime
if err := ParseEvent(data, &event); err != nil {
return nil, err
}
return event, nil
default:
return nil, UnknownEventTypeError{EventType: string(eventType)}
}
@@ -192,3 +201,38 @@ func (e HBAISeats) Fields() map[string]any {
"count": e.Count,
}
}
// HBAgentRuntime is the event associated with hb_agent_runtime_v1. RuntimeMs
// is the total agent-loop runtime in milliseconds consumed by Coder Agents
// (chats) in one UTC hour. Each measured step spans model streaming (including
// provider-executed tools) and stream retries, and ends when the model stream
// finishes. Time spent executing local tools between steps, including
// sub-agents that bill their own model calls, is excluded.
//
// This measures the new Coder Agents (the `chats` tables), not the deprecated
// Tasks counted by dc_managed_agents_v1.
type HBAgentRuntime struct {
RuntimeMs int64 `json:"runtime_ms"`
}
var _ HeartbeatEvent = HBAgentRuntime{}
func (HBAgentRuntime) usageEvent() {}
func (HBAgentRuntime) heartbeatUsageEvent() {}
func (HBAgentRuntime) EventType() UsageEventType {
return UsageEventTypeHBAgentRuntimeV1
}
func (e HBAgentRuntime) Valid() error {
if e.RuntimeMs < 0 {
return xerrors.New("runtime_ms cannot be negative")
}
// The runtime can be 0 (idle hour).
return nil
}
func (e HBAgentRuntime) Fields() map[string]any {
return map[string]any{
"runtime_ms": e.RuntimeMs,
}
}
+21
View File
@@ -76,4 +76,25 @@ func TestParseEventWithType(t *testing.T) {
require.Equal(t, eventType, event.EventType())
require.Equal(t, map[string]any{"count": int64(1)}, event.Fields())
})
t.Run("HBAgentRuntimeV1", func(t *testing.T) {
t.Parallel()
eventType := usagetypes.UsageEventTypeHBAgentRuntimeV1
event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1234}`))
require.NoError(t, err)
require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 1234}, event)
require.Equal(t, eventType, event.EventType())
require.Equal(t, map[string]any{"runtime_ms": int64(1234)}, event.Fields())
event, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 0}`))
require.NoError(t, err)
require.Equal(t, usagetypes.HBAgentRuntime{RuntimeMs: 0}, event)
_, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": -1}`))
require.ErrorContains(t, err, "runtime_ms cannot be negative")
_, err = usagetypes.ParseEventWithType(eventType, []byte(`{"runtime_ms": 1, "extra": "field"}`))
require.ErrorContains(t, err, "unmarshal *usagetypes.HBAgentRuntime event")
})
}
+9
View File
@@ -4,6 +4,7 @@
- number of agent workspace builds consumed
- number of AI Governance seats consumed
- total Coder Agent runtime consumed per hour, in milliseconds
No user-identifiable information or additional metrics are sent to Tallyman. This information is also shared with [Metronome](https://metronome.com), a Stripe product and Coder partner for usage-based billing and reporting.
@@ -33,6 +34,14 @@ Coder-Deployment-ID: 8a4e92f1-3b7c-4d5e-9f12-abc123def456 # your deployment ID
"count": 1
},
"created_at": "2025-01-15T14:30:00Z"
},
{
"id": "hb_agent_runtime_v1:2025-01-15_14:00:00", # unique event ID generated by Coder
"event_type": "hb_agent_runtime_v1", # aka. Coder Agent runtime
"event_data": {
"runtime_ms": 3600000
},
"created_at": "2025-01-15T14:00:00Z" # start of the hour the runtime was consumed in
}
]
}
+6
View File
@@ -154,6 +154,12 @@ func (r *RootCmd) Server(_ func()) *serpent.Command {
usageCron.Start(ctx)
closers.Add(usageCron)
// Usage generation is deliberately not license-gated; the
// publish_usage_data license flag only gates publishing to Tallyman.
usageGenerator := usage.NewGenerator(quartz.NewReal(), options.Logger.Named("usage-event-generator"), options.Database, *options.UsageInserter.Load())
usageGenerator.Start(ctx)
closers.Add(usageGenerator)
// In-memory AI Bridge Proxy daemon. The bridge daemon itself is
// started unconditionally by AGPL cli/server.go (chatd uses its
// in-memory roundtripper regardless of license); only the proxy
+6 -3
View File
@@ -23,7 +23,10 @@ import (
var epoch = time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)
const (
cronDateFormat = "2006-01-02_15:04:05"
// usageEventIDTimeFormat is the timestamp layout used in every
// deterministic usage event ID, both the cron's boundary IDs and the
// generator's bucket IDs.
usageEventIDTimeFormat = "2006-01-02_15:04:05"
)
// HeartbeatFunc generates a heartbeat event and its stable ID.
@@ -145,7 +148,7 @@ func (c *Cron) run(ctx context.Context, job CronJob) {
// 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)
stableID := string(job.EventType) + ":" + boundary.UTC().Format(usageEventIDTimeFormat)
// Skip if this bucket was already recorded — avoids running
// the potentially expensive heartbeat function for a
@@ -184,7 +187,7 @@ func (c *Cron) run(ctx context.Context, job CronJob) {
continue
}
if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, event); err != nil {
if err := c.ins.InsertHeartbeatUsageEvent(ctx, c.db, stableID, c.clock.Now(), event); err != nil {
c.log.Warn(ctx, "cron heartbeat insert failed",
slog.F("job", job.Name),
slog.Error(err),
+242
View File
@@ -0,0 +1,242 @@
package usage
import (
"context"
"math/rand"
"sync"
"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"
)
const (
// AgentRuntimeInterval is the bucket size of hb_agent_runtime_v1 events.
AgentRuntimeInterval = time.Hour
// AgentRuntimeWindow is the trailing window scanned for missing buckets.
// Buckets still missing beyond this window (e.g. because the deployment
// was down for longer) are forfeited, which can only ever undercount
// usage.
AgentRuntimeWindow = 7 * 24 * time.Hour
// AgentRuntimeEligibilityLag is how long after a bucket closes before it
// becomes eligible for generation, giving replicas time to commit
// in-flight chat messages with timestamps inside the bucket. This
// assumes chat message inserts commit within the lag of their
// statement-time created_at; a message committing later than the lag
// after its bucket closes lands in an already-sealed bucket and its
// runtime is dropped (undercount-only).
AgentRuntimeEligibilityLag = 5 * time.Minute
// agentRuntimeJitter staggers replicas after each hour boundary so one
// is likely to complete the work before others attempt it.
agentRuntimeJitter = 4 * time.Minute
// agentRuntimeStartupDelay is the floor on the first pass after start,
// giving the deployment time to finish booting before the generator
// competes for database work. Jitter is added on top of it.
agentRuntimeStartupDelay = time.Minute
// generatorTimerName tags the quartz timer so tests can trap it.
generatorTimerName = "agent-runtime-generator"
)
// Generator reconciles hb_agent_runtime_v1 heartbeat usage events. Unlike
// Cron jobs, which sample live state when they fire, the Generator derives
// events from data already persisted in the database, so it can
// deterministically backfill hours missed while the deployment was down,
// zero-filling idle hours. Deterministic event IDs plus the database's
// ON CONFLICT (id) DO NOTHING make concurrent replicas safe without locking.
//
// Events are generated unconditionally in enterprise builds; the
// publish_usage_data license flag only gates publishing to Tallyman.
type Generator struct {
clock quartz.Clock
log slog.Logger
db database.Store
ins agplusage.Inserter
cancel context.CancelFunc
wg sync.WaitGroup
startOnce sync.Once
}
// NewGenerator creates an unstarted Generator.
func NewGenerator(clock quartz.Clock, log slog.Logger, db database.Store, ins agplusage.Inserter) *Generator {
return &Generator{
clock: clock,
log: log,
db: db,
ins: ins,
}
}
// Start launches the reconciliation goroutine. Subsequent calls are no-ops;
// a closed Generator cannot be restarted.
func (g *Generator) Start(ctx context.Context) {
g.startOnce.Do(func() {
ctx, g.cancel = context.WithCancel(ctx)
g.wg.Add(1)
pproflabel.Go(ctx, pproflabel.Service(pproflabel.ServiceUsageEventGenerator), func(ctx context.Context) {
g.run(ctx)
})
})
}
// Close stops the Generator and waits for its goroutine to exit.
// It always returns nil; the error return exists to satisfy io.Closer, as the
// Generator is registered with the server's closer list.
func (g *Generator) Close() error {
if g.cancel != nil {
g.cancel()
}
g.wg.Wait()
return nil
}
func (g *Generator) run(ctx context.Context) {
//nolint:gocritic // We are a publisher in this function.
ctx = dbauthz.AsUsagePublisher(ctx)
defer g.wg.Done()
// The random initial delay staggers replicas that start simultaneously.
//nolint:gosec // Jitter does not need cryptographic randomness.
delay := agentRuntimeStartupDelay + time.Duration(rand.Int63n(int64(agentRuntimeJitter)))
for {
timer := g.clock.NewTimer(delay, generatorTimerName)
select {
case <-ctx.Done():
if !timer.Stop() {
<-timer.C
}
return
case <-timer.C:
}
err := g.generateAgentRuntimeEvents(ctx)
if ctx.Err() != nil {
return
}
if err != nil {
g.log.Warn(ctx, "generate agent runtime usage events", slog.Error(err))
}
// Wake at the next eligibility instant (hour boundary + lag), not
// the next hour boundary. Computing the tick against the
// lag-shifted clock keeps a bucket whose eligibility is still
// pending in this hour (e.g. a pass that ran just before HH:05)
// from waiting a whole extra hour.
_, delay = nextTick(g.clock.Now().Add(-AgentRuntimeEligibilityLag), AgentRuntimeInterval, agentRuntimeJitter)
}
}
// generateAgentRuntimeEvents inserts one hb_agent_runtime_v1 event per
// missing hourly bucket in the trailing window. Per-bucket errors skip only
// that bucket; the next tick rescans the whole window, so transient failures
// self-heal.
func (g *Generator) generateAgentRuntimeEvents(ctx context.Context) error {
now := g.clock.Now().UTC()
// Bucket [H, H+1) becomes eligible at H + interval + lag.
latestEligible := now.Add(-AgentRuntimeInterval - AgentRuntimeEligibilityLag).Truncate(AgentRuntimeInterval)
earliest := now.Truncate(AgentRuntimeInterval).Add(-AgentRuntimeWindow)
if latestEligible.Before(earliest) {
return nil
}
existingTimes, err := g.db.ListUsageEventCreatedAtsByTypeSince(ctx, database.ListUsageEventCreatedAtsByTypeSinceParams{
EventType: string(usagetypes.UsageEventTypeHBAgentRuntimeV1),
Since: earliest,
})
if err != nil {
return xerrors.Errorf("list existing agent runtime events: %w", err)
}
// A row marks its bucket complete regardless of publish outcome, so a
// bucket whose event Tallyman permanently rejected is never
// regenerated (re-inserting under the deterministic ID is a no-op via
// ON CONFLICT (id) DO NOTHING).
//
// The runtime is not lost locally: the row still holds it, and the
// event can be re-queued for publishing with
//
// UPDATE usage_events
// SET published_at = NULL, publish_started_at = NULL, failure_message = NULL
// WHERE id = 'hb_agent_runtime_v1:<bucket start, e.g. 2026-07-15_14:00:00>';
//
// That re-arm only has an effect while the bucket is inside the
// publisher's 30-day cutoff: SelectUsageEventsForPublishing also
// filters created_at > now - INTERVAL '30 days', and created_at is the
// bucket start, so past that the UPDATE reports success but the row is
// never picked up again. The release gate (Tallyman must accept this
// event type before coderd ships it) is what keeps permanent
// rejections exceptional.
existing := make(map[time.Time]struct{}, len(existingTimes))
for _, ts := range existingTimes {
// created_at is always the exact bucket start for this event type;
// truncation just normalizes timezone and precision.
existing[ts.UTC().Truncate(AgentRuntimeInterval)] = struct{}{}
}
var filled, failed int
for bucket := earliest; !bucket.After(latestEligible); bucket = bucket.Add(AgentRuntimeInterval) {
if _, ok := existing[bucket]; ok {
continue
}
if ctx.Err() != nil {
return ctx.Err()
}
err := g.generateBucket(ctx, bucket)
if err != nil {
if ctx.Err() != nil {
// A cancel landing mid-query surfaces as a bucket error.
// Report the cancellation instead of logging shutdown as
// a bucket failure.
return ctx.Err()
}
// Skip only the failed bucket so one bad bucket (e.g. an
// invalid runtime sum) cannot stall every later bucket until
// it ages out of the window.
g.log.Warn(ctx, "generate agent runtime usage event for bucket",
slog.F("bucket", bucket),
slog.Error(err),
)
failed++
continue
}
filled++
}
if filled > 0 || failed > 0 {
g.log.Info(ctx, "generated agent runtime usage events",
slog.F("buckets_filled", filled),
slog.F("buckets_failed", failed),
slog.F("window_start", earliest),
slog.F("latest_eligible", latestEligible),
)
}
return nil
}
// generateBucket computes and inserts the event for a single hourly bucket.
func (g *Generator) generateBucket(ctx context.Context, bucket time.Time) error {
runtimeMs, err := g.db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{
StartTime: bucket,
EndTime: bucket.Add(AgentRuntimeInterval),
})
if err != nil {
return xerrors.Errorf("sum chat message runtime: %w", err)
}
// The deterministic ID makes concurrent inserts of the same bucket
// idempotent, and created_at is the bucket start (not the insertion
// time) so daily rollups attribute backfilled hours to the correct day.
stableID := string(usagetypes.UsageEventTypeHBAgentRuntimeV1) + ":" + bucket.Format(usageEventIDTimeFormat)
err = g.ins.InsertHeartbeatUsageEvent(ctx, g.db, stableID, bucket, usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs})
if err != nil {
return xerrors.Errorf("insert usage event: %w", err)
}
return nil
}
+481
View File
@@ -0,0 +1,481 @@
package usage_test
import (
"context"
"database/sql"
"sync/atomic"
"testing"
"time"
"github.com/google/uuid"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/require"
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/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/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/rbac"
"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"
)
// generatorTimerName must match the tag the Generator passes to
// clock.NewTimer so tests can trap its timers.
const generatorTimerName = "agent-runtime-generator"
// warnSink counts log entries at Warn or above. The generator downgrades
// per-bucket failures to Warn logs (which slogtest tolerates), so tests that
// must prove the error paths stayed quiet assert on this counter instead.
type warnSink struct{ count atomic.Int64 }
func (s *warnSink) LogEntry(_ context.Context, e slog.SinkEntry) {
if e.Level >= slog.LevelWarn {
s.count.Add(1)
}
}
func (*warnSink) Sync() {}
// generatorHarness runs the generator against a dbauthz-wrapped store so the
// tests also verify that the usage publisher subject holds the permissions
// the generator's queries require.
type generatorHarness struct {
db database.Store
authzDB database.Store
rawDB *sql.DB
user database.User
modelConfig database.ChatModelConfig
chat database.Chat
chat2 database.Chat
}
func newGeneratorHarness(t *testing.T) *generatorHarness {
t.Helper()
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t)
log := slogtest.Make(t, nil)
authzDB := dbauthz.New(db, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), log, coderdtest.AccessControlStorePointer())
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
_ = dbgen.ChatProvider(t, db, database.ChatProvider{
Provider: "openai",
DisplayName: "OpenAI",
})
mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
Model: "test-model",
ContextLimit: 8192,
})
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: mc.ID,
})
chat2 := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: mc.ID,
})
return &generatorHarness{
db: db,
authzDB: authzDB,
rawDB: rawDB,
user: user,
modelConfig: mc,
chat: chat,
chat2: chat2,
}
}
func (h *generatorHarness) insertRuntimeMessage(ctx context.Context, t *testing.T, chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) {
t.Helper()
msg := dbgen.ChatMessage(t, h.db, database.ChatMessage{
ChatID: chatID,
CreatedBy: uuid.NullUUID{UUID: h.user.ID, Valid: true},
ModelConfigID: uuid.NullUUID{UUID: h.modelConfig.ID, Valid: true},
Role: database.ChatMessageRoleAssistant,
RuntimeMs: sql.NullInt64{Int64: runtimeMs, Valid: true},
})
_, err := h.rawDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1, deleted = $2 WHERE id = $3", createdAt, deleted, msg.ID)
require.NoError(t, err)
}
// fetchRuntimeEvents fails the test if any bucket has more than one event.
func (h *generatorHarness) fetchRuntimeEvents(ctx context.Context, t *testing.T) (map[time.Time]int64, map[time.Time]string) {
t.Helper()
rows, err := h.rawDB.QueryContext(ctx, `
SELECT id, (event_data->>'runtime_ms')::bigint, created_at
FROM usage_events
WHERE event_type = 'hb_agent_runtime_v1'
`)
require.NoError(t, err)
defer rows.Close()
runtimes := make(map[time.Time]int64)
ids := make(map[time.Time]string)
for rows.Next() {
var (
id string
runtimeMs int64
createdAt time.Time
)
require.NoError(t, rows.Scan(&id, &runtimeMs, &createdAt))
bucket := createdAt.UTC()
_, ok := runtimes[bucket]
require.False(t, ok, "duplicate event for bucket %s", bucket)
runtimes[bucket] = runtimeMs
ids[bucket] = id
}
require.NoError(t, rows.Err())
return runtimes, ids
}
func expectedBuckets(first, last time.Time, overrides map[time.Time]int64) map[time.Time]int64 {
expected := make(map[time.Time]int64)
for bucket := first; !bucket.After(last); bucket = bucket.Add(time.Hour) {
expected[bucket] = 0
}
for bucket, runtimeMs := range overrides {
expected[bucket] = runtimeMs
}
return expected
}
func TestGenerator(t *testing.T) {
t.Parallel()
// startTime is exactly on an hour boundary, so the first tick (which
// fires 1-5 minutes later) always lands before the just-closed bucket
// [13:00, 14:00) becomes eligible at 14:05.
startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC)
ctx := testutil.Context(t, testutil.WaitLong)
log := slogtest.Make(t, nil)
h := newGeneratorHarness(t)
clock := quartz.NewMock(t)
clock.Set(startTime)
var (
bucketA = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC)
bucketB = time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC)
// The most recent closed bucket; not eligible at the first tick.
bucketC = time.Date(2025, 3, 10, 13, 0, 0, 0, time.UTC)
// Window bounds at the first tick.
windowFirst = startTime.Add(-usage.AgentRuntimeWindow) // 2025-03-03 14:00
windowLast = startTime.Add(-2 * time.Hour) // 2025-03-10 12:00
)
// Bucket A: a message on the bucket start boundary, a message from a
// second chat, and a soft-deleted message. All must be counted.
h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA, false)
h.insertRuntimeMessage(ctx, t, h.chat2.ID, 2000, bucketA.Add(15*time.Minute), false)
h.insertRuntimeMessage(ctx, t, h.chat.ID, 4000, bucketA.Add(30*time.Minute), true)
// Bucket B: a message exactly on the A/B boundary belongs to B.
h.insertRuntimeMessage(ctx, t, h.chat.ID, 8000, bucketB, false)
// Older than the window: must never be generated.
h.insertRuntimeMessage(ctx, t, h.chat.ID, 16000, windowFirst.Add(-30*time.Minute), false)
// Bucket C: only becomes eligible at 14:05, after the first tick.
h.insertRuntimeMessage(ctx, t, h.chat.ID, 32000, bucketC.Add(30*time.Minute), false)
trap := clock.Trap().NewTimer(generatorTimerName)
defer trap.Close()
gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter())
gen.Start(ctx)
defer gen.Close()
call := trap.MustWait(ctx)
call.MustRelease(ctx)
require.GreaterOrEqual(t, call.Duration, time.Minute)
require.Less(t, call.Duration, 5*time.Minute)
clock.Advance(call.Duration).MustWait(ctx)
// The generator creates the next timer only after the tick completes,
// so trapping it synchronizes with the end of the pass.
call = trap.MustWait(ctx)
call.MustRelease(ctx)
// The first pass fills every bucket in [windowFirst, windowLast]: bucket
// C is not yet eligible, the pre-window message is excluded, and idle
// hours are zero-filled.
runtimes, ids := h.fetchRuntimeEvents(ctx, t)
require.Equal(t, expectedBuckets(windowFirst, windowLast, map[time.Time]int64{
bucketA: 7000,
bucketB: 8000,
}), runtimes)
require.Equal(t, "hb_agent_runtime_v1:2025-03-10_10:00:00", ids[bucketA])
// The next tick fires at bucket C's eligibility instant (14:05 plus
// jitter), in the same hour the first pass ran, rather than waiting for
// the next hour boundary.
fireTime := clock.Now().Add(call.Duration)
require.Equal(t, startTime, fireTime.Truncate(usage.AgentRuntimeInterval))
require.GreaterOrEqual(t, fireTime.Sub(fireTime.Truncate(usage.AgentRuntimeInterval)), usage.AgentRuntimeEligibilityLag)
clock.Advance(call.Duration).MustWait(ctx)
call = trap.MustWait(ctx)
call.MustRelease(ctx)
// The second tick fills only the newly-eligible bucket C (13:00). The
// window's start advanced by an hour, but the bucket at the old
// windowFirst is kept (rows are never deleted).
runtimes, _ = h.fetchRuntimeEvents(ctx, t)
require.Equal(t, expectedBuckets(windowFirst, windowLast.Add(time.Hour), map[time.Time]int64{
bucketA: 7000,
bucketB: 8000,
bucketC: 32000,
}), runtimes)
// The third tick fires after the next hour boundary (15:05 plus jitter)
// and fills the idle 14:00 bucket.
fireTime = clock.Now().Add(call.Duration)
require.Equal(t, startTime.Add(time.Hour), fireTime.Truncate(usage.AgentRuntimeInterval))
require.GreaterOrEqual(t, fireTime.Sub(fireTime.Truncate(usage.AgentRuntimeInterval)), usage.AgentRuntimeEligibilityLag)
clock.Advance(call.Duration).MustWait(ctx)
call = trap.MustWait(ctx)
call.MustRelease(ctx)
runtimes, _ = h.fetchRuntimeEvents(ctx, t)
require.Equal(t, expectedBuckets(windowFirst, windowLast.Add(2*time.Hour), map[time.Time]int64{
bucketA: 7000,
bucketB: 8000,
bucketC: 32000,
}), runtimes)
// A separate generator started later (e.g. another replica restarting)
// finds nothing to do: all buckets in its window already exist.
gen2 := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter())
gen2.Start(ctx)
defer gen2.Close()
call = trap.MustWait(ctx)
call.MustRelease(ctx)
clock.Advance(call.Duration).MustWait(ctx)
call = trap.MustWait(ctx)
call.MustRelease(ctx)
runtimes2, _ := h.fetchRuntimeEvents(ctx, t)
require.Equal(t, runtimes, runtimes2)
}
// TestGeneratorBackfillAfterDowntime simulates a deployment that was down
// for several hours: a fresh generator's first pass backfills exactly the
// missing buckets.
func TestGeneratorBackfillAfterDowntime(t *testing.T) {
t.Parallel()
startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC)
ctx := testutil.Context(t, testutil.WaitLong)
log := slogtest.Make(t, nil)
h := newGeneratorHarness(t)
windowFirst := startTime.Add(-usage.AgentRuntimeWindow)
gapBucket := time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC)
h.insertRuntimeMessage(ctx, t, h.chat.ID, 5000, gapBucket.Add(45*time.Minute), false)
// Simulate events generated before the deployment went down at 09:00:
// every bucket in [windowFirst, 08:00] exists with runtime 1.
inserter := usage.NewDBInserter()
for bucket := windowFirst; !bucket.After(time.Date(2025, 3, 10, 8, 0, 0, 0, time.UTC)); bucket = bucket.Add(time.Hour) {
err := inserter.InsertHeartbeatUsageEvent(ctx, h.db, "hb_agent_runtime_v1:"+bucket.Format("2006-01-02_15:04:05"), bucket, usagetypes.HBAgentRuntime{RuntimeMs: 1})
require.NoError(t, err)
}
clock := quartz.NewMock(t)
clock.Set(startTime)
trap := clock.Trap().NewTimer(generatorTimerName)
defer trap.Close()
gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter())
gen.Start(ctx)
defer gen.Close()
call := trap.MustWait(ctx)
call.MustRelease(ctx)
clock.Advance(call.Duration).MustWait(ctx)
call = trap.MustWait(ctx)
call.MustRelease(ctx)
// The pass must fill exactly the gap [09:00, 12:00] and leave the
// pre-existing rows untouched.
runtimes, _ := h.fetchRuntimeEvents(ctx, t)
expected := expectedBuckets(windowFirst, startTime.Add(-2*time.Hour), map[time.Time]int64{
gapBucket: 5000,
})
for bucket := windowFirst; !bucket.After(time.Date(2025, 3, 10, 8, 0, 0, 0, time.UTC)); bucket = bucket.Add(time.Hour) {
expected[bucket] = 1
}
require.Equal(t, expected, runtimes)
}
func TestGeneratorInserterArguments(t *testing.T) {
t.Parallel()
startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC)
ctx := testutil.Context(t, testutil.WaitLong)
log := slogtest.Make(t, nil)
h := newGeneratorHarness(t)
clock := quartz.NewMock(t)
clock.Set(startTime)
var (
bucketA = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC)
windowFirst = startTime.Add(-usage.AgentRuntimeWindow)
// The first pass stops at 12:00; the second pass fires after 14:05
// and adds the newly eligible 13:00 bucket.
windowLast = startTime.Add(-time.Hour)
)
h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA.Add(10*time.Minute), false)
ins := coderdtest.NewUsageInserter()
trap := clock.Trap().NewTimer(generatorTimerName)
defer trap.Close()
gen := usage.NewGenerator(clock, log, h.authzDB, ins)
gen.Start(ctx)
defer gen.Close()
// Each pass requests its next timer only after finishing, so trapping
// that request synchronizes with the end of the pass.
for range 2 {
call := trap.MustWait(ctx)
call.MustRelease(ctx)
clock.Advance(call.Duration).MustWait(ctx)
}
call := trap.MustWait(ctx)
call.MustRelease(ctx)
var expected []coderdtest.HeartbeatEvent
for bucket := windowFirst; !bucket.After(windowLast); bucket = bucket.Add(usage.AgentRuntimeInterval) {
var runtimeMs int64
if bucket.Equal(bucketA) {
runtimeMs = 1000
}
expected = append(expected, coderdtest.HeartbeatEvent{
ID: "hb_agent_runtime_v1:" + bucket.Format("2006-01-02_15:04:05"),
CreatedAt: bucket,
Event: usagetypes.HBAgentRuntime{RuntimeMs: runtimeMs},
})
}
require.Equal(t, expected, ins.GetHeartbeatEvents())
}
// TestGeneratorConcurrentReplicas runs two generators against the same
// database concurrently and verifies exactly one event is produced per
// bucket. Each replica gets its own mock clock (as real replicas have their
// own wall clocks) so their first passes can be fired independently and run
// at the same time.
func TestGeneratorConcurrentReplicas(t *testing.T) {
t.Parallel()
startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC)
ctx := testutil.Context(t, testutil.WaitLong)
// If insert idempotency regressed (e.g. ON CONFLICT DO NOTHING was
// dropped), the losing replica's inserts would error and surface as
// Warn logs, which slogtest alone would not catch.
sink := &warnSink{}
log := slogtest.Make(t, nil).AppendSinks(sink)
h := newGeneratorHarness(t)
bucketA := time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC)
h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, bucketA.Add(10*time.Minute), false)
// Both replicas' first passes fire between 14:01 and 14:05, so they
// compute identical windows regardless of their random startup jitter.
var traps []*quartz.Trap
for range 2 {
clock := quartz.NewMock(t)
clock.Set(startTime)
trap := clock.Trap().NewTimer(generatorTimerName)
t.Cleanup(trap.Close)
traps = append(traps, trap)
gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter())
gen.Start(ctx)
// Cleanups run LIFO, so each generator is closed before its trap.
t.Cleanup(func() { _ = gen.Close() })
call := trap.MustWait(ctx)
call.MustRelease(ctx)
// Fire the initial timer without waiting for the pass to complete so
// both replicas' passes overlap.
clock.Advance(call.Duration).MustWait(ctx)
}
// Wait for both passes to complete (each requests its next timer only
// after the pass finishes).
for _, trap := range traps {
call := trap.MustWait(ctx)
call.MustRelease(ctx)
}
// fetchRuntimeEvents fails on duplicate buckets; the expected map proves
// both replicas raced without double-inserting, and the warn counter
// proves the duplicate inserts were deduplicated rather than rejected
// with errors.
runtimes, _ := h.fetchRuntimeEvents(ctx, t)
require.Equal(t, expectedBuckets(
startTime.Add(-usage.AgentRuntimeWindow),
startTime.Add(-2*time.Hour),
map[time.Time]int64{bucketA: 1000},
), runtimes)
require.Zero(t, sink.count.Load(), "no replica may log Warn or above during the race")
}
// TestGeneratorPoisonBucket verifies that a bucket which fails
// deterministically on every tick does not stall generation of later
// buckets.
func TestGeneratorPoisonBucket(t *testing.T) {
t.Parallel()
startTime := time.Date(2025, 3, 10, 14, 0, 0, 0, time.UTC)
ctx := testutil.Context(t, testutil.WaitLong)
sink := &warnSink{}
log := slogtest.Make(t, nil).AppendSinks(sink)
h := newGeneratorHarness(t)
clock := quartz.NewMock(t)
clock.Set(startTime)
var (
poisonBucket = time.Date(2025, 3, 10, 10, 0, 0, 0, time.UTC)
laterBucket = time.Date(2025, 3, 10, 11, 0, 0, 0, time.UTC)
)
// Nothing prevents a negative runtime_ms row in the database, and a
// negative sum fails event validation on every tick.
h.insertRuntimeMessage(ctx, t, h.chat.ID, -5, poisonBucket.Add(10*time.Minute), false)
h.insertRuntimeMessage(ctx, t, h.chat.ID, 1000, laterBucket.Add(10*time.Minute), false)
trap := clock.Trap().NewTimer(generatorTimerName)
defer trap.Close()
gen := usage.NewGenerator(clock, log, h.authzDB, usage.NewDBInserter())
gen.Start(ctx)
defer gen.Close()
call := trap.MustWait(ctx)
call.MustRelease(ctx)
clock.Advance(call.Duration).MustWait(ctx)
call = trap.MustWait(ctx)
call.MustRelease(ctx)
// Every bucket except the poison bucket is generated, including buckets
// after it, and the failure is logged.
runtimes, _ := h.fetchRuntimeEvents(ctx, t)
expected := expectedBuckets(
startTime.Add(-usage.AgentRuntimeWindow),
startTime.Add(-2*time.Hour),
map[time.Time]int64{laterBucket: 1000},
)
delete(expected, poisonBucket)
require.Equal(t, expected, runtimes)
require.NotZero(t, sink.count.Load(), "the poison bucket failure must be logged")
}
+9 -2
View File
@@ -3,6 +3,7 @@ package usage
import (
"context"
"encoding/json"
"time"
"github.com/google/uuid"
"golang.org/x/xerrors"
@@ -68,10 +69,16 @@ func (i *dbInserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.S
}
// InsertHeartbeatUsageEvent implements agplusage.Inserter.
func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error {
func (*dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, createdAt time.Time, event usagetypes.HeartbeatEvent) error {
if !event.EventType().IsHeartbeat() {
return xerrors.Errorf("event type %q is not a heartbeat event", event.EventType())
}
// A zero createdAt stores the row at year 1, where bucket reconciliation
// can never match it again while the deterministic id turns every retry
// into a no-op, silently forfeiting the bucket's usage.
if createdAt.IsZero() {
return xerrors.Errorf("createdAt must be set for %q event", event.EventType())
}
if err := event.Valid(); err != nil {
return xerrors.Errorf("invalid %q event: %w", event.EventType(), err)
}
@@ -87,6 +94,6 @@ func (i *dbInserter) InsertHeartbeatUsageEvent(ctx context.Context, tx database.
ID: id,
EventType: string(event.EventType()),
EventData: jsonData,
CreatedAt: dbtime.Time(i.clock.Now()),
CreatedAt: dbtime.Time(createdAt),
})
}
+50
View File
@@ -68,6 +68,35 @@ func TestInserter(t *testing.T) {
}
})
t.Run("Heartbeat", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
inserter := usage.NewDBInserter()
// Heartbeat inserts must store the provided id and createdAt
// verbatim.
event := usagetypes.HBAgentRuntime{RuntimeMs: 1234}
eventJSON := jsoninate(t, event)
id := "hb_agent_runtime_v1:2025-01-02_03:00:00"
createdAt := time.Date(2025, 1, 2, 3, 0, 0, 0, time.UTC)
db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn(
func(ctx interface{}, params database.InsertUsageEventParams) error {
assert.Equal(t, id, params.ID)
assert.Equal(t, event.EventType(), usagetypes.UsageEventType(params.EventType))
assert.JSONEq(t, eventJSON, string(params.EventData))
assert.Equal(t, dbtime.Time(createdAt), params.CreatedAt)
return nil
},
).Times(1)
err := inserter.InsertHeartbeatUsageEvent(ctx, db, id, createdAt, event)
require.NoError(t, err)
})
t.Run("InvalidEvent", func(t *testing.T) {
t.Parallel()
@@ -81,5 +110,26 @@ func TestInserter(t *testing.T) {
Count: 0, // invalid
})
assert.ErrorContains(t, err, `invalid "dc_managed_agents_v1" event: count must be greater than 0`)
err = inserter.InsertHeartbeatUsageEvent(ctx, db, "some-id", time.Now(), usagetypes.HBAgentRuntime{
RuntimeMs: -1, // invalid
})
assert.ErrorContains(t, err, `invalid "hb_agent_runtime_v1" event: runtime_ms cannot be negative`)
})
t.Run("ZeroCreatedAt", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
ctrl := gomock.NewController(t)
// The mock store fails the test on any unexpected call, so no insert
// may reach the database.
db := dbmock.NewMockStore(ctrl)
inserter := usage.NewDBInserter()
err := inserter.InsertHeartbeatUsageEvent(ctx, db, "some-id", time.Time{}, usagetypes.HBAgentRuntime{
RuntimeMs: 1,
})
assert.ErrorContains(t, err, `createdAt must be set for "hb_agent_runtime_v1" event`)
})
}