feat: make ServerTailnet set peers lost when it reconnects to the coordinator

spikecurtis · spikecurtis · commit 135bd481e3e3 · 2024-01-19T16:36:08.000+04:00
diff --git a/Makefile b/Makefile
@@ -476,7 +476,9 @@ gen: \
 	site/e2e/provisionerGenerated.ts \
 	site/src/theme/icons.json \
 	examples/examples.gen.json \
-	tailnet/tailnettest/coordinatormock.go
+	tailnet/tailnettest/coordinatormock.go \
+	tailnet/tailnettest/coordinateemock.go \
+	tailnet/tailnettest/multiagentmock.go
 .PHONY: gen
 
 # Mark all generated files as fresh so make thinks they're up-to-date. This is
@@ -504,6 +506,8 @@ gen/mark-fresh:
 		site/src/theme/icons.json \
 		examples/examples.gen.json \
 		tailnet/tailnettest/coordinatormock.go \
+		tailnet/tailnettest/coordinateemock.go \
+		tailnet/tailnettest/multiagentmockmock.go \
 	"
 	for file in $$files; do
 		echo "$$file"
@@ -531,7 +535,7 @@ coderd/database/querier.go: coderd/database/sqlc.yaml coderd/database/dump.sql $
 coderd/database/dbmock/dbmock.go: coderd/database/db.go coderd/database/querier.go
 	go generate ./coderd/database/dbmock/
 
-tailnet/tailnettest/coordinatormock.go: tailnet/coordinator.go
+tailnet/tailnettest/coordinatormock.go tailnet/tailnettest/multiagentmock.go tailnet/tailnettest/coordinateemock.go: tailnet/coordinator.go tailnet/multiagent.go
 	go generate ./tailnet/tailnettest/
 
 tailnet/proto/tailnet.pb.go: tailnet/proto/tailnet.proto
diff --git a/coderd/tailnet.go b/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
diff --git a/coderd/tailnet_internal_test.go b/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)
+}
diff --git a/tailnet/tailnettest/coordinateemock.go b/tailnet/tailnettest/coordinateemock.go
diff --git a/tailnet/tailnettest/multiagentmock.go b/tailnet/tailnettest/multiagentmock.go
diff --git a/tailnet/tailnettest/tailnettest.go b/tailnet/tailnettest/tailnettest.go
@@ -21,7 +21,9 @@ import (
 	"github.com/coder/coder/v2/tailnet"
 )
 
+//go:generate mockgen -destination ./multiagentmock.go -package tailnettest github.com/coder/coder/v2/tailnet MultiAgentConn
 //go:generate mockgen -destination ./coordinatormock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinator
+//go:generate mockgen -destination ./coordinateemock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinatee
 
 // RunDERPAndSTUN creates a DERP mapping for tests.
 func RunDERPAndSTUN(t *testing.T) (*tailcfg.DERPMap, *derp.Server) {

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,9 @@ import (`
`21`	`21`	`"github.com/coder/coder/v2/tailnet"`
`22`	`22`	`)`
`23`	`23`
	`24`	`+//go:generate mockgen -destination ./multiagentmock.go -package tailnettest github.com/coder/coder/v2/tailnet MultiAgentConn`
`24`	`25`	`//go:generate mockgen -destination ./coordinatormock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinator`
	`26`	`+//go:generate mockgen -destination ./coordinateemock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinatee`
`25`	`27`
`26`	`28`	`// RunDERPAndSTUN creates a DERP mapping for tests.`
`27`	`29`	`func RunDERPAndSTUN(t testing.T) (tailcfg.DERPMap, *derp.Server) {`