multipool -> mpool

pool.go

+// Package mpool provides a sync.Pool equivalent that buckets incoming
+// requests to one of 32 sub-pools, one for each power of 2, 0-32.
+//
+//	import "github.com/jbenet/go-msgio/mpool"
+//	var p mpool.Pool
+//
+//	small := make([]byte, 1024)
+//	large := make([]byte, 4194304)
+//	p.Put(1024, small)
+//	p.Put(4194304, large)
+//
+//	small2 := p.Get(1024).([]byte)
+//	large2 := p.Get(4194304).([]byte)
+//	fmt.Println("small2 len:", len(small2))
+//	fmt.Println("large2 len:", len(large2))
+//
+//	// Output:
+//	// small2 len: 1024
+//	// large2 len: 4194304
+//
+package mpool
+
+import (
+	"fmt"
+	"sync"
+)
+
+// ByteSlicePool is a static Pool for reusing byteslices of various sizes.
+var ByteSlicePool Pool
+
+func init() {
+	ByteSlicePool.New = func(length int) interface{} {
+		return make([]byte, length)
+	}
+}
+
+// MaxLength is the maximum length of an element that can be added to the Pool.
+const MaxLength = 1 << 32
+
+// Pool is a pool to handle cases of reusing elements of varying sizes.
+// It maintains up to  32 internal pools, for each power of 2 in 0-32.
+type Pool struct {
+	small      int            // the size of the first pool
+	pools      [32]*sync.Pool // a list of singlePools
+	sync.Mutex                // protecting list
+
+	// New is a function that constructs a new element in the pool, with given len
+	New func(len int) interface{}
+}
+
+func (p *Pool) getPool(idx uint32) *sync.Pool {
+	if idx > uint32(len(p.pools)) {
+		panic(fmt.Errorf("index too large: %d", idx))
+	}
+
+	p.Lock()
+	defer p.Unlock()
+
+	sp := p.pools[idx]
+	if sp == nil {
+		sp = new(sync.Pool)
+		p.pools[idx] = sp
+	}
+	return sp
+}
+
+// Get selects an arbitrary item from the Pool, removes it from the Pool,
+// and returns it to the caller. Get may choose to ignore the pool and
+// treat it as empty. Callers should not assume any relation between values
+// passed to Put and the values returned by Get.
+//
+// If Get would otherwise return nil and p.New is non-nil, Get returns the
+// result of calling p.New.
+func (p *Pool) Get(length uint32) interface{} {
+	idx := largerPowerOfTwo(length)
+	sp := p.getPool(idx)
+	val := sp.Get()
+	if val == nil && p.New != nil {
+		val = p.New(0x1 << idx)
+	}
+	return val
+}
+
+// Put adds x to the pool.
+func (p *Pool) Put(length uint32, val interface{}) {
+	idx := smallerPowerOfTwo(length)
+	sp := p.getPool(idx)
+	sp.Put(val)
+}
+
+func largerPowerOfTwo(num uint32) uint32 {
+	for p := uint32(0); p < 32; p++ {
+		if (0x1 << p) >= num {
+			return p
+		}
+	}
+
+	panic("unreachable")
+}
+
+func smallerPowerOfTwo(num uint32) uint32 {
+	for p := uint32(1); p < 32; p++ {
+		if (0x1 << p) > num {
+			return p - 1
+		}
+	}
+
+	panic("unreachable")
+}

pool_test.go

