improve buffer pool

lxzan · May 3, 2024 · 2e9ba55 · 2e9ba55
1 parent bf09c37
commit 2e9ba55
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 173 deletions.
diff --git a/assets/performance-compress-disabled.png b/assets/performance-compress-disabled.png
diff --git a/init.go b/init.go
@@ -3,7 +3,7 @@ package gws
 import "github.com/lxzan/gws/internal"
 
 var (
-	framePadding  = frameHeader{}            // 帧头填充物
-	binaryPool    = internal.NewBufferPool() // 缓冲池
-	defaultLogger = new(stdLogger)           // 默认日志工具
+	framePadding  = frameHeader{}                         // 帧头填充物
+	binaryPool    = internal.NewBufferPool(128, 256*1024) // 缓冲池
+	defaultLogger = new(stdLogger)                        // 默认日志工具
 )
diff --git a/internal/pool.go b/internal/pool.go
@@ -5,96 +5,60 @@ import (
 	"sync"
 )
 
-const (
-	poolSize = 10
-
-	Lv1 = 128
-	Lv2 = 1024
-	Lv3 = 2 * 1024
-	Lv4 = 4 * 1024
-	Lv5 = 8 * 1024
-	Lv6 = 16 * 1024
-	Lv7 = 32 * 1024
-	Lv8 = 64 * 1024
-	Lv9 = 128 * 1024
-)
-
 type BufferPool struct {
-	pools  []*sync.Pool
-	limits []int
+	begin      uint32
+	shards     []*sync.Pool
+	size2index map[int]int
 }
 
