github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/bufalloc/byte_allocator.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package bufalloc
    12  
    13  // ByteAllocator provides chunk allocation of []byte, amortizing the overhead
    14  // of each allocation. Because the underlying storage for the slices is shared,
    15  // they should share a similar lifetime in order to avoid pinning large amounts
    16  // of memory unnecessarily. The allocator itself is a []byte where cap()
    17  // indicates the total amount of memory and len() is the amount already
    18  // allocated. The size of the buffer to allocate from is grown exponentially
    19  // when it runs out of room up to a maximum size (chunkAllocMaxSize).
    20  type ByteAllocator []byte
    21  
    22  const chunkAllocMinSize = 512
    23  const chunkAllocMaxSize = 16384
    24  
    25  func (a ByteAllocator) reserve(n int) ByteAllocator {
    26  	allocSize := cap(a) * 2
    27  	if allocSize < chunkAllocMinSize {
    28  		allocSize = chunkAllocMinSize
    29  	} else if allocSize > chunkAllocMaxSize {
    30  		allocSize = chunkAllocMaxSize
    31  	}
    32  	if allocSize < n {
    33  		allocSize = n
    34  	}
    35  	return make([]byte, 0, allocSize)
    36  }
    37  
    38  // Alloc allocates a new chunk of memory with the specified length. extraCap
    39  // indicates additional zero bytes that will be present in the returned []byte,
    40  // but not part of the length.
    41  func (a ByteAllocator) Alloc(n int, extraCap int) (ByteAllocator, []byte) {
    42  	if cap(a)-len(a) < n+extraCap {
    43  		a = a.reserve(n + extraCap)
    44  	}
    45  	p := len(a)
    46  	r := a[p : p+n : p+n+extraCap]
    47  	a = a[:p+n+extraCap]
    48  	return a, r
    49  }
    50  
    51  // Copy allocates a new chunk of memory, initializing it from src.
    52  // extraCap indicates additional zero bytes that will be present in the returned
    53  // []byte, but not part of the length.
    54  func (a ByteAllocator) Copy(src []byte, extraCap int) (ByteAllocator, []byte) {
    55  	var alloc []byte
    56  	a, alloc = a.Alloc(len(src), extraCap)
    57  	copy(alloc, src)
    58  	return a, alloc
    59  }