package bufferpool: Strategy

pkg/bufferpool/bufferpool.go

@@ -19,24 +19,29 @@ var (
 )
 
 type bufferPool struct {
-	status          Status
+	strategy        Strategy
+	singleBucket    Bucket
 	capToFreeBucket map[int]Bucket
+	status          Status
 }
 
 func (bp *bufferPool) Get(size int) *bytes.Buffer {
 	atomic.AddInt64(&bp.status.getCount, 1)
-	ss := up2power(size)
-	if ss < minSize {
-		ss = minSize
+
+	var bucket Bucket
+	if bp.strategy == StategyMultiStdPoolBucket || bp.strategy == StategyMultiSlicePoolBucket {
+		ss := up2power(size)
+		if ss < minSize {
+			ss = minSize
+		}
+		bucket = bp.capToFreeBucket[ss]
+	} else {
+		bucket = bp.singleBucket
 	}
 
-	bucket := bp.capToFreeBucket[ss]
 	buf := bucket.Get()
 	if buf == nil {
-		var buf bytes.Buffer
-		buf.Grow(ss)
-		atomic.AddInt64(&bp.status.mallocCount, 1)
-		return &buf
+		return &bytes.Buffer{}
 	}
 
 	atomic.AddInt64(&bp.status.hitCount, 1)
@@ -48,22 +53,28 @@ func (bp *bufferPool) Put(buf *bytes.Buffer) {
 	c := buf.Cap()
 	atomic.AddInt64(&bp.status.putCount, 1)
 	atomic.AddInt64(&bp.status.sizeBytes, int64(c))
-	size := down2power(c)
-	if size < minSize {
-		size = minSize
+
+	var bucket Bucket
+	if bp.strategy == StategyMultiStdPoolBucket || bp.strategy == StategyMultiSlicePoolBucket {
+		size := down2power(c)
+		if size < minSize {
+			size = minSize
+		}
+
+		bucket = bp.capToFreeBucket[size]
+	} else {
+		bucket = bp.singleBucket
 	}
 
-	bucket := bp.capToFreeBucket[size]
 	bucket.Put(buf)
 }
 
 func (bp *bufferPool) RetrieveStatus() Status {
 	return Status{
-		getCount:    atomic.LoadInt64(&bp.status.getCount),
-		putCount:    atomic.LoadInt64(&bp.status.putCount),
-		hitCount:    atomic.LoadInt64(&bp.status.hitCount),
-		mallocCount: atomic.LoadInt64(&bp.status.mallocCount),
-		sizeBytes:   atomic.LoadInt64(&bp.status.sizeBytes),
+		getCount:  atomic.LoadInt64(&bp.status.getCount),
+		putCount:  atomic.LoadInt64(&bp.status.putCount),
+		hitCount:  atomic.LoadInt64(&bp.status.hitCount),
+		sizeBytes: atomic.LoadInt64(&bp.status.sizeBytes),
 	}
 }
 

pkg/bufferpool/bufferpool_benchmark_test.go

@@ -59,7 +59,7 @@ func BenchmarkOrigin(b *testing.B) {
 }
 
 func BenchmarkBufferPool(b *testing.B) {
-	bp = NewBufferPool()
+	bp = NewBufferPool(StategyMultiStdPoolBucket)
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		bufferPoolFunc()

pkg/bufferpool/bufferpool_test.go

@@ -17,16 +17,26 @@ import (
 
 func TestBufferPool(t *testing.T) {
 	// TODO chef: assert result
-	bp := NewBufferPool()
-	buf := &bytes.Buffer{}
-	bp.Get(128)
-	bp.Put(buf)
-	buf = bp.Get(128)
-	buf.Grow(4096)
-	bp.Put(buf)
-	buf = bp.Get(4096)
-	bp.Put(buf)
-	bp.RetrieveStatus()
+
+	strategyList := []Strategy{
+		StrategySingleStdPoolBucket,
+		StrategySingleSlicePoolBucket,
+		StategyMultiStdPoolBucket,
+		StategyMultiSlicePoolBucket,
+	}
+
+	for _, s := range strategyList {
+		bp := NewBufferPool(s)
+		buf := &bytes.Buffer{}
+		bp.Get(128)
+		bp.Put(buf)
+		buf = bp.Get(128)
+		buf.Grow(4096)
+		bp.Put(buf)
+		buf = bp.Get(4096)
+		bp.Put(buf)
+		bp.RetrieveStatus()
+	}
 }
 
 func TestGlobal(t *testing.T) {

pkg/bufferpool/global.go

@@ -24,6 +24,10 @@ func RetrieveStatus() Status {
 	return global.RetrieveStatus()
 }
 
+func Init(strategy Strategy) {
+	global = NewBufferPool(strategy)
+}
+
 func init() {
-	global = NewBufferPool()
+	Init(StategyMultiStdPoolBucket)
 }

pkg/bufferpool/interface.go

@@ -10,10 +10,8 @@ package bufferpool
 
 import "bytes"
 
-// 使用sync.Pool作为底层bucket实现时，池内自由选择合适的时机，自动释放空闲Buffer
-
 type BufferPool interface {
-	// 获取一个已经预申请大于<size>大小的Buffer对象，如果池内没有满足条件的空闲Buffer，会向Go内存管理模块申请
+	// 获取一个预估容量为<size>的Buffer对象，不同的策略返回的Buffer对象实际容量可能会有不同，但是不会返回nil
 	Get(size int) *bytes.Buffer
 
 	// 将Buffer对象放回池中
@@ -24,21 +22,61 @@ type BufferPool interface {
 }
 
 type Status struct {
-	getCount    int64 // 调用Get方法的次数
-	putCount    int64 // 调用Put方法的次数
-	hitCount    int64 // 调用Get方法时，池内存在满足条件的空闲Buffer，这种情况的计数
-	mallocCount int64 // 调用Get方法时，池内不存在满足条件的空闲Buffer，向Go内存管理模块申请，这种情况的计数
-	sizeBytes   int64 // 池内所有空闲Buffer占用的内存大小，单位字节
+	getCount  int64 // 调用Get方法的次数
+	putCount  int64 // 调用Put方法的次数
+	hitCount  int64 // 调用Get方法时，池内存在满足条件的空闲Buffer，这种情况的计数
+	sizeBytes int64 // 池内所有空闲Buffer占用的内存大小，单位字节
 }
 
-func NewBufferPool() BufferPool {
-	capToFreeBucket := make(map[int]Bucket)
-	for i := minSize; i <= maxSize; i <<= 1 {
-		capToFreeBucket[i] = NewStdPoolBucket()
-		//capToFreeBucket[i] = NewSliceBucket()
+type Strategy int
+
+const (
+	// 直接使用一个标准库中的sync.Pool
+	StrategySingleStdPoolBucket Strategy = iota + 1
+
+	// 直接使用一个切片存储所有的Buffer对象
+	StrategySingleSlicePoolBucket
+
+	// 按Buffer对象的容量哈希到不同的桶中，每个桶是一个sync.Pool
+	StategyMultiStdPoolBucket
+
+	// 按Buffer对象的容量哈希到不同的桶中，每个桶是一个切片
+	StategyMultiSlicePoolBucket
+)
+
+type Bucket interface {
+	// 桶内没有Buffer对象时，返回nil
+	Get() *bytes.Buffer
+
+	Put(buf *bytes.Buffer)
+}
+
+func NewBufferPool(strategy Strategy) BufferPool {
+	var (
+		singleBucket    Bucket
+		capToFreeBucket map[int]Bucket
+	)
+
+	switch strategy {
+	case StrategySingleStdPoolBucket:
+		singleBucket = NewStdPoolBucket()
+	case StrategySingleSlicePoolBucket:
+		singleBucket = NewSliceBucket()
+	case StategyMultiStdPoolBucket:
+		capToFreeBucket = make(map[int]Bucket)
+		for i := minSize; i <= maxSize; i <<= 1 {
+			capToFreeBucket[i] = NewStdPoolBucket()
+		}
+	case StategyMultiSlicePoolBucket:
+		capToFreeBucket = make(map[int]Bucket)
+		for i := minSize; i <= maxSize; i <<= 1 {
+			capToFreeBucket[i] = NewSliceBucket()
+		}
 	}
 
 	return &bufferPool{
+		strategy:        strategy,
+		singleBucket:    singleBucket,
 		capToFreeBucket: capToFreeBucket,
 	}
 }

pkg/bufferpool/stdpool_bucket.go

@@ -13,11 +13,6 @@ import (
 	"sync"
 )
 
-type Bucket interface {
-	Get() *bytes.Buffer
-	Put(buf *bytes.Buffer)
-}
-
 type StdPoolBucket struct {
 	core *sync.Pool
 }