Use the stack clock everywhere

Updates #5939.
Updates #6012.

RELNOTES: n/a
PiperOrigin-RevId: 375931554

1 files changed, 157 insertions, 57 deletions

diff --git a/pkg/tcpip/tests/integration/link_resolution_test.go b/pkg/tcpip/tests/integration/link_resolution_test.go
index 4c06460e4..16fd7a99a 100644
--- a/pkg/tcpip/tests/integration/link_resolution_test.go
+++ b/pkg/tcpip/tests/integration/link_resolution_test.go
@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"fmt"
 	"net"
+	"runtime"
 	"testing"
 	"time"
 
@@ -287,9 +288,11 @@ func TestTCPLinkResolutionFailure(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
+			clock := faketime.NewManualClock()
 			stackOpts := stack.Options{
 				NetworkProtocols:   []stack.NetworkProtocolFactory{arp.NewProtocol, ipv4.NewProtocol, ipv6.NewProtocol},
 				TransportProtocols: []stack.TransportProtocolFactory{tcp.NewProtocol},
+				Clock:              clock,
 			}
 
 			host1Stack, host2Stack := setupStack(t, stackOpts, host1NICID, host2NICID)
@@ -333,7 +336,17 @@ func TestTCPLinkResolutionFailure(t *testing.T) {
 
 			// Wait for an error due to link resolution failing, or the endpoint to be
 			// writable.
+			if test.expectedWriteErr != nil {
+				nudConfigs, err := host1Stack.NUDConfigurations(host1NICID, test.netProto)
+				if err != nil {
+					t.Fatalf("host1Stack.NUDConfigurations(%d, %d): %s", host1NICID, test.netProto, err)
+				}
+				clock.Advance(time.Duration(nudConfigs.MaxMulticastProbes) * nudConfigs.RetransmitTimer)
+			} else {
+				clock.RunImmediatelyScheduledJobs()
+			}
 			<-ch
+
 			{
 				var r bytes.Reader
 				r.Reset([]byte{0})
@@ -693,8 +706,10 @@ func TestGetLinkAddress(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
+			clock := faketime.NewManualClock()
 			stackOpts := stack.Options{
 				NetworkProtocols: []stack.NetworkProtocolFactory{arp.NewProtocol, ipv4.NewProtocol, ipv6.NewProtocol},
+				Clock:            clock,
 			}
 
 			host1Stack, _ := setupStack(t, stackOpts, host1NICID, host2NICID)
@@ -710,8 +725,20 @@ func TestGetLinkAddress(t *testing.T) {
 			if test.expectedErr == nil {
 				wantRes.LinkAddress = utils.LinkAddr2
 			}
-			if diff := cmp.Diff(wantRes, <-ch); diff != "" {
-				t.Fatalf("link resolution result mismatch (-want +got):\n%s", diff)
+
+			nudConfigs, err := host1Stack.NUDConfigurations(host1NICID, test.netProto)
+			if err != nil {
+				t.Fatalf("host1Stack.NUDConfigurations(%d, %d): %s", host1NICID, test.netProto, err)
+			}
+
+			clock.Advance(time.Duration(nudConfigs.MaxMulticastProbes) * nudConfigs.RetransmitTimer)
+			select {
+			case got := <-ch:
+				if diff := cmp.Diff(wantRes, got); diff != "" {
+					t.Fatalf("link resolution result mismatch (-want +got):\n%s", diff)
+				}
+			default:
+				t.Fatal("event didn't arrive")
 			}
 		})
 	}
@@ -788,8 +815,10 @@ func TestRouteResolvedFields(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
+			clock := faketime.NewManualClock()
 			stackOpts := stack.Options{
 				NetworkProtocols: []stack.NetworkProtocolFactory{arp.NewProtocol, ipv4.NewProtocol, ipv6.NewProtocol},
+				Clock:            clock,
 			}
 
 			host1Stack, _ := setupStack(t, stackOpts, host1NICID, host2NICID)
@@ -819,8 +848,20 @@ func TestRouteResolvedFields(t *testing.T) {
 				if _, ok := err.(*tcpip.ErrWouldBlock); !ok {
 					t.Errorf("got r.ResolvedFields(_) = %s, want = %s", err, &tcpip.ErrWouldBlock{})
 				}
-				if diff := cmp.Diff(stack.ResolvedFieldsResult{RouteInfo: wantRouteInfo, Err: test.expectedErr}, <-ch, cmp.AllowUnexported(stack.RouteInfo{})); diff != "" {
-					t.Errorf("route resolve result mismatch (-want +got):\n%s", diff)
+
+				nudConfigs, err := host1Stack.NUDConfigurations(host1NICID, test.netProto)
+				if err != nil {
+					t.Fatalf("host1Stack.NUDConfigurations(%d, %d): %s", host1NICID, test.netProto, err)
+				}
+				clock.Advance(time.Duration(nudConfigs.MaxMulticastProbes) * nudConfigs.RetransmitTimer)
+
+				select {
+				case got := <-ch:
+					if diff := cmp.Diff(stack.ResolvedFieldsResult{RouteInfo: wantRouteInfo, Err: test.expectedErr}, got, cmp.AllowUnexported(stack.RouteInfo{})); diff != "" {
+						t.Errorf("route resolve result mismatch (-want +got):\n%s", diff)
+					}
+				default:
+					t.Fatalf("event didn't arrive")
 				}
 
 				if test.expectedErr != nil {
@@ -1033,11 +1074,16 @@ func (d *nudDispatcher) OnNeighborRemoved(nicID tcpip.NICID, entry stack.Neighbo
 	d.c <- e
 }
 
-func (d *nudDispatcher) waitForEvent(want eventInfo) error {
-	if diff := cmp.Diff(want, <-d.c, cmp.AllowUnexported(eventInfo{}), cmpopts.IgnoreFields(stack.NeighborEntry{}, "UpdatedAt")); diff != "" {
-		return fmt.Errorf("got invalid event (-want +got):\n%s", diff)
+func (d *nudDispatcher) expectEvent(want eventInfo) error {
+	select {
+	case got := <-d.c:
+		if diff := cmp.Diff(want, got, cmp.AllowUnexported(eventInfo{}), cmpopts.IgnoreFields(stack.NeighborEntry{}, "UpdatedAt")); diff != "" {
+			return fmt.Errorf("got invalid event (-want +got):\n%s", diff)
+		}
+		return nil
+	default:
+		return fmt.Errorf("event didn't arrive")
 	}
-	return nil
 }
 
 // TestTCPConfirmNeighborReachability tests that TCP informs layers beneath it
@@ -1048,7 +1094,7 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 		netProto        tcpip.NetworkProtocolNumber
 		remoteAddr      tcpip.Address
 		neighborAddr    tcpip.Address
-		getEndpoints    func(*testing.T, *stack.Stack, *stack.Stack, *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{})
+		getEndpoints    func(*testing.T, *stack.Stack, *stack.Stack, *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{})
 		isHost1Listener bool
 	}{
 		{
@@ -1056,23 +1102,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv4.ProtocolNumber,
 			remoteAddr:   utils.Host2IPv4Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv4Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := host2Stack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("host2Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 		},
 		{
@@ -1080,23 +1128,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv6.ProtocolNumber,
 			remoteAddr:   utils.Host2IPv6Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv6Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := host2Stack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("host2Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 		},
 		{
@@ -1104,23 +1154,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv4.ProtocolNumber,
 			remoteAddr:   utils.RouterNIC1IPv4Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv4Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := routerStack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("routerStack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 		},
 		{
@@ -1128,23 +1180,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv6.ProtocolNumber,
 			remoteAddr:   utils.RouterNIC1IPv6Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv6Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := routerStack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("routerStack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 		},
 		{
@@ -1152,23 +1206,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv4.ProtocolNumber,
 			remoteAddr:   utils.Host1IPv4Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv4Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := routerStack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("routerStack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 			isHost1Listener: true,
 		},
@@ -1177,23 +1233,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv6.ProtocolNumber,
 			remoteAddr:   utils.Host1IPv6Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv6Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, routerStack, _ *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := routerStack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("routerStack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 			isHost1Listener: true,
 		},
@@ -1202,23 +1260,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv4.ProtocolNumber,
 			remoteAddr:   utils.Host1IPv4Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv4Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := host2Stack.NewEndpoint(tcp.ProtocolNumber, ipv4.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("host2Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv4.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 			isHost1Listener: true,
 		},
@@ -1227,23 +1287,25 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			netProto:     ipv6.ProtocolNumber,
 			remoteAddr:   utils.Host1IPv6Addr.AddressWithPrefix.Address,
 			neighborAddr: utils.RouterNIC1IPv6Addr.AddressWithPrefix.Address,
-			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, tcpip.Endpoint, <-chan struct{}) {
+			getEndpoints: func(t *testing.T, host1Stack, _, host2Stack *stack.Stack) (tcpip.Endpoint, <-chan struct{}, tcpip.Endpoint, <-chan struct{}) {
 				var listenerWQ waiter.Queue
+				listenerWE, listenerCH := waiter.NewChannelEntry(nil)
+				listenerWQ.EventRegister(&listenerWE, waiter.EventIn)
 				listenerEP, err := host1Stack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &listenerWQ)
 				if err != nil {
 					t.Fatalf("host1Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
+				t.Cleanup(listenerEP.Close)
 
 				var clientWQ waiter.Queue
 				clientWE, clientCH := waiter.NewChannelEntry(nil)
-				clientWQ.EventRegister(&clientWE, waiter.WritableEvents)
+				clientWQ.EventRegister(&clientWE, waiter.ReadableEvents|waiter.WritableEvents)
 				clientEP, err := host2Stack.NewEndpoint(tcp.ProtocolNumber, ipv6.ProtocolNumber, &clientWQ)
 				if err != nil {
-					listenerEP.Close()
 					t.Fatalf("host2Stack.NewEndpoint(%d, %d, _): %s", tcp.ProtocolNumber, ipv6.ProtocolNumber, err)
 				}
 
-				return listenerEP, clientEP, clientCH
+				return listenerEP, listenerCH, clientEP, clientCH
 			},
 			isHost1Listener: true,
 		},
@@ -1281,14 +1343,14 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 					t.Fatalf("link resolution mismatch (-want +got):\n%s", diff)
 				}
 			}
-			if err := nudDisp.waitForEvent(eventInfo{
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryAdded,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Incomplete, Addr: test.neighborAddr},
 			}); err != nil {
 				t.Fatalf("error waiting for initial NUD event: %s", err)
 			}
-			if err := nudDisp.waitForEvent(eventInfo{
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryChanged,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Reachable, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},
@@ -1308,7 +1370,7 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			// See NUDConfigurations.BaseReachableTime for more information.
 			maxReachableTime := time.Duration(float32(nudConfigs.BaseReachableTime) * nudConfigs.MaxRandomFactor)
 			clock.Advance(maxReachableTime)
-			if err := nudDisp.waitForEvent(eventInfo{
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryChanged,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Stale, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},
@@ -1316,8 +1378,7 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 				t.Fatalf("error waiting for stale NUD event: %s", err)
 			}
 
-			listenerEP, clientEP, clientCH := test.getEndpoints(t, host1Stack, routerStack, host2Stack)
-			defer listenerEP.Close()
+			listenerEP, listenerCH, clientEP, clientCH := test.getEndpoints(t, host1Stack, routerStack, host2Stack)
 			defer clientEP.Close()
 			listenerAddr := tcpip.FullAddress{Addr: test.remoteAddr, Port: 1234}
 			if err := listenerEP.Bind(listenerAddr); err != nil {
@@ -1338,14 +1399,15 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 			// with confirmation that the neighbor is reachable (indicated by a
 			// successful 3-way handshake).
 			<-clientCH
-			if err := nudDisp.waitForEvent(eventInfo{
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryChanged,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Delay, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},
 			}); err != nil {
 				t.Fatalf("error waiting for delay NUD event: %s", err)
 			}
-			if err := nudDisp.waitForEvent(eventInfo{
+			<-listenerCH
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryChanged,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Reachable, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},
@@ -1353,26 +1415,55 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 				t.Fatalf("error waiting for reachable NUD event: %s", err)
 			}
 
+			peerEP, peerWQ, err := listenerEP.Accept(nil)
+			if err != nil {
+				t.Fatalf("listenerEP.Accept(): %s", err)
+			}
+			defer peerEP.Close()
+			peerWE, peerCH := waiter.NewChannelEntry(nil)
+			peerWQ.EventRegister(&peerWE, waiter.ReadableEvents)
+
 			// Wait for the neighbor to be stale again then send data to the remote.
 			//
 			// On successful transmission, the neighbor should become reachable
 			// without probing the neighbor as a TCP ACK would be received which is an
 			// indication of the neighbor being reachable.
 			clock.Advance(maxReachableTime)
-			if err := nudDisp.waitForEvent(eventInfo{
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryChanged,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Stale, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},
 			}); err != nil {
 				t.Fatalf("error waiting for stale NUD event: %s", err)
 			}
-			var r bytes.Reader
-			r.Reset([]byte{0})
-			var wOpts tcpip.WriteOptions
-			if _, err := clientEP.Write(&r, wOpts); err != nil {
-				t.Errorf("clientEP.Write(_, %#v): %s", wOpts, err)
+			{
+				var r bytes.Reader
+				r.Reset([]byte{0})
+				var wOpts tcpip.WriteOptions
+				if _, err := clientEP.Write(&r, wOpts); err != nil {
+					t.Errorf("clientEP.Write(_, %#v): %s", wOpts, err)
+				}
+			}
+			// Heads up, there is a race here.
+			//
+			// Incoming TCP segments are handled in
+			// tcp.(*endpoint).handleSegmentLocked:
+			//
+			// - tcp.(*endpoint).rcv.handleRcvdSegment puts the segment on the
+			// segment queue and notifies waiting readers (such as this channel)
+			//
+			// - tcp.(*endpoint).snd.handleRcvdSegment sends an ACK for the segment
+			// and notifies the NUD machinery that the peer is reachable
+			//
+			// Thus we must permit a delay between the readable signal and the
+			// expected NUD event.
+			//
+			// At the time of writing, this race is reliably hit with gotsan.
+			<-peerCH
+			for len(nudDisp.c) == 0 {
+				runtime.Gosched()
 			}
-			if err := nudDisp.waitForEvent(eventInfo{
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryChanged,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Delay, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},
@@ -1385,7 +1476,7 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 				// TCP should not mark the route reachable and NUD should go through the
 				// probe state.
 				clock.Advance(nudConfigs.DelayFirstProbeTime)
-				if err := nudDisp.waitForEvent(eventInfo{
+				if err := nudDisp.expectEvent(eventInfo{
 					eventType: entryChanged,
 					nicID:     utils.Host1NICID,
 					entry:     stack.NeighborEntry{State: stack.Probe, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},
@@ -1393,7 +1484,16 @@ func TestTCPConfirmNeighborReachability(t *testing.T) {
 					t.Fatalf("error waiting for probe NUD event: %s", err)
 				}
 			}
-			if err := nudDisp.waitForEvent(eventInfo{
+			{
+				var r bytes.Reader
+				r.Reset([]byte{0})
+				var wOpts tcpip.WriteOptions
+				if _, err := peerEP.Write(&r, wOpts); err != nil {
+					t.Errorf("peerEP.Write(_, %#v): %s", wOpts, err)
+				}
+			}
+			<-clientCH
+			if err := nudDisp.expectEvent(eventInfo{
 				eventType: entryChanged,
 				nicID:     utils.Host1NICID,
 				entry:     stack.NeighborEntry{State: stack.Reachable, Addr: test.neighborAddr, LinkAddr: utils.LinkAddr2},