golang实现一个简易的内存池

浮尘丶若梦

已于 2022-11-17 16:02:18 修改

阅读量499

点赞数

CC 4.0 BY-SA版权

分类专栏： golang基础文章标签： golang 开发语言后端

于 2022-11-17 15:57:47 首次发布

本文链接：https://blog.youkuaiyun.com/rbin_2009/article/details/127905631

golang基础专栏收录该内容

9 篇文章

订阅专栏

本文介绍了一种内存池管理方案的实现细节，包括如何通过 Go 语言创建和维护不同大小的内存池来提高程序效率。该方案能够根据所需内存大小动态调整并复用内存块，减少频繁分配与释放内存带来的开销。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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)
}