package bufferpool: new package, pool of bytes.Buffer

pkg/bufferpool/bufferpool.go

@@ -0,0 +1,112 @@
+package bufferpool
+
+import (
+	"bytes"
+	"sync"
+	"sync/atomic"
+)
+
+var minSize = 1024
+
+type BufferPool struct {
+	getCount uint32
+	putCount uint32
+	hitCount uint32
+	mallocCount uint32
+
+	m sync.Mutex
+	// TODO chef: 这个map可以预申请，做成fixed size的
+	sizeToList map[int]*[]*bytes.Buffer
+}
+
+func NewBufferPool() *BufferPool {
+	return &BufferPool{
+		sizeToList: make(map[int]*[]*bytes.Buffer),
+	}
+}
+
+func (bp *BufferPool) Get(size int) *bytes.Buffer {
+	atomic.AddUint32(&bp.getCount, 1)
+	ss := up2power(size)
+	if ss < minSize {
+		ss = minSize
+	}
+
+	bp.m.Lock()
+	l, ok := bp.sizeToList[ss]
+	if !ok {
+		bp.m.Unlock()
+		return bp.newBuffer(ss)
+	} else {
+		if len(*l) == 0 {
+			bp.m.Unlock()
+			return bp.newBuffer(ss)
+		}
+		buf := (*l)[len(*l)-1]
+		*l = (*l)[:len(*l)-1]
+		bp.m.Unlock()
+		buf.Reset()
+		atomic.AddUint32(&bp.hitCount, 1)
+		return buf
+	}
+}
+
+func (bp *BufferPool) Put(buf *bytes.Buffer) {
+	atomic.AddUint32(&bp.putCount, 1)
+	size := down2power(buf.Cap())
+	if size < minSize {
+		size = minSize
+	}
+
+	bp.m.Lock()
+	l, ok := bp.sizeToList[size]
+	if !ok {
+		l = new([]*bytes.Buffer)
+		*l = append(*l, buf)
+		// TODO
+		bp.sizeToList[size] = l
+	} else {
+		*l = append(*l, buf)
+	}
+	bp.m.Unlock()
+}
+
+func (bp *BufferPool) newBuffer(n int) *bytes.Buffer {
+	var buf bytes.Buffer
+	buf.Grow(n)
+	atomic.AddUint32(&bp.mallocCount, 1)
+	return &buf
+}
+
+// @return 范围为 [2, 4, 8, 16, ..., 1073741824]，如果大于等于1073741824，则直接返回n
+func up2power(n int) int {
+	if n >= 1073741824 {
+		return n
+	}
+
+	var i uint32
+	for ; n > (2 << i); i++ {
+	}
+	return 2 << i
+}
+
+// @return 范围为 [2, 4, 8, 16, ..., 1073741824]
+func down2power(n int) int {
+	if n < 2 {
+		return 2
+	} else if n >= 1073741824 {
+		return 1073741824
+	}
+
+	var i uint32
+	for {
+		nn := 2 << i
+		if n > nn {
+			i++
+		} else if n == nn {
+			return n
+		} else if n < nn {
+			return 2 << (i - 1)
+		}
+	}
+}

pkg/bufferpool/bufferpool_test.go

@@ -0,0 +1,133 @@
+package bufferpool
+
+import (
+	"bytes"
+	"github.com/q191201771/naza/pkg/assert"
+	"math/rand"
+	"testing"
+	"time"
+)
+
+var bp *BufferPool
+var count int
+
+func TestBufferPool(t *testing.T) {
+	bp := NewBufferPool()
+	buf := &bytes.Buffer{}
+	bp.Put(buf)
+	buf = bp.Get(4096)
+	buf.Grow(4096)
+	bp.Put(buf)
+	buf = bp.Get(4096)
+	bp.Put(buf)
+}
+
+func size() int {
+	//return 1024
+
+	//ss := []int{1000, 2000, 5000}
+	////ss := []int{128, 1024, 4096, 16384}
+	//count++
+	//return ss[count % 3]
+
+	return random(0, 128 * 1024)
+}
+
+func random(l, r int) int {
+	return l + (rand.Int() % (r - l))
+}
+
+func origin() {
+	var buf bytes.Buffer
+	size := size()
+	buf.Grow(size)
+}
+
+func bufferPool() {
+	size := size()
+	buf := bp.Get(size)
+	buf.Grow(size)
+	bp.Put(buf)
+}
+
+func BenchmarkOrigin(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		origin()
+	}
+}
+
+func BenchmarkBufferPool(b *testing.B) {
+	bp = NewBufferPool()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		bufferPool()
+	}
+}
+
+func BenchmarkOriginParallel(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				origin()
+			}
+		})
+	}
+}
+
+func BenchmarkBufferPoolParallel(b *testing.B) {
+	bp = NewBufferPool()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				bufferPool()
+			}
+		})
+	}
+}
+
+func TestUp2power(t *testing.T) {
+	assert.Equal(t, 2, up2power(0))
+	assert.Equal(t, 2, up2power(1))
+	assert.Equal(t, 2, up2power(2))
+	assert.Equal(t, 4, up2power(3))
+	assert.Equal(t, 4, up2power(4))
+	assert.Equal(t, 8, up2power(5))
+	assert.Equal(t, 8, up2power(6))
+	assert.Equal(t, 8, up2power(7))
+	assert.Equal(t, 8, up2power(8))
+	assert.Equal(t, 16, up2power(9))
+	assert.Equal(t, 1024, up2power(1023))
+	assert.Equal(t, 1024, up2power(1024))
+	assert.Equal(t, 2048, up2power(1025))
+	assert.Equal(t, 1073741824, up2power(1073741824-1))
+	assert.Equal(t, 1073741824, up2power(1073741824))
+	assert.Equal(t, 1073741824+1, up2power(1073741824+1))
+	assert.Equal(t, 2047483647-1, up2power(2047483647-1))
+	assert.Equal(t, 2047483647, up2power(2047483647))
+}
+
+func TestDown2power(t *testing.T) {
+	assert.Equal(t, 2, down2power(0))
+	assert.Equal(t, 2, down2power(1))
+	assert.Equal(t, 2, down2power(2))
+	assert.Equal(t, 2, down2power(3))
+	assert.Equal(t, 4, down2power(4))
+	assert.Equal(t, 4, down2power(5))
+	assert.Equal(t, 4, down2power(6))
+	assert.Equal(t, 4, down2power(7))
+	assert.Equal(t, 8, down2power(8))
+	assert.Equal(t, 8, down2power(9))
+	assert.Equal(t, 512, down2power(1023))
+	assert.Equal(t, 1024, down2power(1024))
+	assert.Equal(t, 1024, down2power(1025))
+	assert.Equal(t, 1073741824 >> 1, down2power(1073741824-1))
+	assert.Equal(t, 1073741824, down2power(1073741824))
+	assert.Equal(t, 1073741824, down2power(1073741824+1))
+	assert.Equal(t, 1073741824, down2power(2047483647-1))
+	assert.Equal(t, 1073741824, down2power(2047483647))
+}
+
+func init() {
+	rand.Seed(time.Now().Unix())
+}

test.sh

@@ -38,4 +38,4 @@ done
 
 # go test -race -coverprofile=profile.out -covermode=atomic && go tool cover -html=profile.out -o coverage.html && open coverage.html
 # go test -test.bench=".*"
-# go test -bench=. -benchtime=10s
+# go test -bench=. -benchmem -benchtime=10s