github.com/pingcap/badger@v1.5.1-0.20230103063557-828f39b09b6d/buffer/buffer.go (about)

     1  // Package buffer implements a variable-sized bytes pool.
     2  package buffer
     3  
     4  import (
     5  	"sync"
     6  )
     7  
     8  const (
     9  	numBuckets = 256
    10  	pageSize   = 1024
    11  )
    12  
    13  // Buffers contains buckets for sharding.
    14  type Buffers struct {
    15  	pageSize int
    16  	buckets  [numBuckets]bucket
    17  }
    18  
    19  type bucket struct {
    20  	lock  sync.RWMutex
    21  	pools map[int]*Pool
    22  }
    23  
    24  // Pool represents a fixed size bytes pool.
    25  type Pool struct {
    26  	size int
    27  	pool *sync.Pool
    28  }
    29  
    30  // GetBuffer returns a bytes from the pool with the given size.
    31  func (p *Pool) GetBuffer(size int) []byte {
    32  	return p.pool.Get().([]byte)[:size]
    33  }
    34  
    35  // PutBuffer frees the bytes to the pool.
    36  func (p *Pool) PutBuffer(buf []byte) {
    37  	buf = buf[:cap(buf)]
    38  	if cap(buf) >= p.size {
    39  		p.pool.Put(buf)
    40  	}
    41  }
    42  
    43  // NewBuffers creates a new Buffers with the given page size.
    44  func NewBuffers(pageSize int) *Buffers {
    45  	b := new(Buffers)
    46  	for i := range b.buckets {
    47  		b.buckets[i].pools = make(map[int]*Pool)
    48  	}
    49  	if pageSize < 1 {
    50  		b.pageSize = 1
    51  	} else {
    52  		b.pageSize = pageSize
    53  	}
    54  	return b
    55  }
    56  
    57  // AssignPool assigns a fixed size bytes pool with the given size.
    58  func (b *Buffers) AssignPool(size int) (p *Pool) {
    59  	var alignedSize = size
    60  	if size%b.pageSize > 0 {
    61  		alignedSize = size/b.pageSize*b.pageSize + b.pageSize
    62  	}
    63  	m := &b.buckets[alignedSize/b.pageSize%numBuckets]
    64  	var ok bool
    65  	m.lock.RLock()
    66  	if p, ok = m.pools[alignedSize]; ok {
    67  		m.lock.RUnlock()
    68  		return
    69  	}
    70  	m.lock.RUnlock()
    71  	m.lock.Lock()
    72  	if p, ok = m.pools[alignedSize]; !ok {
    73  		p = &Pool{
    74  			pool: &sync.Pool{New: func() interface{} {
    75  				return make([]byte, alignedSize)
    76  			}},
    77  			size: alignedSize,
    78  		}
    79  		m.pools[alignedSize] = p
    80  	}
    81  	m.lock.Unlock()
    82  	return
    83  }
    84  
    85  // GetBuffer returns a bytes from the pool with the given size.
    86  func (b *Buffers) GetBuffer(size int) []byte {
    87  	return b.AssignPool(size).GetBuffer(size)
    88  }
    89  
    90  // PutBuffer frees the bytes to the pool.
    91  func (b *Buffers) PutBuffer(buf []byte) {
    92  	b.AssignPool(cap(buf)).PutBuffer(buf)
    93  }
    94  
    95  // defaultBuffers is the default instance of *Buffers.
    96  var defaultBuffers = NewBuffers(pageSize)
    97  
    98  // AssignPool assigns a fixed size bytes pool with the given size.
    99  func AssignPool(size int) *Pool {
   100  	return defaultBuffers.AssignPool(size)
   101  }
   102  
   103  // GetBuffer returns a bytes from the pool with the given size.
   104  func GetBuffer(size int) []byte {
   105  	return defaultBuffers.GetBuffer(size)
   106  }
   107  
   108  // PutBuffer frees the bytes to the pool.
   109  func PutBuffer(buf []byte) {
   110  	defaultBuffers.PutBuffer(buf)
   111  }