From 66da66de302bbb1d0fe5da853055e76369abba61 Mon Sep 17 00:00:00 2001
From: Lucas Manning <lucasmanning@google.com>
Date: Wed, 2 Feb 2022 21:48:57 -0800
Subject: [PATCH] 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       | 29 ++++++++-------
 pkg/tcpip/link/qdisc/fifo/qdisc_test.go | 49 ++++++++++++++++++++++---
 2 files changed, 59 insertions(+), 19 deletions(-)

diff --git a/pkg/tcpip/link/qdisc/fifo/fifo.go b/pkg/tcpip/link/qdisc/fifo/fifo.go
index d9b57a914..607012723 100644
--- a/pkg/tcpip/link/qdisc/fifo/fifo.go
+++ b/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()
 	}
 }
 
diff --git a/pkg/tcpip/link/qdisc/fifo/qdisc_test.go b/pkg/tcpip/link/qdisc/fifo/qdisc_test.go
index f4f155f7c..f707ced64 100644
--- a/pkg/tcpip/link/qdisc/fifo/qdisc_test.go
+++ b/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()