github.com/grafana/pyroscope@v1.18.0/pkg/block/writer.go (about)

     1  package block
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"os"
     7  	"path/filepath"
     8  
     9  	"github.com/grafana/pyroscope/pkg/objstore"
    10  )
    11  
    12  type Writer struct {
    13  	path string
    14  	f    *os.File
    15  	w    *bufio.Writer
    16  	off  uint64
    17  }
    18  
    19  func NewBlockWriter(tmpdir string) (*Writer, error) {
    20  	var err error
    21  	if err = os.MkdirAll(tmpdir, 0755); err != nil {
    22  		return nil, err
    23  	}
    24  	w := &Writer{path: filepath.Join(tmpdir, FileNameDataObject)}
    25  	if w.f, err = os.Create(w.path); err != nil {
    26  		return nil, err
    27  	}
    28  	w.w = bufio.NewWriter(w.f)
    29  	return w, nil
    30  }
    31  
    32  func (w *Writer) Write(p []byte) (n int, err error) {
    33  	n, err = w.w.Write(p)
    34  	w.off += uint64(n)
    35  	return n, err
    36  }
    37  
    38  func (w *Writer) Offset() uint64 { return w.off }
    39  
    40  func (w *Writer) Upload(ctx context.Context, bucket objstore.Bucket, path string) error {
    41  	if err := w.w.Flush(); err != nil {
    42  		return err
    43  	}
    44  	if _, err := w.f.Seek(0, 0); err != nil {
    45  		return err
    46  	}
    47  	return bucket.Upload(ctx, path, w.f)
    48  }
    49  
    50  func (w *Writer) Close() error {
    51  	err := w.f.Close()
    52  	w.f = nil
    53  	w.w = nil
    54  	return err
    55  }