golang实现一个简易的内存池
package base
import (
"fmt"
m "math"
"sync"
)
type Pool struct {
sync.RWMutex
pool sync.Pool
size int
idle int32
active int32
maxIdle int32
}
func NewPool(size int, maxIdle int) *Pool {
return &Pool{
size: size,
maxIdle: int32(maxIdle),
pool: sync.Pool{New: func() any {
arr := make([]byte, size)
return &arr
}},
}
}
func (p *Pool) Get() *[]byte {
p.Lock()
defer p.Unlock()
p.active++
if p.idle > 0 {
p.idle--
}
if p.active > 1000 {
fmt.Printf("size=%d, idle=%d, active=%d\n", p.size, p.idle, p.active)
}
return p.pool.Get().(*[]byte)
}
func (p *Pool) Put(x any) {
p.Lock()
defer p.Unlock()
p.active--
if p.idle >= p.maxIdle {
return
}
p.idle++
p.pool.Put(x)
}
type MemoryPool struct {
sync.RWMutex
pools []*Pool
BaseSize int
MaxIdle int
}
func (p *MemoryPool) Get(size int) *[]byte {
var bytes *[]byte
p.RLock()
if p.BaseSize <= 0 {
p.RUnlock()
return nil
}
for _, pool := range p.pools {
if size <= pool.size {
bytes = pool.Get()
break
}
}
p.RUnlock()
if bytes != nil {
return bytes
}
p.Lock()
defer p.Unlock()
var pool *Pool
for _, pool = range p.pools {
if size <= pool.size {
return pool.Get()
}
}
var maxSize int
if pool != nil {
maxSize = 2 * pool.size
} else {
maxSize = p.BaseSize
}
if size < p.BaseSize {
size = p.BaseSize
}
for sz := maxSize; sz <= size; sz *= 2 {
pool = NewPool(sz, p.MaxIdle)
p.pools = append(p.pools, pool)
}
return pool.Get()
}
func (p *MemoryPool) Put(bytes *[]byte) {
c := cap(*bytes)
p.RLock()
index := int(m.Log2(float64(c / p.BaseSize)))
if index > len(p.pools) {
p.RUnlock()
return
}
pool := p.pools[index]
p.RUnlock()
pool.Put(bytes)
}