Progress

coder · ethanndickson · Jun 26, 2024 · Jun 26, 2024 · Jun 26, 2024 · Jun 26, 2024
commit 305fd2b4c939b0e33de2535477a18352053e1f05
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/schedule"
 	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/coderd/workspacestats"
@@ -62,7 +61,6 @@ type Options struct {
 	Pubsub                            pubsub.Pubsub
 	DerpMapFn                         func() *tailcfg.DERPMap
 	TailnetCoordinator                *atomic.Pointer[tailnet.Coordinator]
-	TemplateScheduleStore             *atomic.Pointer[schedule.TemplateScheduleStore]
 	StatsReporter                     *workspacestats.Reporter
 	AppearanceFetcher                 *atomic.Pointer[appearance.Fetcher]
 	PublishWorkspaceUpdateFn          func(ctx context.Context, workspaceID uuid.UUID)

diff --git a/coderd/agentapi/api_test.go b/coderd/agentapi/api_test.go
@@ -0,0 +1,106 @@
+package agentapi_test
+
+import (
+	"context"
+	"net/url"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+	"tailscale.com/tailcfg"
+
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/agent/proto"
+	"github.com/coder/coder/v2/coderd/agentapi"
+	"github.com/coder/coder/v2/coderd/appearance"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"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"
+	"github.com/coder/coder/v2/tailnet/tailnettest"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
+func Test_APIClose(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitMedium)
+	log := slogtest.Make(t, nil)
+
+	db, pubsub := dbtestutil.NewDB(t)
+	fCoord := tailnettest.NewFakeCoordinator()
+	var coord tailnet.Coordinator = fCoord
+	coordPtr := atomic.Pointer[tailnet.Coordinator]{}
+	coordPtr.Store(&coord)
+
+	mockTemplateScheduleStore := schedule.MockTemplateScheduleStore{}
+	var templateScheduleStore schedule.TemplateScheduleStore = mockTemplateScheduleStore
+	templateScheduleStorePtr := atomic.Pointer[schedule.TemplateScheduleStore]{}
+	templateScheduleStorePtr.Store(&templateScheduleStore)
+
+	statsBatcher, closeBatcher, err := workspacestats.NewBatcher(ctx, workspacestats.BatcherWithStore(db))
+	require.NoError(t, err)
+	t.Cleanup(closeBatcher)
+	statsTracker := workspacestats.NewTracker(db)
+	t.Cleanup(func() {
+		_ = statsTracker.Close()
+	})
+	statsReporter := workspacestats.NewReporter(workspacestats.ReporterOptions{
+		Database:              db,
+		Logger:                log,
+		Pubsub:                pubsub,
+		TemplateScheduleStore: &templateScheduleStorePtr,
+		StatsBatcher:          statsBatcher,
+		UsageTracker:          statsTracker,
+		UpdateAgentMetricsFn:  func(_ context.Context, _ prometheusmetrics.AgentMetricLabels, _ []*proto.Stats_Metric) {},
+		AppStatBatchSize:      0,
+	})
+
+	appearanceFetcherPtr := atomic.Pointer[appearance.Fetcher]{}
+	appearanceFetcherPtr.Store(&appearance.DefaultFetcher)
+
+	api := agentapi.New(agentapi.Options{
+		AgentID:  uuid.New(),
+		Ctx:      ctx,
+		Log:      log,
+		Database: db,
+		Pubsub:   pubsub,
+		DerpMapFn: func() *tailcfg.DERPMap {
+			return &tailcfg.DERPMap{Regions: map[int]*tailcfg.DERPRegion{999: {RegionCode: "test"}}}
+		},
+		TailnetCoordinator:       &coordPtr,
+		StatsReporter:            statsReporter,
+		AppearanceFetcher:        &appearanceFetcherPtr,
+		PublishWorkspaceUpdateFn: func(_ context.Context, _ uuid.UUID) {},
+		NetworkTelemetryBatchFn:  func(_ []telemetry.NetworkEvent) {},
+		AccessURL: &url.URL{
+			Scheme: "http",
+			Host:   "localhost",
+		},
+		AppHostname:                    "",
+		AgentStatsRefreshInterval:      time.Second,
+		DisableDirectConnections:       false,
+		DerpForceWebSockets:            false,
+		DerpMapUpdateFrequency:         time.Second,
+		NetworkTelemetryBatchFrequency: time.Second,
+		NetworkTelemetryBatchMaxSize:   1,
+		ExternalAuthConfigs:            []*externalauth.Config{},
+		Experiments:                    codersdk.Experiments{},
+		WorkspaceID:                    uuid.New(),
+		UpdateAgentMetricsFn:           func(_ context.Context, _ prometheusmetrics.AgentMetricLabels, _ []*proto.Stats_Metric) {},
+	})
+
+	err = api.Close()
+	require.NoError(t, err)
+}
diff --git a/coderd/telemetry/telemetry.go b/coderd/telemetry/telemetry.go
@@ -1240,8 +1240,8 @@ type NetworkEvent struct {
 	Status              string                          `json:"status"` // connected, disconnected
 	DisconnectionReason string                          `json:"disconnection_reason"`
 	ClientType          string                          `json:"client_type"` // cli, agent, coderd, wsproxy
-	NodeIDSelf          uuid.UUID                       `json:"node_id_self"`
-	NodeIDRemote        uuid.UUID                       `json:"node_id_remote"`
+	NodeIDSelf          uint64                          `json:"node_id_self"`
+	NodeIDRemote        uint64                          `json:"node_id_remote"`
 	P2PEndpoint         NetworkEventP2PEndpoint         `json:"p2p_endpoint"`
 	LogIPHashes         map[string]NetworkEventIPFields `json:"log_ip_hashes"`
 	HomeDERP            string                          `json:"home_derp"`
@@ -1272,14 +1272,6 @@ func NetworkEventFromProto(proto *tailnetproto.TelemetryEvent) (NetworkEvent, er
 	if err != nil {
 		return NetworkEvent{}, xerrors.Errorf("parse id %q: %w", proto.Id, err)
 	}
-	nodeIDSelf, err := uuid.ParseBytes(proto.NodeIdSelf)
-	if err != nil {
-		return NetworkEvent{}, xerrors.Errorf("parse node_id_self %q: %w", proto.NodeIdSelf, err)
-	}
-	nodeIDRemote, err := uuid.ParseBytes(proto.NodeIdRemote)
-	if err != nil {
-		return NetworkEvent{}, xerrors.Errorf("parse node_id_remote %q: %w", proto.NodeIdRemote, err)
-	}
 
 	logIPHashes := make(map[string]NetworkEventIPFields, len(proto.LogIpHashes))
 	for k, v := range proto.LogIpHashes {
@@ -1293,8 +1285,8 @@ func NetworkEventFromProto(proto *tailnetproto.TelemetryEvent) (NetworkEvent, er
 		Status:              strings.ToLower(proto.Status.String()),
 		DisconnectionReason: proto.DisconnectionReason,
 		ClientType:          strings.ToLower(proto.ClientType.String()),
-		NodeIDSelf:          nodeIDSelf,
-		NodeIDRemote:        nodeIDRemote,
+		NodeIDSelf:          proto.NodeIdSelf,
+		NodeIDRemote:        proto.NodeIdRemote,
 		P2PEndpoint:         p2pEndpointFromProto(proto.P2PEndpoint),
 		LogIPHashes:         logIPHashes,
 		HomeDERP:            proto.HomeDerp,

diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go
@@ -131,7 +131,6 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		Pubsub:                            api.Pubsub,
 		DerpMapFn:                         api.DERPMap,
 		TailnetCoordinator:                &api.TailnetCoordinator,
-		TemplateScheduleStore:             api.TemplateScheduleStore,
 		AppearanceFetcher:                 &api.AppearanceFetcher,
 		StatsReporter:                     api.statsReporter,
 		PublishWorkspaceUpdateFn:          api.publishWorkspaceUpdate,
@@ -157,6 +156,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		WorkspaceID:          build.WorkspaceID, // saves the extra lookup later
 		UpdateAgentMetricsFn: api.UpdateAgentMetrics,
 	})
+	defer agentAPI.Close()
 
 	streamID := tailnet.StreamID{
 		Name: fmt.Sprintf("%s-%s-%s", owner.Username, workspace.Name, workspaceAgent.Name),

diff --git a/tailnet/proto/tailnet.pb.go b/tailnet/proto/tailnet.pb.go
diff --git a/tailnet/proto/tailnet.proto b/tailnet/proto/tailnet.proto
@@ -170,8 +170,8 @@ message TelemetryEvent {
 	Status status = 4;
 	string disconnection_reason = 5;
 	ClientType client_type = 6;
-	bytes node_id_self = 7;
-	bytes node_id_remote = 8;
+	uint64 node_id_self = 7;
+	uint64 node_id_remote = 8;
 	P2PEndpoint p2p_endpoint = 9;
 	map<string, IPFields> log_ip_hashes = 10;
 	string home_derp = 11;