-
Notifications
You must be signed in to change notification settings - Fork 1
/
base2bufpool.go
169 lines (152 loc) · 3.5 KB
/
base2bufpool.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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
package base2bufpool
import (
"fmt"
"math/bits"
"sync"
)
type pool struct {
mtx sync.Mutex
bufs [][]byte
shift uint
}
func (p *pool) Put(buf []byte) {
p.mtx.Lock()
defer p.mtx.Unlock()
if len(buf) != 1<<p.shift {
panic(fmt.Sprintf("implementation error: %v %v", len(buf), 1<<p.shift))
}
if len(p.bufs) > 10 { // FIXME constant
return
}
p.bufs = append(p.bufs, buf)
}
func (p *pool) Get() []byte {
p.mtx.Lock()
defer p.mtx.Unlock()
if len(p.bufs) > 0 {
ret := p.bufs[len(p.bufs)-1]
p.bufs = p.bufs[0 : len(p.bufs)-1]
return ret
}
return make([]byte, 1<<p.shift)
}
type Pool struct {
minShift uint
maxShift uint
pools []pool
onNoFit NoFitBehavior
}
type Buffer struct {
// always power of 2, from Pool.pools
shiftBuf []byte
// presentedLen
payloadLen uint
// backref too pool for Free
pool *Pool
}
func (b Buffer) Bytes() []byte {
return b.shiftBuf[0:b.payloadLen]
}
func (b *Buffer) Shrink(newPayloadLen uint) {
if newPayloadLen > b.payloadLen {
panic(fmt.Sprintf("shrink is actually an expand, invalid: %v %v", newPayloadLen, b.payloadLen))
}
b.payloadLen = newPayloadLen
}
func (b Buffer) Free() {
if b.pool != nil {
b.pool.put(b)
}
}
//go:generate enumer -type NoFitBehavior
type NoFitBehavior uint
const (
AllocateSmaller NoFitBehavior = 1 << iota
AllocateLarger
Allocate NoFitBehavior = AllocateSmaller | AllocateLarger
Panic NoFitBehavior = 0
)
func New(minShift, maxShift uint, noFitBehavior NoFitBehavior) *Pool {
if minShift > 63 || maxShift > 63 {
panic(fmt.Sprintf("{min|max}Shift are the _exponent_, got minShift=%v maxShift=%v and limit of 63, which amounts to %v bits", minShift, maxShift, uint64(1)<<63))
}
pools := make([]pool, maxShift-minShift+1)
for i := uint(0); i < uint(len(pools)); i++ {
i := i // the closure below must copy i
pools[i] = pool{
shift: minShift + i,
bufs: make([][]byte, 0, 10),
}
}
return &Pool{
minShift: minShift,
maxShift: maxShift,
pools: pools,
onNoFit: noFitBehavior,
}
}
func fittingShift(x uint) uint {
if x == 0 {
return 0
}
blen := uint(bits.Len(x))
if 1<<(blen-1) == x {
return blen - 1
}
return blen
}
func (p *Pool) handlePotentialNoFit(reqShift uint) (buf Buffer, didHandle bool) {
if reqShift == 0 {
if p.onNoFit&AllocateSmaller != 0 {
return Buffer{[]byte{}, 0, nil}, true
} else {
goto doPanic
}
}
if reqShift < p.minShift {
if p.onNoFit&AllocateSmaller != 0 {
goto alloc
} else {
goto doPanic
}
}
if reqShift > p.maxShift {
if p.onNoFit&AllocateLarger != 0 {
goto alloc
} else {
goto doPanic
}
}
return Buffer{}, false
alloc:
return Buffer{make([]byte, 1<<reqShift), 1 << reqShift, nil}, true
doPanic:
panic(fmt.Sprintf("base2bufpool: configured to panic on shift=%v (minShift=%v maxShift=%v)", reqShift, p.minShift, p.maxShift))
}
func (p *Pool) Get(minSize uint) Buffer {
shift := fittingShift(minSize)
buf, didHandle := p.handlePotentialNoFit(shift)
if didHandle {
buf.Shrink(minSize)
return buf
}
idx := int64(shift) - int64(p.minShift)
return Buffer{p.pools[idx].Get(), minSize, p}
}
func (p *Pool) put(buffer Buffer) {
if buffer.pool != p {
panic("putting buffer to pool where it didn't originate from")
}
buf := buffer.shiftBuf
if bits.OnesCount(uint(len(buf))) > 1 {
panic(fmt.Sprintf("putting buffer that is not power of two len: %v", len(buf)))
}
if len(buf) == 0 {
return
}
shift := fittingShift(uint(len(buf)))
if shift < p.minShift || shift > p.maxShift {
return // drop it
}
p.pools[shift-p.minShift].Put(buf)
}