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 }