github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/pool/bytesbuffer.go (about)

     1  package pool
     2  
     3  import (
     4  	"bytes"
     5  	"sync"
     6  )
     7  
     8  // BufferPool is a bucketed pool for variably bytes buffers.
     9  type BufferPool struct {
    10  	buckets []sync.Pool
    11  	sizes   []int
    12  }
    13  
    14  // NewBuffer a new Pool with size buckets for minSize to maxSize
    15  // increasing by the given factor.
    16  func NewBuffer(minSize, maxSize int, factor float64) *BufferPool {
    17  	if minSize < 1 {
    18  		panic("invalid minimum pool size")
    19  	}
    20  	if maxSize < 1 {
    21  		panic("invalid maximum pool size")
    22  	}
    23  	if factor < 1 {
    24  		panic("invalid factor")
    25  	}
    26  
    27  	var sizes []int
    28  
    29  	for s := minSize; s <= maxSize; s = int(float64(s) * factor) {
    30  		sizes = append(sizes, s)
    31  	}
    32  
    33  	return &BufferPool{
    34  		buckets: make([]sync.Pool, len(sizes)),
    35  		sizes:   sizes,
    36  	}
    37  }
    38  
    39  // Get returns a byte buffer that fits the given size.
    40  func (p *BufferPool) Get(sz int) *bytes.Buffer {
    41  	for i, bktSize := range p.sizes {
    42  		if sz > bktSize {
    43  			continue
    44  		}
    45  		b := p.buckets[i].Get()
    46  		if b == nil {
    47  			b = bytes.NewBuffer(make([]byte, 0, bktSize))
    48  		}
    49  		buf := b.(*bytes.Buffer)
    50  		buf.Reset()
    51  		return b.(*bytes.Buffer)
    52  	}
    53  	return bytes.NewBuffer(make([]byte, 0, sz))
    54  }
    55  
    56  // Put adds a byte buffer to the right bucket in the pool.
    57  func (p *BufferPool) Put(s *bytes.Buffer) {
    58  	if s == nil {
    59  		return
    60  	}
    61  	cap := s.Cap()
    62  	for i, size := range p.sizes {
    63  		if cap > size {
    64  			continue
    65  		}
    66  		p.buckets[i].Put(s)
    67  		return
    68  	}
    69  }