github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/symdb/block_writer.go (about)

     1  package symdb
     2  
     3  import (
     4  	"bufio"
     5  	"io"
     6  	"os"
     7  	"path/filepath"
     8  
     9  	"github.com/grafana/pyroscope/pkg/phlaredb/block"
    10  )
    11  
    12  type blockWriter interface {
    13  	writePartitions(partitions []*PartitionWriter) error
    14  	meta() []block.File
    15  }
    16  
    17  type fileWriter struct {
    18  	path string
    19  	buf  *bufio.Writer
    20  	f    *os.File
    21  	w    *writerOffset
    22  }
    23  
    24  func newFileWriter(path string) (*fileWriter, error) {
    25  	f, err := os.Create(path)
    26  	if err != nil {
    27  		return nil, err
    28  	}
    29  	// There is no particular reason to use
    30  	// a buffer larger than the default 4K.
    31  	b := bufio.NewWriterSize(f, 4096)
    32  	w := withWriterOffset(b)
    33  	fw := fileWriter{
    34  		path: path,
    35  		buf:  b,
    36  		f:    f,
    37  		w:    w,
    38  	}
    39  	return &fw, nil
    40  }
    41  
    42  func (f *fileWriter) Write(p []byte) (n int, err error) {
    43  	return f.w.Write(p)
    44  }
    45  
    46  func (f *fileWriter) Close() (err error) {
    47  	if err = f.buf.Flush(); err != nil {
    48  		return err
    49  	}
    50  	return f.f.Close()
    51  }
    52  
    53  func (f *fileWriter) meta() (m block.File) {
    54  	m.RelPath = filepath.Base(f.path)
    55  	if stat, err := os.Stat(f.path); err == nil {
    56  		m.SizeBytes = uint64(stat.Size())
    57  	}
    58  	return m
    59  }
    60  
    61  type writerOffset struct {
    62  	io.Writer
    63  	offset int64
    64  	err    error
    65  }
    66  
    67  func withWriterOffset(w io.Writer) *writerOffset {
    68  	return &writerOffset{Writer: w}
    69  }
    70  
    71  func (w *writerOffset) write(p []byte) {
    72  	if w.err == nil {
    73  		n, err := w.Writer.Write(p)
    74  		w.offset += int64(n)
    75  		w.err = err
    76  	}
    77  }
    78  
    79  func (w *writerOffset) Write(p []byte) (n int, err error) {
    80  	n, err = w.Writer.Write(p)
    81  	w.offset += int64(n)
    82  	return n, err
    83  }