Skip to content

Commit 6b2f72f

Browse files
deansheatherdeansheather
committed
chore: move usage types to new package
1 parent 7c36cc0 commit 6b2f72f

File tree

12 files changed

+470
-225
lines changed

12 files changed

+470
-225
lines changed

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,14 +28,6 @@ import (
2828
protobuf "google.golang.org/protobuf/proto"
2929

3030
"cdr.dev/slog"
31-
32-
"github.com/coder/coder/v2/coderd/usage"
33-
"github.com/coder/coder/v2/coderd/util/slice"
34-
35-
"github.com/coder/coder/v2/codersdk/drpcsdk"
36-
37-
"github.com/coder/quartz"
38-
3931
"github.com/coder/coder/v2/coderd/apikey"
4032
"github.com/coder/coder/v2/coderd/audit"
4133
"github.com/coder/coder/v2/coderd/database"
@@ -49,13 +41,18 @@ import (
4941
"github.com/coder/coder/v2/coderd/schedule"
5042
"github.com/coder/coder/v2/coderd/telemetry"
5143
"github.com/coder/coder/v2/coderd/tracing"
44+
"github.com/coder/coder/v2/coderd/usage"
45+
"github.com/coder/coder/v2/coderd/usage/usagetypes"
46+
"github.com/coder/coder/v2/coderd/util/slice"
5247
"github.com/coder/coder/v2/coderd/wspubsub"
5348
"github.com/coder/coder/v2/codersdk"
5449
"github.com/coder/coder/v2/codersdk/agentsdk"
50+
"github.com/coder/coder/v2/codersdk/drpcsdk"
5551
"github.com/coder/coder/v2/provisioner"
5652
"github.com/coder/coder/v2/provisionerd/proto"
5753
"github.com/coder/coder/v2/provisionersdk"
5854
sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
55+
"github.com/coder/quartz"
5956
)
6057

6158
const (
@@ -2037,7 +2034,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
20372034
// Insert usage event for managed agents.
20382035
usageInserter := s.UsageInserter.Load()
20392036
if usageInserter != nil {
2040-
event := usage.DCManagedAgentsV1{
2037+
event := usagetypes.DCManagedAgentsV1{
20412038
Count: 1,
20422039
}
20432040
err = (*usageInserter).InsertDiscreteUsageEvent(ctx, db, event)

‎coderd/provisionerdserver/provisionerdserver_test.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import (
4545
"github.com/coder/coder/v2/coderd/schedule/cron"
4646
"github.com/coder/coder/v2/coderd/telemetry"
4747
"github.com/coder/coder/v2/coderd/usage"
48+
"github.com/coder/coder/v2/coderd/usage/usagetypes"
4849
"github.com/coder/coder/v2/coderd/wspubsub"
4950
"github.com/coder/coder/v2/codersdk"
5051
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -2924,7 +2925,7 @@ func TestCompleteJob(t *testing.T) {
29242925

29252926
// Check that a usage event was collected.
29262927
require.Len(t, fakeUsageInserter.collectedEvents, 1)
2927-
require.Equal(t, usage.DCManagedAgentsV1{
2928+
require.Equal(t, usagetypes.DCManagedAgentsV1{
29282929
Count: 1,
29292930
}, fakeUsageInserter.collectedEvents[0])
29302931
} else {
@@ -4102,7 +4103,7 @@ func (s *fakeStream) cancel() {
41024103
}
41034104

41044105
type fakeUsageInserter struct {
4105-
collectedEvents []usage.Event
4106+
collectedEvents []usagetypes.Event
41064107
}
41074108

41084109
var _ usage.Inserter = &fakeUsageInserter{}
@@ -4115,7 +4116,7 @@ func newFakeUsageInserter() (*fakeUsageInserter, *atomic.Pointer[usage.Inserter]
41154116
return fake, ptr
41164117
}
41174118

4118-
func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usage.DiscreteEvent) error {
4119+
func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
41194120
f.collectedEvents = append(f.collectedEvents, event)
41204121
return nil
41214122
}

‎coderd/usage/events.go

Lines changed: 0 additions & 82 deletions
This file was deleted.

‎coderd/usage/inserter.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,14 @@ import (
44
"context"
55

66
"github.com/coder/coder/v2/coderd/database"
7+
"github.com/coder/coder/v2/coderd/usage/usagetypes"
78
)
89

910
// Inserter accepts usage events generated by the product.
1011
type Inserter interface {
1112
// InsertDiscreteUsageEvent writes a discrete usage event to the database
1213
// within the given transaction.
13-
InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event DiscreteEvent) error
14+
InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event usagetypes.DiscreteEvent) error
1415
}
1516

1617
// AGPLInserter is a no-op implementation of Inserter.
@@ -24,6 +25,6 @@ func NewAGPLInserter() Inserter {
2425

2526
// InsertDiscreteUsageEvent is a no-op implementation of
2627
// InsertDiscreteUsageEvent.
27-
func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ DiscreteEvent) error {
28+
func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ usagetypes.DiscreteEvent) error {
2829
return nil
2930
}

‎coderd/usage/usagetypes/events.go

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
// Package usagetypes contains the types for usage events. These are kept in
2+
// their own package to avoid importing any real code from coderd.
3+
//
4+
// Imports in this package should be limited to the standard library and the
5+
// following packages ONLY:
6+
// - github.com/google/uuid
7+
// - golang.org/x/xerrors
8+
//
9+
// This package is imported by the Tallyman codebase.
10+
package usagetypes
11+
12+
// Please read the package documentation before adding imports.
13+
import (
14+
"bytes"
15+
"encoding/json"
16+
"strings"
17+
18+
"golang.org/x/xerrors"
19+
)
20+
21+
// UsageEventType is an enum of all usage event types. It mirrors the database
22+
// type `usage_event_type`.
23+
type UsageEventType string
24+
25+
const (
26+
UsageEventTypeDCManagedAgentsV1 UsageEventType = "dc_managed_agents_v1"
27+
)
28+
29+
func (e UsageEventType) Valid() bool {
30+
switch e {
31+
case UsageEventTypeDCManagedAgentsV1:
32+
return true
33+
default:
34+
return false
35+
}
36+
}
37+
38+
func (e UsageEventType) IsDiscrete() bool {
39+
return e.Valid() && strings.HasPrefix(string(e), "dc_")
40+
}
41+
42+
func (e UsageEventType) IsHeartbeat() bool {
43+
return e.Valid() && strings.HasPrefix(string(e), "hb_")
44+
}
45+
46+
// ParseEvent parses the raw event data into the specified Go type. It fails if
47+
// there is any unknown fields or extra data after the event. The returned event
48+
// is validated.
49+
func ParseEvent[T Event](data json.RawMessage) (T, error) {
50+
dec := json.NewDecoder(bytes.NewReader(data))
51+
dec.DisallowUnknownFields()
52+
53+
var event T
54+
err := dec.Decode(&event)
55+
if err != nil {
56+
return event, xerrors.Errorf("unmarshal %T event: %w", event, err)
57+
}
58+
if dec.More() {
59+
return event, xerrors.Errorf("extra data after %T event", event)
60+
}
61+
err = event.Valid()
62+
if err != nil {
63+
return event, xerrors.Errorf("invalid %T event: %w", event, err)
64+
}
65+
66+
return event, nil
67+
}
68+
69+
// ParseEventWithType parses the raw event data into the specified Go type. It
70+
// fails if there is any unknown fields or extra data after the event. The
71+
// returned event is validated.
72+
func ParseEventWithType(eventType UsageEventType, data json.RawMessage) (Event, error) {
73+
switch eventType {
74+
case UsageEventTypeDCManagedAgentsV1:
75+
return ParseEvent[DCManagedAgentsV1](data)
76+
default:
77+
return nil, xerrors.Errorf("unknown event type: %s", eventType)
78+
}
79+
}
80+
81+
// Event is a usage event that can be collected by the usage collector.
82+
//
83+
// Note that the following event types should not be updated once they are
84+
// merged into the product. Please consult Dean before making any changes.
85+
//
86+
// This type cannot be implemented outside of this package as it this package
87+
// is the source of truth for the coder/tallyman repo.
88+
type Event interface {
89+
usageEvent() // to prevent external types from implementing this interface
90+
EventType() UsageEventType
91+
Valid() error
92+
Fields() map[string]any // fields to be marshaled and sent to tallyman/Metronome
93+
}
94+
95+
// DiscreteEvent is a usage event that is collected as a discrete event.
96+
type DiscreteEvent interface {
97+
Event
98+
discreteUsageEvent() // marker method, also prevents external types from implementing this interface
99+
}
100+
101+
// DCManagedAgentsV1 is a discrete usage event for the number of managed agents.
102+
// This event is sent in the following situations:
103+
// - Once on first startup after usage tracking is added to the product with
104+
// the count of all existing managed agents (count=N)
105+
// - A new managed agent is created (count=1)
106+
type DCManagedAgentsV1 struct {
107+
Count uint64 `json:"count"`
108+
}
109+
110+
var _ DiscreteEvent = DCManagedAgentsV1{}
111+
112+
func (DCManagedAgentsV1) usageEvent() {}
113+
func (DCManagedAgentsV1) discreteUsageEvent() {}
114+
func (DCManagedAgentsV1) EventType() UsageEventType {
115+
return UsageEventTypeDCManagedAgentsV1
116+
}
117+
118+
func (e DCManagedAgentsV1) Valid() error {
119+
if e.Count == 0 {
120+
return xerrors.New("count must be greater than 0")
121+
}
122+
return nil
123+
}
124+
125+
func (e DCManagedAgentsV1) Fields() map[string]any {
126+
return map[string]any{
127+
"count": e.Count,
128+
}
129+
}

‎coderd/usage/usagetypes/events_test.go

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
package usagetypes_test
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/require"
7+
8+
"github.com/coder/coder/v2/coderd/usage/usagetypes"
9+
)
10+
11+
func TestParseEvent(t *testing.T) {
12+
t.Parallel()
13+
14+
t.Run("ExtraFields", func(t *testing.T) {
15+
t.Parallel()
16+
_, err := usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": 1, "extra": "field"}`))
17+
require.ErrorContains(t, err, "unmarshal usagetypes.DCManagedAgentsV1 event")
18+
})
19+
20+
t.Run("ExtraData", func(t *testing.T) {
21+
t.Parallel()
22+
_, err := usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": 1}{"count": 2}`))
23+
require.ErrorContains(t, err, "extra data after usagetypes.DCManagedAgentsV1 event")
24+
})
25+
26+
t.Run("DCManagedAgentsV1", func(t *testing.T) {
27+
t.Parallel()
28+
29+
event, err := usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": 1}`))
30+
require.NoError(t, err)
31+
require.Equal(t, usagetypes.DCManagedAgentsV1{Count: 1}, event)
32+
require.Equal(t, map[string]any{"count": uint64(1)}, event.Fields())
33+
34+
_, err = usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": "invalid"}`))
35+
require.ErrorContains(t, err, "unmarshal usagetypes.DCManagedAgentsV1 event")
36+
37+
_, err = usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{}`))
38+
require.ErrorContains(t, err, "invalid usagetypes.DCManagedAgentsV1 event: count must be greater than 0")
39+
})
40+
}
41+
42+
func TestParseEventWithType(t *testing.T) {
43+
t.Parallel()
44+
45+
t.Run("UnknownEvent", func(t *testing.T) {
46+
t.Parallel()
47+
_, err := usagetypes.ParseEventWithType(usagetypes.UsageEventType("fake"), []byte(`{}`))
48+
require.ErrorContains(t, err, "unknown event type: fake")
49+
})
50+
51+
t.Run("DCManagedAgentsV1", func(t *testing.T) {
52+
t.Parallel()
53+
54+
eventType := usagetypes.UsageEventTypeDCManagedAgentsV1
55+
event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"count": 1}`))
56+
require.NoError(t, err)
57+
require.Equal(t, usagetypes.DCManagedAgentsV1{Count: 1}, event)
58+
require.Equal(t, eventType, event.EventType())
59+
require.Equal(t, map[string]any{"count": uint64(1)}, event.Fields())
60+
})
61+
}

0 commit comments

Comments
 (0)