Skip to content

Commit eb12fd7

Browse files
spikecurtisspikecurtis
authored
feat: make ServerTailnet set peers lost when it reconnects to the coordinator (coder#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.
1 parent f86186e commit eb12fd7

File tree

7 files changed

+315
-7
lines changed

7 files changed

+315
-7
lines changed

‎Makefile

Lines changed: 6 additions & 2 deletions
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/multiagentmock.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

Lines changed: 11 additions & 4 deletions
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

Lines changed: 75 additions & 0 deletions
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/service_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ func TestClientService_ServeClient_V2(t *testing.T) {
102102
err = c.Close()
103103
require.NoError(t, err)
104104
err = testutil.RequireRecvCtx(ctx, t, errCh)
105-
require.ErrorIs(t, err, io.EOF)
105+
require.True(t, xerrors.Is(err, io.EOF) || xerrors.Is(err, io.ErrClosedPipe))
106106
}
107107

108108
func TestClientService_ServeClient_V1(t *testing.T) {

‎tailnet/tailnettest/coordinateemock.go

Lines changed: 79 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tailnet/tailnettest/multiagentmock.go

Lines changed: 141 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tailnet/tailnettest/tailnettest.go

Lines changed: 2 additions & 0 deletions
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)