Skip to content

Commit 1a601c3

Browse files
deansheatherdeansheather
authored
chore: move usage types to new package (#19103)
1 parent 6eb02d1 commit 1a601c3

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 (
@@ -2041,7 +2038,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
20412038
// Insert usage event for managed agents.
20422039
usageInserter := s.UsageInserter.Load()
20432040
if usageInserter != nil {
2044-
event := usage.DCManagedAgentsV1{
2041+
event := usagetypes.DCManagedAgentsV1{
20452042
Count: 1,
20462043
}
20472044
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
@@ -48,6 +48,7 @@ import (
4848
"github.com/coder/coder/v2/coderd/schedule/cron"
4949
"github.com/coder/coder/v2/coderd/telemetry"
5050
"github.com/coder/coder/v2/coderd/usage"
51+
"github.com/coder/coder/v2/coderd/usage/usagetypes"
5152
"github.com/coder/coder/v2/coderd/wspubsub"
5253
"github.com/coder/coder/v2/codersdk"
5354
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -3044,7 +3045,7 @@ func TestCompleteJob(t *testing.T) {
30443045
if tc.expectUsageEvent {
30453046
// Check that a usage event was collected.
30463047
require.Len(t, fakeUsageInserter.collectedEvents, 1)
3047-
require.Equal(t, usage.DCManagedAgentsV1{
3048+
require.Equal(t, usagetypes.DCManagedAgentsV1{
30483049
Count: 1,
30493050
}, fakeUsageInserter.collectedEvents[0])
30503051
} else {
@@ -4226,7 +4227,7 @@ func (s *fakeStream) cancel() {
42264227
}
42274228

42284229
type fakeUsageInserter struct {
4229-
collectedEvents []usage.Event
4230+
collectedEvents []usagetypes.Event
42304231
}
42314232

42324233
var _ usage.Inserter = &fakeUsageInserter{}
@@ -4239,7 +4240,7 @@ func newFakeUsageInserter() (*fakeUsageInserter, *atomic.Pointer[usage.Inserter]
42394240
return fake, ptr
42404241
}
42414242

4242-
func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usage.DiscreteEvent) error {
4243+
func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
42434244
f.collectedEvents = append(f.collectedEvents, event)
42444245
return nil
42454246
}

‎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,6 +4,7 @@ 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.
@@ -12,7 +13,7 @@ type Inserter interface {
1213
// within the given transaction.
1314
// The caller context must be authorized to create usage events in the
1415
// database.
15-
InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event DiscreteEvent) error
16+
InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event usagetypes.DiscreteEvent) error
1617
}
1718

1819
// AGPLInserter is a no-op implementation of Inserter.
@@ -26,6 +27,6 @@ func NewAGPLInserter() Inserter {
2627

2728
// InsertDiscreteUsageEvent is a no-op implementation of
2829
// InsertDiscreteUsageEvent.
29-
func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ DiscreteEvent) error {
30+
func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ usagetypes.DiscreteEvent) error {
3031
return nil
3132
}

‎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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy