From 7d8db149b68145099a12a72367a32909316d84fb Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 19 Apr 2023 09:58:54 +0200
Subject: [PATCH] introduce a buffer pool for large buffers (20k)

---
 buffer_pool.go                | 29 +++++++++++++++++++++--------
 buffer_pool_test.go           | 12 ++++++++----
 internal/protocol/protocol.go |  5 ++++-
 sys_conn_oob_test.go          |  2 +-
 sys_conn_test.go              |  2 +-
 5 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/buffer_pool.go b/buffer_pool.go
index f6745b08..7d676c84 100644
--- a/buffer_pool.go
+++ b/buffer_pool.go
@@ -56,13 +56,18 @@ func (b *packetBuffer) Len() protocol.ByteCount {
 }
 
 func (b *packetBuffer) putBack() {
-	if cap(b.Data) != int(protocol.MaxPacketBufferSize) {
-		panic("putPacketBuffer called with packet of wrong size!")
+	if cap(b.Data) == protocol.MaxPacketBufferSize {
+		bufferPool.Put(b)
+		return
 	}
-	bufferPool.Put(b)
+	if cap(b.Data) == protocol.MaxLargePacketBufferSize {
+		largeBufferPool.Put(b)
+		return
+	}
+	panic("putPacketBuffer called with packet of wrong size!")
 }
 
-var bufferPool sync.Pool
+var bufferPool, largeBufferPool sync.Pool
 
 func getPacketBuffer() *packetBuffer {
 	buf := bufferPool.Get().(*packetBuffer)
@@ -71,10 +76,18 @@ func getPacketBuffer() *packetBuffer {
 	return buf
 }
 
+func getLargePacketBuffer() *packetBuffer {
+	buf := largeBufferPool.Get().(*packetBuffer)
+	buf.refCount = 1
+	buf.Data = buf.Data[:0]
+	return buf
+}
+
 func init() {
-	bufferPool.New = func() interface{} {
-		return &packetBuffer{
-			Data: make([]byte, 0, protocol.MaxPacketBufferSize),
-		}
+	bufferPool.New = func() any {
+		return &packetBuffer{Data: make([]byte, 0, protocol.MaxPacketBufferSize)}
+	}
+	largeBufferPool.New = func() any {
+		return &packetBuffer{Data: make([]byte, 0, protocol.MaxLargePacketBufferSize)}
 	}
 }
diff --git a/buffer_pool_test.go b/buffer_pool_test.go
index 7e7a28ea..c565d4a4 100644
--- a/buffer_pool_test.go
+++ b/buffer_pool_test.go
@@ -9,13 +9,17 @@ import (
 
 var _ = Describe("Buffer Pool", func() {
 	It("returns buffers of cap", func() {
-		buf := getPacketBuffer()
-		Expect(buf.Data).To(HaveCap(int(protocol.MaxPacketBufferSize)))
+		buf1 := getPacketBuffer()
+		Expect(buf1.Data).To(HaveCap(protocol.MaxPacketBufferSize))
+		buf2 := getLargePacketBuffer()
+		Expect(buf2.Data).To(HaveCap(protocol.MaxLargePacketBufferSize))
 	})
 
 	It("releases buffers", func() {
-		buf := getPacketBuffer()
-		buf.Release()
+		buf1 := getPacketBuffer()
+		buf1.Release()
+		buf2 := getLargePacketBuffer()
+		buf2.Release()
 	})
 
 	It("gets the length", func() {
diff --git a/internal/protocol/protocol.go b/internal/protocol/protocol.go
index 93ad96cc..98bc9ffb 100644
--- a/internal/protocol/protocol.go
+++ b/internal/protocol/protocol.go
@@ -59,7 +59,10 @@ type StatelessResetToken [16]byte
 // ethernet's max size, minus the IP and UDP headers. IPv6 has a 40 byte header,
 // UDP adds an additional 8 bytes.  This is a total overhead of 48 bytes.
 // Ethernet's max packet size is 1500 bytes,  1500 - 48 = 1452.
-const MaxPacketBufferSize ByteCount = 1452
+const MaxPacketBufferSize = 1452
+
+// MaxLargePacketBufferSize is used when using GSO
+const MaxLargePacketBufferSize = 20 * 1024
 
 // MinInitialPacketSize is the minimum size an Initial packet is required to have.
 const MinInitialPacketSize = 1200
diff --git a/sys_conn_oob_test.go b/sys_conn_oob_test.go
index eab73df4..7e0bf7c5 100644
--- a/sys_conn_oob_test.go
+++ b/sys_conn_oob_test.go
@@ -216,7 +216,7 @@ var _ = Describe("OOB Conn Test", func() {
 				Expect(ms).To(HaveLen(batchSize))
 				for i := 0; i < numMsgRead; i++ {
 					Expect(ms[i].Buffers).To(HaveLen(1))
-					Expect(ms[i].Buffers[0]).To(HaveLen(int(protocol.MaxPacketBufferSize)))
+					Expect(ms[i].Buffers[0]).To(HaveLen(protocol.MaxPacketBufferSize))
 					data := []byte(fmt.Sprintf("message %d", counter))
 					counter++
 					ms[i].Buffers[0] = data
diff --git a/sys_conn_test.go b/sys_conn_test.go
index f3bbdacf..418e2c31 100644
--- a/sys_conn_test.go
+++ b/sys_conn_test.go
@@ -18,7 +18,7 @@ var _ = Describe("Basic Conn Test", func() {
 		addr := &net.UDPAddr{IP: net.IPv4(1, 2, 3, 4), Port: 1234}
 		c.EXPECT().ReadFrom(gomock.Any()).DoAndReturn(func(b []byte) (int, net.Addr, error) {
 			data := []byte("foobar")
-			Expect(b).To(HaveLen(int(protocol.MaxPacketBufferSize)))
+			Expect(b).To(HaveLen(protocol.MaxPacketBufferSize))
 			return copy(b, data), addr, nil
 		})