Skip to content

Commit 042a5e2

Browse files
deansheatherdeansheather
committed
chore: move usage types to new package
1 parent 89b1757 commit 042a5e2

File tree

13 files changed

+453
-188
lines changed

13 files changed

+453
-188
lines changed

‎coderd/database/modelmethods.go

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"encoding/hex"
55
"sort"
66
"strconv"
7-
"strings"
87
"time"
98

109
"github.com/google/uuid"
@@ -629,11 +628,3 @@ func (m WorkspaceAgentVolumeResourceMonitor) Debounce(
629628

630629
return m.DebouncedUntil, false
631630
}
632-
633-
func (e UsageEventType) IsDiscrete() bool {
634-
return e.Valid() && strings.HasPrefix(string(e), "dc_")
635-
}
636-
637-
func (e UsageEventType) IsHeartbeat() bool {
638-
return e.Valid() && strings.HasPrefix(string(e), "hb_")
639-
}

‎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 (
@@ -1903,7 +1900,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
19031900
// Collect usage event for managed agents.
19041901
usageCollector := s.UsageCollector.Load()
19051902
if usageCollector != nil {
1906-
event := usage.DCManagedAgentsV1{
1903+
event := usagetypes.DCManagedAgentsV1{
19071904
Count: 1,
19081905
}
19091906
err = (*usageCollector).CollectDiscreteUsageEvent(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"
@@ -2678,7 +2679,7 @@ func TestCompleteJob(t *testing.T) {
26782679

26792680
// Check that a usage event was collected.
26802681
require.Len(t, fakeUsageCollector.collectedEvents, 1)
2681-
require.Equal(t, usage.DCManagedAgentsV1{
2682+
require.Equal(t, usagetypes.DCManagedAgentsV1{
26822683
Count: 1,
26832684
}, fakeUsageCollector.collectedEvents[0])
26842685
} else {
@@ -3850,7 +3851,7 @@ func (s *fakeStream) cancel() {
38503851
}
38513852

38523853
type fakeUsageCollector struct {
3853-
collectedEvents []usage.Event
3854+
collectedEvents []usagetypes.Event
38543855
}
38553856

38563857
var _ usage.Collector = &fakeUsageCollector{}
@@ -3863,7 +3864,7 @@ func newFakeUsageCollector() (*fakeUsageCollector, *atomic.Pointer[usage.Collect
38633864
return fake, ptr
38643865
}
38653866

3866-
func (f *fakeUsageCollector) CollectDiscreteUsageEvent(_ context.Context, _ database.Store, event usage.DiscreteEvent) error {
3867+
func (f *fakeUsageCollector) CollectDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
38673868
f.collectedEvents = append(f.collectedEvents, event)
38683869
return nil
38693870
}

‎coderd/usage/collector.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
// Collector is a sink for usage events generated by the product.
1011
type Collector interface {
1112
// CollectDiscreteUsageEvent writes a discrete usage event to the database
1213
// with the given database or transaction.
13-
CollectDiscreteUsageEvent(ctx context.Context, db database.Store, event DiscreteEvent) error
14+
CollectDiscreteUsageEvent(ctx context.Context, db database.Store, event usagetypes.DiscreteEvent) error
1415
}
1516

1617
// AGPLCollector is a no-op implementation of Collector.
@@ -24,6 +25,6 @@ func NewAGPLCollector() Collector {
2425

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

‎coderd/usage/events.go

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

‎coderd/usage/usagetypes/events.go

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
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+
type Event interface {
86+
usageEvent() // to prevent external types from implementing this interface
87+
EventType() UsageEventType
88+
Valid() error
89+
}
90+
91+
// DiscreteEvent is a usage event that is collected as a discrete event.
92+
type DiscreteEvent interface {
93+
Event
94+
discreteUsageEvent() // marker method, also prevents external types from implementing this interface
95+
}
96+
97+
// DCManagedAgentsV1 is a discrete usage event for the number of managed agents.
98+
// This event is sent in the following situations:
99+
// - Once on first startup after usage tracking is added to the product with
100+
// the count of all existing managed agents (count=N)
101+
// - A new managed agent is created (count=1)
102+
type DCManagedAgentsV1 struct {
103+
Count uint64 `json:"count"`
104+
}
105+
106+
var _ DiscreteEvent = DCManagedAgentsV1{}
107+
108+
func (DCManagedAgentsV1) usageEvent() {}
109+
func (DCManagedAgentsV1) discreteUsageEvent() {}
110+
func (DCManagedAgentsV1) EventType() UsageEventType {
111+
return UsageEventTypeDCManagedAgentsV1
112+
}
113+
114+
func (e DCManagedAgentsV1) Valid() error {
115+
if e.Count == 0 {
116+
return xerrors.New("count must be greater than 0")
117+
}
118+
return nil
119+
}

‎coderd/usage/usagetypes/events_test.go

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
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+
33+
_, err = usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": "invalid"}`))
34+
require.ErrorContains(t, err, "unmarshal usagetypes.DCManagedAgentsV1 event")
35+
36+
_, err = usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{}`))
37+
require.ErrorContains(t, err, "invalid usagetypes.DCManagedAgentsV1 event: count must be greater than 0")
38+
})
39+
}
40+
41+
func TestParseEventWithType(t *testing.T) {
42+
t.Parallel()
43+
44+
t.Run("UnknownEvent", func(t *testing.T) {
45+
t.Parallel()
46+
_, err := usagetypes.ParseEventWithType(usagetypes.UsageEventType("fake"), []byte(`{}`))
47+
require.ErrorContains(t, err, "unknown event type: fake")
48+
})
49+
50+
t.Run("DCManagedAgentsV1", func(t *testing.T) {
51+
t.Parallel()
52+
53+
eventType := usagetypes.UsageEventTypeDCManagedAgentsV1
54+
event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"count": 1}`))
55+
require.NoError(t, err)
56+
require.Equal(t, usagetypes.DCManagedAgentsV1{Count: 1}, event)
57+
require.Equal(t, eventType, event.EventType())
58+
})
59+
}

0 commit comments

Comments
 (0)