add ByteProfileMerger

Author: threadedstream

merge.go

@@ -28,6 +28,46 @@ type locationKey struct {
 	isFolded           bool
 }
 
+// ByteProfileUnPacker is the unpacker for MergedByteProfile
+type ByteProfileUnPacker struct {
+	mergedProfile *MergedByteProfile
+}
+
+// NewByteProfileUnPacker returns new ByteProfileUnPacker instance
+func NewByteProfileUnPacker(mergedProfile *MergedByteProfile) *ByteProfileUnPacker {
+	return &ByteProfileUnPacker{
+		mergedProfile: mergedProfile,
+	}
+}
+
+func (pu *ByteProfileUnPacker) UnpackRaw(compressedRawProfile []byte, idx uint64) ([]byte, error) {
+	bb := bytes.NewBuffer(compressedRawProfile)
+
+	gzReader, err := gzip.NewReader(bb)
+	if err != nil {
+		return nil, err
+	}
+
+	rawProfile, err := io.ReadAll(gzReader)
+	if err != nil {
+		return nil, err
+	}
+
+	if pu.mergedProfile == nil {
+		pu.mergedProfile = new(MergedByteProfile)
+	}
+
+	if err = proto.Unmarshal(rawProfile, pu.mergedProfile); err != nil {
+		return nil, err
+	}
+
+	return pu.Unpack(idx)
+}
+
+func (pu *ByteProfileUnPacker) Unpack(idx uint64) ([]byte, error) {
+	return pu.mergedProfile.Profiles[idx], nil
+}
+
 // ProfileUnPacker recovers any of the profiles stored inside mergedProfile
 type ProfileUnPacker struct {
 	mergedProfile *MergedProfile
@@ -287,6 +327,47 @@ func (pu *ProfileUnPacker) unpackMapping(p *profile.Profile, id uint64) *profile
 	return profileMapping
 }
 
+type ByteProfileMerger struct {
+	mergedProfile *MergedByteProfile
+}
+
+func NewByteProfileMerger() *ByteProfileMerger {
+	return &ByteProfileMerger{
+		mergedProfile: new(MergedByteProfile),
+	}
+}
+
+func (bm *ByteProfileMerger) Merge(profiles ...[]byte) *MergedByteProfile {
+	bm.mergedProfile.Profiles = make([][]byte, len(profiles))
+	for i, p := range profiles {
+		bm.mergedProfile.Profiles[i] = p
+	}
+
+	return bm.mergedProfile
+}
+
+func (bm *ByteProfileMerger) 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(bm.mergedProfile)
+	if err != nil {
+		return err
+	}
+
+	_, err = zw.Write(serialized)
+	return err
+}
+
+func (bm *ByteProfileMerger) WriteUncompressed(w io.Writer) error {
+	serialized, err := proto.Marshal(bm.mergedProfile)
+	if err != nil {
+		return err
+	}
+	_, err = w.Write(serialized)
+	return err
+}
+
 // ProfileMerger merges several profiles into a single one
 type ProfileMerger struct {
 	mergedProfile *MergedProfile

merge_test.go

@@ -2,6 +2,7 @@ package ppmerge
 
 import (
 	"bytes"
+	"io"
 	"os"
 	"testing"
 
@@ -119,20 +120,73 @@ func TestMergeWrite(t *testing.T) {
 }
 
 func TestMergeUnpack(t *testing.T) {
-	profiles := getProfiles(t, "hprof1", "hprof2", "hprof3", "hprof4")
+	t.Run("general merge unpack", func(t *testing.T) {
+		profiles := getProfiles(t, "hprof1", "hprof2", "hprof3", "hprof4")
 
-	profileMerger := NewProfileMerger()
-	mergedProfile := profileMerger.Merge(profiles...)
-	require.NotNil(t, mergedProfile)
+		profileMerger := NewProfileMerger()
+		mergedProfile := profileMerger.Merge(profiles...)
+		require.NotNil(t, mergedProfile)
 
-	compressedBB := bytes.NewBuffer(nil)
-	require.NoError(t, profileMerger.WriteCompressed(compressedBB))
-	require.Greater(t, compressedBB.Len(), 0)
+		compressedBB := bytes.NewBuffer(nil)
+		require.NoError(t, profileMerger.WriteCompressed(compressedBB))
+		require.Greater(t, compressedBB.Len(), 0)
+
+		unpacker := NewProfileUnPacker(nil)
+		p, err := unpacker.UnpackRaw(compressedBB.Bytes(), 0)
+		require.NoError(t, err)
+		require.NotNil(t, p)
+	})
+
+	t.Run("merge unpack debug goroutine profiles", func(t *testing.T) {
+		profiles := getDebugProfiles(t, "parca_goroutine_debug_1_1", "parca_goroutine_debug_1_2", "parca_goroutine_debug_1_2")
+
+		profileMerger := NewByteProfileMerger()
+		mergedProfile := profileMerger.Merge(profiles...)
+		require.NotNil(t, mergedProfile)
+
+		unpacker := NewByteProfileUnPacker(mergedProfile)
+		p, err := unpacker.Unpack(0)
+		require.NoError(t, err)
+		require.NotNil(t, p)
+		require.Equal(t, profiles[0], p)
+
+		p, err = unpacker.Unpack(1)
+		require.NoError(t, err)
+		require.NotNil(t, p)
+		require.Equal(t, profiles[1], p)
+
+		p, err = unpacker.Unpack(2)
+		require.NoError(t, err)
+		require.NotNil(t, p)
+		require.Equal(t, profiles[2], p)
+	})
 
-	unpacker := NewProfileUnPacker(nil)
-	p, err := unpacker.UnpackRaw(compressedBB.Bytes(), 0)
-	require.NoError(t, err)
-	require.NotNil(t, p)
+	t.Run("merge unpack raw debug goroutine profiles", func(t *testing.T) {
+		profiles := getDebugProfiles(t, "parca_goroutine_debug_1_1", "parca_goroutine_debug_1_2", "parca_goroutine_debug_1_2")
+
+		profileMerger := NewByteProfileMerger()
+		mergedProfile := profileMerger.Merge(profiles...)
+		require.NotNil(t, mergedProfile)
+
+		bb := bytes.NewBuffer(nil)
+		require.NoError(t, profileMerger.WriteCompressed(bb))
+
+		unpacker := NewByteProfileUnPacker(mergedProfile)
+		p, err := unpacker.UnpackRaw(bb.Bytes(), 0)
+		require.NoError(t, err)
+		require.NotNil(t, p)
+		require.Equal(t, profiles[0], p)
+
+		p, err = unpacker.UnpackRaw(bb.Bytes(), 1)
+		require.NoError(t, err)
+		require.NotNil(t, p)
+		require.Equal(t, profiles[1], p)
+
+		p, err = unpacker.UnpackRaw(bb.Bytes(), 2)
+		require.NoError(t, err)
+		require.NotNil(t, p)
+		require.Equal(t, profiles[2], p)
+	})
 }
 
 func BenchmarkProfileMerger(b *testing.B) {
@@ -173,3 +227,17 @@ func getProfiles(t require.TestingT, paths ...string) []*profile.Profile {
 
 	return profiles
 }
+
+func getDebugProfiles(t require.TestingT, paths ...string) [][]byte {
+	dir := "./testdata/"
+	profiles := make([][]byte, len(paths))
+	for i, profileName := range paths {
+		file, err := os.OpenFile(dir+profileName, os.O_RDONLY, 0666)
+		require.NoError(t, err)
+		p, err := io.ReadAll(file)
+		require.NoError(t, err)
+		profiles[i] = p
+	}
+
+	return profiles
+}