From 8bccff393cd1fb124c6e2bc4f28d7c334ec2fb90 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Mon, 31 Jul 2023 16:39:06 -0700
Subject: [PATCH] netstack: fix flaky forwarding test

PiperOrigin-RevId: 552626894
---
 pkg/tcpip/transport/tcp/test/e2e/BUILD        |  1 +
 .../transport/tcp/test/e2e/forwarder_test.go  | 20 ++++++++++++++-----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/pkg/tcpip/transport/tcp/test/e2e/BUILD b/pkg/tcpip/transport/tcp/test/e2e/BUILD
index fc3d98127..b3ae490cf 100644
--- a/pkg/tcpip/transport/tcp/test/e2e/BUILD
+++ b/pkg/tcpip/transport/tcp/test/e2e/BUILD
@@ -76,6 +76,7 @@ go_test(
     srcs = ["forwarder_test.go"],
     deps = [
         ":e2e",
+        "//pkg/atomicbitops",
         "//pkg/refs",
         "//pkg/tcpip",
         "//pkg/tcpip/checker",
diff --git a/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go b/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go
index 610e40909..5ef83c08a 100644
--- a/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go
+++ b/pkg/tcpip/transport/tcp/test/e2e/forwarder_test.go
@@ -19,6 +19,7 @@ import (
 	"testing"
 	"time"
 
+	"gvisor.dev/gvisor/pkg/atomicbitops"
 	"gvisor.dev/gvisor/pkg/refs"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/checker"
@@ -175,15 +176,22 @@ func TestForwarderDroppedStats(t *testing.T) {
 	c := context.New(t, mtu)
 	defer c.Cleanup()
 
-	s := c.Stack()
 	const maxInFlight = 2
+	iters := atomicbitops.FromInt64(maxInFlight)
+	s := c.Stack()
+	checkedStats := make(chan struct{})
+	done := make(chan struct{})
 	f := tcp.NewForwarder(s, 65536, maxInFlight, func(r *tcp.ForwarderRequest) {
+		<-checkedStats
 		// Complete all requests without doing anything
 		r.Complete(false)
+		if iter := iters.Add(-1); iter == 0 {
+			close(done)
+		}
 	})
 	s.SetTransportProtocolHandler(tcp.ProtocolNumber, f.HandlePacket)
 
-	for i := 0; i < maxInFlight*10; i++ {
+	for i := 0; i < maxInFlight+1; i++ {
 		iss := seqnum.Value(context.TestInitialSequenceNumber + i)
 		c.SendPacket(nil, &context.Headers{
 			SrcPort: uint16(context.TestPort + i),
@@ -194,10 +202,12 @@ func TestForwarderDroppedStats(t *testing.T) {
 		})
 	}
 
-	// Verify that we got some ignored packets
-	if curr := s.Stats().TCP.ForwardMaxInFlightDrop.Value(); curr == 0 {
-		t.Errorf("Expected at least one dropped connection")
+	// Verify that we got one ignored packet.
+	if curr := s.Stats().TCP.ForwardMaxInFlightDrop.Value(); curr != 1 {
+		t.Errorf("Expected one dropped connection, but got %d", curr)
 	}
+	close(checkedStats)
+	<-done
 }
 
 func TestMain(m *testing.M) {