gitee.com/sy_183/go-common@v1.0.5-0.20231205030221-958cfe129b47/pool/ring-buffer-pool.go (about)

     1  package pool
     2  
     3  import "sync/atomic"
     4  
     5  type RingBufferPool struct {
     6  	buffers   []*Buffer
     7  	allocated []atomic.Bool
     8  	index     int
     9  }
    10  
    11  func NewRingBufferPool(count uint, size uint, reversed uint) *RingBufferPool {
    12  	if count == 0 {
    13  		count = 1
    14  	}
    15  	p := &RingBufferPool{
    16  		buffers:   make([]*Buffer, count),
    17  		allocated: make([]atomic.Bool, count),
    18  	}
    19  	for i := range p.buffers {
    20  		p.buffers[i] = NewBuffer(size, reversed).SetOnReleased(func(*Buffer) {
    21  			p.allocated[i].Store(false)
    22  		})
    23  	}
    24  	p.allocated[0].Store(true)
    25  	p.buffers[0].AddRef()
    26  	return p
    27  }
    28  
    29  func (p *RingBufferPool) Get() []byte {
    30  	if buf := p.buffers[p.index].Get(); buf != nil {
    31  		return buf
    32  	}
    33  	p.buffers[p.index].Release()
    34  	index := p.index + 1
    35  	if index == len(p.buffers) {
    36  		index = 0
    37  	}
    38  	if !p.allocated[index].CompareAndSwap(false, true) {
    39  		return nil
    40  	}
    41  	p.index = index
    42  	return p.buffers[index].Use().Get()
    43  }
    44  
    45  func (p *RingBufferPool) Alloc(size uint) *Data {
    46  	return p.buffers[p.index].Alloc(size)
    47  }