feat: use tailnet v2 API for coordination (#11638)

This one is huge, and I'm sorry.

The problem is that once I change `tailnet.Conn` to start doing v2 behavior, I kind of have to change it everywhere, including in CoderSDK (CLI), the agent, wsproxy, and ServerTailnet.

There is still a bit more cleanup to do, and I need to add code so that when we lose connection to the Coordinator, we mark all peers as LOST, but that will be in a separate PR since this is big enough!

coderd/coderd_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/provisioner/echo"
 	"github.com/coder/coder/v2/tailnet"
+	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
 	"github.com/coder/coder/v2/testutil"
 )
 
@@ -98,14 +99,32 @@ func TestDERP(t *testing.T) {
 
 	w2Ready := make(chan struct{})
 	w2ReadyOnce := sync.Once{}
+	w1ID := uuid.New()
 	w1.SetNodeCallback(func(node *tailnet.Node) {
-		w2.UpdateNodes([]*tailnet.Node{node}, false)
+		pn, err := tailnet.NodeToProto(node)
+		if !assert.NoError(t, err) {
+			return
+		}
+		w2.UpdatePeers([]*tailnetproto.CoordinateResponse_PeerUpdate{{
+			Id:   w1ID[:],
+			Node: pn,
+			Kind: tailnetproto.CoordinateResponse_PeerUpdate_NODE,
+		}})
 		w2ReadyOnce.Do(func() {
 			close(w2Ready)
 		})
 	})
+	w2ID := uuid.New()
 	w2.SetNodeCallback(func(node *tailnet.Node) {
-		w1.UpdateNodes([]*tailnet.Node{node}, false)
+		pn, err := tailnet.NodeToProto(node)
+		if !assert.NoError(t, err) {
+			return
+		}
+		w1.UpdatePeers([]*tailnetproto.CoordinateResponse_PeerUpdate{{
+			Id:   w2ID[:],
+			Node: pn,
+			Kind: tailnetproto.CoordinateResponse_PeerUpdate_NODE,
+		}})
 	})
 
 	conn := make(chan struct{})
@@ -199,7 +218,11 @@ func TestDERPForceWebSockets(t *testing.T) {
 	defer cancel()
 
 	resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
-	conn, err := client.DialWorkspaceAgent(ctx, resources[0].Agents[0].ID, nil)
+	conn, err := client.DialWorkspaceAgent(ctx, resources[0].Agents[0].ID,
+		&codersdk.DialWorkspaceAgentOptions{
+			Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("client"),
+		},
+	)
 	require.NoError(t, err)
 	defer func() {
 		_ = conn.Close()

coderd/tailnet.go

@@ -121,12 +121,23 @@ func NewServerTailnet(
 	}
 	tn.agentConn.Store(&agentConn)
 
-	err = tn.getAgentConn().UpdateSelf(conn.Node())
+	pn, err := tailnet.NodeToProto(conn.Node())
 	if err != nil {
-		tn.logger.Warn(context.Background(), "server tailnet update self", slog.Error(err))
+		tn.logger.Critical(context.Background(), "failed to convert node", slog.Error(err))
+	} else {
+		err = tn.getAgentConn().UpdateSelf(pn)
+		if err != nil {
+			tn.logger.Warn(context.Background(), "server tailnet update self", slog.Error(err))
+		}
 	}
+
 	conn.SetNodeCallback(func(node *tailnet.Node) {
-		err := tn.getAgentConn().UpdateSelf(node)
+		pn, err := tailnet.NodeToProto(node)
+		if err != nil {
+			tn.logger.Critical(context.Background(), "failed to convert node", slog.Error(err))
+			return
+		}
+		err = tn.getAgentConn().UpdateSelf(pn)
 		if err != nil {
 			tn.logger.Warn(context.Background(), "broadcast server node to agents", slog.Error(err))
 		}
@@ -191,21 +202,9 @@ func (s *ServerTailnet) doExpireOldAgents(cutoff time.Duration) {
 		// If no one has connected since the cutoff and there are no active
 		// connections, remove the agent.
 		if time.Since(lastConnection) > cutoff && len(s.agentTickets[agentID]) == 0 {
-			deleted, err := s.conn.RemovePeer(tailnet.PeerSelector{
-				ID: tailnet.NodeID(agentID),
-				IP: netip.PrefixFrom(tailnet.IPFromUUID(agentID), 128),
-			})
-			if err != nil {
-				s.logger.Warn(ctx, "failed to remove peer from server tailnet", slog.Error(err))
-				continue
-			}
-			if !deleted {
-				s.logger.Warn(ctx, "peer didn't exist in tailnet", slog.Error(err))
-			}
-
 			deletedCount++
 			delete(s.agentConnectionTimes, agentID)
-			err = agentConn.UnsubscribeAgent(agentID)
+			err := agentConn.UnsubscribeAgent(agentID)
 			if err != nil {
 				s.logger.Error(ctx, "unsubscribe expired agent", slog.Error(err), slog.F("agent_id", agentID))
 			}
@@ -221,7 +220,7 @@ func (s *ServerTailnet) doExpireOldAgents(cutoff time.Duration) {
 func (s *ServerTailnet) watchAgentUpdates() {
 	for {
 		conn := s.getAgentConn()
-		nodes, ok := conn.NextUpdate(s.ctx)
+		resp, ok := conn.NextUpdate(s.ctx)
 		if !ok {
 			if conn.IsClosed() && s.ctx.Err() == nil {
 				s.logger.Warn(s.ctx, "multiagent closed, reinitializing")
@@ -231,7 +230,7 @@ func (s *ServerTailnet) watchAgentUpdates() {
 			return
 		}
 
-		err := s.conn.UpdateNodes(nodes, false)
+		err := s.conn.UpdatePeers(resp.GetPeerUpdates())
 		if err != nil {
 			if xerrors.Is(err, tailnet.ErrConnClosed) {
 				s.logger.Warn(context.Background(), "tailnet conn closed, exiting watchAgentUpdates", slog.Error(err))

coderd/tailnet_test.go

@@ -3,7 +3,6 @@ package coderd_test
 import (
 	"context"
 	"fmt"
-	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/netip"
@@ -204,22 +203,20 @@ func setupAgent(t *testing.T, agentAddresses []netip.Prefix) (uuid.UUID, agent.A
 			Logger:    logger.Named("client"),
 		})
 		require.NoError(t, err)
-		clientConn, serverConn := net.Pipe()
-		serveClientDone := make(chan struct{})
 		t.Cleanup(func() {
-			_ = clientConn.Close()
-			_ = serverConn.Close()
 			_ = conn.Close()
-			<-serveClientDone
 		})
-		go func() {
-			defer close(serveClientDone)
-			coord.ServeClient(serverConn, uuid.New(), manifest.AgentID)
-		}()
-		sendNode, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
-			return conn.UpdateNodes(node, false)
+		clientID := uuid.New()
+		testCtx, testCtxCancel := context.WithCancel(context.Background())
+		t.Cleanup(testCtxCancel)
+		coordination := tailnet.NewInMemoryCoordination(
+			testCtx, logger,
+			clientID, manifest.AgentID,
+			coord, conn,
+		)
+		t.Cleanup(func() {
+			_ = coordination.Close()
 		})
-		conn.SetNodeCallback(sendNode)
 		return codersdk.NewWorkspaceAgentConn(conn, codersdk.WorkspaceAgentConnOptions{
 			AgentID:   manifest.AgentID,
 			AgentIP:   codersdk.WorkspaceAgentIP,

coderd/util/apiversion/apiversion.go

@@ -30,6 +30,10 @@ func (v *APIVersion) WithBackwardCompat(majs ...int) *APIVersion {
 	return v
 }
 
+func (v *APIVersion) String() string {
+	return fmt.Sprintf("%d.%d", v.supportedMajor, v.supportedMinor)
+}
+
 // Validate validates the given version against the given constraints:
 // A given major.minor version is valid iff:
 //  1. The requested major version is contained within v.supportedMajors
@@ -42,10 +46,6 @@ func (v *APIVersion) WithBackwardCompat(majs ...int) *APIVersion {
 // - 1.x is supported,
 // - 2.0, 2.1, and 2.2 are supported,
 // - 2.3+ is not supported.
-func (v *APIVersion) String() string {
-	return fmt.Sprintf("%d.%d", v.supportedMajor, v.supportedMinor)
-}
-
 func (v *APIVersion) Validate(version string) error {
 	major, minor, err := Parse(version)
 	if err != nil {

coderd/workspaceagents.go

@@ -857,8 +857,6 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
 // Deprecated: use api.tailnet.AgentConn instead.
 // See: https://github.com/coder/coder/issues/8218
 func (api *API) _dialWorkspaceAgentTailnet(agentID uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
-	clientConn, serverConn := net.Pipe()
-
 	derpMap := api.DERPMap()
 	conn, err := tailnet.NewConn(&tailnet.Options{
 		Addresses:           []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
@@ -868,8 +866,6 @@ func (api *API) _dialWorkspaceAgentTailnet(agentID uuid.UUID) (*codersdk.Workspa
 		BlockEndpoints:      api.DeploymentValues.DERP.Config.BlockDirect.Value(),
 	})
 	if err != nil {
-		_ = clientConn.Close()
-		_ = serverConn.Close()
 		return nil, xerrors.Errorf("create tailnet conn: %w", err)
 	}
 	ctx, cancel := context.WithCancel(api.ctx)
@@ -887,10 +883,10 @@ func (api *API) _dialWorkspaceAgentTailnet(agentID uuid.UUID) (*codersdk.Workspa
 		return left
 	})
 
-	sendNodes, _ := tailnet.ServeCoordinator(clientConn, func(nodes []*tailnet.Node) error {
-		return conn.UpdateNodes(nodes, true)
-	})
-	conn.SetNodeCallback(sendNodes)
+	clientID := uuid.New()
+	coordination := tailnet.NewInMemoryCoordination(ctx, api.Logger,
+		clientID, agentID,
+		*(api.TailnetCoordinator.Load()), conn)
 
 	// Check for updated DERP map every 5 seconds.
 	go func() {
@@ -920,27 +916,13 @@ func (api *API) _dialWorkspaceAgentTailnet(agentID uuid.UUID) (*codersdk.Workspa
 		AgentID: agentID,
 		AgentIP: codersdk.WorkspaceAgentIP,
 		CloseFunc: func() error {
+			_ = coordination.Close()
 			cancel()
-			_ = clientConn.Close()
-			_ = serverConn.Close()
 			return nil
 		},
 	})
-	go func() {
-		err := (*api.TailnetCoordinator.Load()).ServeClient(serverConn, uuid.New(), agentID)
-		if err != nil {
-			// Sometimes, we get benign closed pipe errors when the server is
-			// shutting down.
-			if api.ctx.Err() == nil {
-				api.Logger.Warn(ctx, "tailnet coordinator client error", slog.Error(err))
-			}
-			_ = agentConn.Close()
-		}
-	}()
 	if !agentConn.AwaitReachable(ctx) {
 		_ = agentConn.Close()
-		_ = serverConn.Close()
-		_ = clientConn.Close()
 		cancel()
 		return nil, xerrors.Errorf("agent not reachable")
 	}

coderd/workspaceagents_test.go

@@ -535,7 +535,6 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) {
 	})
 	require.NoError(t, err)
 	defer conn.Close()
-	require.True(t, conn.BlockEndpoints())
 
 	require.True(t, conn.AwaitReachable(ctx))
 	_, p2p, _, err := conn.Ping(ctx)

coderd/wsconncache/wsconncache_test.go

@@ -12,14 +12,19 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"go.uber.org/atomic"
 	"go.uber.org/goleak"
+	"golang.org/x/xerrors"
+	"storj.io/drpc"
+	"storj.io/drpc/drpcmux"
+	"storj.io/drpc/drpcserver"
+	"tailscale.com/tailcfg"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
@@ -27,7 +32,9 @@ import (
 	"github.com/coder/coder/v2/coderd/wsconncache"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
+	drpcsdk "github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/tailnet"
+	"github.com/coder/coder/v2/tailnet/proto"
 	"github.com/coder/coder/v2/tailnet/tailnettest"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -41,7 +48,7 @@ func TestCache(t *testing.T) {
 	t.Run("Same", func(t *testing.T) {
 		t.Parallel()
 		cache := wsconncache.New(func(id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
-			return setupAgent(t, agentsdk.Manifest{}, 0), nil
+			return setupAgent(t, agentsdk.Manifest{}, 0)
 		}, 0)
 		defer func() {
 			_ = cache.Close()
@@ -54,10 +61,10 @@ func TestCache(t *testing.T) {
 	})
 	t.Run("Expire", func(t *testing.T) {
 		t.Parallel()
-		called := atomic.NewInt32(0)
+		called := int32(0)
 		cache := wsconncache.New(func(id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
-			called.Add(1)
-			return setupAgent(t, agentsdk.Manifest{}, 0), nil
+			atomic.AddInt32(&called, 1)
+			return setupAgent(t, agentsdk.Manifest{}, 0)
 		}, time.Microsecond)
 		defer func() {
 			_ = cache.Close()
@@ -70,12 +77,12 @@ func TestCache(t *testing.T) {
 		require.NoError(t, err)
 		release()
 		<-conn.Closed()
-		require.Equal(t, int32(2), called.Load())
+		require.Equal(t, int32(2), called)
 	})
 	t.Run("NoExpireWhenLocked", func(t *testing.T) {
 		t.Parallel()
 		cache := wsconncache.New(func(id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
-			return setupAgent(t, agentsdk.Manifest{}, 0), nil
+			return setupAgent(t, agentsdk.Manifest{}, 0)
 		}, time.Microsecond)
 		defer func() {
 			_ = cache.Close()
@@ -108,7 +115,7 @@ func TestCache(t *testing.T) {
 		go server.Serve(random)
 
 		cache := wsconncache.New(func(id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
-			return setupAgent(t, agentsdk.Manifest{}, 0), nil
+			return setupAgent(t, agentsdk.Manifest{}, 0)
 		}, time.Microsecond)
 		defer func() {
 			_ = cache.Close()
@@ -154,7 +161,7 @@ func TestCache(t *testing.T) {
 	})
 }
 
-func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Duration) *codersdk.WorkspaceAgentConn {
+func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Duration) (*codersdk.WorkspaceAgentConn, error) {
 	t.Helper()
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 	manifest.DERPMap, _ = tailnettest.RunDERPAndSTUN(t)
@@ -184,18 +191,25 @@ func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Durati
 		DERPForceWebSockets: manifest.DERPForceWebSockets,
 		Logger:              slogtest.Make(t, nil).Named("tailnet").Leveled(slog.LevelDebug),
 	})
-	require.NoError(t, err)
-	clientConn, serverConn := net.Pipe()
+	// setupAgent is called by wsconncache Dialer, so we can't use require here as it will end the
+	// test, which in turn closes the wsconncache, which in turn waits for the Dialer and deadlocks.
+	if !assert.NoError(t, err) {
+		return nil, err
+	}
 	t.Cleanup(func() {
-		_ = clientConn.Close()
-		_ = serverConn.Close()
 		_ = conn.Close()
 	})
-	go coordinator.ServeClient(serverConn, uuid.New(), manifest.AgentID)
-	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(nodes []*tailnet.Node) error {
-		return conn.UpdateNodes(nodes, false)
+	clientID := uuid.New()
+	testCtx, testCtxCancel := context.WithCancel(context.Background())
+	t.Cleanup(testCtxCancel)
+	coordination := tailnet.NewInMemoryCoordination(
+		testCtx, logger,
+		clientID, manifest.AgentID,
+		coordinator, conn,
+	)
+	t.Cleanup(func() {
+		_ = coordination.Close()
 	})
-	conn.SetNodeCallback(sendNode)
 	agentConn := codersdk.NewWorkspaceAgentConn(conn, codersdk.WorkspaceAgentConnOptions{
 		AgentID: manifest.AgentID,
 		AgentIP: codersdk.WorkspaceAgentIP,
@@ -206,16 +220,20 @@ func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Durati
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
 	defer cancel()
 	if !agentConn.AwaitReachable(ctx) {
-		t.Fatal("agent not reachable")
+		// setupAgent is called by wsconncache Dialer, so we can't use t.Fatal here as it will end
+		// the test, which in turn closes the wsconncache, which in turn waits for the Dialer and
+		// deadlocks.
+		t.Error("agent not reachable")
+		return nil, xerrors.New("agent not reachable")
 	}
-	return agentConn
+	return agentConn, nil
 }
 
 type client struct {
 	t           *testing.T
 	agentID     uuid.UUID
 	manifest    agentsdk.Manifest
-	coordinator tailnet.CoordinatorV1
+	coordinator tailnet.Coordinator
 }
 
 func (c *client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
@@ -240,19 +258,53 @@ func (*client) DERPMapUpdates(_ context.Context) (<-chan agentsdk.DERPMapUpdate,
 	}, nil
 }
 
-func (c *client) Listen(_ context.Context) (net.Conn, error) {
-	clientConn, serverConn := net.Pipe()
+func (c *client) Listen(_ context.Context) (drpc.Conn, error) {
+	logger := slogtest.Make(c.t, nil).Leveled(slog.LevelDebug).Named("drpc")
+	conn, lis := drpcsdk.MemTransportPipe()
 	closed := make(chan struct{})
 	c.t.Cleanup(func() {
-		_ = serverConn.Close()
-		_ = clientConn.Close()
+		_ = conn.Close()
+		_ = lis.Close()
 		<-closed
 	})
+	coordPtr := atomic.Pointer[tailnet.Coordinator]{}
+	coordPtr.Store(&c.coordinator)
+	mux := drpcmux.New()
+	drpcService := &tailnet.DRPCService{
+		CoordPtr: &coordPtr,
+		Logger:   logger,
+		// TODO: handle DERPMap too!
+		DerpMapUpdateFrequency: time.Hour,
+		DerpMapFn:              func() *tailcfg.DERPMap { panic("not implemented") },
+	}
+	err := proto.DRPCRegisterTailnet(mux, drpcService)
+	if err != nil {
+		return nil, xerrors.Errorf("register DRPC service: %w", err)
+	}
+	server := drpcserver.NewWithOptions(mux, drpcserver.Options{
+		Log: func(err error) {
+			if xerrors.Is(err, io.EOF) ||
+				xerrors.Is(err, context.Canceled) ||
+				xerrors.Is(err, context.DeadlineExceeded) {
+				return
+			}
+			logger.Debug(context.Background(), "drpc server error", slog.Error(err))
+		},
+	})
+	serveCtx, cancel := context.WithCancel(context.Background())
+	c.t.Cleanup(cancel)
+	auth := tailnet.AgentTunnelAuth{}
+	streamID := tailnet.StreamID{
+		Name: "wsconncache_test-agent",
+		ID:   c.agentID,
+		Auth: auth,
+	}
+	serveCtx = tailnet.WithStreamID(serveCtx, streamID)
 	go func() {
-		_ = c.coordinator.ServeAgent(serverConn, c.agentID, "")
+		server.Serve(serveCtx, lis)
 		close(closed)
 	}()
-	return clientConn, nil
+	return conn, nil
 }
 
 func (*client) ReportStats(_ context.Context, _ slog.Logger, _ <-chan *agentsdk.Stats, _ func(time.Duration)) (io.Closer, error) {