From 463f6b43a581fe8476a87b0ca62f237d573a96e2 Mon Sep 17 00:00:00 2001
From: gildarov <gildarov@ozon.ru>
Date: Tue, 21 May 2024 18:24:38 +0300
Subject: [PATCH] use UnpackRaw method to unpack profile from compressed merged
 profile

---
 merge.go      | 25 +++++++++++++++++++++++++
 merge_test.go | 17 +++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/merge.go b/merge.go
index 09fb519..973f8d9 100644
--- a/merge.go
+++ b/merge.go
@@ -1,6 +1,7 @@
 package ppmerge
 
 import (
+	"bytes"
 	"compress/gzip"
 	"io"
 	"math"
@@ -46,6 +47,30 @@ func NewProfileUnPacker(mergedProfile *MergedProfile) *ProfileUnPacker {
 	}
 }
 
+func (pu *ProfileUnPacker) UnpackRaw(compressedRawProfile []byte, idx uint64) (*profile.Profile, 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(MergedProfile)
+	}
+
+	if err = proto.Unmarshal(rawProfile, pu.mergedProfile); err != nil {
+		return nil, err
+	}
+
+	return pu.Unpack(idx)
+}
+
 func (pu *ProfileUnPacker) Unpack(idx uint64) (*profile.Profile, error) {
 	var p profile.Profile
 	if err := pu.unpackSampleTypes(&p, idx); err != nil {
diff --git a/merge_test.go b/merge_test.go
index b774b56..50e9ce4 100644
--- a/merge_test.go
+++ b/merge_test.go
@@ -118,6 +118,23 @@ func TestMergeWrite(t *testing.T) {
 	require.Less(t, compressedBB.Len(), noCompactBB.Len())
 }
 
+func TestMergeUnpack(t *testing.T) {
+	profiles := getProfiles(t, "hprof1", "hprof2", "hprof3", "hprof4")
+
+	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)
+
+	unpacker := NewProfileUnPacker(nil)
+	p, err := unpacker.UnpackRaw(compressedBB.Bytes(), 0)
+	require.NoError(t, err)
+	require.NotNil(t, p)
+}
+
 func BenchmarkProfileMerger(b *testing.B) {
 	profiles := getProfiles(b, "hprof1", "hprof2", "hprof3", "hprof4")