-
Notifications
You must be signed in to change notification settings - Fork 0
/
buffer_pool.go
135 lines (120 loc) · 2.61 KB
/
buffer_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
126
127
128
129
130
131
132
133
134
135
package buf
import (
"os"
"runtime"
"strconv"
"sync"
)
// Pool provides functionality to generate and recycle buffers on demand.
type Pool interface {
// Allocate either returns a unused buffer from the pool, or generates a new one from system.
Allocate() *Buffer
// Free recycles the given buffer.
Free(*Buffer)
}
// SyncPool is a buffer pool based on sync.Pool
type SyncPool struct {
allocator *sync.Pool
}
// NewSyncPool creates a SyncPool with given buffer size.
func NewSyncPool(bufferSize uint32) *SyncPool {
pool := &SyncPool{
allocator: &sync.Pool{
New: func() interface{} { return make([]byte, bufferSize) },
},
}
return pool
}
// Allocate implements Pool.Allocate().
func (p *SyncPool) Allocate() *Buffer {
return &Buffer{
v: p.allocator.Get().([]byte),
pool: p,
}
}
// Free implements Pool.Free().
func (p *SyncPool) Free(buffer *Buffer) {
rawBuffer := buffer.v
if rawBuffer == nil {
return
}
p.allocator.Put(rawBuffer)
}
// BufferPool is a Pool that utilizes an internal cache.
type BufferPool struct {
chain chan []byte
allocator *sync.Pool
}
// NewBufferPool creates a new BufferPool with given buffer size, and internal cache size.
func NewBufferPool(bufferSize, poolSize uint32) *BufferPool {
pool := &BufferPool{
chain: make(chan []byte, poolSize),
allocator: &sync.Pool{
New: func() interface{} { return make([]byte, bufferSize) },
},
}
for i := uint32(0); i < poolSize; i++ {
pool.chain <- make([]byte, bufferSize)
}
return pool
}
// Allocate implements Pool.Allocate().
func (p *BufferPool) Allocate() *Buffer {
var b []byte
select {
case b = <-p.chain:
default:
b = p.allocator.Get().([]byte)
}
return &Buffer{
v: b,
pool: p,
}
}
// Free implements Pool.Free().
func (p *BufferPool) Free(buffer *Buffer) {
rawBuffer := buffer.v
if rawBuffer == nil {
return
}
select {
case p.chain <- rawBuffer:
default:
p.allocator.Put(rawBuffer)
}
}
const (
// Size of a regular buffer.
Size = 8 * 1024
// SizeSmall is the size of a small buffer.
SizeSmall = 2 * 1024
poolSizeEnvKey = "v2ray.buffer.size"
)
var (
mediumPool Pool
smallPool = NewSyncPool(SizeSmall)
)
func getDefaultPoolSize() uint32 {
switch runtime.GOARCH {
case "amd64", "386":
return 20
default:
return 5
}
}
func init() {
size := getDefaultPoolSize()
sizeStr := os.Getenv(poolSizeEnvKey)
if len(sizeStr) > 0 {
customSize, err := strconv.ParseUint(sizeStr, 10, 32)
if err == nil {
size = uint32(customSize)
}
}
if size > 0 {
totalByteSize := size * 1024 * 1024
mediumPool = NewBufferPool(Size, totalByteSize/Size)
} else {
mediumPool = NewSyncPool(Size)
}
}