feat: make ServerTailnet set peers lost when it reconnects to the coordinator (#11682)

Adds support to `ServerTailnet` to set all peers lost before attempting to reconnect to the coordinator. In practice, this only really affects `wsproxy` since coderd has a local connection to the coordinator that only goes down if we're shutting down or change licenses.

coderd/tailnet.go

@@ -95,6 +95,7 @@ func NewServerTailnet(
 		logger:               logger,
 		tracer:               traceProvider.Tracer(tracing.TracerName),
 		conn:                 conn,
+		coordinatee:          conn,
 		getMultiAgent:        getMultiAgent,
 		cache:                cache,
 		agentConnectionTimes: map[uuid.UUID]time.Time{},
@@ -224,13 +225,14 @@ func (s *ServerTailnet) watchAgentUpdates() {
 		if !ok {
 			if conn.IsClosed() && s.ctx.Err() == nil {
 				s.logger.Warn(s.ctx, "multiagent closed, reinitializing")
+				s.coordinatee.SetAllPeersLost()
 				s.reinitCoordinator()
 				continue
 			}
 			return
 		}
 
-		err := s.conn.UpdatePeers(resp.GetPeerUpdates())
+		err := s.coordinatee.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))
@@ -280,9 +282,14 @@ type ServerTailnet struct {
 	cancel               func()
 	derpMapUpdaterClosed chan struct{}
 
-	logger        slog.Logger
-	tracer        trace.Tracer
-	conn          *tailnet.Conn
+	logger slog.Logger
+	tracer trace.Tracer
+
+	// in prod, these are the same, but coordinatee is a subset of Conn's
+	// methods which makes some tests easier.
+	conn        *tailnet.Conn
+	coordinatee tailnet.Coordinatee
+
 	getMultiAgent func(context.Context) (tailnet.MultiAgentConn, error)
 	agentConn     atomic.Pointer[tailnet.MultiAgentConn]
 	cache         *wsconncache.Cache

coderd/tailnet_internal_test.go

@@ -0,0 +1,75 @@
+package coderd
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"go.uber.org/mock/gomock"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/tailnet"
+	"github.com/coder/coder/v2/tailnet/tailnettest"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// TestServerTailnet_Reconnect tests that ServerTailnet calls SetAllPeersLost on the Coordinatee
+// (tailnet.Conn in production) when it disconnects from the Coordinator (via MultiAgentConn) and
+// reconnects.
+func TestServerTailnet_Reconnect(t *testing.T) {
+	t.Parallel()
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
+	ctrl := gomock.NewController(t)
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	mMultiAgent0 := tailnettest.NewMockMultiAgentConn(ctrl)
+	mMultiAgent1 := tailnettest.NewMockMultiAgentConn(ctrl)
+	mac := make(chan tailnet.MultiAgentConn, 2)
+	mac <- mMultiAgent0
+	mac <- mMultiAgent1
+	mCoord := tailnettest.NewMockCoordinatee(ctrl)
+
+	uut := &ServerTailnet{
+		ctx:         ctx,
+		logger:      logger,
+		coordinatee: mCoord,
+		getMultiAgent: func(ctx context.Context) (tailnet.MultiAgentConn, error) {
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			case m := <-mac:
+				return m, nil
+			}
+		},
+		agentConn:            atomic.Pointer[tailnet.MultiAgentConn]{},
+		agentConnectionTimes: make(map[uuid.UUID]time.Time),
+	}
+	// reinit the Coordinator once, to load mMultiAgent0
+	uut.reinitCoordinator()
+
+	mMultiAgent0.EXPECT().NextUpdate(gomock.Any()).
+		Times(1).
+		Return(nil, false) // this indicates there are no more updates
+	closed0 := mMultiAgent0.EXPECT().IsClosed().
+		Times(1).
+		Return(true) // this triggers reconnect
+	setLost := mCoord.EXPECT().SetAllPeersLost().Times(1).After(closed0)
+	mMultiAgent1.EXPECT().NextUpdate(gomock.Any()).
+		Times(1).
+		After(setLost).
+		Return(nil, false)
+	mMultiAgent1.EXPECT().IsClosed().
+		Times(1).
+		Return(false) // this causes us to exit and not reconnect
+
+	done := make(chan struct{})
+	go func() {
+		uut.watchAgentUpdates()
+		close(done)
+	}()
+
+	testutil.RequireRecvCtx(ctx, t, done)
+}