-func NewBufferPool() *BufferPool {
-	var p = &BufferPool{
-		pools:  make([]*sync.Pool, poolSize),
-		limits: []int{0, Lv1, Lv2, Lv3, Lv4, Lv5, Lv6, Lv7, Lv8, Lv9},
-	}
-	for i := 1; i < poolSize; i++ {
-		var capacity = p.limits[i]
-		p.pools[i] = &sync.Pool{New: func() any {
-			return bytes.NewBuffer(make([]byte, 0, capacity))
-		}}
+// NewBufferPool Creating a memory pool
+// Left, right indicate the interval range of the memory pool, they will be transformed into pow(2,n)。
+// Below left, Get method will return at least left bytes; above right, Put method will not reclaim the buffer.
+func NewBufferPool(left, right uint32) *BufferPool {
+	var begin, end = int(binaryCeil(left)), int(binaryCeil(right))
+	var p = &BufferPool{begin: uint32(begin), size2index: map[int]int{}}
+	for i, j := begin, 0; i <= end; i *= 2 {
+		capacity := i
+		pool := &sync.Pool{New: func() any { return bytes.NewBuffer(make([]byte, 0, capacity)) }}
+		p.shards = append(p.shards, pool)
+		p.size2index[i] = j
+		j++
 	}
 	return p
 }
 
+// Put Return buffer to memory pool
 func (p *BufferPool) Put(b *bytes.Buffer) {
-	if b == nil || b.Cap() == 0 {
-		return
-	}
-	if index := p.getIndex(uint32(b.Cap())); index > 0 {
-		p.pools[index].Put(b)
+	if b != nil {
+		if index, ok := p.size2index[b.Cap()]; ok {
+			p.shards[index].Put(b)
+		}
 	}
 }
 
+// Get Fetch a buffer from the memory pool, of at least n bytes
 func (p *BufferPool) Get(n int) *bytes.Buffer {
-	var index = p.getIndex(uint32(n))
-	if index == 0 {
-		return bytes.NewBuffer(make([]byte, 0, n))
-	}
-
-	b := p.pools[index].Get().(*bytes.Buffer)
-	if b.Cap() < n {
-		b.Grow(p.limits[index])
+	var size = int(Max(binaryCeil(uint32(n)), p.begin))
+	if index, ok := p.size2index[size]; ok {
+		b := p.shards[index].Get().(*bytes.Buffer)
+		if b.Cap() < size {
+			b.Grow(size)
+		}
+		b.Reset()
+		return b
 	}
-	b.Reset()
-	return b
+	return bytes.NewBuffer(make([]byte, 0, n))
 }
 
-func (p *BufferPool) getIndex(v uint32) int {
-	if v > Lv9 {
-		return 0
-	}
-	if v <= 128 {
-		return 1
-	}
-
+func binaryCeil(v uint32) uint32 {
 	v--
 	v |= v >> 1
 	v |= v >> 2
 	v |= v >> 4
 	v |= v >> 8
 	v |= v >> 16
 	v++
-
-	switch v {
-	case Lv3:
-		return 3
-	case Lv4:
-		return 4
-	case Lv5:
-		return 5
-	case Lv6:
-		return 6
-	case Lv7:
-		return 7
-	case Lv8:
-		return 8
-	case Lv9:
-		return 9
-	default:
-		return 2
-	}
+	return v
 }
 
 func NewPool[T any](f func() T) *Pool[T] {

diff --git a/internal/pool_test.go b/internal/pool_test.go
@@ -9,7 +9,7 @@ import (
 
 func TestBufferPool(t *testing.T) {
 	var as = assert.New(t)
-	var pool = NewBufferPool()
+	var pool = NewBufferPool(128, 128*1024)
 
 	for i := 0; i < 10; i++ {
 		var n = AlphabetNumeric.Intn(126)
@@ -19,17 +19,17 @@ func TestBufferPool(t *testing.T) {
 	}
 	for i := 0; i < 10; i++ {
 		var buf = pool.Get(500)
-		as.Equal(Lv2, buf.Cap())
+		as.Equal(512, buf.Cap())
 		as.Equal(0, buf.Len())
 	}
 	for i := 0; i < 10; i++ {
 		var buf = pool.Get(2000)
-		as.Equal(Lv3, buf.Cap())
+		as.Equal(2048, buf.Cap())
 		as.Equal(0, buf.Len())
 	}
 	for i := 0; i < 10; i++ {
 		var buf = pool.Get(5000)
-		as.Equal(Lv5, buf.Cap())
+		as.Equal(8192, buf.Cap())
 		as.Equal(0, buf.Len())
 	}
 
@@ -56,58 +56,3 @@ func TestPool(t *testing.T) {
 	assert.Equal(t, 0, p.Get())
 	p.Put(1)
 }
-
-func TestBufferPool_GetIndex(t *testing.T) {
-	var p = NewBufferPool()
-	assert.Equal(t, p.getIndex(200*1024), 0)
-
-	assert.Equal(t, p.getIndex(0), 1)
-	assert.Equal(t, p.getIndex(1), 1)
-	assert.Equal(t, p.getIndex(10), 1)
-	assert.Equal(t, p.getIndex(100), 1)
-	assert.Equal(t, p.getIndex(128), 1)
-
-	assert.Equal(t, p.getIndex(200), 2)
-	assert.Equal(t, p.getIndex(1000), 2)
-	assert.Equal(t, p.getIndex(500), 2)
-	assert.Equal(t, p.getIndex(1024), 2)
-
-	assert.Equal(t, p.getIndex(2*1024), 3)
-	assert.Equal(t, p.getIndex(2000), 3)
-	assert.Equal(t, p.getIndex(1025), 3)
-
-	assert.Equal(t, p.getIndex(4*1024), 4)
-	assert.Equal(t, p.getIndex(3000), 4)
-	assert.Equal(t, p.getIndex(2*1024+1), 4)
-
-	assert.Equal(t, p.getIndex(8*1024), 5)
-	assert.Equal(t, p.getIndex(5000), 5)
-	assert.Equal(t, p.getIndex(4*1024+1), 5)
-
-	assert.Equal(t, p.getIndex(16*1024), 6)
-	assert.Equal(t, p.getIndex(10000), 6)
-	assert.Equal(t, p.getIndex(8*1024+1), 6)
-
-	assert.Equal(t, p.getIndex(32*1024), 7)
-	assert.Equal(t, p.getIndex(20000), 7)
-	assert.Equal(t, p.getIndex(16*1024+1), 7)
-
-	assert.Equal(t, p.getIndex(64*1024), 8)
-	assert.Equal(t, p.getIndex(40000), 8)
-	assert.Equal(t, p.getIndex(32*1024+1), 8)
-
-	assert.Equal(t, p.getIndex(128*1024), 9)
-	assert.Equal(t, p.getIndex(100000), 9)
-	assert.Equal(t, p.getIndex(64*1024+1), 9)
-}
-
-func BenchmarkPool_GetIndex(b *testing.B) {
-	var p = NewBufferPool()
-
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		for j := 0; j < 1000000; j++ {
-			p.getIndex(uint32(j))
-		}
-	}
-}
diff --git a/internal/utils.go b/internal/utils.go
@@ -171,13 +171,6 @@ func SelectValue[T any](ok bool, a, b T) T {
 	return b
 }
 
-func IsNil(v any) bool {
-	if v == nil {
-		return true
-	}
-	return reflect.ValueOf(v).IsNil()
-}
-
 func ToBinaryNumber[T Integer](n T) T {
 	var x T = 1
 	for x < n {
@@ -212,14 +205,14 @@ func WithDefault[T comparable](rawValue, newValue T) T {
 	return rawValue
 }
 
-func Min[T int | int64](a, b T) T {
+func Min(a, b int) int {
 	if a < b {
 		return a
 	}
 	return b
 }
 
-func Max[T int | int64](a, b T) T {
+func Max(a, b uint32) uint32 {
 	if a > b {
 		return a
 	}
@@ -237,12 +230,3 @@ func IsSameSlice[T comparable](a, b []T) bool {
 	}
 	return true
 }
-
-func CheckErrors(errs ...error) error {
-	for _, item := range errs {
-		if item != nil {
-			return item
-		}
-	}
-	return nil
-}
diff --git a/internal/utils_test.go b/internal/utils_test.go
@@ -4,10 +4,8 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/hex"
-	"errors"
 	"hash/fnv"
 	"io"
-	"net/http"
 	"reflect"
 	"strings"
 	"testing"
@@ -147,18 +145,6 @@ func TestSelectInt(t *testing.T) {
 	assert.Equal(t, 2, SelectValue(false, 1, 2))
 }
 
-func TestIsNil(t *testing.T) {
-	{
-		var v io.Reader
-		assert.True(t, IsNil(v))
-	}
-	{
-		var v *http.Request
-		var v1 any = v
-		assert.True(t, IsNil(v1))
-	}
-}
-
 func TestToBinaryNumber(t *testing.T) {
 	assert.Equal(t, 8, ToBinaryNumber(7))
 	assert.Equal(t, 1, ToBinaryNumber(0))
@@ -227,8 +213,8 @@ func TestMin(t *testing.T) {
 }
 
 func TestMax(t *testing.T) {
-	assert.Equal(t, Max(1, 2), 2)
-	assert.Equal(t, Max(4, 3), 4)
+	assert.Equal(t, Max(1, 2), uint32(2))
+	assert.Equal(t, Max(4, 3), uint32(4))
 }
 
 func TestIsSameSlice(t *testing.T) {
@@ -247,12 +233,3 @@ func TestIsSameSlice(t *testing.T) {
 		[]int{1, 2, 4},
 	))
 }
-
-func TestCheckErrors(t *testing.T) {
-	var err0 error
-	var err1 error
-	var err2 = errors.New("1")
-	assert.NoError(t, CheckErrors(err0, err1))
-	assert.Error(t, CheckErrors(err0, err1, err2))
-	assert.True(t, errors.Is(CheckErrors(err0, err1, err2), err2))
-}