Merge pull request #4 from threadedstream/use_vtprotobuf

use vtprotobuf client for optimizing memory usage

Author: threadedstream

README.md

@@ -18,7 +18,7 @@ input profiles.
 In order to understand how merger works, we need to dig a bit deeper into a profile structure
 
 ```go
-// Profile is an in-memory representation of profile.proto.
+// Profile is an in-memory representation of merged_profile.proto.
 type Profile struct {
 	SampleType        []*ValueType
 	DefaultSampleType string

go.mod

@@ -5,8 +5,9 @@ go 1.21
 require (
 	github.com/google/pprof v0.0.0-20220829040838-70bd9ae97f40
 	github.com/pkg/errors v0.9.1
+	github.com/planetscale/vtprotobuf v0.6.0
 	github.com/stretchr/testify v1.9.0
-	google.golang.org/protobuf v1.28.1
+	google.golang.org/protobuf v1.31.0
 )
 
 require (

go.sum

@@ -7,15 +7,17 @@ github.com/google/pprof v0.0.0-20220829040838-70bd9ae97f40 h1:ykKxL12NZd3JmWZnyq
 github.com/google/pprof v0.0.0-20220829040838-70bd9ae97f40/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/planetscale/vtprotobuf v0.6.0 h1:nBeETjudeJ5ZgBHUz1fVHvbqUKnYOXNhsIEabROxmNA=
+github.com/planetscale/vtprotobuf v0.6.0/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
-google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w=
-google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
+google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

merge.go

@@ -57,7 +57,7 @@ func (pu *ByteProfileUnPacker) UnpackRaw(compressedRawProfile []byte, idx uint64
 		pu.mergedProfile = new(MergedByteProfile)
 	}
 
-	if err = proto.Unmarshal(rawProfile, pu.mergedProfile); err != nil {
+	if err = pu.mergedProfile.UnmarshalVT(rawProfile); err != nil {
 		return nil, err
 	}
 
@@ -68,6 +68,7 @@ func (pu *ByteProfileUnPacker) Unpack(idx uint64) ([]byte, error) {
 	return pu.mergedProfile.Profiles[idx], nil
 }
 
+// ProfileUnPacker recovers any of the profiles stored inside mergedProfile
 // ProfileUnPacker recovers any of the profiles stored inside mergedProfile
 type ProfileUnPacker struct {
 	mergedProfile *MergedProfile
@@ -266,7 +267,7 @@ func (pu *ProfileUnPacker) unpackLocation(p *profile.Profile, id uint64) *profil
 	return loc
 }
 
-func (pu *ProfileUnPacker) unpackLine(p *profile.Profile, line *Line) profile.Line {
+func (pu *ProfileUnPacker) unpackLine(p *profile.Profile, line *MergeLine) profile.Line {
 	return profile.Line{
 		Line:     line.Line,
 		Function: pu.unpackFunction(p, line.FunctionId),
@@ -277,7 +278,7 @@ func (pu *ProfileUnPacker) getString(id int) string {
 	if id < 0 || id > len(pu.mergedProfile.StringTable) {
 		return ""
 	}
-	return pu.mergedProfile.StringTable[id-1]
+	return pu.mergedProfile.StringTable[id]
 }
 
 func (pu *ProfileUnPacker) unpackFunction(p *profile.Profile, id uint64) *profile.Function {
@@ -360,7 +361,7 @@ func (bm *ByteProfileMerger) WriteCompressed(w io.Writer) error {
 }
 
 func (bm *ByteProfileMerger) WriteUncompressed(w io.Writer) error {
-	serialized, err := proto.Marshal(bm.mergedProfile)
+	serialized, err := bm.mergedProfile.MarshalVT()
 	if err != nil {
 		return err
 	}
@@ -392,7 +393,7 @@ func (pw *ProfileMerger) WriteCompressed(w io.Writer) error {
 	// Write writes the profile as a gzip-compressed marshaled protobuf.
 	zw := gzip.NewWriter(w)
 	defer zw.Close()
-	serialized, err := proto.Marshal(pw.mergedProfile)
+	serialized, err := pw.mergedProfile.MarshalVT()
 	if err != nil {
 		return err
 	}
@@ -402,15 +403,15 @@ func (pw *ProfileMerger) WriteCompressed(w io.Writer) error {
 }
 
 func (pw *ProfileMerger) WriteUncompressed(w io.Writer) error {
-	serialized, err := proto.Marshal(pw.mergedProfile)
+	serialized, err := pw.mergedProfile.MarshalVT()
 	if err != nil {
 		return err
 	}
 	_, err = w.Write(serialized)
 	return err
 }
 
-func (pw *ProfileMerger) Merge(ps ...*profile.Profile) *MergedProfile {
+func (pw *ProfileMerger) Merge(ps ...*Profile) *MergedProfile {
 	pw.mergedProfile.NumFunctions = make([]uint64, 0, len(ps))
 	pw.mergedProfile.NumLocations = make([]uint64, 0, len(ps))
 	pw.mergedProfile.NumSampleTypes = make([]uint64, 0, len(ps))
@@ -432,65 +433,66 @@ func (pw *ProfileMerger) Merge(ps ...*profile.Profile) *MergedProfile {
 	pw.mergePeriods(ps...)
 	pw.mergePeriodTypes(ps...)
 
-	pw.mergedProfile.StringTable = make([]string, len(pw.stringTable), len(pw.stringTable))
+	pw.mergedProfile.StringTable = make([]string, len(pw.stringTable)+1)
+	pw.mergedProfile.StringTable[0] = ""
 	for st, id := range pw.stringTable {
-		pw.mergedProfile.StringTable[id-1] = st
+		pw.mergedProfile.StringTable[id] = st
 	}
 
 	return pw.mergedProfile
 }
 
-func (pw *ProfileMerger) mergeSamples(ps ...*profile.Profile) {
+func (pw *ProfileMerger) mergeSamples(ps ...*Profile) {
 	// allocate samples slice beforehand
 	size := 0
 	for _, p := range ps {
 		size += len(p.Sample)
 	}
-	pw.mergedProfile.Samples = make([]*Sample, 0, size)
+	pw.mergedProfile.Samples = make([]*MergeSample, 0, size)
 
 	for _, p := range ps {
 		for _, s := range p.Sample {
-			pw.mergedProfile.Samples = append(pw.mergedProfile.Samples, pw.asMergedSample(s))
+			pw.mergedProfile.Samples = append(pw.mergedProfile.Samples, pw.asMergedSample(s, p))
 		}
 	}
 }
 
-func (pw *ProfileMerger) mergePeriodTypes(ps ...*profile.Profile) {
+func (pw *ProfileMerger) mergePeriodTypes(ps ...*Profile) {
 	pw.mergedProfile.PeriodTypes = make([]int64, 0, len(ps)*2)
 
 	for _, p := range ps {
 		pw.mergedProfile.PeriodTypes = append(pw.mergedProfile.PeriodTypes,
-			int64(pw.putString(p.PeriodType.Type)),
-			int64(pw.putString(p.PeriodType.Unit)),
+			p.PeriodType.Type,
+			p.PeriodType.Unit,
 		)
 	}
 }
 
-func (pw *ProfileMerger) mergeTimeNanos(ps ...*profile.Profile) {
+func (pw *ProfileMerger) mergeTimeNanos(ps ...*Profile) {
 	pw.mergedProfile.TimesNanos = make([]int64, 0, len(ps))
 
 	for _, p := range ps {
 		pw.mergedProfile.TimesNanos = append(pw.mergedProfile.TimesNanos, p.TimeNanos)
 	}
 }
 
-func (pw *ProfileMerger) mergeDurationNanos(ps ...*profile.Profile) {
+func (pw *ProfileMerger) mergeDurationNanos(ps ...*Profile) {
 	pw.mergedProfile.DurationsNanos = make([]int64, 0, len(ps))
 
 	for _, p := range ps {
 		pw.mergedProfile.DurationsNanos = append(pw.mergedProfile.DurationsNanos, p.DurationNanos)
 	}
 }
 
-func (pw *ProfileMerger) mergePeriods(ps ...*profile.Profile) {
+func (pw *ProfileMerger) mergePeriods(ps ...*Profile) {
 	pw.mergedProfile.Periods = make([]int64, 0, len(ps))
 
 	for _, p := range ps {
 		pw.mergedProfile.Periods = append(pw.mergedProfile.Periods, p.Period)
 	}
 }
 
-func (pw *ProfileMerger) mergeSampleTypes(ps ...*profile.Profile) {
+func (pw *ProfileMerger) mergeSampleTypes(ps ...*Profile) {
 	size := 0
 	for _, p := range ps {
 		size += len(p.SampleType)
@@ -501,24 +503,24 @@ func (pw *ProfileMerger) mergeSampleTypes(ps ...*profile.Profile) {
 	for _, p := range ps {
 		for _, vt := range p.SampleType {
 			pw.mergedProfile.SampleType = append(pw.mergedProfile.SampleType,
-				int64(pw.putString(vt.Type)),
-				int64(pw.putString(vt.Unit)),
+				int64(pw.getStringRef(uint64(vt.Type), p)),
+				int64(pw.getStringRef(uint64(vt.Unit), p)),
 			)
 		}
 	}
 }
 
-func (pw *ProfileMerger) putMapping(src *profile.Mapping) uint64 {
+func (pw *ProfileMerger) putMapping(src *Mapping, p *Profile) uint64 {
 	if src == nil {
 		return math.MaxUint64
 	}
 
-	mapping := &Mapping{
-		MemoryStart:     src.Start,
-		MemoryLimit:     src.Limit,
-		FileOffset:      src.Offset,
-		Filename:        int64(pw.putString(src.File)),
-		BuildId:         int64(pw.putString(src.BuildID)),
+	mapping := &MergeMapping{
+		MemoryStart:     src.MemoryStart,
+		MemoryLimit:     src.MemoryLimit,
+		FileOffset:      src.FileOffset,
+		Filename:        int64(pw.getStringRef(uint64(src.Filename), p)),
+		BuildId:         int64(pw.getStringRef(uint64(src.BuildId), p)),
 		HasFilenames:    src.HasFilenames,
 		HasFunctions:    src.HasFunctions,
 		HasInlineFrames: src.HasInlineFrames,
@@ -537,51 +539,52 @@ func (pw *ProfileMerger) putMapping(src *profile.Mapping) uint64 {
 	return mapping.Id
 }
 
-func (pw *ProfileMerger) asMergedSample(s *profile.Sample) *Sample {
-	mergedProfileSample := &Sample{
-		LocationId: make([]int64, 0, len(s.Location)),
+func (pw *ProfileMerger) asMergedSample(s *Sample, p *Profile) *MergeSample {
+	mergedProfileSample := &MergeSample{
+		LocationId: make([]int64, 0, len(s.LocationId)),
 		Value:      s.Value,
 	}
 
-	for _, loc := range s.Location {
-		mergedProfileSample.LocationId = append(mergedProfileSample.LocationId, int64(pw.putLocation(loc)))
+	for _, locId := range s.LocationId {
+		mergedProfileSample.LocationId = append(mergedProfileSample.LocationId, int64(pw.putLocation(p.Location[locId-1], p)))
 	}
 
 	return mergedProfileSample
 }
 
-func (pw *ProfileMerger) asMergedValueType(vt *profile.ValueType) *ValueType {
-	return &ValueType{
-		Type: int64(pw.putString(vt.Type)),
-		Unit: int64(pw.putString(vt.Unit)),
+func (pw *ProfileMerger) asMergedValueType(vt *ValueType) *MergeValueType {
+	return &MergeValueType{
+		Type: vt.Type,
+		Unit: vt.Unit,
 	}
 }
 
-func (pw *ProfileMerger) asMergedProfileLines(lines []profile.Line) []*Line {
-	mergedProfileLines := make([]*Line, 0, len(lines))
+func (pw *ProfileMerger) asMergedProfileLines(lines []*Line, p *Profile) []*MergeLine {
+	mergedProfileLines := make([]*MergeLine, 0, len(lines))
 	for _, ln := range lines {
-		mergedProfileLines = append(mergedProfileLines, pw.asMergedProfileLine(ln))
+		mergedProfileLines = append(mergedProfileLines, pw.asMergedProfileLine(ln, p))
 	}
 	return mergedProfileLines
 }
 
-func (pw *ProfileMerger) asMergedProfileLine(line profile.Line) *Line {
-	return &Line{
-		FunctionId: pw.putFunction(line.Function),
+func (pw *ProfileMerger) asMergedProfileLine(line *Line, p *Profile) *MergeLine {
+	return &MergeLine{
+		FunctionId: pw.putFunction(p.Function[line.FunctionId-1], p),
 		Line:       line.Line,
 	}
 }
 
-func (pw *ProfileMerger) putString(val string) int {
-	id, ok := pw.stringTable[val]
-	if !ok {
-		id = len(pw.stringTable) + 1
-		pw.stringTable[val] = id
+func (pw *ProfileMerger) getStringRef(id uint64, p *Profile) int {
+	strVal := p.StringTable[id]
+	if localId, ok := pw.stringTable[strVal]; ok {
+		return localId
 	}
-	return id
+	newId := len(pw.stringTable) + 1
+	pw.stringTable[strVal] = newId
+	return newId
 }
 
-func (pw *ProfileMerger) getFunctionKey(fn *Function) functionKey {
+func (pw *ProfileMerger) getFunctionKey(fn *MergeFunction) functionKey {
 	return functionKey{
 		name:       fn.Name,
 		systemName: fn.SystemName,
@@ -590,7 +593,7 @@ func (pw *ProfileMerger) getFunctionKey(fn *Function) functionKey {
 	}
 }
 
-func (pw *ProfileMerger) getMappingKey(m *Mapping) mappingKey {
+func (pw *ProfileMerger) getMappingKey(m *MergeMapping) mappingKey {
 	key := mappingKey{
 		start:  m.MemoryStart,
 		limit:  m.MemoryLimit,
@@ -607,7 +610,7 @@ func (pw *ProfileMerger) getMappingKey(m *Mapping) mappingKey {
 	return key
 }
 
-func (pw *ProfileMerger) getLocationKey(loc *Location) locationKey {
+func (pw *ProfileMerger) getLocationKey(loc *MergeLocation) locationKey {
 	key := locationKey{
 		mappingID: loc.MappingId,
 		address:   loc.Address,
@@ -626,27 +629,27 @@ func (pw *ProfileMerger) getLocationKey(loc *Location) locationKey {
 	return key
 }
 
-func (pw *ProfileMerger) putLine(src profile.Line) *Line {
-	return &Line{
-		FunctionId: pw.putFunction(src.Function),
+func (pw *ProfileMerger) putLine(src *Line, p *Profile) *MergeLine {
+	return &MergeLine{
+		FunctionId: pw.putFunction(p.Function[src.FunctionId-1], p),
 		Line:       src.Line,
 	}
 }
 
-func (pw *ProfileMerger) putLocation(src *profile.Location) uint64 {
+func (pw *ProfileMerger) putLocation(src *Location, p *Profile) uint64 {
 	if src == nil {
 		return math.MaxUint64
 	}
 
-	loc := &Location{
-		MappingId: pw.putMapping(src.Mapping),
+	loc := &MergeLocation{
+		MappingId: pw.putMapping(p.Mapping[src.MappingId-1], p),
 		Address:   src.Address,
-		Line:      make([]*Line, len(src.Line), len(src.Line)),
+		Line:      make([]*MergeLine, len(src.Line), len(src.Line)),
 		IsFolded:  src.IsFolded,
 	}
 
 	for i, line := range src.Line {
-		loc.Line[i] = pw.putLine(line)
+		loc.Line[i] = pw.putLine(line, p)
 	}
 
 	key := pw.getLocationKey(loc)
@@ -660,15 +663,15 @@ func (pw *ProfileMerger) putLocation(src *profile.Location) uint64 {
 	return loc.Id
 }
 
-func (pw *ProfileMerger) putFunction(src *profile.Function) uint64 {
+func (pw *ProfileMerger) putFunction(src *Function, p *Profile) uint64 {
 	if src == nil {
 		return math.MaxUint64
 	}
 
-	f := &Function{
-		Name:       int64(pw.putString(src.Name)),
-		SystemName: int64(pw.putString(src.SystemName)),
-		Filename:   int64(pw.putString(src.Filename)),
+	f := &MergeFunction{
+		Name:       int64(pw.getStringRef(uint64(src.Name), p)),
+		SystemName: int64(pw.getStringRef(uint64(src.SystemName), p)),
+		Filename:   int64(pw.getStringRef(uint64(src.Filename), p)),
 		StartLine:  src.StartLine,
 	}