chore: move usage types to new package (#19103)

coderd/provisionerdserver/provisionerdserver.go

@@ -28,14 +28,6 @@ import (
 	protobuf "google.golang.org/protobuf/proto"
 
 	"cdr.dev/slog"
-
-	"github.com/coder/coder/v2/coderd/usage"
-	"github.com/coder/coder/v2/coderd/util/slice"
-
-	"github.com/coder/coder/v2/codersdk/drpcsdk"
-
-	"github.com/coder/quartz"
-
 	"github.com/coder/coder/v2/coderd/apikey"
 	"github.com/coder/coder/v2/coderd/audit"
 	"github.com/coder/coder/v2/coderd/database"
@@ -49,13 +41,18 @@ import (
 	"github.com/coder/coder/v2/coderd/schedule"
 	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/coderd/usage"
+	"github.com/coder/coder/v2/coderd/usage/usagetypes"
+	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/coderd/wspubsub"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
 	"github.com/coder/coder/v2/provisioner"
 	"github.com/coder/coder/v2/provisionerd/proto"
 	"github.com/coder/coder/v2/provisionersdk"
 	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+	"github.com/coder/quartz"
 )
 
 const (
@@ -2041,7 +2038,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 			// Insert usage event for managed agents.
 			usageInserter := s.UsageInserter.Load()
 			if usageInserter != nil {
-				event := usage.DCManagedAgentsV1{
+				event := usagetypes.DCManagedAgentsV1{
 					Count: 1,
 				}
 				err = (*usageInserter).InsertDiscreteUsageEvent(ctx, db, event)

coderd/provisionerdserver/provisionerdserver_test.go

@@ -48,6 +48,7 @@ import (
 	"github.com/coder/coder/v2/coderd/schedule/cron"
 	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/coderd/usage"
+	"github.com/coder/coder/v2/coderd/usage/usagetypes"
 	"github.com/coder/coder/v2/coderd/wspubsub"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -3044,7 +3045,7 @@ func TestCompleteJob(t *testing.T) {
 					if tc.expectUsageEvent {
 						// Check that a usage event was collected.
 						require.Len(t, fakeUsageInserter.collectedEvents, 1)
-						require.Equal(t, usage.DCManagedAgentsV1{
+						require.Equal(t, usagetypes.DCManagedAgentsV1{
 							Count: 1,
 						}, fakeUsageInserter.collectedEvents[0])
 					} else {
@@ -4226,7 +4227,7 @@ func (s *fakeStream) cancel() {
 }
 
 type fakeUsageInserter struct {
-	collectedEvents []usage.Event
+	collectedEvents []usagetypes.Event
 }
 
 var _ usage.Inserter = &fakeUsageInserter{}
@@ -4239,7 +4240,7 @@ func newFakeUsageInserter() (*fakeUsageInserter, *atomic.Pointer[usage.Inserter]
 	return fake, ptr
 }
 
-func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usage.DiscreteEvent) error {
+func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
 	f.collectedEvents = append(f.collectedEvents, event)
 	return nil
 }

coderd/usage/events.go

@@ -1,82 +0,0 @@
-package usage
-
-import (
-	"strings"
-
-	"golang.org/x/xerrors"
-)
-
-// EventType is an enum of all usage event types. It mirrors the check
-// constraint on the `event_type` column in the `usage_events` table.
-type EventType string //nolint:revive
-
-const (
-	UsageEventTypeDCManagedAgentsV1 EventType = "dc_managed_agents_v1"
-)
-
-func (e EventType) Valid() bool {
-	switch e {
-	case UsageEventTypeDCManagedAgentsV1:
-		return true
-	default:
-		return false
-	}
-}
-
-func (e EventType) IsDiscrete() bool {
-	return e.Valid() && strings.HasPrefix(string(e), "dc_")
-}
-
-func (e EventType) IsHeartbeat() bool {
-	return e.Valid() && strings.HasPrefix(string(e), "hb_")
-}
-
-// Event is a usage event that can be collected by the usage collector.
-//
-// Note that the following event types should not be updated once they are
-// merged into the product. Please consult Dean before making any changes.
-//
-// Event types cannot be implemented outside of this package, as they are
-// imported by the coder/tallyman repository.
-type Event interface {
-	usageEvent() // to prevent external types from implementing this interface
-	EventType() EventType
-	Valid() error
-	Fields() map[string]any // fields to be marshaled and sent to tallyman/Metronome
-}
-
-// DiscreteEvent is a usage event that is collected as a discrete event.
-type DiscreteEvent interface {
-	Event
-	discreteUsageEvent() // marker method, also prevents external types from implementing this interface
-}
-
-// DCManagedAgentsV1 is a discrete usage event for the number of managed agents.
-// This event is sent in the following situations:
-//   - Once on first startup after usage tracking is added to the product with
-//     the count of all existing managed agents (count=N)
-//   - A new managed agent is created (count=1)
-type DCManagedAgentsV1 struct {
-	Count uint64 `json:"count"`
-}
-
-var _ DiscreteEvent = DCManagedAgentsV1{}
-
-func (DCManagedAgentsV1) usageEvent()         {}
-func (DCManagedAgentsV1) discreteUsageEvent() {}
-func (DCManagedAgentsV1) EventType() EventType {
-	return UsageEventTypeDCManagedAgentsV1
-}
-
-func (e DCManagedAgentsV1) Valid() error {
-	if e.Count == 0 {
-		return xerrors.New("count must be greater than 0")
-	}
-	return nil
-}
-
-func (e DCManagedAgentsV1) Fields() map[string]any {
-	return map[string]any{
-		"count": e.Count,
-	}
-}

coderd/usage/inserter.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/usage/usagetypes"
 )
 
 // Inserter accepts usage events generated by the product.
@@ -12,7 +13,7 @@ type Inserter interface {
 	// within the given transaction.
 	// The caller context must be authorized to create usage events in the
 	// database.
-	InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event DiscreteEvent) error
+	InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event usagetypes.DiscreteEvent) error
 }
 
 // AGPLInserter is a no-op implementation of Inserter.
@@ -26,6 +27,6 @@ func NewAGPLInserter() Inserter {
 
 // InsertDiscreteUsageEvent is a no-op implementation of
 // InsertDiscreteUsageEvent.
-func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ DiscreteEvent) error {
+func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ usagetypes.DiscreteEvent) error {
 	return nil
 }

coderd/usage/usagetypes/events.go

@@ -0,0 +1,129 @@
+// Package usagetypes contains the types for usage events. These are kept in
+// their own package to avoid importing any real code from coderd.
+//
+// Imports in this package should be limited to the standard library and the
+// following packages ONLY:
+//   - github.com/google/uuid
+//   - golang.org/x/xerrors
+//
+// This package is imported by the Tallyman codebase.
+package usagetypes
+
+// Please read the package documentation before adding imports.
+import (
+	"bytes"
+	"encoding/json"
+	"strings"
+
+	"golang.org/x/xerrors"
+)
+
+// UsageEventType is an enum of all usage event types. It mirrors the database
+// type `usage_event_type`.
+type UsageEventType string
+
+const (
+	UsageEventTypeDCManagedAgentsV1 UsageEventType = "dc_managed_agents_v1"
+)
+
+func (e UsageEventType) Valid() bool {
+	switch e {
+	case UsageEventTypeDCManagedAgentsV1:
+		return true
+	default:
+		return false
+	}
+}
+
+func (e UsageEventType) IsDiscrete() bool {
+	return e.Valid() && strings.HasPrefix(string(e), "dc_")
+}
+
+func (e UsageEventType) IsHeartbeat() bool {
+	return e.Valid() && strings.HasPrefix(string(e), "hb_")
+}
+
+// ParseEvent parses the raw event data into the specified Go type. It fails if
+// there is any unknown fields or extra data after the event. The returned event
+// is validated.
+func ParseEvent[T Event](data json.RawMessage) (T, error) {
+	dec := json.NewDecoder(bytes.NewReader(data))
+	dec.DisallowUnknownFields()
+
+	var event T
+	err := dec.Decode(&event)
+	if err != nil {
+		return event, xerrors.Errorf("unmarshal %T event: %w", event, err)
+	}
+	if dec.More() {
+		return event, xerrors.Errorf("extra data after %T event", event)
+	}
+	err = event.Valid()
+	if err != nil {
+		return event, xerrors.Errorf("invalid %T event: %w", event, err)
+	}
+
+	return event, nil
+}
+
+// ParseEventWithType parses the raw event data into the specified Go type. It
+// fails if there is any unknown fields or extra data after the event. The
+// returned event is validated.
+func ParseEventWithType(eventType UsageEventType, data json.RawMessage) (Event, error) {
+	switch eventType {
+	case UsageEventTypeDCManagedAgentsV1:
+		return ParseEvent[DCManagedAgentsV1](data)
+	default:
+		return nil, xerrors.Errorf("unknown event type: %s", eventType)
+	}
+}
+
+// Event is a usage event that can be collected by the usage collector.
+//
+// Note that the following event types should not be updated once they are
+// merged into the product. Please consult Dean before making any changes.
+//
+// This type cannot be implemented outside of this package as it this package
+// is the source of truth for the coder/tallyman repo.
+type Event interface {
+	usageEvent() // to prevent external types from implementing this interface
+	EventType() UsageEventType
+	Valid() error
+	Fields() map[string]any // fields to be marshaled and sent to tallyman/Metronome
+}
+
+// DiscreteEvent is a usage event that is collected as a discrete event.
+type DiscreteEvent interface {
+	Event
+	discreteUsageEvent() // marker method, also prevents external types from implementing this interface
+}
+
+// DCManagedAgentsV1 is a discrete usage event for the number of managed agents.
+// This event is sent in the following situations:
+//   - Once on first startup after usage tracking is added to the product with
+//     the count of all existing managed agents (count=N)
+//   - A new managed agent is created (count=1)
+type DCManagedAgentsV1 struct {
+	Count uint64 `json:"count"`
+}
+
+var _ DiscreteEvent = DCManagedAgentsV1{}
+
+func (DCManagedAgentsV1) usageEvent()         {}
+func (DCManagedAgentsV1) discreteUsageEvent() {}
+func (DCManagedAgentsV1) EventType() UsageEventType {
+	return UsageEventTypeDCManagedAgentsV1
+}
+
+func (e DCManagedAgentsV1) Valid() error {
+	if e.Count == 0 {
+		return xerrors.New("count must be greater than 0")
+	}
+	return nil
+}
+
+func (e DCManagedAgentsV1) Fields() map[string]any {
+	return map[string]any{
+		"count": e.Count,
+	}
+}

coderd/usage/usagetypes/events_test.go

@@ -0,0 +1,61 @@
+package usagetypes_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/usage/usagetypes"
+)
+
+func TestParseEvent(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ExtraFields", func(t *testing.T) {
+		t.Parallel()
+		_, err := usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": 1, "extra": "field"}`))
+		require.ErrorContains(t, err, "unmarshal usagetypes.DCManagedAgentsV1 event")
+	})
+
+	t.Run("ExtraData", func(t *testing.T) {
+		t.Parallel()
+		_, err := usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": 1}{"count": 2}`))
+		require.ErrorContains(t, err, "extra data after usagetypes.DCManagedAgentsV1 event")
+	})
+
+	t.Run("DCManagedAgentsV1", func(t *testing.T) {
+		t.Parallel()
+
+		event, err := usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": 1}`))
+		require.NoError(t, err)
+		require.Equal(t, usagetypes.DCManagedAgentsV1{Count: 1}, event)
+		require.Equal(t, map[string]any{"count": uint64(1)}, event.Fields())
+
+		_, err = usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{"count": "invalid"}`))
+		require.ErrorContains(t, err, "unmarshal usagetypes.DCManagedAgentsV1 event")
+
+		_, err = usagetypes.ParseEvent[usagetypes.DCManagedAgentsV1]([]byte(`{}`))
+		require.ErrorContains(t, err, "invalid usagetypes.DCManagedAgentsV1 event: count must be greater than 0")
+	})
+}
+
+func TestParseEventWithType(t *testing.T) {
+	t.Parallel()
+
+	t.Run("UnknownEvent", func(t *testing.T) {
+		t.Parallel()
+		_, err := usagetypes.ParseEventWithType(usagetypes.UsageEventType("fake"), []byte(`{}`))
+		require.ErrorContains(t, err, "unknown event type: fake")
+	})
+
+	t.Run("DCManagedAgentsV1", func(t *testing.T) {
+		t.Parallel()
+
+		eventType := usagetypes.UsageEventTypeDCManagedAgentsV1
+		event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"count": 1}`))
+		require.NoError(t, err)
+		require.Equal(t, usagetypes.DCManagedAgentsV1{Count: 1}, event)
+		require.Equal(t, eventType, event.EventType())
+		require.Equal(t, map[string]any{"count": uint64(1)}, event.Fields())
+	})
+}

coderd/usage/usagetypes/tallyman.go

@@ -0,0 +1,70 @@
+package usagetypes
+
+// Please read the package documentation before adding imports.
+import (
+	"encoding/json"
+	"time"
+
+	"golang.org/x/xerrors"
+)
+
+const (
+	TallymanCoderLicenseKeyHeader   = "Coder-License-Key"
+	TallymanCoderDeploymentIDHeader = "Coder-Deployment-ID"
+)
+
+// TallymanV1Response is a generic response with a message from the Tallyman
+// API. It is typically returned when there is an error.
+type TallymanV1Response struct {
+	Message string `json:"message"`
+}
+
+// TallymanV1IngestRequest is a request to the Tallyman API to ingest usage
+// events.
+type TallymanV1IngestRequest struct {
+	Events []TallymanV1IngestEvent `json:"events"`
+}
+
+// TallymanV1IngestEvent is an event to be ingested into the Tallyman API.
+type TallymanV1IngestEvent struct {
+	ID        string          `json:"id"`
+	EventType UsageEventType  `json:"event_type"`
+	EventData json.RawMessage `json:"event_data"`
+	CreatedAt time.Time       `json:"created_at"`
+}
+
+// Valid validates the TallymanV1IngestEvent. It does not validate the event
+// body.
+func (e TallymanV1IngestEvent) Valid() error {
+	if e.ID == "" {
+		return xerrors.New("id is required")
+	}
+	if !e.EventType.Valid() {
+		return xerrors.Errorf("event_type %q is invalid", e.EventType)
+	}
+	if e.CreatedAt.IsZero() {
+		return xerrors.New("created_at cannot be zero")
+	}
+	return nil
+}
+
+// TallymanV1IngestResponse is a response from the Tallyman API to ingest usage
+// events.
+type TallymanV1IngestResponse struct {
+	AcceptedEvents []TallymanV1IngestAcceptedEvent `json:"accepted_events"`
+	RejectedEvents []TallymanV1IngestRejectedEvent `json:"rejected_events"`
+}
+
+// TallymanV1IngestAcceptedEvent is an event that was accepted by the Tallyman
+// API.
+type TallymanV1IngestAcceptedEvent struct {
+	ID string `json:"id"`
+}
+
+// TallymanV1IngestRejectedEvent is an event that was rejected by the Tallyman
+// API.
+type TallymanV1IngestRejectedEvent struct {
+	ID        string `json:"id"`
+	Message   string `json:"message"`
+	Permanent bool   `json:"permanent"`
+}

coderd/usage/usagetypes/tallyman_test.go

@@ -0,0 +1,85 @@
+package usagetypes_test
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/usage/usagetypes"
+)
+
+func TestTallymanV1UsageEvent(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name         string
+		event        usagetypes.TallymanV1IngestEvent
+		errorMessage string
+	}{
+		{
+			name: "OK",
+			event: usagetypes.TallymanV1IngestEvent{
+				ID:        "123",
+				EventType: usagetypes.UsageEventTypeDCManagedAgentsV1,
+				// EventData is not validated.
+				EventData: json.RawMessage{},
+				CreatedAt: time.Now(),
+			},
+			errorMessage: "",
+		},
+		{
+			name: "NoID",
+			event: usagetypes.TallymanV1IngestEvent{
+				EventType: usagetypes.UsageEventTypeDCManagedAgentsV1,
+				EventData: json.RawMessage{},
+				CreatedAt: time.Now(),
+			},
+			errorMessage: "id is required",
+		},
+		{
+			name: "NoEventType",
+			event: usagetypes.TallymanV1IngestEvent{
+				ID:        "123",
+				EventType: usagetypes.UsageEventType(""),
+				EventData: json.RawMessage{},
+				CreatedAt: time.Now(),
+			},
+			errorMessage: `event_type "" is invalid`,
+		},
+		{
+			name: "UnknownEventType",
+			event: usagetypes.TallymanV1IngestEvent{
+				ID:        "123",
+				EventType: usagetypes.UsageEventType("unknown"),
+				EventData: json.RawMessage{},
+				CreatedAt: time.Now(),
+			},
+			errorMessage: `event_type "unknown" is invalid`,
+		},
+		{
+			name: "NoCreatedAt",
+			event: usagetypes.TallymanV1IngestEvent{
+				ID:        "123",
+				EventType: usagetypes.UsageEventTypeDCManagedAgentsV1,
+				EventData: json.RawMessage{},
+				CreatedAt: time.Time{},
+			},
+			errorMessage: "created_at cannot be zero",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			err := tc.event.Valid()
+			if tc.errorMessage == "" {
+				require.NoError(t, err)
+			} else {
+				require.ErrorContains(t, err, tc.errorMessage)
+			}
+		})
+	}
+}