Pool PacketBuffers.

New PacketBuffers are now allocated from a sync.Pool, and PacketBuffers
are returned to the pool when their reference count is zero.

**Without pooling**
goos: linux
goarch: amd64
pkg: gvisor/test/benchmarks/network/network
cpu: Intel(R) Xeon(R) CPU @ 2.20GHz
BenchmarkIperf/operation.Download-16
28458945
1249 ns/op
52458.82 MB/s
863949824 bandwidth.bytes_per_second

Alloc'd mem:      32.6GB
Alloc'd objects:  17,619,980

**With pooling**
goos: linux
goarch: amd64
pkg: gvisor/test/benchmarks/network/network
cpu: Intel(R) Xeon(R) CPU @ 2.20GHz
BenchmarkIperf/operation.Download-16
30767242
1211 ns/op (-3%)
54107.47 MB/s (+3%)
887847936 bandwidth.bytes_per_second (+2.5%)

Alloc'd mem:     29.7GB (-9%)
Allod'd objects: 15,947,073 (-9.5%)

PiperOrigin-RevId: 408729962

pkg/tcpip/stack/packet_buffer.go

@@ -32,6 +32,12 @@ const (
 	numHeaderType
 )
 
+var pkPool = sync.Pool{
+	New: func() interface{} {
+		return &PacketBuffer{}
+	},
+}
+
 // PacketBufferOptions specifies options for PacketBuffer creation.
 type PacketBufferOptions struct {
 	// ReserveHeaderBytes is the number of bytes to reserve for headers. Total
@@ -157,9 +163,9 @@ type PacketBuffer struct {
 
 // NewPacketBuffer creates a new PacketBuffer with opts.
 func NewPacketBuffer(opts PacketBufferOptions) *PacketBuffer {
-	pk := &PacketBuffer{
-		buf: &buffer.Buffer{},
-	}
+	pk := pkPool.Get().(*PacketBuffer)
+	pk.reset()
+	pk.buf = &buffer.Buffer{}
 	if opts.ReserveHeaderBytes != 0 {
 		pk.buf.AppendOwned(make([]byte, opts.ReserveHeaderBytes))
 		pk.reserved = opts.ReserveHeaderBytes
@@ -174,17 +180,27 @@ func NewPacketBuffer(opts PacketBufferOptions) *PacketBuffer {
 	return pk
 }
 
-// DecRef overrides refsvfs2 DecRef and passes a nil destroy function.
-func (pk *PacketBuffer) DecRef() {
-	pk.packetBufferRefs.DecRef(nil)
-}
-
 // PreserveObject marks this PacketBuffer so it is not recycled by internal
 // pooling.
 func (pk *PacketBuffer) PreserveObject() {
 	pk.preserveObject = true
 }
 
+// DecRef decrements the PacketBuffer's refcount. If the refcount is
+// decremented to zero, the PacketBuffer is returned to the PacketBuffer
+// pool.
+func (pk *PacketBuffer) DecRef() {
+	pk.packetBufferRefs.DecRef(func() {
+		if pk.packetBufferRefs.refCount == 0 && !pk.preserveObject {
+			pkPool.Put(pk)
+		}
+	})
+}
+
+func (pk *PacketBuffer) reset() {
+	*pk = PacketBuffer{}
+}
+
 // ReservedHeaderBytes returns the number of bytes initially reserved for
 // headers.
 func (pk *PacketBuffer) ReservedHeaderBytes() int {
@@ -307,26 +323,25 @@ func (pk *PacketBuffer) headerView(typ headerType) tcpipbuffer.View {
 // Clone makes a semi-deep copy of pk. The underlying packet payload is
 // shared. Hence, no modifications is done to underlying packet payload.
 func (pk *PacketBuffer) Clone() *PacketBuffer {
-	newPk := &PacketBuffer{
-		PacketBufferEntry:            pk.PacketBufferEntry,
-		buf:                          pk.buf.Clone(),
-		reserved:                     pk.reserved,
-		pushed:                       pk.pushed,
-		consumed:                     pk.consumed,
-		headers:                      pk.headers,
-		Hash:                         pk.Hash,
-		Owner:                        pk.Owner,
-		GSOOptions:                   pk.GSOOptions,
-		NetworkProtocolNumber:        pk.NetworkProtocolNumber,
-		DNATDone:                     pk.DNATDone,
-		SNATDone:                     pk.SNATDone,
-		TransportProtocolNumber:      pk.TransportProtocolNumber,
-		PktType:                      pk.PktType,
-		NICID:                        pk.NICID,
-		RXTransportChecksumValidated: pk.RXTransportChecksumValidated,
-		NetworkPacketInfo:            pk.NetworkPacketInfo,
-		tuple:                        pk.tuple,
-	}
+	newPk := pkPool.Get().(*PacketBuffer)
+	newPk.PacketBufferEntry = pk.PacketBufferEntry
+	newPk.buf = pk.buf.Clone()
+	newPk.reserved = pk.reserved
+	newPk.pushed = pk.pushed
+	newPk.consumed = pk.consumed
+	newPk.headers = pk.headers
+	newPk.Hash = pk.Hash
+	newPk.Owner = pk.Owner
+	newPk.GSOOptions = pk.GSOOptions
+	newPk.NetworkProtocolNumber = pk.NetworkProtocolNumber
+	newPk.DNATDone = pk.DNATDone
+	newPk.SNATDone = pk.SNATDone
+	newPk.TransportProtocolNumber = pk.TransportProtocolNumber
+	newPk.PktType = pk.PktType
+	newPk.NICID = pk.NICID
+	newPk.RXTransportChecksumValidated = pk.RXTransportChecksumValidated
+	newPk.NetworkPacketInfo = pk.NetworkPacketInfo
+	newPk.tuple = pk.tuple
 	newPk.InitRefs()
 	return newPk
 }
@@ -351,13 +366,13 @@ func (pk *PacketBuffer) Network() header.Network {
 // See PacketBuffer.Data for details about how a packet buffer holds an inbound
 // packet.
 func (pk *PacketBuffer) CloneToInbound() *PacketBuffer {
-	newPk := &PacketBuffer{
-		buf: pk.buf.Clone(),
-		// Treat unfilled header portion as reserved.
-		reserved: pk.AvailableHeaderBytes(),
-		tuple:    pk.tuple,
-	}
+	newPk := pkPool.Get().(*PacketBuffer)
+	newPk.reset()
+	newPk.buf = pk.buf.Clone()
 	newPk.InitRefs()
+	// Treat unfilled header portion as reserved.
+	newPk.reserved = pk.AvailableHeaderBytes()
+	newPk.tuple = pk.tuple
 	return newPk
 }
 
@@ -405,8 +420,14 @@ func (pk *PacketBufferList) IncRef() {
 // DecRef decreases the reference count on each PacketBuffer
 // stored in the PacketBufferList.
 func (pk *PacketBufferList) DecRef() {
-	for pb := pk.Front(); pb != nil; pb = pb.Next() {
+	// Using a while-loop here (instead of for-loop) because DecRef() can cause
+	// the pb to be recycled. If it is recycled during execution of this loop,
+	// there is a possibility of a data race during a call to pb.Next().
+	pb := pk.Front()
+	for pb != nil {
+		next := pb.Next()
 		pb.DecRef()
+		pb = next
 	}
 }