From 3c088b33a234395e2fdd5cb7c99fb8689b549976 Mon Sep 17 00:00:00 2001
From: H1JK <hell0jack@protonmail.com>
Date: Sat, 18 Nov 2023 21:40:50 +0800
Subject: [PATCH] chore: Shrink allocator pool range

---
 common/pool/alloc.go      | 85 ++++++++++++++-------------------------
 common/pool/alloc_test.go |  4 +-
 2 files changed, 32 insertions(+), 57 deletions(-)

diff --git a/common/pool/alloc.go b/common/pool/alloc.go
index ee3fa1a1..0a629403 100644
--- a/common/pool/alloc.go
+++ b/common/pool/alloc.go
@@ -12,7 +12,7 @@ var defaultAllocator = NewAllocator()
 
 // Allocator for incoming frames, optimized to prevent overwriting after zeroing
 type Allocator struct {
-	buffers [17]sync.Pool
+	buffers [11]sync.Pool
 }
 
 // NewAllocator initiates a []byte allocator for frames less than 65536 bytes,
@@ -20,13 +20,7 @@ type Allocator struct {
 // no more than 50%.
 func NewAllocator() *Allocator {
 	return &Allocator{
-		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) }},
@@ -52,46 +46,38 @@ func (alloc *Allocator) Get(size int) []byte {
 	case size > 65536:
 		return make([]byte, size)
 	default:
-		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")
@@ -110,44 +96,33 @@ func (alloc *Allocator) Put(buf []byte) error {
 	if 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")
diff --git a/common/pool/alloc_test.go b/common/pool/alloc_test.go
index eb66b808..0ac72ee0 100644
--- a/common/pool/alloc_test.go
+++ b/common/pool/alloc_test.go
@@ -13,8 +13,8 @@ func TestAllocGet(t *testing.T) {
 	assert.Equal(t, 1, len(alloc.Get(1)))
 	assert.Equal(t, 2, len(alloc.Get(2)))
 	assert.Equal(t, 3, len(alloc.Get(3)))
-	assert.Equal(t, 4, cap(alloc.Get(3)))
-	assert.Equal(t, 4, cap(alloc.Get(4)))
+	assert.Equal(t, 64, cap(alloc.Get(3)))
+	assert.Equal(t, 64, cap(alloc.Get(4)))
 	assert.Equal(t, 1023, len(alloc.Get(1023)))
 	assert.Equal(t, 1024, cap(alloc.Get(1023)))
 	assert.Equal(t, 1024, len(alloc.Get(1024)))