/
pool.go
125 lines (109 loc) · 2.92 KB
/
pool.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
// 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 := nextPowerOfTwo(length)
sp := p.getPool(idx)
// fmt.Printf("Get(%d) idx(%d)\n", length, 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 := prevPowerOfTwo(length)
// fmt.Printf("Put(%d, -) idx(%d)\n", length, idx)
sp := p.getPool(idx)
sp.Put(val)
}
func nextPowerOfTwo(v uint32) uint32 {
// fmt.Printf("nextPowerOfTwo(%d) ", v)
v--
v |= v >> 1
v |= v >> 2
v |= v >> 4
v |= v >> 8
v |= v >> 16
v++
// fmt.Printf("-> %d", v)
i := uint32(0)
for i = 0; v > 1; i++ {
v = v >> 1
}
// fmt.Printf("-> %d\n", i)
return i
}
func prevPowerOfTwo(num uint32) uint32 {
next := nextPowerOfTwo(num)
// fmt.Printf("prevPowerOfTwo(%d) next: %d", num, next)
switch {
case num == (1 << next): // num is a power of 2
case next == 0:
default:
next = next - 1 // smaller
}
// fmt.Printf(" = %d\n", next)
return next
}