github.com/XiaoMi/Gaea@v1.2.5/util/bucketpool/bucketpool.go (about)

     1  package bucketpool
     2  
     3  import (
     4  	"math"
     5  	"sync"
     6  )
     7  
     8  type sizedPool struct {
     9  	size int
    10  	pool sync.Pool
    11  }
    12  
    13  func newSizedPool(size int) *sizedPool {
    14  	return &sizedPool{
    15  		size: size,
    16  		pool: sync.Pool{
    17  			New: func() interface{} { return makeSlicePointer(size) },
    18  		},
    19  	}
    20  }
    21  
    22  // Pool is actually multiple pools which store buffers of specific size.
    23  // i.e. it can be three pools which return buffers 32K, 64K and 128K.
    24  type Pool struct {
    25  	minSize int
    26  	maxSize int
    27  	pools   []*sizedPool
    28  }
    29  
    30  // New returns Pool which has buckets from minSize to maxSize.
    31  // Buckets increase with the power of two, i.e with multiplier 2: [2b, 4b, 16b, ... , 1024b]
    32  // Last pool will always be capped to maxSize.
    33  func New(minSize, maxSize int) *Pool {
    34  	if maxSize < minSize {
    35  		panic("maxSize can't be less than minSize")
    36  	}
    37  	const multiplier = 2
    38  	var pools []*sizedPool
    39  	curSize := minSize
    40  	for curSize < maxSize {
    41  		pools = append(pools, newSizedPool(curSize))
    42  		curSize *= multiplier
    43  	}
    44  	pools = append(pools, newSizedPool(maxSize))
    45  	return &Pool{
    46  		minSize: minSize,
    47  		maxSize: maxSize,
    48  		pools:   pools,
    49  	}
    50  }
    51  
    52  func (p *Pool) findPool(size int) *sizedPool {
    53  	if size > p.maxSize {
    54  		return nil
    55  	}
    56  	idx := int(math.Ceil(math.Log2(float64(size) / float64(p.minSize))))
    57  	if idx < 0 {
    58  		idx = 0
    59  	}
    60  	if idx > len(p.pools)-1 {
    61  		return nil
    62  	}
    63  	return p.pools[idx]
    64  }
    65  
    66  // Get returns pointer to []byte which has len size.
    67  // If there is no bucket with buffers >= size, slice will be allocated.
    68  func (p *Pool) Get(size int) *[]byte {
    69  	sp := p.findPool(size)
    70  	if sp == nil {
    71  		return makeSlicePointer(size)
    72  	}
    73  	buf := sp.pool.Get().(*[]byte)
    74  	*buf = (*buf)[:size]
    75  	return buf
    76  }
    77  
    78  // Put returns pointer to slice to some bucket. Discards slice for which there is no bucket
    79  func (p *Pool) Put(b *[]byte) {
    80  	sp := p.findPool(cap(*b))
    81  	if sp == nil {
    82  		return
    83  	}
    84  	*b = (*b)[:cap(*b)]
    85  	sp.pool.Put(b)
    86  }
    87  
    88  func makeSlicePointer(size int) *[]byte {
    89  	data := make([]byte, size)
    90  	return &data
    91  }