pkgs/bytes: init bytes pool

pkgs/bytes/pool.go

@@ -0,0 +1,31 @@
+package bytes
+
+import "sync"
+
+type BytesPool struct {
+	size uint
+	pool sync.Pool
+}
+
+func NewBytesPool(size uint) *BytesPool {
+	return &BytesPool{
+		size: size,
+		pool: sync.Pool{
+			New: func() any {
+				return make([]byte, size)
+			},
+		},
+	}
+}
+
+func (p *BytesPool) Get() []byte {
+	return p.pool.Get().([]byte)
+}
+
+func (p *BytesPool) Put(data []byte) {
+	if uint(len(data)) != p.size {
+		panic("buf size not match pool size")
+	}
+
+	p.pool.Put(data)
+}

pkgs/bytes/pool_test.go

@@ -0,0 +1,33 @@
+package bytes
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPoolGet(t *testing.T) {
+	pool := NewBytesPool(200)
+	assert.Equal(t, 200, len(pool.Get()))
+
+	pool = NewBytesPool(1025)
+	assert.Equal(t, 1025, len(pool.Get()))
+
+	pool = NewBytesPool(2 * 1024)
+	assert.Equal(t, 2*1024, len(pool.Get()))
+
+	pool = NewBytesPool(5 * 2000)
+	assert.Equal(t, 5*2000, len(pool.Get()))
+}
+
+func TestPoolPut(t *testing.T) {
+	pool := NewBytesPool(200)
+
+	assert.Panics(t, func() {
+		pool.Put(make([]byte, 2000))
+	})
+
+	assert.NotPanics(t, func() {
+		pool.Put(make([]byte, 200))
+	})
+}