Skip to content

Commit e96ab0e

Browse files
spikecurtisspikecurtis
authored
test: fix PingDirect test tear down (#20687)
Fixes coder/internal#66 The problem is a race during test tear down where the node callback can fire after the destination tailnet.Conn has already closed, causing an error. The fix I have employed is to remove the callback in a t.Cleanup() and also refactor some tests to ensure they close the tailnet.Conn in a Cleanup deeper in the stack.
1 parent c21b3e4 commit e96ab0e

File tree

1 file changed

+21
-29
lines changed

1 file changed

+21
-29
lines changed

‎tailnet/conn_test.go‎

Lines changed: 21 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -95,9 +95,6 @@ func TestTailnet(t *testing.T) {
9595
node := testutil.TryReceive(ctx, t, nodes)
9696
// Ensure this connected over raw (not websocket) DERP!
9797
require.Len(t, node.DERPForcedWebsocket, 0)
98-
99-
w1.Close()
100-
w2.Close()
10198
})
10299

103100
t.Run("ForcesWebSockets", func(t *testing.T) {
@@ -163,9 +160,6 @@ func TestTailnet(t *testing.T) {
163160
require.Len(t, node.DERPForcedWebsocket, 1)
164161
// Ensure the reason is valid!
165162
require.Equal(t, `GET failed with status code 400 (a proxy could be disallowing the use of 'Upgrade: derp'): Invalid "Upgrade" header: DERP`, node.DERPForcedWebsocket[derpMap.RegionIDs()[0]])
166-
167-
w1.Close()
168-
w2.Close()
169163
})
170164

171165
t.Run("PingDirect", func(t *testing.T) {
@@ -206,9 +200,6 @@ func TestTailnet(t *testing.T) {
206200
}
207201
return true
208202
}, testutil.WaitShort, testutil.IntervalFast)
209-
210-
w1.Close()
211-
w2.Close()
212203
})
213204

214205
t.Run("PingDERPOnly", func(t *testing.T) {
@@ -251,9 +242,6 @@ func TestTailnet(t *testing.T) {
251242
}
252243
return true
253244
}, testutil.WaitShort, testutil.IntervalFast)
254-
255-
w1.Close()
256-
w2.Close()
257245
})
258246
}
259247

@@ -302,10 +290,10 @@ func TestConn_UpdateDERP(t *testing.T) {
302290
BlockEndpoints: true,
303291
})
304292
require.NoError(t, err)
305-
defer func() {
293+
t.Cleanup(func() {
306294
err := conn.Close()
307295
assert.NoError(t, err)
308-
}()
296+
})
309297

310298
// Buffer channel so callback doesn't block
311299
nodes := make(chan *tailnet.Node, 50)
@@ -314,7 +302,7 @@ func TestConn_UpdateDERP(t *testing.T) {
314302
})
315303

