github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/metrics/encoding/protobuf/buffer.go (about)

     1  // Copyright (c) 2018 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package protobuf
    22  
    23  import (
    24  	"math"
    25  
    26  	"github.com/m3db/m3/src/x/pool"
    27  )
    28  
    29  // PoolReleaseFn is a function used to release underlying slice back to bytes pool.
    30  type PoolReleaseFn func([]byte)
    31  
    32  // Buffer contains a byte slice backed by an optional bytes pool.
    33  type Buffer struct {
    34  	finalizer PoolReleaseFn
    35  	buf       []byte
    36  }
    37  
    38  // NewBuffer create a new buffer.
    39  func NewBuffer(buf []byte, p PoolReleaseFn) Buffer {
    40  	return Buffer{buf: buf, finalizer: p}
    41  }
    42  
    43  // Bytes returns the raw byte slice.
    44  func (b Buffer) Bytes() []byte { return b.buf }
    45  
    46  // Truncate truncates the raw byte slice.
    47  func (b *Buffer) Truncate(n int) { b.buf = b.buf[:n] }
    48  
    49  // Close closes the buffer.
    50  func (b *Buffer) Close() {
    51  	if b.finalizer != nil && b.buf != nil {
    52  		b.finalizer(b.buf)
    53  	}
    54  	b.buf = nil
    55  	b.finalizer = nil
    56  }
    57  
    58  type copyDataMode int
    59  
    60  const (
    61  	dontCopyData copyDataMode = iota
    62  	copyData
    63  )
    64  
    65  // ensureBufferSize returns a buffer with at least the specified target size.
    66  // If the specified buffer has enough capacity, it is returned as is. Otherwise,
    67  // a new buffer is allocated with at least the specified target size, and the
    68  // specified buffer is returned to pool if possible.
    69  func ensureBufferSize(
    70  	buf []byte,
    71  	p pool.BytesPool,
    72  	targetSize int,
    73  	copyDataMode copyDataMode,
    74  ) []byte {
    75  	bufSize := len(buf)
    76  	if bufSize >= targetSize {
    77  		return buf
    78  	}
    79  	newSize := int(math.Max(float64(targetSize), float64(bufSize*2)))
    80  	newBuf := allocate(p, newSize)
    81  	if copyDataMode == copyData {
    82  		copy(newBuf, buf)
    83  	}
    84  	if p != nil && buf != nil {
    85  		p.Put(buf)
    86  	}
    87  	return newBuf
    88  }
    89  
    90  // allocate allocates a byte slice with at least the specified size.
    91  func allocate(p pool.BytesPool, targetSize int) []byte {
    92  	if p == nil {
    93  		return make([]byte, targetSize)
    94  	}
    95  	b := p.Get(targetSize)
    96  	return b[:cap(b)]
    97  }