Skip to content

Commit 0b0b1a4

Browse files
spikecurtisspikecurtis
committed
feat: make ServerTailnet set peers lost when it reconnects to the coordinator
1 parent 53b97df commit 0b0b1a4

File tree

6 files changed

+314
-6
lines changed

6 files changed

+314
-6
lines changed

‎Makefile

+6-2
Original file line numberDiff line numberDiff line change
@@ -476,7 +476,9 @@ gen: \
476476
site/e2e/provisionerGenerated.ts \
477477
site/src/theme/icons.json \
478478
examples/examples.gen.json \
479-
tailnet/tailnettest/coordinatormock.go
479+
tailnet/tailnettest/coordinatormock.go \
480+
tailnet/tailnettest/coordinateemock.go \
481+
tailnet/tailnettest/multiagentmock.go
480482
.PHONY: gen
481483

482484
# Mark all generated files as fresh so make thinks they're up-to-date. This is
@@ -504,6 +506,8 @@ gen/mark-fresh:
504506
site/src/theme/icons.json \
505507
examples/examples.gen.json \
506508
tailnet/tailnettest/coordinatormock.go \
509+
tailnet/tailnettest/coordinateemock.go \
510+
tailnet/tailnettest/multiagentmockmock.go \
507511
"
508512
for file in $$files; do
509513
echo "$$file"
@@ -531,7 +535,7 @@ coderd/database/querier.go: coderd/database/sqlc.yaml coderd/database/dump.sql $
531535
coderd/database/dbmock/dbmock.go: coderd/database/db.go coderd/database/querier.go
532536
go generate ./coderd/database/dbmock/
533537

534-
tailnet/tailnettest/coordinatormock.go: tailnet/coordinator.go
538+
tailnet/tailnettest/coordinatormock.go tailnet/tailnettest/multiagentmock.go tailnet/tailnettest/coordinateemock.go: tailnet/coordinator.go tailnet/multiagent.go
535539
go generate ./tailnet/tailnettest/
536540

537541
tailnet/proto/tailnet.pb.go: tailnet/proto/tailnet.proto

‎coderd/tailnet.go

