Correct sharedmem stress test flakiness.

This change increases test size to decrease the likelihood of a timeout,
adds a closed switch to qdisc to stop packets from being queued, and moves
queue clearing to when then dispatcher actually receives a close notification
rather than after it returns to make it more clear to the programmer that we
are clearing the queue during a "close" event. The latter change is purely
stylistic.

PiperOrigin-RevId: 423871729

pkg/tcpip/link/qdisc/fifo/fifo.go

@@ -18,6 +18,8 @@
 package fifo
 
 import (
+	"sync/atomic"
+
 	"gvisor.dev/gvisor/pkg/sleep"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -34,8 +36,13 @@ var _ stack.QueueingDiscipline = (*discipline)(nil)
 type discipline struct {
 	wg          sync.WaitGroup
 	dispatchers []queueDispatcher
+
+	// +checkatomic
+	closed int32
 }
 
+const qDiscClosed = 1
+
 // queueDispatcher is responsible for dispatching all outbound packets in its
 // queue. It will also smartly batch packets when possible and write them
 // through the lower LinkWriter.
@@ -69,13 +76,6 @@ func New(lower stack.LinkWriter, n int, queueLen int) stack.QueueingDiscipline {
 		go func() {
 			defer d.wg.Done()
 			qd.dispatchLoop()
-			qd.mu.Lock()
-			for qd.queue.Front() != nil {
-				p := qd.queue.Front()
-				qd.queue.Remove(p)
-				p.DecRef()
-			}
-			qd.mu.Unlock()
 		}()
 	}
 	return d
@@ -93,6 +93,14 @@ func (qd *queueDispatcher) dispatchLoop() {
 		switch w := s.Fetch(true); w {
 		case &qd.newPacketWaker:
 		case &qd.closeWaker:
+			qd.mu.Lock()
+			for p := qd.queue.Front(); p != nil; p = qd.queue.Front() {
+				qd.queue.Remove(p)
+				p.DecRef()
+				qd.used--
+			}
+			qd.queue.DecRef()
+			qd.mu.Unlock()
 			return
 		default:
 			panic("unknown waker")
@@ -123,6 +131,9 @@ func (qd *queueDispatcher) dispatchLoop() {
 //  - pkt.GSOOptions
 //  - pkt.NetworkProtocolNumber
 func (d *discipline) WritePacket(pkt *stack.PacketBuffer) tcpip.Error {
+	if atomic.LoadInt32(&d.closed) == qDiscClosed {
+		return &tcpip.ErrClosedForSend{}
+	}
 	qd := &d.dispatchers[int(pkt.Hash)%len(d.dispatchers)]
 	qd.mu.Lock()
 	haveSpace := qd.used < qd.limit
@@ -140,6 +151,7 @@ func (d *discipline) WritePacket(pkt *stack.PacketBuffer) tcpip.Error {
 }
 
 func (d *discipline) Close() {
+	atomic.StoreInt32(&d.closed, qDiscClosed)
 	for i := range d.dispatchers {
 		d.dispatchers[i].closeWaker.Assert()
 	}

pkg/tcpip/link/qdisc/fifo/qdisc_test.go

@@ -67,6 +67,17 @@ func TestFastSimultaneousWrites(t *testing.T) {
 	}
 }
 
+func TestWriteRefusedAfterClosed(t *testing.T) {
+	linkEp := fifo.New(nil, 1, 2)
+
+	linkEp.Close()
+	err := linkEp.WritePacket(nil)
+	_, ok := err.(*tcpip.ErrClosedForSend)
+	if !ok {
+		t.Errorf("got err = %s, want %s", err, &tcpip.ErrClosedForSend{})
+	}
+}
+
 func TestMain(m *testing.M) {
 	refs.SetLeakMode(refs.LeaksPanic)
 	code := m.Run()

pkg/tcpip/link/sharedmem/BUILD

@@ -54,7 +54,7 @@ go_test(
 
 go_test(
     name = "sharedmem_server_test",
-    size = "small",
+    size = "medium",
     srcs = ["sharedmem_server_test.go"],
     deps = [
         ":sharedmem",

pkg/tcpip/link/sharedmem/sharedmem_server_test.go

@@ -189,12 +189,32 @@ func newTestContext(t *testing.T) *testContext {
 }
 
 func (ctx *testContext) cleanup() {
+	ctx.clientStk.RemoveNIC(tcpip.NICID(1))
+	ctx.serverStk.RemoveNIC(tcpip.NICID(1))
 	unix.Close(ctx.peerFDs[0])
 	unix.Close(ctx.peerFDs[1])
 	ctx.clientStk.Close()
 	ctx.serverStk.Close()
 }
 
+func makeRequest(ctx *testContext) (*http.Response, error) {
+	listenAddr := tcpip.FullAddress{Addr: remoteIPv4Address, Port: serverPort}
+	dialFunc := func(address, protocol string) (net.Conn, error) {
+		return gonet.DialTCP(ctx.clientStk, listenAddr, ipv4.ProtocolNumber)
+	}
+	httpClient := &http.Client{
+		Transport: &http.Transport{
+			Dial: dialFunc,
+		},
+	}
+	// Close idle "keep alive" connections. If any connections remain open after
+	// a test ends, DoLeakCheck() will erroneously detect leaked packets.
+	defer httpClient.CloseIdleConnections()
+	serverURL := fmt.Sprintf("http://[%s]:%d/", net.IP(remoteIPv4Address), serverPort)
+	response, err := httpClient.Get(serverURL)
+	return response, err
+}
+
 func TestServerRoundTrip(t *testing.T) {
 	ctx := newTestContext(t)
 	defer ctx.cleanup()
@@ -211,17 +231,7 @@ func TestServerRoundTrip(t *testing.T) {
 		}))
 	}()
 
-	dialFunc := func(address, protocol string) (net.Conn, error) {
-		return gonet.DialTCP(ctx.clientStk, listenAddr, ipv4.ProtocolNumber)
-	}
-
-	httpClient := &http.Client{
-		Transport: &http.Transport{
-			Dial: dialFunc,
-		},
-	}
-	serverURL := fmt.Sprintf("http://[%s]:%d/", net.IP(remoteIPv4Address), serverPort)
-	response, err := httpClient.Get(serverURL)
+	response, err := makeRequest(ctx)
 	if err != nil {
 		t.Fatalf("httpClient.Get(\"/\") failed: %s", err)
 	}
@@ -254,20 +264,10 @@ func TestServerRoundTripStress(t *testing.T) {
 		}))
 	}()
 
-	dialFunc := func(address, protocol string) (net.Conn, error) {
-		return gonet.DialTCP(ctx.clientStk, listenAddr, ipv4.ProtocolNumber)
-	}
-
-	serverURL := fmt.Sprintf("http://[%s]:%d/", net.IP(remoteIPv4Address), serverPort)
 	var errs errgroup.Group
 	for i := 0; i < 1000; i++ {
 		errs.Go(func() error {
-			httpClient := &http.Client{
-				Transport: &http.Transport{
-					Dial: dialFunc,
-				},
-			}
-			response, err := httpClient.Get(serverURL)
+			response, err := makeRequest(ctx)
 			if err != nil {
 				return fmt.Errorf("httpClient.Get(\"/\") failed: %s", err)
 			}