From 5728f719d5fad052364bc9aa06316d20b0023e9d Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Fri, 29 Mar 2024 15:49:28 -0700
Subject: [PATCH] netstack: more small alloc optimizations

PiperOrigin-RevId: 620360076
---
 pkg/tcpip/link/fdbased/packet_dispatchers.go |  9 +++++----
 pkg/tcpip/tcpip.go                           | 12 ++++++++----
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/pkg/tcpip/link/fdbased/packet_dispatchers.go b/pkg/tcpip/link/fdbased/packet_dispatchers.go
index 0c2539808..f8caca86f 100644
--- a/pkg/tcpip/link/fdbased/packet_dispatchers.go
+++ b/pkg/tcpip/link/fdbased/packet_dispatchers.go
@@ -237,6 +237,9 @@ type recvMMsgDispatcher struct {
 	// array is passed as the parameter to recvmmsg call to retrieve
 	// potentially more than 1 packet per unix.
 	msgHdrs []rawfile.MMsgHdr
+
+	// pkts is reused to avoid allocations.
+	pkts stack.PacketBufferList
 }
 
 const (
@@ -290,20 +293,18 @@ func (d *recvMMsgDispatcher) dispatch() (bool, tcpip.Error) {
 		return false, err
 	}
 	// Process each of received packets.
-	// Keep a list of packets so we can DecRef outside of the loop.
-	var pkts stack.PacketBufferList
 
 	d.e.mu.RLock()
 	dsp := d.e.dispatcher
 	d.e.mu.RUnlock()
 
-	defer func() { pkts.DecRef() }()
+	defer d.pkts.Reset()
 	for k := 0; k < nMsgs; k++ {
 		n := int(d.msgHdrs[k].Len)
 		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
 			Payload: d.bufs[k].pullBuffer(n),
 		})
-		pkts.PushBack(pkt)
+		d.pkts.PushBack(pkt)
 
 		// Mark that this iovec has been processed.
 		d.msgHdrs[k].Msg.Iovlen = 0
diff --git a/pkg/tcpip/tcpip.go b/pkg/tcpip/tcpip.go
index d2d233b03..7c21c43c1 100644
--- a/pkg/tcpip/tcpip.go
+++ b/pkg/tcpip/tcpip.go
@@ -2668,15 +2668,19 @@ func (a AddressWithPrefix) Subnet() Subnet {
 	addrLen := a.Address.length
 	if a.PrefixLen <= 0 {
 		return Subnet{
-			address: AddrFromSlice(bytes.Repeat([]byte{0}, addrLen)),
-			mask:    MaskFromBytes(bytes.Repeat([]byte{0}, addrLen)),
+			address: Address{length: addrLen},
+			mask:    AddressMask{length: addrLen},
 		}
 	}
 	if a.PrefixLen >= addrLen*8 {
-		return Subnet{
+		sub := Subnet{
 			address: a.Address,
-			mask:    MaskFromBytes(bytes.Repeat([]byte{0xff}, addrLen)),
+			mask:    AddressMask{length: addrLen},
 		}
+		for i := 0; i < addrLen; i++ {
+			sub.mask.mask[i] = 0xff
+		}
+		return sub
 	}
 
 	sa := Address{length: addrLen}