316304
ctx1, cancel1 := context.WithTimeout(context.Background(), testutil.WaitShort)
317-
defer cancel1()
305+
t.Cleanup(cancel1)
318306
select {
319307
case node := <-nodes:
320308
require.Equal(t, 1, node.PreferredDERP)
@@ -330,10 +318,10 @@ func TestConn_UpdateDERP(t *testing.T) {
330318
BlockEndpoints: true,
331319
})
332320
require.NoError(t, err)
333-
defer func() {
321+
t.Cleanup(func() {
334322
err := client1.Close()
335323
assert.NoError(t, err)
336-
}()
324+
})
337325
stitch(t, conn, client1)
338326
pn, err := tailnet.NodeToProto(conn.Node())
339327
require.NoError(t, err)
@@ -346,7 +334,7 @@ func TestConn_UpdateDERP(t *testing.T) {
346334
require.NoError(t, err)
347335

348336
awaitReachableCtx1, awaitReachableCancel1 := context.WithTimeout(context.Background(), testutil.WaitShort)
349-
defer awaitReachableCancel1()
337+
t.Cleanup(awaitReachableCancel1)
350338
require.True(t, client1.AwaitReachable(awaitReachableCtx1, ip))
351339

352340
// Update the DERP map and wait for the preferred DERP server to change.
@@ -361,7 +349,7 @@ func TestConn_UpdateDERP(t *testing.T) {
361349
conn.SetDERPMap(derpMap2)
362350

363351
ctx2, cancel2 := context.WithTimeout(context.Background(), testutil.WaitShort)
364-
defer cancel2()
352+
t.Cleanup(cancel2)
365353
parentLoop:
366354
for {
367355
select {
@@ -379,7 +367,7 @@ parentLoop:
379367

380368
// Client1 should be dropped...
381369
awaitReachableCtx2, awaitReachableCancel2 := context.WithTimeout(context.Background(), testutil.WaitShort)
382-
defer awaitReachableCancel2()
370+
t.Cleanup(awaitReachableCancel2)
383371
require.False(t, client1.AwaitReachable(awaitReachableCtx2, ip))
384372

385373
// ... unless the client updates it's derp map and nodes.
@@ -392,7 +380,7 @@ parentLoop:
392380
Kind: proto.CoordinateResponse_PeerUpdate_NODE,
393381
}})
394382
awaitReachableCtx3, awaitReachableCancel3 := context.WithTimeout(context.Background(), testutil.WaitShort)
395-
defer awaitReachableCancel3()
383+
t.Cleanup(awaitReachableCancel3)
396384
require.True(t, client1.AwaitReachable(awaitReachableCtx3, ip))
397385

398386
// Connect from a different different client with up-to-date derp map and
@@ -404,10 +392,10 @@ parentLoop:
404392
BlockEndpoints: true,
405393
})
406394
require.NoError(t, err)
407-
defer func() {
395+
t.Cleanup(func() {
408396
err := client2.Close()
409397
assert.NoError(t, err)
410-
}()
398+
})
411399
stitch(t, conn, client2)
412400
pn, err = tailnet.NodeToProto(conn.Node())
413401
require.NoError(t, err)
@@ -418,7 +406,7 @@ parentLoop:
418406
}})
419407

420408
awaitReachableCtx4, awaitReachableCancel4 := context.WithTimeout(context.Background(), testutil.WaitShort)
421-
defer awaitReachableCancel4()
409+
t.Cleanup(awaitReachableCancel4)
422410
require.True(t, client2.AwaitReachable(awaitReachableCtx4, ip))
423411
}
424412

@@ -437,10 +425,10 @@ func TestConn_BlockEndpoints(t *testing.T) {
437425
BlockEndpoints: true,
438426
})
439427
require.NoError(t, err)
440-
defer func() {
428+
t.Cleanup(func() {
441429
err := conn1.Close()
442430
assert.NoError(t, err)
443-
}()
431+
})
444432

445433
// Setup conn 2.
446434
ip2 := tailnet.TailscaleServicePrefix.RandomAddr()
@@ -451,16 +439,16 @@ func TestConn_BlockEndpoints(t *testing.T) {
451439
BlockEndpoints: true,
452440
})
453441
require.NoError(t, err)
454-
defer func() {
442+
t.Cleanup(func() {
455443
err := conn2.Close()
456444
assert.NoError(t, err)
457-
}()
445+
})
458446

459447
// Connect them together and wait for them to be reachable.
460448
stitch(t, conn2, conn1)
461449
stitch(t, conn1, conn2)
462450
awaitReachableCtx, awaitReachableCancel := context.WithTimeout(context.Background(), testutil.WaitShort)
463-
defer awaitReachableCancel()
451+
t.Cleanup(awaitReachableCancel)
464452
require.True(t, conn1.AwaitReachable(awaitReachableCtx, ip2))
465453

466454
// Wait 10s for endpoints to potentially be sent over Disco. There's no way
@@ -495,6 +483,10 @@ func stitch(t *testing.T, dst, src *tailnet.Conn) {
495483
}})
496484
assert.NoError(t, err)
497485
})
486+
// ensures we don't send callbacks after the test ends and connections are closed.
487+
t.Cleanup(func() {
488+
src.SetNodeCallback(nil)
489+
})
498490
}
499491

500492
func TestTailscaleServicePrefix(t *testing.T) {

0 commit comments

Comments
 (0)