Avoid slim database import

deansheather · ethanndickson · commit 2148b6263388 · 2024-06-28T07:21:14.000Z
diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go
@@ -22,7 +22,6 @@ import (
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/externalauth"
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
-	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/codersdk"
@@ -65,7 +64,7 @@ type Options struct {
 	AppearanceFetcher                 *atomic.Pointer[appearance.Fetcher]
 	PublishWorkspaceUpdateFn          func(ctx context.Context, workspaceID uuid.UUID)
 	PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
-	NetworkTelemetryBatchFn           func(batch []telemetry.NetworkEvent)
+	NetworkTelemetryBatchFn           func(batch []*tailnetproto.TelemetryEvent)
 
 	AccessURL                      *url.URL
 	AppHostname                    string
diff --git a/coderd/agentapi/api_test.go b/coderd/agentapi/api_test.go
@@ -20,10 +20,10 @@ import (
 	"github.com/coder/coder/v2/coderd/externalauth"
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/schedule"
-	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/tailnet"
+	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
 	"github.com/coder/coder/v2/tailnet/tailnettest"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -83,7 +83,7 @@ func Test_APIClose(t *testing.T) {
 		StatsReporter:            statsReporter,
 		AppearanceFetcher:        &appearanceFetcherPtr,
 		PublishWorkspaceUpdateFn: func(_ context.Context, _ uuid.UUID) {},
-		NetworkTelemetryBatchFn:  func(_ []telemetry.NetworkEvent) {},
+		NetworkTelemetryBatchFn:  func(_ []*tailnetproto.TelemetryEvent) {},
 		AccessURL: &url.URL{
 			Scheme: "http",
 			Host:   "localhost",
diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go
@@ -28,6 +28,7 @@ import (
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/tailnet"
+	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
 )
 
 // @Summary Workspace agent RPC API
@@ -135,9 +136,19 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		StatsReporter:                     api.statsReporter,
 		PublishWorkspaceUpdateFn:          api.publishWorkspaceUpdate,
 		PublishWorkspaceAgentLogsUpdateFn: api.publishWorkspaceAgentLogsUpdate,
-		NetworkTelemetryBatchFn: func(batch []telemetry.NetworkEvent) {
+		NetworkTelemetryBatchFn: func(batch []*tailnetproto.TelemetryEvent) {
+			telemetryEvents := make([]telemetry.NetworkEvent, 0, len(batch))
+			for _, pEvent := range batch {
+				tEvent, err := telemetry.NetworkEventFromProto(pEvent)
+				if err != nil {
+					// Events that fail to be converted get discarded for now.
+					continue
+				}
+				telemetryEvents = append(telemetryEvents, tEvent)
+			}
+
 			api.Telemetry.Report(&telemetry.Snapshot{
-				NetworkEvents: batch,
+				NetworkEvents: telemetryEvents,
 			})
 		},
 
diff --git a/tailnet/service.go b/tailnet/service.go
@@ -10,16 +10,14 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/hashicorp/yamux"
+	"golang.org/x/xerrors"
 	"storj.io/drpc/drpcmux"
 	"storj.io/drpc/drpcserver"
 	"tailscale.com/tailcfg"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/apiversion"
-	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/tailnet/proto"
-
-	"golang.org/x/xerrors"
 )
 
 type streamIDContextKey struct{}
@@ -125,13 +123,14 @@ type DRPCService struct {
 	DerpMapFn                      func() *tailcfg.DERPMap
 	NetworkTelemetryBatchFrequency time.Duration
 	NetworkTelemetryBatchMaxSize   int
-	NetworkTelemetryBatchFn        func(batch []telemetry.NetworkEvent)
+	NetworkTelemetryBatchFn        func(batch []*proto.TelemetryEvent)
 
-	mu                   sync.Mutex
-	pendingNetworkEvents []telemetry.NetworkEvent
+	mu                      sync.Mutex
+	networkEventBatchTicker *time.Ticker
+	pendingNetworkEvents    []*proto.TelemetryEvent
 }
 
-func (s *DRPCService) writeTelemetryEvents(events []telemetry.NetworkEvent) {
+func (s *DRPCService) writeTelemetryEvents(events []*proto.TelemetryEvent) {
 	if s.NetworkTelemetryBatchFn == nil {
 		return
 	}
@@ -142,13 +141,14 @@ func (s *DRPCService) sendTelemetryBatch() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	events := s.pendingNetworkEvents
-	s.pendingNetworkEvents = []telemetry.NetworkEvent{}
+	s.pendingNetworkEvents = []*proto.TelemetryEvent{}
 	go s.writeTelemetryEvents(events)
 }
 
 // PeriodicTelemetryBatcher starts a goroutine to periodically send telemetry
 // events to the telemetry backend. The returned function is a cleanup function
-// that should be called when the service is no longer needed.
+// that should be called when the service is no longer needed. Calling this more
+// than once will panic.
 //
 // Note: calling the returned function does not unblock any in-flight calls to
 // the underlying telemetry backend that come from PostTelemetry due to
@@ -162,9 +162,16 @@ func (s *DRPCService) PeriodicTelemetryBatcher() func() {
 		return func() {}
 	}
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.networkEventBatchTicker != nil {
+		panic("PeriodicTelemetryBatcher called more than once")
+	}
+	ticker := time.NewTicker(s.NetworkTelemetryBatchFrequency)
+	s.networkEventBatchTicker = ticker
+
 	go func() {
 		defer close(done)
-		ticker := time.NewTicker(s.NetworkTelemetryBatchFrequency)
 		defer ticker.Stop()
 
 		for {
@@ -189,19 +196,16 @@ func (s *DRPCService) PostTelemetry(_ context.Context, req *proto.TelemetryReque
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	for _, pEvent := range req.Events {
-		tEvent, err := telemetry.NetworkEventFromProto(pEvent)
-		if err != nil {
-			// TODO(@deansheather): log? return an error?
-			continue
-		}
-
-		s.pendingNetworkEvents = append(s.pendingNetworkEvents, tEvent)
+	for _, event := range req.Events {
+		s.pendingNetworkEvents = append(s.pendingNetworkEvents, event)
 
 		if len(s.pendingNetworkEvents) >= s.NetworkTelemetryBatchMaxSize {
 			events := s.pendingNetworkEvents
-			s.pendingNetworkEvents = []telemetry.NetworkEvent{}
+			s.pendingNetworkEvents = []*proto.TelemetryEvent{}
 			// Perform the send in a goroutine to avoid blocking the DRPC call.
+			if s.networkEventBatchTicker != nil {
+				s.networkEventBatchTicker.Reset(s.NetworkTelemetryBatchFrequency)
+			}
 			go s.writeTelemetryEvents(events)
 		}
 	}
