From 01f8f2db2f6e791b121c3ab139cfbdd726c74158 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Tue, 10 Jun 2025 10:54:08 +0800
Subject: [PATCH] chore: cleanup allocator code

---
 common/pool/alloc.go             | 19 ++++++++++++-------
 common/pool/buffer_low_memory.go |  7 +++----
 common/pool/buffer_standard.go   |  9 ++++-----
 common/pool/pool.go              |  4 ++--
 common/pool/sing.go              |  2 +-
 5 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/common/pool/alloc.go b/common/pool/alloc.go
index 80927a2c..cb586442 100644
--- a/common/pool/alloc.go
+++ b/common/pool/alloc.go
@@ -8,18 +8,23 @@ import (
 	"sync"
 )
 
-var defaultAllocator = NewAllocator()
+var DefaultAllocator = NewAllocator()
 
-// Allocator for incoming frames, optimized to prevent overwriting after zeroing
-type Allocator struct {
+type Allocator interface {
+	Get(size int) []byte
+	Put(buf []byte) error
+}
+
+// defaultAllocator for incoming frames, optimized to prevent overwriting after zeroing
+type defaultAllocator struct {
 	buffers [11]sync.Pool
 }
 
 // NewAllocator initiates a []byte allocator for frames less than 65536 bytes,
 // the waste(memory fragmentation) of space allocation is guaranteed to be
 // no more than 50%.
-func NewAllocator() *Allocator {
-	return &Allocator{
+func NewAllocator() Allocator {
+	return &defaultAllocator{
 		buffers: [...]sync.Pool{ // 64B -> 64K
 			{New: func() any { return new([1 << 6]byte) }},
 			{New: func() any { return new([1 << 7]byte) }},
@@ -37,7 +42,7 @@ func NewAllocator() *Allocator {
 }
 
 // Get a []byte from pool with most appropriate cap
-func (alloc *Allocator) Get(size int) []byte {
+func (alloc *defaultAllocator) Get(size int) []byte {
 	switch {
 	case size < 0:
 		panic("alloc.Get: len out of range")
@@ -87,7 +92,7 @@ func (alloc *Allocator) Get(size int) []byte {
 
 // Put returns a []byte to pool for future use,
 // which the cap must be exactly 2^n
-func (alloc *Allocator) Put(buf []byte) error {
+func (alloc *defaultAllocator) Put(buf []byte) error {
 	if cap(buf) == 0 || cap(buf) > 65536 {
 		return nil
 	}
diff --git a/common/pool/buffer_low_memory.go b/common/pool/buffer_low_memory.go
index 24e18a75..a165a37c 100644
--- a/common/pool/buffer_low_memory.go
+++ b/common/pool/buffer_low_memory.go
@@ -3,13 +3,12 @@
 package pool
 
 const (
+	// RelayBufferSize using for tcp
 	// io.Copy default buffer size is 32 KiB
-	// but the maximum packet size of vmess/shadowsocks is about 16 KiB
-	// so define a buffer of 20 KiB to reduce the memory of each TCP relay
 	RelayBufferSize = 16 * 1024
 
-	// RelayBufferSize uses 20KiB, but due to the allocator it will actually
-	// request 32Kib. Most UDPs are smaller than the MTU, and the TUN's MTU
+	// UDPBufferSize using for udp
+	// Most UDPs are smaller than the MTU, and the TUN's MTU
 	// set to 9000, so the UDP Buffer size set to 16Kib
 	UDPBufferSize = 8 * 1024
 )
diff --git a/common/pool/buffer_standard.go b/common/pool/buffer_standard.go
index ff758700..e0ff1195 100644
--- a/common/pool/buffer_standard.go
+++ b/common/pool/buffer_standard.go
@@ -3,13 +3,12 @@
 package pool
 
 const (
+	// RelayBufferSize using for tcp
 	// io.Copy default buffer size is 32 KiB
-	// but the maximum packet size of vmess/shadowsocks is about 16 KiB
-	// so define a buffer of 20 KiB to reduce the memory of each TCP relay
-	RelayBufferSize = 20 * 1024
+	RelayBufferSize = 32 * 1024
 
-	// RelayBufferSize uses 20KiB, but due to the allocator it will actually
-	// request 32Kib. Most UDPs are smaller than the MTU, and the TUN's MTU
+	// UDPBufferSize using for udp
+	// Most UDPs are smaller than the MTU, and the TUN's MTU
 	// set to 9000, so the UDP Buffer size set to 16Kib
 	UDPBufferSize = 16 * 1024
 )
diff --git a/common/pool/pool.go b/common/pool/pool.go
index 288ea467..3d324120 100644
--- a/common/pool/pool.go
+++ b/common/pool/pool.go
@@ -1,9 +1,9 @@
 package pool
 
 func Get(size int) []byte {
-	return defaultAllocator.Get(size)
+	return DefaultAllocator.Get(size)
 }
 
 func Put(buf []byte) error {
-	return defaultAllocator.Put(buf)
+	return DefaultAllocator.Put(buf)
 }
diff --git a/common/pool/sing.go b/common/pool/sing.go
index a2c2b060..5506233d 100644
--- a/common/pool/sing.go
+++ b/common/pool/sing.go
@@ -3,5 +3,5 @@ package pool
 import "github.com/metacubex/sing/common/buf"
 
 func init() {
-	buf.DefaultAllocator = defaultAllocator
+	buf.DefaultAllocator = DefaultAllocator
 }