mihomo/common/pool/alloc.go
2023-11-30 20:30:05 +08:00

161 lines
4.2 KiB
Go

package pool
// Inspired by https://github.com/xtaci/smux/blob/master/alloc.go
import (
"errors"
"math/bits"
"sync"
)
var defaultAllocator = NewAllocator()
// Allocator for incoming frames, optimized to prevent overwriting after zeroing
type Allocator struct {
buffers [17]sync.Pool
}
// NewAllocator initiates a []byte allocator for frames less than 65536 bytes,
// the waste(memory fragmentation) of space allocation is guaranteed to be
// no more than 50%.
func NewAllocator() *Allocator {
return &Allocator{
buffers: [...]sync.Pool{ // 1B -> 64K
{New: func() any { return new([1]byte) }},
{New: func() any { return new([1 << 1]byte) }},
{New: func() any { return new([1 << 2]byte) }},
{New: func() any { return new([1 << 3]byte) }},
{New: func() any { return new([1 << 4]byte) }},
{New: func() any { return new([1 << 5]byte) }},
{New: func() any { return new([1 << 6]byte) }},
{New: func() any { return new([1 << 7]byte) }},
{New: func() any { return new([1 << 8]byte) }},
{New: func() any { return new([1 << 9]byte) }},
{New: func() any { return new([1 << 10]byte) }},
{New: func() any { return new([1 << 11]byte) }},
{New: func() any { return new([1 << 12]byte) }},
{New: func() any { return new([1 << 13]byte) }},
{New: func() any { return new([1 << 14]byte) }},
{New: func() any { return new([1 << 15]byte) }},
{New: func() any { return new([1 << 16]byte) }},
},
}
}
// Get a []byte from pool with most appropriate cap
func (alloc *Allocator) Get(size int) []byte {
switch {
case size < 0:
panic("alloc.Get: len out of range")
case size == 0:
return nil
case size > 65536:
return make([]byte, size)
default:
index := msb(size)
if size != 1<<index {
index += 1
}
buffer := alloc.buffers[index].Get()
switch index {
case 0:
return buffer.(*[1]byte)[:size]
case 1:
return buffer.(*[1 << 1]byte)[:size]
case 2:
return buffer.(*[1 << 2]byte)[:size]
case 3:
return buffer.(*[1 << 3]byte)[:size]
case 4:
return buffer.(*[1 << 4]byte)[:size]
case 5:
return buffer.(*[1 << 5]byte)[:size]
case 6:
return buffer.(*[1 << 6]byte)[:size]
case 7:
return buffer.(*[1 << 7]byte)[:size]
case 8:
return buffer.(*[1 << 8]byte)[:size]
case 9:
return buffer.(*[1 << 9]byte)[:size]
case 10:
return buffer.(*[1 << 10]byte)[:size]
case 11:
return buffer.(*[1 << 11]byte)[:size]
case 12:
return buffer.(*[1 << 12]byte)[:size]
case 13:
return buffer.(*[1 << 13]byte)[:size]
case 14:
return buffer.(*[1 << 14]byte)[:size]
case 15:
return buffer.(*[1 << 15]byte)[:size]
case 16:
return buffer.(*[1 << 16]byte)[:size]
default:
panic("invalid pool index")
}
}
}
// Put returns a []byte to pool for future use,
// which the cap must be exactly 2^n
func (alloc *Allocator) Put(buf []byte) error {
if cap(buf) == 0 || cap(buf) > 65536 {
return nil
}
bits := msb(cap(buf))
if cap(buf) != 1<<bits {
return errors.New("allocator Put() incorrect buffer size")
}
buf = buf[:cap(buf)]
//nolint
//lint:ignore SA6002 ignore temporarily
switch bits {
case 0:
alloc.buffers[bits].Put((*[1]byte)(buf))
case 1:
alloc.buffers[bits].Put((*[1 << 1]byte)(buf))
case 2:
alloc.buffers[bits].Put((*[1 << 2]byte)(buf))
case 3:
alloc.buffers[bits].Put((*[1 << 3]byte)(buf))
case 4:
alloc.buffers[bits].Put((*[1 << 4]byte)(buf))
case 5:
alloc.buffers[bits].Put((*[1 << 5]byte)(buf))
case 6:
alloc.buffers[bits].Put((*[1 << 6]byte)(buf))
case 7:
alloc.buffers[bits].Put((*[1 << 7]byte)(buf))
case 8:
alloc.buffers[bits].Put((*[1 << 8]byte)(buf))
case 9:
alloc.buffers[bits].Put((*[1 << 9]byte)(buf))
case 10:
alloc.buffers[bits].Put((*[1 << 10]byte)(buf))
case 11:
alloc.buffers[bits].Put((*[1 << 11]byte)(buf))
case 12:
alloc.buffers[bits].Put((*[1 << 12]byte)(buf))
case 13:
alloc.buffers[bits].Put((*[1 << 13]byte)(buf))
case 14:
alloc.buffers[bits].Put((*[1 << 14]byte)(buf))
case 15:
alloc.buffers[bits].Put((*[1 << 15]byte)(buf))
case 16:
alloc.buffers[bits].Put((*[1 << 16]byte)(buf))
default:
panic("invalid pool index")
}
return nil
}
// msb return the pos of most significant bit
func msb(size int) uint16 {
return uint16(bits.Len32(uint32(size)) - 1)
}