Fix incorrect behavior for qdisc wakeups.

Without this fix, the qdisc dispatch loop writes at most one batch,
then waits for another wakeup before writing any more packets. All
packets in the queue should be written to the link endpoint after
a wakeup.

PiperOrigin-RevId: 426053587

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

@@ -28,6 +28,13 @@ import (
 
 var _ stack.QueueingDiscipline = (*discipline)(nil)
 
+const (
+	// BatchSize represents the number of packets written to the
+	// lower link endpoint during calls to WritePackets.
+	BatchSize   = 32
+	qDiscClosed = 1
+)
+
 // discipline represents a QueueingDiscipline which implements a FIFO queue for
 // all outgoing packets. discipline can have 1 or more underlying
 // queueDispatchers. All outgoing packets are consistenly hashed to a single
@@ -41,8 +48,6 @@ type discipline struct {
 	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.
@@ -87,7 +92,6 @@ func (qd *queueDispatcher) dispatchLoop() {
 	s.AddWaker(&qd.closeWaker)
 	defer s.Done()
 
-	const batchSize = 32
 	var batch stack.PacketBufferList
 	for {
 		switch w := s.Fetch(true); w {
@@ -106,21 +110,20 @@ func (qd *queueDispatcher) dispatchLoop() {
 			panic("unknown waker")
 		}
 		qd.mu.Lock()
-		for batch.Len() < batchSize {
-			pkt := qd.queue.Front()
-			if pkt == nil {
-				break
-			}
-
+		for pkt := qd.queue.Front(); pkt != nil; pkt = qd.queue.Front() {
 			qd.queue.Remove(pkt)
 			qd.used--
 			batch.PushBack(pkt)
+			if batch.Len() < BatchSize && qd.used != 0 {
+				continue
+			}
+			qd.mu.Unlock()
+			_, _ = qd.lower.WritePackets(batch)
+			batch.DecRef()
+			batch.Reset()
+			qd.mu.Lock()
 		}
 		qd.mu.Unlock()
-
-		_, _ = qd.lower.WritePackets(batch)
-		batch.DecRef()
-		batch.Reset()
 	}
 }
 

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

@@ -18,6 +18,7 @@ import (
 	"math/rand"
 	"os"
 	"testing"
+	"time"
 
 	"gvisor.dev/gvisor/pkg/refs"
 	"gvisor.dev/gvisor/pkg/refsvfs2"
@@ -28,20 +29,31 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
-var _ stack.LinkWriter = (*discardWriter)(nil)
+var _ stack.LinkWriter = (*countWriter)(nil)
 
-// discardWriter implements LinkWriter.
-type discardWriter struct {
+// countWriter implements LinkWriter.
+type countWriter struct {
+	mu             sync.Mutex
+	packetsWritten int
+	packetsWanted  int
+	done           chan struct{}
 }
 
-func (*discardWriter) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
+func (cw *countWriter) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
+	cw.mu.Lock()
+	defer cw.mu.Unlock()
+	cw.packetsWritten += pkts.Len()
+	// Opt out of using the done channel if packetsWanted is not set.
+	if cw.packetsWanted > 0 && cw.packetsWritten == cw.packetsWanted {
+		close(cw.done)
+	}
 	return pkts.Len(), nil
 }
 
 // In b/209690936, fast simultaneous writes on qdisc will cause panics. This test
 // reproduces the behavior shown in that bug.
 func TestFastSimultaneousWrites(t *testing.T) {
-	lower := &discardWriter{}
+	lower := &countWriter{}
 	linkEP := fifo.New(lower, 16, 1000)
 
 	v := make(buffer.View, 1)
@@ -50,7 +62,6 @@ func TestFastSimultaneousWrites(t *testing.T) {
 	nWriters := 100
 	nWrites := 100
 	var wg sync.WaitGroup
-	defer wg.Done()
 	for i := 0; i < nWriters; i++ {
 		wg.Add(1)
 		go func() {
@@ -65,6 +76,8 @@ func TestFastSimultaneousWrites(t *testing.T) {
 			}
 		}()
 	}
+	wg.Wait()
+	linkEP.Close()
 }
 
 func TestWriteRefusedAfterClosed(t *testing.T) {
@@ -78,6 +91,30 @@ func TestWriteRefusedAfterClosed(t *testing.T) {
 	}
 }
 
+func TestWriteMorePacketsThanBatchSize(t *testing.T) {
+	tc := []int{fifo.BatchSize + 1, fifo.BatchSize*2 + 1}
+	v := make(buffer.View, 1)
+
+	for _, want := range tc {
+		done := make(chan struct{})
+		lower := &countWriter{done: done, packetsWanted: want}
+		linkEp := fifo.New(lower, 1, 1000)
+		for i := 0; i < want; i++ {
+			pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
+				Data: v.ToVectorisedView(),
+			})
+			linkEp.WritePacket(pkt)
+			pkt.DecRef()
+		}
+		select {
+		case <-done:
+		case <-time.After(1 * time.Second):
+			t.Fatalf("expected %d packets, but got only %d", want, lower.packetsWritten)
+		}
+		linkEp.Close()
+	}
+}
+
 func TestMain(m *testing.M) {
 	refs.SetLeakMode(refs.LeaksPanic)
 	code := m.Run()