aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--weed/util/mem/slot_pool.go44
-rw-r--r--weed/util/mem/slot_pool_test.go37
2 files changed, 81 insertions, 0 deletions
diff --git a/weed/util/mem/slot_pool.go b/weed/util/mem/slot_pool.go
new file mode 100644
index 000000000..e6680d3cb
--- /dev/null
+++ b/weed/util/mem/slot_pool.go
@@ -0,0 +1,44 @@
+package mem
+
+import "sync"
+
+var pools []*sync.Pool
+
+const (
+ min_size = 1024
+)
+
+func bitCount(size int) (count int) {
+ for ; size > min_size; count++ {
+ size = size >> 1
+ }
+ return
+}
+
+func init() {
+ // 1KB ~ 256MB
+ pools = make([]*sync.Pool, bitCount(1024*1024*256))
+ for i := 0; i < len(pools); i++ {
+ slotSize := 1024 << i
+ pools[i] = &sync.Pool{
+ New: func() interface{} {
+ buffer := make([]byte, slotSize)
+ return &buffer
+ },
+ }
+ }
+}
+
+func getSlotPool(size int) *sync.Pool {
+ index := bitCount(size)
+ return pools[index]
+}
+
+func Allocate(size int) []byte {
+ slab := *getSlotPool(size).Get().(*[]byte)
+ return slab[:size]
+}
+
+func Free(buf []byte) {
+ getSlotPool(cap(buf)).Put(&buf)
+}
diff --git a/weed/util/mem/slot_pool_test.go b/weed/util/mem/slot_pool_test.go
new file mode 100644
index 000000000..114951522
--- /dev/null
+++ b/weed/util/mem/slot_pool_test.go
@@ -0,0 +1,37 @@
+package mem
+
+import (
+ "testing"
+)
+
+func TestAllocateFree(t *testing.T) {
+ buf := Allocate(12)
+ Free(buf)
+ if cap(buf) != min_size {
+ t.Errorf("min size error allocated capacity=%d", cap(buf))
+ }
+ if len(buf) != 12 {
+ t.Errorf("size error")
+ }
+
+ buf = Allocate(4883)
+ Free(buf)
+ if cap(buf) != 1024<<bitCount(4883) {
+ t.Errorf("min size error allocated capacity=%d", cap(buf))
+ }
+ if len(buf) != 4883 {
+ t.Errorf("size error")
+ }
+
+}
+
+func TestBitCount(t *testing.T) {
+ count := bitCount(12)
+ if count != 0 {
+ t.Errorf("bitCount error count=%d", count)
+ }
+ if count != bitCount(min_size) {
+ t.Errorf("bitCount error count=%d", count)
+ }
+
+}