+// Copyright 2013 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Pool is no-op under race detector, so all these tests do not work.
+// +build !race
+
+package mpool
+
+import (
+	"fmt"
+	"runtime"
+	"runtime/debug"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestPool(t *testing.T) {
+	// disable GC so we can control when it happens.
+	defer debug.SetGCPercent(debug.SetGCPercent(-1))
+	var p Pool
+	if p.Get(10) != nil {
+		t.Fatal("expected empty")
+	}
+	p.Put(16, "a")
+	p.Put(2048, "b")
+	if g := p.Get(16); g != "a" {
+		t.Fatalf("got %#v; want a", g)
+	}
+	if g := p.Get(2048); g != "b" {
+		t.Fatalf("got %#v; want b", g)
+	}
+	if g := p.Get(16); g != nil {
+		t.Fatalf("got %#v; want nil", g)
+	}
+	if g := p.Get(2048); g != nil {
+		t.Fatalf("got %#v; want nil", g)
+	}
+	if g := p.Get(1); g != nil {
+		t.Fatalf("got %#v; want nil", g)
+	}
+	p.Put(1023, "d")
+	if g := p.Get(1024); g != nil {
+		t.Fatalf("got %#v; want nil", g)
+	}
+	if g := p.Get(512); g != "d" {
+		t.Fatalf("got %#v; want d", g)
+	}
+
+	debug.SetGCPercent(100) // to allow following GC to actually run
+	runtime.GC()
+	if g := p.Get(10); g != nil {
+		t.Fatalf("got %#v; want nil after GC", g)
+	}
+}
+
+func TestPoolNew(t *testing.T) {
+	// disable GC so we can control when it happens.
+	defer debug.SetGCPercent(debug.SetGCPercent(-1))
+
+	s := [32]int{}
+	p := Pool{
+		New: func(length int) interface{} {
+			idx := largerPowerOfTwo(uint32(length))
+			s[idx]++
+			return s[idx]
+		},
+	}
+	if v := p.Get(1 << 5); v != 1 {
+		t.Fatalf("got %v; want 1", v)
+	}
+	if v := p.Get(1 << 2); v != 1 {
+		t.Fatalf("got %v; want 1", v)
+	}
+	if v := p.Get(1 << 2); v != 2 {
+		t.Fatalf("got %v; want 2", v)
+	}
+	if v := p.Get(1 << 5); v != 2 {
+		t.Fatalf("got %v; want 2", v)
+	}
+	p.Put(1<<2, 42)
+	p.Put(1<<5, 42)
+	if v := p.Get(1 << 2); v != 42 {
+		t.Fatalf("got %v; want 42", v)
+	}
+	if v := p.Get(1 << 2); v != 3 {
+		t.Fatalf("got %v; want 3", v)
+	}
+	if v := p.Get(1 << 5); v != 42 {
+		t.Fatalf("got %v; want 42", v)
+	}
+	if v := p.Get(1 << 5); v != 3 {
+		t.Fatalf("got %v; want 3", v)
+	}
+}
+
+// Test that Pool does not hold pointers to previously cached
+// resources
+func TestPoolGC(t *testing.T) {
+	var p Pool
+	var fin uint32
+	const N = 100
+	for i := 0; i < N; i++ {
+		v := new(string)
+		runtime.SetFinalizer(v, func(vv *string) {
+			atomic.AddUint32(&fin, 1)
+		})
+		p.Put(uint32(i), v)
+	}
+	for i := 0; i < N; i++ {
+		p.Get(uint32(i))
+	}
+	for i := 0; i < 5; i++ {
+		runtime.GC()
+		time.Sleep(time.Duration(i*100+10) * time.Millisecond)
+		// 1 pointer can remain on stack or elsewhere
+		if atomic.LoadUint32(&fin) >= N-1 {
+			return
+		}
+	}
+	t.Fatalf("only %v out of %v resources are finalized",
+		atomic.LoadUint32(&fin), N)
+}
+
+func TestPoolStress(t *testing.T) {
+	const P = 10
+	N := int(1e6)
+	if testing.Short() {
+		N /= 100
+	}
+	var p Pool
+	done := make(chan bool)
+	for i := 0; i < P; i++ {
+		go func() {
+			var v interface{} = 0
+			for j := 0; j < N; j++ {
+				if v == nil {
+					v = 0
+				}
+				p.Put(uint32(j), v)
+				v = p.Get(uint32(j))
+				if v != nil && v.(int) != 0 {
+					t.Fatalf("expect 0, got %v", v)
+				}
+			}
+			done <- true
+		}()
+	}
+	for i := 0; i < P; i++ {
+		<-done
+	}
+}
+
+func BenchmarkPool(b *testing.B) {
+	var p Pool
+	b.RunParallel(func(pb *testing.PB) {
+		i := 0
+		for pb.Next() {
+			i = i << 1
+			p.Put(uint32(i), 1)
+			p.Get(uint32(i))
+		}
+	})
+}
+
+func BenchmarkPoolOverlflow(b *testing.B) {
+	var p Pool
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			for pow := uint32(0); pow < 32; pow++ {
+				for b := 0; b < 100; b++ {
+					p.Put(uint32(1<<pow), 1)
+				}
+			}
+			for pow := uint32(0); pow < 32; pow++ {
+				for b := 0; b < 100; b++ {
+					p.Get(uint32(1 << pow))
+				}
+			}
+		}
+	})
+}
+
+func ExamplePool() {
+	var p Pool
+
+	small := make([]byte, 1024)
+	large := make([]byte, 4194304)
+	p.Put(uint32(len(small)), small)
+	p.Put(uint32(len(large)), large)
+
+	small2 := p.Get(uint32(len(small))).([]byte)
+	large2 := p.Get(uint32(len(large))).([]byte)
+	fmt.Println("small2 len:", len(small2))
+	fmt.Println("large2 len:", len(large2))
+	// Output:
+	// small2 len: 1024
+	// large2 len: 4194304
+}