mirror of
https://github.com/v2fly/v2ray-core.git
synced 2024-09-12 06:54:23 -04:00
106 lines
1.9 KiB
Go
106 lines
1.9 KiB
Go
package alloc
|
|
|
|
import (
|
|
"os"
|
|
"strconv"
|
|
"sync"
|
|
)
|
|
|
|
type Pool interface {
|
|
Allocate() *Buffer
|
|
Free(*Buffer)
|
|
}
|
|
|
|
type SyncPool struct {
|
|
allocator *sync.Pool
|
|
}
|
|
|
|
func NewSyncPool(bufferSize uint32) *SyncPool {
|
|
pool := &SyncPool{
|
|
allocator: &sync.Pool{
|
|
New: func() interface{} { return make([]byte, bufferSize) },
|
|
},
|
|
}
|
|
return pool
|
|
}
|
|
|
|
func (p *SyncPool) Allocate() *Buffer {
|
|
return CreateBuffer(p.allocator.Get().([]byte), p)
|
|
}
|
|
|
|
func (p *SyncPool) Free(buffer *Buffer) {
|
|
rawBuffer := buffer.head
|
|
if rawBuffer == nil {
|
|
return
|
|
}
|
|
p.allocator.Put(rawBuffer)
|
|
}
|
|
|
|
type BufferPool struct {
|
|
chain chan []byte
|
|
allocator *sync.Pool
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
func (p *BufferPool) Allocate() *Buffer {
|
|
var b []byte
|
|
select {
|
|
case b = <-p.chain:
|
|
default:
|
|
b = p.allocator.Get().([]byte)
|
|
}
|
|
return CreateBuffer(b, p)
|
|
}
|
|
|
|
func (p *BufferPool) Free(buffer *Buffer) {
|
|
rawBuffer := buffer.head
|
|
if rawBuffer == nil {
|
|
return
|
|
}
|
|
select {
|
|
case p.chain <- rawBuffer:
|
|
default:
|
|
p.allocator.Put(rawBuffer)
|
|
}
|
|
}
|
|
|
|
const (
|
|
mediumBufferByteSize = 8 * 1024
|
|
BufferSize = mediumBufferByteSize - defaultOffset
|
|
|
|
smallBufferByteSize = 2 * 1024
|
|
SmallBufferSize = smallBufferByteSize - defaultOffset
|
|
|
|
PoolSizeEnvKey = "v2ray.buffer.size"
|
|
)
|
|
|
|
var (
|
|
mediumPool *BufferPool
|
|
smallPool = NewSyncPool(2048)
|
|
)
|
|
|
|
func init() {
|
|
var size uint32 = 20
|
|
sizeStr := os.Getenv(PoolSizeEnvKey)
|
|
if len(sizeStr) > 0 {
|
|
customSize, err := strconv.ParseUint(sizeStr, 10, 32)
|
|
if err == nil {
|
|
size = uint32(customSize)
|
|
}
|
|
}
|
|
totalByteSize := size * 1024 * 1024
|
|
mediumPool = NewBufferPool(mediumBufferByteSize, totalByteSize/mediumBufferByteSize)
|
|
}
|