symdb: add WritePartitionV3 function (#3503)

* symdb: add WritePartitionV3

* Update block_writer_v3.go

Co-authored-by: Anton Kolesnikov <[email protected]>

* add test

* fmt

---------

Co-authored-by: Anton Kolesnikov <[email protected]>

Author: korniltsev

pkg/phlaredb/symdb/block_writer_v3.go

@@ -19,6 +19,10 @@ type writerV3 struct {
 	files    []block.File
 	footer   Footer
 
+	encodersV3
+}
+
+type encodersV3 struct {
 	stringsEncoder   *symbolsEncoder[string]
 	mappingsEncoder  *symbolsEncoder[v1.InMemoryMapping]
 	functionsEncoder *symbolsEncoder[v1.InMemoryFunction]
@@ -27,22 +31,10 @@ type writerV3 struct {
 
 func newWriterV3(c *Config) *writerV3 {
 	return &writerV3{
-		config: c,
-		index: IndexFile{
-			Header: IndexHeader{
-				Magic:   symdbMagic,
-				Version: FormatV3,
-			},
-		},
-		footer: Footer{
-			Magic:   symdbMagic,
-			Version: FormatV3,
-		},
-
-		stringsEncoder:   newStringsEncoder(),
-		mappingsEncoder:  newMappingsEncoder(),
-		functionsEncoder: newFunctionsEncoder(),
-		locationsEncoder: newLocationsEncoder(),
+		config:     c,
+		index:      newIndexFileV3(),
+		footer:     newFooterV3(),
+		encodersV3: newEncodersV3(),
 	}
 }
 
@@ -58,7 +50,7 @@ func (w *writerV3) writePartitions(partitions []*PartitionWriter) (err error) {
 		w.files = []block.File{w.dataFile.meta()}
 	}()
 	for _, p := range partitions {
-		if err = writePartitionV3(w, p); err != nil {
+		if err = writePartitionV3(w.dataFile.w, &w.encodersV3, p); err != nil {
 			return fmt.Errorf("failed to write partition: %w", err)
 		}
 		w.index.PartitionHeaders = append(w.index.PartitionHeaders, &p.header)
@@ -83,17 +75,17 @@ func (w *writerV3) newFile(path string) (f *fileWriter, err error) {
 	return f, err
 }
 
-func writePartitionV3(w *writerV3, p *PartitionWriter) (err error) {
-	if p.header.V3.Strings, err = writeSymbolsBlock(w.dataFile, p.strings.slice, w.stringsEncoder); err != nil {
+func writePartitionV3(w *writerOffset, e *encodersV3, p *PartitionWriter) (err error) {
+	if p.header.V3.Strings, err = writeSymbolsBlock(w, p.strings.slice, e.stringsEncoder); err != nil {
 		return err
 	}
-	if p.header.V3.Mappings, err = writeSymbolsBlock(w.dataFile, p.mappings.slice, w.mappingsEncoder); err != nil {
+	if p.header.V3.Mappings, err = writeSymbolsBlock(w, p.mappings.slice, e.mappingsEncoder); err != nil {
 		return err
 	}
-	if p.header.V3.Functions, err = writeSymbolsBlock(w.dataFile, p.functions.slice, w.functionsEncoder); err != nil {
+	if p.header.V3.Functions, err = writeSymbolsBlock(w, p.functions.slice, e.functionsEncoder); err != nil {
 		return err
 	}
-	if p.header.V3.Locations, err = writeSymbolsBlock(w.dataFile, p.locations.slice, w.locationsEncoder); err != nil {
+	if p.header.V3.Locations, err = writeSymbolsBlock(w, p.locations.slice, e.locationsEncoder); err != nil {
 		return err
 	}
 	for ci, c := range p.stacktraces.chunks {
@@ -102,7 +94,7 @@ func writePartitionV3(w *writerV3, p *PartitionWriter) (err error) {
 			stacks = uint32(len(p.stacktraces.hashToIdx))
 		}
 		h := StacktraceBlockHeader{
-			Offset:             w.dataFile.w.offset,
+			Offset:             w.offset,
 			Partition:          p.header.Partition,
 			BlockIndex:         uint16(ci),
 			Encoding:           StacktraceEncodingGroupVarint,
@@ -111,7 +103,7 @@ func writePartitionV3(w *writerV3, p *PartitionWriter) (err error) {
 			StacktraceMaxNodes: c.partition.maxNodesPerChunk,
 		}
 		crc := crc32.New(castagnoli)
-		if h.Size, err = c.WriteTo(io.MultiWriter(crc, w.dataFile)); err != nil {
+		if h.Size, err = c.WriteTo(io.MultiWriter(crc, w)); err != nil {
 			return fmt.Errorf("writing stacktrace chunk data: %w", err)
 		}
 		h.CRC = crc.Sum32()
@@ -120,18 +112,63 @@ func writePartitionV3(w *writerV3, p *PartitionWriter) (err error) {
 	return nil
 }
 
-func writeSymbolsBlock[T any](w *fileWriter, s []T, e *symbolsEncoder[T]) (h SymbolsBlockHeader, err error) {
-	h.Offset = uint64(w.w.offset)
+func writeSymbolsBlock[T any](w *writerOffset, s []T, e *symbolsEncoder[T]) (h SymbolsBlockHeader, err error) {
+	h.Offset = uint64(w.offset)
 	crc := crc32.New(castagnoli)
-	mw := io.MultiWriter(crc, w.w)
+	mw := io.MultiWriter(crc, w)
 	if err = e.encode(mw, s); err != nil {
 		return h, err
 	}
-	h.Size = uint32(w.w.offset) - uint32(h.Offset)
+	h.Size = uint32(w.offset) - uint32(h.Offset)
 	h.CRC = crc.Sum32()
 	h.Length = uint32(len(s))
 	h.BlockSize = uint32(e.blockSize)
 	h.BlockHeaderSize = uint16(e.blockEncoder.headerSize())
 	h.Format = e.blockEncoder.format()
 	return h, nil
 }
+
+func WritePartition(p *PartitionWriter, dst io.Writer) error {
+	index := newIndexFileV3()
+	footer := newFooterV3()
+	encoders := newEncodersV3()
+	w := withWriterOffset(dst)
+
+	if err := writePartitionV3(w, &encoders, p); err != nil {
+		return fmt.Errorf("failed to write partition: %w", err)
+	}
+	index.PartitionHeaders = append(index.PartitionHeaders, &p.header)
+	footer.IndexOffset = uint64(w.offset)
+	if _, err := index.WriteTo(w); err != nil {
+		return fmt.Errorf("failed to write index: %w", err)
+	}
+	if _, err := w.Write(footer.MarshalBinary()); err != nil {
+		return fmt.Errorf("failed to write footer: %w", err)
+	}
+	return nil
+}
+
+func newEncodersV3() encodersV3 {
+	return encodersV3{
+		stringsEncoder:   newStringsEncoder(),
+		mappingsEncoder:  newMappingsEncoder(),
+		functionsEncoder: newFunctionsEncoder(),
+		locationsEncoder: newLocationsEncoder(),
+	}
+}
+
+func newFooterV3() Footer {
+	return Footer{
+		Magic:   symdbMagic,
+		Version: FormatV3,
+	}
+}
+
+func newIndexFileV3() IndexFile {
+	return IndexFile{
+		Header: IndexHeader{
+			Magic:   symdbMagic,
+			Version: FormatV3,
+		},
+	}
+}

pkg/phlaredb/symdb/symdb.go

@@ -188,18 +188,18 @@ func (s *SymDB) PartitionWriter(partition uint64) *PartitionWriter {
 		s.m.Unlock()
 		return p
 	}
-	p = s.newPartition(partition)
+	p = NewPartitionWriter(partition, &s.config)
 	s.partitions[partition] = p
 	s.m.Unlock()
 	return p
 }
 
-func (s *SymDB) newPartition(partition uint64) *PartitionWriter {
+func NewPartitionWriter(partition uint64, config *Config) *PartitionWriter {
 	p := PartitionWriter{
 		header:      PartitionHeader{Partition: partition},
-		stacktraces: newStacktracesPartition(s.config.Stacktraces.MaxNodesPerChunk),
+		stacktraces: newStacktracesPartition(config.Stacktraces.MaxNodesPerChunk),
 	}
-	switch s.config.Version {
+	switch config.Version {
 	case FormatV2:
 		p.header.V2 = new(PartitionHeaderV2)
 	case FormatV3:

pkg/phlaredb/symdb/symdb_test.go

@@ -1,11 +1,17 @@
 package symdb
 
 import (
+	"bytes"
 	"context"
 	"io"
 	"sort"
 	"sync/atomic"
 	"testing"
+	"time"
+
+	phlareobj "github.com/grafana/pyroscope/pkg/objstore"
+	"github.com/grafana/pyroscope/pkg/objstore/providers/memory"
+	pprofth "github.com/grafana/pyroscope/pkg/pprof/testhelper"
 
 	"github.com/cespare/xxhash/v2"
 	"github.com/stretchr/testify/require"
@@ -101,8 +107,8 @@ func (b *testBucket) GetRange(ctx context.Context, name string, off, length int6
 	return b.Bucket.GetRange(ctx, name, off, length)
 }
 
-func newTestFileWriter(w io.Writer) *fileWriter {
-	return &fileWriter{w: &writerOffset{Writer: w}}
+func newTestFileWriter(w io.Writer) *writerOffset {
+	return &writerOffset{Writer: w}
 }
 
 //nolint:unparam
@@ -180,3 +186,44 @@ func Test_Stats(t *testing.T) {
 	}
 	require.Equal(t, expected, actual)
 }
+
+func TestWritePartition(t *testing.T) {
+	p := NewPartitionWriter(0, &Config{
+		Version: FormatV3,
+		Stacktraces: StacktracesConfig{
+			MaxNodesPerChunk: 4 << 20,
+		},
+		Parquet: ParquetConfig{
+			MaxBufferRowCount: 100 << 10,
+		},
+	})
+	profile := pprofth.NewProfileBuilder(time.Now().UnixNano()).
+		CPUProfile().
+		WithLabels(phlaremodel.LabelNameServiceName, "svc").
+		ForStacktraceString("foo", "bar").
+		AddSamples(1).
+		ForStacktraceString("qwe", "foo", "bar").
+		AddSamples(2)
+
+	profiles := p.WriteProfileSymbols(profile.Profile)
+	symdbBlob := bytes.NewBuffer(nil)
+	err := WritePartition(p, symdbBlob)
+	require.NoError(t, err)
+
+	bucket := phlareobj.NewBucket(memory.NewInMemBucket())
+	require.NoError(t, bucket.Upload(context.Background(), DefaultFileName, bytes.NewReader(symdbBlob.Bytes())))
+	reader, err := Open(context.Background(), bucket, testBlockMeta)
+	require.NoError(t, err)
+
+	r := NewResolver(context.Background(), reader)
+	defer r.Release()
+	r.AddSamples(0, profiles[0].Samples)
+	resolved, err := r.Tree()
+	require.NoError(t, err)
+	expected := `.
+└── bar: self 0 total 3
+    └── foo: self 1 total 3
+        └── qwe: self 2 total 2
+`
+	require.Equal(t, expected, resolved.String())
+}