From 978ca2e01d8388abf7ab38f7431e14d9b47e54fb Mon Sep 17 00:00:00 2001
From: H1JK <hell0jack@protonmail.com>
Date: Sat, 18 Nov 2023 21:37:58 +0800
Subject: [PATCH] Shrink buf pool range

---
 common/buf/alloc.go | 85 ++++++++++++++++-----------------------------
 1 file changed, 30 insertions(+), 55 deletions(-)

diff --git a/common/buf/alloc.go b/common/buf/alloc.go
index 69048c4..b556d93 100644
--- a/common/buf/alloc.go
+++ b/common/buf/alloc.go
@@ -17,7 +17,7 @@ type Allocator interface {
 
 // defaultAllocator for incoming frames, optimized to prevent overwriting after zeroing
 type defaultAllocator struct {
-	buffers [17]sync.Pool
+	buffers [11]sync.Pool
 }
 
 // NewAllocator initiates a []byte allocator for frames less than 65536 bytes,
@@ -25,13 +25,7 @@ type defaultAllocator struct {
 // no more than 50%.
 func newDefaultAllocator() Allocator {
 	return &defaultAllocator{
-		buffers: [...]sync.Pool{ // 1B -> 64K
-			{New: func() any { return new([1]byte) }},
-			{New: func() any { return new([1 << 1]byte) }},
-			{New: func() any { return new([1 << 2]byte) }},
-			{New: func() any { return new([1 << 3]byte) }},
-			{New: func() any { return new([1 << 4]byte) }},
-			{New: func() any { return new([1 << 5]byte) }},
+		buffers: [...]sync.Pool{ // 64B -> 64K
 			{New: func() any { return new([1 << 6]byte) }},
 			{New: func() any { return new([1 << 7]byte) }},
 			{New: func() any { return new([1 << 8]byte) }},
@@ -53,46 +47,38 @@ func (alloc *defaultAllocator) Get(size int) []byte {
 		return nil
 	}
 
-	index := msb(size)
-	if size != 1<<index {
-		index += 1
+	var index uint16
+	if size > 64 {
+		index = msb(size)
+		if size != 1<<index {
+			index += 1
+		}
+		index -= 6
 	}
 
 	buffer := alloc.buffers[index].Get()
 	switch index {
 	case 0:
-		return buffer.(*[1]byte)[:size]
-	case 1:
-		return buffer.(*[1 << 1]byte)[:size]
-	case 2:
-		return buffer.(*[1 << 2]byte)[:size]
-	case 3:
-		return buffer.(*[1 << 3]byte)[:size]
-	case 4:
-		return buffer.(*[1 << 4]byte)[:size]
-	case 5:
-		return buffer.(*[1 << 5]byte)[:size]
-	case 6:
 		return buffer.(*[1 << 6]byte)[:size]
-	case 7:
+	case 1:
 		return buffer.(*[1 << 7]byte)[:size]
-	case 8:
+	case 2:
 		return buffer.(*[1 << 8]byte)[:size]
-	case 9:
+	case 3:
 		return buffer.(*[1 << 9]byte)[:size]
-	case 10:
+	case 4:
 		return buffer.(*[1 << 10]byte)[:size]
-	case 11:
+	case 5:
 		return buffer.(*[1 << 11]byte)[:size]
-	case 12:
+	case 6:
 		return buffer.(*[1 << 12]byte)[:size]
-	case 13:
+	case 7:
 		return buffer.(*[1 << 13]byte)[:size]
-	case 14:
+	case 8:
 		return buffer.(*[1 << 14]byte)[:size]
-	case 15:
+	case 9:
 		return buffer.(*[1 << 15]byte)[:size]
-	case 16:
+	case 10:
 		return buffer.(*[1 << 16]byte)[:size]
 	default:
 		panic("invalid pool index")
@@ -106,44 +92,33 @@ func (alloc *defaultAllocator) Put(buf []byte) error {
 	if cap(buf) == 0 || cap(buf) > 65536 || cap(buf) != 1<<bits {
 		return errors.New("allocator Put() incorrect buffer size")
 	}
+	bits -= 6
 	buf = buf[:cap(buf)]
 
 	//nolint
 	//lint:ignore SA6002 ignore temporarily
 	switch bits {
 	case 0:
-		alloc.buffers[bits].Put((*[1]byte)(buf))
-	case 1:
-		alloc.buffers[bits].Put((*[1 << 1]byte)(buf))
-	case 2:
-		alloc.buffers[bits].Put((*[1 << 2]byte)(buf))
-	case 3:
-		alloc.buffers[bits].Put((*[1 << 3]byte)(buf))
-	case 4:
-		alloc.buffers[bits].Put((*[1 << 4]byte)(buf))
-	case 5:
-		alloc.buffers[bits].Put((*[1 << 5]byte)(buf))
-	case 6:
 		alloc.buffers[bits].Put((*[1 << 6]byte)(buf))
-	case 7:
+	case 1:
 		alloc.buffers[bits].Put((*[1 << 7]byte)(buf))
-	case 8:
+	case 2:
 		alloc.buffers[bits].Put((*[1 << 8]byte)(buf))
-	case 9:
+	case 3:
 		alloc.buffers[bits].Put((*[1 << 9]byte)(buf))
-	case 10:
+	case 4:
 		alloc.buffers[bits].Put((*[1 << 10]byte)(buf))
-	case 11:
+	case 5:
 		alloc.buffers[bits].Put((*[1 << 11]byte)(buf))
-	case 12:
+	case 6:
 		alloc.buffers[bits].Put((*[1 << 12]byte)(buf))
-	case 13:
+	case 7:
 		alloc.buffers[bits].Put((*[1 << 13]byte)(buf))
-	case 14:
+	case 8:
 		alloc.buffers[bits].Put((*[1 << 14]byte)(buf))
-	case 15:
+	case 9:
 		alloc.buffers[bits].Put((*[1 << 15]byte)(buf))
-	case 16:
+	case 10:
 		alloc.buffers[bits].Put((*[1 << 16]byte)(buf))
 	default:
 		panic("invalid pool index")