bytes pool (#14)

* bytes pool

bytes/bytes_pool.go

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package gxbytes
+
+import (
+	"sync"
+)
+
+// BytesPool hold specific size []byte
+type BytesPool struct {
+	sizes  []int // sizes declare the cap of each slot
+	slots  []sync.Pool
+	length int
+}
+
+var defaultBytesPool = NewBytesPool([]int{512, 1 << 10, 4 << 10, 16 << 10, 64 << 10})
+
+// NewBytesPool ...
+func NewBytesPool(slotSize []int) *BytesPool {
+	bp := &BytesPool{}
+	bp.sizes = slotSize
+	bp.length = len(bp.sizes)
+
+	bp.slots = make([]sync.Pool, bp.length)
+	for i, size := range bp.sizes {
+		size := size
+		bp.slots[i] = sync.Pool{New: func() interface{} {
+			return make([]byte, 0, size)
+		}}
+	}
+	return bp
+}
+
+// SetDefaultBytesPool change default pool options
+func SetDefaultBytesPool(bp *BytesPool) {
+	defaultBytesPool = bp
+}
+
+func (bp *BytesPool) findIndex(size int) int {
+	for i := 0; i < bp.length; i++ {
+		if bp.sizes[i] >= size {
+			return i
+		}
+	}
+	return bp.length
+}
+
+// AcquireBytes get specific make([]byte, 0, size)
+func (bp *BytesPool) AcquireBytes(size int) []byte {
+	idx := bp.findIndex(size)
+	if idx >= bp.length {
+		return make([]byte, 0, size)
+	}
+
+	return bp.slots[idx].Get().([]byte)[:size]
+}
+
+// ReleaseBytes ...
+func (bp *BytesPool) ReleaseBytes(buf []byte) {
+	idx := bp.findIndex(cap(buf))
+	if idx >= bp.length {
+		return
+	}
+
+	bp.slots[idx].Put(buf)
+}
+
+// AcquireBytes called by defaultBytesPool
+func AcquireBytes(size int) []byte { return defaultBytesPool.AcquireBytes(size) }
+
+// ReleaseBytes called by defaultBytesPool
+func ReleaseBytes(buf []byte) { defaultBytesPool.ReleaseBytes(buf) }

bytes/bytes_pool_test.go

+package gxbytes
+
+import (
+	"fmt"
+	"testing"
+)
+
+func Test_findIndex(t *testing.T) {
+	bp := NewBytesPool([]int{16, 4 << 10, 16 << 10, 32 << 10, 64 << 10})
+
+	type args struct {
+		size int
+	}
+	tests := []struct {
+		args args
+		want int
+	}{
+		{args{1}, 0},
+		{args{15}, 0},
+		{args{16}, 0},
+		{args{17}, 1},
+		{args{4095}, 1},
+		{args{4096}, 1},
+		{args{4097}, 2},
+		{args{16 << 10}, 2},
+		{args{64 << 10}, 4},
+		{args{64 << 11}, 5},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprint(tt.args.size), func(t *testing.T) {
+			if got := bp.findIndex(tt.args.size); got != tt.want {
+				t.Errorf("[%v] findIndex() = %v, want %v", tt.args.size, got, tt.want)
+			}
+		})
+	}
+}
+
+func BenchmarkAcquireBytesSize32(b *testing.B)  { benchmarkAcquireBytes(b, 32) }
+func BenchmarkAcquireBytesSize10k(b *testing.B) { benchmarkAcquireBytes(b, 10000) }
+func BenchmarkAcquireBytesSize60k(b *testing.B) { benchmarkAcquireBytes(b, 60000) }
+func BenchmarkAcquireBytesSize70k(b *testing.B) { benchmarkAcquireBytes(b, 70000) }
+
+func benchmarkAcquireBytes(b *testing.B, size int) {
+	for i := 0; i < b.N; i++ {
+		bytes := AcquireBytes(size)
+		ReleaseBytes(bytes)
+	}
+}
+
+func BenchmarkFindIndexSize8(b *testing.B)   { benchmarkfindIndex(b, 8) }
+func BenchmarkFindIndexSize60k(b *testing.B) { benchmarkfindIndex(b, 60000) }
+
+func benchmarkfindIndex(b *testing.B, size int) {
+	for i := 0; i < b.N; i++ {
+		defaultBytesPool.findIndex(size)
+	}
+}

