github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/bytearena.go (about)

     1  package util
     2  
     3  // ByteArena for reduce GC pressure, not concurrent safe
     4  type ByteArena struct {
     5  	alloc             []byte
     6  	chunkAllocMinSize int
     7  	chunkAllocMaxSize int
     8  }
     9  
    10  // NewByteArena is ctor for ByteArena
    11  func NewByteArena(chunkAllocMinSize, chunkAllocMaxSize int) *ByteArena {
    12  	return &ByteArena{chunkAllocMinSize: chunkAllocMinSize, chunkAllocMaxSize: chunkAllocMaxSize}
    13  }
    14  
    15  // AllocBytes for allocate bytes
    16  func (a *ByteArena) AllocBytes(n int) (bytes []byte) {
    17  	if cap(a.alloc)-len(a.alloc) < n {
    18  		bytes = a.reserveOrAlloc(n)
    19  		if bytes != nil {
    20  			return
    21  		}
    22  	}
    23  
    24  	pos := len(a.alloc)
    25  	bytes = a.alloc[pos : pos+n : pos+n]
    26  	a.alloc = a.alloc[:pos+n]
    27  
    28  	return
    29  }
    30  
    31  // UnsafeReset for reuse
    32  func (a *ByteArena) UnsafeReset() {
    33  	a.alloc = a.alloc[:0]
    34  }
    35  
    36  func (a *ByteArena) reserveOrAlloc(n int) (bytes []byte) {
    37  
    38  	allocSize := cap(a.alloc) * 2
    39  	if allocSize < a.chunkAllocMinSize {
    40  		allocSize = a.chunkAllocMinSize
    41  	} else if allocSize > a.chunkAllocMaxSize {
    42  		allocSize = a.chunkAllocMaxSize
    43  	}
    44  	if allocSize <= n {
    45  		bytes = make([]byte, 0, n)
    46  		return
    47  	}
    48  
    49  	a.alloc = make([]byte, 0, allocSize)
    50  	return
    51  }