github.com/grafana/pyroscope@v1.18.0/pkg/util/bufferpool/pool.go (about)

     1  package bufferpool
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  // Sized *bytes.Buffer pools: from 2^9 (512b) to 2^30 (1GB).
    10  var pools [maxPool]sync.Pool
    11  
    12  type Buffer struct {
    13  	B []byte
    14  	p int64
    15  }
    16  
    17  const (
    18  	minBits = 9
    19  	maxPool = 22
    20  )
    21  
    22  // GetBuffer returns a buffer from the pool, or creates a new one.
    23  // The returned buffer has at least the requested capacity.
    24  func GetBuffer(size int) *Buffer {
    25  	i := poolIndex(size)
    26  	if i < 0 {
    27  		return &Buffer{B: make([]byte, 0, size)}
    28  	}
    29  	x := pools[i].Get()
    30  	if x != nil {
    31  		return x.(*Buffer)
    32  	}
    33  	c := 2 << (minBits + i - 1)
    34  	c += bytes.MinRead
    35  	return &Buffer{
    36  		B: make([]byte, 0, c),
    37  		p: i,
    38  	}
    39  }
    40  
    41  // Put places the buffer into the pool.
    42  func Put(b *Buffer) {
    43  	if b == nil {
    44  		return
    45  	}
    46  	if p := returnPool(cap(b.B), b.p); p > 0 {
    47  		b.B = b.B[:0]
    48  		pools[p].Put(b)
    49  	}
    50  }
    51  
    52  func returnPool(c int, p int64) int64 {
    53  	// Empty buffers are ignored.
    54  	if c == 0 {
    55  		return -1
    56  	}
    57  	i := poolIndex(c)
    58  	if p == 0 {
    59  		// The buffer does not belong to any pool, or it's
    60  		// of the smallest size. We pick the pool based on
    61  		// its current capacity.
    62  		return i
    63  	}
    64  	d := i - p
    65  	if d < 0 {
    66  		// This buffer was likely obtained outside the pool.
    67  		// For example, an empty one, or with pre-allocated
    68  		// byte slice.
    69  		return i
    70  	}
    71  	if d > 1 {
    72  		// Relocate the buffer, if it's capacity has been
    73  		// grown by more than a power of two.
    74  		return i
    75  	}
    76  	// Otherwise, keep the buffer in the current pool.
    77  	return p
    78  }
    79  
    80  func poolIndex(n int) (i int64) {
    81  	n--
    82  	n >>= minBits
    83  	for n > 0 {
    84  		n >>= 1
    85  		i++
    86  	}
    87  	if i >= maxPool {
    88  		return -1
    89  	}
    90  	return i
    91  }
    92  
    93  func (b *Buffer) ReadFrom(r io.Reader) (int64, error) {
    94  	buf := bytes.NewBuffer(b.B)
    95  	n, err := buf.ReadFrom(r)
    96  	b.B = buf.Bytes()
    97  	return n, err
    98  }