bytes/slice_pool.go

@@ -17,14 +17,9 @@
 
 package gxbytes
 
-import (
-	"sync"
-)
-
 const (
 	minShift = 6
 	maxShift = 18
-	errSlot  = -1
 )
 
 var (
@@ -36,98 +31,46 @@ func init() {
 }
 
 // SlicePool is []byte pools
+// Deprecated
 type SlicePool struct {
-	minShift int
-	minSize  int
-	maxSize  int
-
-	pool []*sliceSlot
-}
-
-type sliceSlot struct {
-	defaultSize int
-	pool        sync.Pool
+	BytesPool
 }
 
 // newSlicePool returns SlicePool
+// Deprecated
+// instead by NewBytesPool
 func NewSlicePool() *SlicePool {
-	p := &SlicePool{
-		minShift: minShift,
-		minSize:  1 << minShift,
-		maxSize:  1 << maxShift,
+	sizes := make([]int, 0, maxShift-minShift+1)
+	for i := minShift; i <= maxShift; i++ {
+		sizes = append(sizes, 1<<uint(i))
 	}
-	for i := 0; i <= maxShift-minShift; i++ {
-		slab := &sliceSlot{
-			defaultSize: 1 << (uint)(i+minShift),
-		}
-		p.pool = append(p.pool, slab)
+	p := &SlicePool{
+		*NewBytesPool(sizes),
 	}
-
 	return p
 }
 
-func (p *SlicePool) slot(size int) int {
-	if size > p.maxSize {
-		return errSlot
-	}
-	slot := 0
-	shift := 0
-	if size > p.minSize {
-		size--
-		for size > 0 {
-			size = size >> 1
-			shift++
-		}
-		slot = shift - p.minShift
-	}
-
-	return slot
-}
-
-func newSlice(size int) []byte {
-	return make([]byte, size)
-}
-
 // Get returns *[]byte from SlicePool
 func (p *SlicePool) Get(size int) *[]byte {
-	slot := p.slot(size)
-	if slot == errSlot {
-		b := newSlice(size)
+	b := p.AcquireBytes(size)
 	return &b
-	}
-	v := p.pool[slot].pool.Get()
-	if v == nil {
-		b := newSlice(p.pool[slot].defaultSize)
-		b = b[0:size]
-		return &b
-	}
-	b := v.(*[]byte)
-	*b = (*b)[0:size]
-	return b
 }
 
 // Put returns *[]byte to SlicePool
 func (p *SlicePool) Put(buf *[]byte) {
-	if buf == nil {
-		return
-	}
-	size := cap(*buf)
-	slot := p.slot(size)
-	if slot == errSlot {
-		return
-	}
-	if size != int(p.pool[slot].defaultSize) {
-		return
-	}
-	p.pool[slot].pool.Put(buf)
+	p.ReleaseBytes(*buf)
 }
 
 // GetBytes returns *[]byte from SlicePool
+// Deprecated
+// instead by AcquireBytes
 func GetBytes(size int) *[]byte {
 	return defaultSlicePool.Get(size)
 }
 
 // PutBytes Put *[]byte to SlicePool
+// Deprecated
+// instead by ReleaseBytes
 func PutBytes(buf *[]byte) {
 	defaultSlicePool.Put(buf)
 }

bytes/slice_pool_test.go

@@ -77,31 +77,3 @@ func TestSlicePoolLargeBytes(t *testing.T) {
 		pool.Put(bp)
 	}
 }
-
-func TestBytesSlot(t *testing.T) {
-	pool := NewSlicePool()
-
-	if pool.slot(pool.minSize-1) != 0 {
-		t.Errorf("Expect get the 0 slot")
-	}
-
-	if pool.slot(pool.minSize) != 0 {
-		t.Errorf("Expect get the 0 slot")
-	}
-
-	if pool.slot(pool.minSize+1) != 1 {
-		t.Errorf("Expect get the 1 slot")
-	}
-
-	if pool.slot(pool.maxSize-1) != maxShift-minShift {
-		t.Errorf("Expect get the %d slot", maxShift-minShift)
-	}
-
-	if pool.slot(pool.maxSize) != maxShift-minShift {
-		t.Errorf("Expect get the %d slot", maxShift-minShift)
-	}
-
-	if pool.slot(pool.maxSize+1) != errSlot {
-		t.Errorf("Expect get errSlot")
-	}
-}