github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/memdb/index/buf.go (about)

     1  package index
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"sync"
     8  
     9  	"github.com/pkg/errors"
    10  )
    11  
    12  type BufferWriter struct {
    13  	buf *bytes.Buffer
    14  	pos uint64
    15  }
    16  
    17  var pool = sync.Pool{
    18  	New: func() interface{} {
    19  		return NewBufferWriter()
    20  	},
    21  }
    22  
    23  func GetBufferWriterFromPool() *BufferWriter {
    24  	res := pool.Get().(*BufferWriter)
    25  	res.Reset()
    26  	return res
    27  }
    28  
    29  func PutBufferWriterToPool(fw *BufferWriter) {
    30  	fw.Reset()
    31  	pool.Put(fw)
    32  }
    33  
    34  // NewBufferWriter returns a new BufferWriter.
    35  // todo: pooling memory
    36  func NewBufferWriter() *BufferWriter {
    37  	return &BufferWriter{
    38  		buf: bytes.NewBuffer(make([]byte, 0, 0x2000)),
    39  		pos: 0,
    40  	}
    41  }
    42  
    43  func (fw *BufferWriter) Pos() uint64 {
    44  	return fw.pos
    45  }
    46  
    47  func (fw *BufferWriter) Write(bufs ...[]byte) error {
    48  	for _, buf := range bufs {
    49  		n, err := fw.buf.Write(buf)
    50  		if err != nil {
    51  			return err
    52  		}
    53  		fw.pos += uint64(n)
    54  	}
    55  	return nil
    56  }
    57  
    58  func (fw *BufferWriter) Flush() error {
    59  	return nil
    60  }
    61  
    62  func (fw *BufferWriter) WriteAt(buf []byte, pos uint64) error {
    63  	if pos > fw.pos {
    64  		return fmt.Errorf("position out of range")
    65  	}
    66  	if pos+uint64(len(buf)) > fw.pos {
    67  		return fmt.Errorf("write exceeds buffer size")
    68  	}
    69  	copy(fw.buf.Bytes()[pos:], buf)
    70  	return nil
    71  }
    72  
    73  func (fw *BufferWriter) Read(buf []byte) (int, error) {
    74  	return fw.buf.Read(buf)
    75  }
    76  
    77  func (fw *BufferWriter) ReadFrom(r io.Reader) (int64, error) {
    78  	n, err := fw.buf.ReadFrom(r)
    79  	if err != nil {
    80  		return n, err
    81  	}
    82  	fw.pos += uint64(n)
    83  	return n, err
    84  }
    85  
    86  func (fw *BufferWriter) AddPadding(size int) error {
    87  	p := fw.pos % uint64(size)
    88  	if p == 0 {
    89  		return nil
    90  	}
    91  	p = uint64(size) - p
    92  
    93  	if err := fw.Write(make([]byte, p)); err != nil {
    94  		return errors.Wrap(err, "add padding")
    95  	}
    96  	return nil
    97  }
    98  
    99  func (fw *BufferWriter) Buffer() ([]byte, io.Closer, error) {
   100  	return fw.buf.Bytes(), io.NopCloser(nil), nil
   101  }
   102  
   103  func (fw *BufferWriter) Close() error {
   104  	return nil
   105  }
   106  
   107  func (fw *BufferWriter) Reset() {
   108  	fw.pos = 0
   109  	fw.buf.Reset()
   110  }
   111  
   112  func (fw *BufferWriter) Remove() error {
   113  	return nil
   114  }