+11-4
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ func NewServerTailnet(
9595
logger: logger,
9696
tracer: traceProvider.Tracer(tracing.TracerName),
9797
conn: conn,
98+
coordinatee: conn,
9899
getMultiAgent: getMultiAgent,
99100
cache: cache,
100101
agentConnectionTimes: map[uuid.UUID]time.Time{},
@@ -224,13 +225,14 @@ func (s *ServerTailnet) watchAgentUpdates() {
224225
if !ok {
225226
if conn.IsClosed() && s.ctx.Err() == nil {
226227
s.logger.Warn(s.ctx, "multiagent closed, reinitializing")
228+
s.coordinatee.SetAllPeersLost()
227229
s.reinitCoordinator()
228230
continue
229231
}
230232
return
231233
}
232234

233-
err := s.conn.UpdatePeers(resp.GetPeerUpdates())
235+
err := s.coordinatee.UpdatePeers(resp.GetPeerUpdates())
234236
if err != nil {
235237
if xerrors.Is(err, tailnet.ErrConnClosed) {
236238
s.logger.Warn(context.Background(), "tailnet conn closed, exiting watchAgentUpdates", slog.Error(err))
@@ -280,9 +282,14 @@ type ServerTailnet struct {
280282
cancel func()
281283
derpMapUpdaterClosed chan struct{}
282284

283-
logger slog.Logger
284-
tracer trace.Tracer
285-
conn *tailnet.Conn
285+
logger slog.Logger
286+
tracer trace.Tracer
287+
288+
// in prod, these are the same, but coordinatee is a subset of Conn's
289+
// methods which makes some tests easier.
290+
conn *tailnet.Conn
291+
coordinatee tailnet.Coordinatee
292+
286293
getMultiAgent func(context.Context) (tailnet.MultiAgentConn, error)
287294
agentConn atomic.Pointer[tailnet.MultiAgentConn]
288295
cache *wsconncache.Cache

‎coderd/tailnet_internal_test.go

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package coderd
2+
3+
import (
4+
"context"
5+
"sync/atomic"
6+
"testing"
7+
"time"
8+
9+
"github.com/google/uuid"
10+
"go.uber.org/mock/gomock"
11+
12+
"cdr.dev/slog"
13+
"cdr.dev/slog/sloggers/slogtest"
14+
"github.com/coder/coder/v2/tailnet"
15+
"github.com/coder/coder/v2/tailnet/tailnettest"
16+
"github.com/coder/coder/v2/testutil"
17+
)
18+
19+
// TestServerTailnet_Reconnect tests that ServerTailnet calls SetAllPeersLost on the Coordinatee
20+
// (tailnet.Conn in production) when it disconnects from the Coordinator (via MultiAgentConn) and
21+
// reconnects.
22+
func TestServerTailnet_Reconnect(t *testing.T) {
23+
t.Parallel()
24+
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
25+
ctrl := gomock.NewController(t)
26+
ctx := testutil.Context(t, testutil.WaitShort)
27+
28+
mMultiAgent0 := tailnettest.NewMockMultiAgentConn(ctrl)
29+
mMultiAgent1 := tailnettest.NewMockMultiAgentConn(ctrl)
30+
mac := make(chan tailnet.MultiAgentConn, 2)
31+
mac <- mMultiAgent0
32+
mac <- mMultiAgent1
33+
mCoord := tailnettest.NewMockCoordinatee(ctrl)
34+
35+
uut := &ServerTailnet{
36+
ctx: ctx,
37+
logger: logger,
38+
coordinatee: mCoord,
39+
getMultiAgent: func(ctx context.Context) (tailnet.MultiAgentConn, error) {
40+
select {
41+
case <-ctx.Done():
42+
return nil, ctx.Err()
43+
case m := <-mac:
44+
return m, nil
45+
}
46+
},
47+
agentConn: atomic.Pointer[tailnet.MultiAgentConn]{},
48+
agentConnectionTimes: make(map[uuid.UUID]time.Time),
49+
}
50+
// reinit the Coordinator once, to load mMultiAgent0
51+
uut.reinitCoordinator()
52+
53+
mMultiAgent0.EXPECT().NextUpdate(gomock.Any()).
54+
Times(1).
55+
Return(nil, false) // this indicates there are no more updates
56+
closed0 := mMultiAgent0.EXPECT().IsClosed().
57+
Times(1).
58+
Return(true) // this triggers reconnect
59+
setLost := mCoord.EXPECT().SetAllPeersLost().Times(1).After(closed0)
60+
mMultiAgent1.EXPECT().NextUpdate(gomock.Any()).
61+
Times(1).
62+
After(setLost).
63+
Return(nil, false)
64+
mMultiAgent1.EXPECT().IsClosed().
65+
Times(1).
66+
Return(false) // this causes us to exit and not reconnect
67+
68+
done := make(chan struct{})
69+
go func() {
70+
uut.watchAgentUpdates()
71+
close(done)
72+
}()
73+
74+
testutil.RequireRecvCtx(ctx, t, done)
75+
}

‎tailnet/tailnettest/coordinateemock.go

+79
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tailnet/tailnettest/multiagentmock.go

+141
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tailnet/tailnettest/tailnettest.go

+2
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ import (
2121
"github.com/coder/coder/v2/tailnet"
2222
)
2323

24+
//go:generate mockgen -destination ./multiagentmock.go -package tailnettest github.com/coder/coder/v2/tailnet MultiAgentConn
2425
//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
2527

2628
// RunDERPAndSTUN creates a DERP mapping for tests.
2729
func RunDERPAndSTUN(t *testing.T) (*tailcfg.DERPMap, *derp.Server) {

0 commit comments

Comments
 (0)