From 6ad1af2b9cbfe5404c1cd7fd3fef5bd55a791955 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Fri, 8 Mar 2024 22:28:08 -0800
Subject: [PATCH] netstack: don't allocate by calling AsSlices() in fdbased
 endpoint

In a redis-benchmark PING_INLINE test, this reduces allocations by 23%.

PiperOrigin-RevId: 614144212
---
 pkg/buffer/BUILD                   |  2 +-
 pkg/buffer/buffer.go               | 10 ++++++++--
 pkg/buffer/view.go                 |  2 +-
 pkg/tcpip/link/fdbased/endpoint.go | 18 ++++++++++++++----
 pkg/tcpip/stack/packet_buffer.go   |  9 +++++++++
 5 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/pkg/buffer/BUILD b/pkg/buffer/BUILD
index 30e10b760..ecf6230bb 100644
--- a/pkg/buffer/BUILD
+++ b/pkg/buffer/BUILD
@@ -21,7 +21,7 @@ go_template_instance(
     name = "view_list",
     out = "view_list.go",
     package = "buffer",
-    prefix = "view",
+    prefix = "View",
     template = "//pkg/ilist:generic_list",
     types = {
         "Element": "*View",
diff --git a/pkg/buffer/buffer.go b/pkg/buffer/buffer.go
index 3e585c008..79151c64b 100644
--- a/pkg/buffer/buffer.go
+++ b/pkg/buffer/buffer.go
@@ -28,7 +28,7 @@ import (
 //
 // +stateify savable
 type Buffer struct {
-	data viewList `state:".([]byte)"`
+	data ViewList `state:".([]byte)"`
 	size int64
 }
 
@@ -398,6 +398,12 @@ func (b *Buffer) Size() int64 {
 	return b.size
 }
 
+// AsViewList returns the ViewList backing b. Users may not save or modify the
+// ViewList returned.
+func (b *Buffer) AsViewList() ViewList {
+	return b.data
+}
+
 // Clone creates a copy-on-write clone of b. The underlying chunks are shared
 // until they are written to.
 func (b *Buffer) Clone() Buffer {
@@ -476,7 +482,7 @@ func (b *Buffer) Checksum(offset int) uint16 {
 // operation completes.
 func (b *Buffer) Merge(other *Buffer) {
 	b.data.PushBackList(&other.data)
-	other.data = viewList{}
+	other.data = ViewList{}
 
 	// Adjust sizes.
 	b.size += other.size
diff --git a/pkg/buffer/view.go b/pkg/buffer/view.go
index d7eb2f118..6c8d17eff 100644
--- a/pkg/buffer/view.go
+++ b/pkg/buffer/view.go
@@ -48,7 +48,7 @@ var viewPool = sync.Pool{
 //
 // +stateify savable
 type View struct {
-	viewEntry `state:"nosave"`
+	ViewEntry `state:"nosave"`
 	read      int
 	write     int
 	chunk     *chunk
diff --git a/pkg/tcpip/link/fdbased/endpoint.go b/pkg/tcpip/link/fdbased/endpoint.go
index faede2b5d..9a51c8ef5 100644
--- a/pkg/tcpip/link/fdbased/endpoint.go
+++ b/pkg/tcpip/link/fdbased/endpoint.go
@@ -642,8 +642,16 @@ func (e *endpoint) sendBatch(batchFDInfo fdInfo, pkts []*stack.PacketBuffer) (in
 				vnetHdrBuf = vnetHdr.marshal()
 			}
 
-			views := pkt.AsSlices()
-			numIovecs := len(views)
+			views, offset := pkt.AsViewList()
+			var skipped int
+			var view *buffer.View
+			for view = views.Front(); view != nil && offset >= view.Size(); view = view.Next() {
+				offset -= view.Size()
+				skipped++
+			}
+
+			// We've made it to the usable views.
+			numIovecs := views.Len() - skipped
 			if len(vnetHdrBuf) != 0 {
 				numIovecs++
 			}
@@ -665,8 +673,10 @@ func (e *endpoint) sendBatch(batchFDInfo fdInfo, pkts []*stack.PacketBuffer) (in
 			// they will escape this loop iteration via mmsgHdrs.
 			iovecs := make([]unix.Iovec, 0, numIovecs)
 			iovecs = rawfile.AppendIovecFromBytes(iovecs, vnetHdrBuf, numIovecs)
-			for _, v := range views {
-				iovecs = rawfile.AppendIovecFromBytes(iovecs, v, numIovecs)
+			// At most one slice has a non-zero offset.
+			iovecs = rawfile.AppendIovecFromBytes(iovecs, view.AsSlice()[offset:], numIovecs)
+			for view = view.Next(); view != nil; view = view.Next() {
+				iovecs = rawfile.AppendIovecFromBytes(iovecs, view.AsSlice(), numIovecs)
 			}
 
 			var mmsgHdr rawfile.MMsgHdr
diff --git a/pkg/tcpip/stack/packet_buffer.go b/pkg/tcpip/stack/packet_buffer.go
index 260dc70e9..82b042b97 100644
--- a/pkg/tcpip/stack/packet_buffer.go
+++ b/pkg/tcpip/stack/packet_buffer.go
@@ -276,6 +276,9 @@ func (pk *PacketBuffer) Data() PacketData {
 }
 
 // AsSlices returns the underlying storage of the whole packet.
+//
+// Note that AsSlices can allocate a lot. In hot paths it may be preferable to
+// iterate over a PacketBuffer's data via AsViewList.
 func (pk *PacketBuffer) AsSlices() [][]byte {
 	var views [][]byte
 	offset := pk.headerOffset()
@@ -285,6 +288,12 @@ func (pk *PacketBuffer) AsSlices() [][]byte {
 	return views
 }
 
+// AsViewList returns the list of Views backing the PacketBuffer along with the
+// header offset into them. Users may not save or modify the ViewList returned.
+func (pk *PacketBuffer) AsViewList() (buffer.ViewList, int) {
+	return pk.buf.AsViewList(), pk.headerOffset()
+}
+
 // ToBuffer returns a caller-owned copy of the underlying storage of the whole
 // packet.
 func (pk *PacketBuffer) ToBuffer() buffer.Buffer {