perf: add gzip writer reuse (#163)

Author: korniltsev

godeltaprof/block.go

@@ -30,6 +30,7 @@ type BlockProfiler struct {
 	runtimeProfile func([]runtime.BlockProfileRecord) (int, bool)
 	scaleProfile   pprof.MutexProfileScaler
 	options        pprof.ProfileBuilderOptions
+	gz             gz
 }
 
 // NewMutexProfiler creates a new BlockProfiler instance for profiling mutex contention.
@@ -115,8 +116,9 @@ func (d *BlockProfiler) Profile(w io.Writer) error {
 
 	sort.Slice(p, func(i, j int) bool { return p[i].Cycles > p[j].Cycles })
 
+	zw := d.gz.get(w)
 	stc := pprof.MutexProfileConfig()
-	b := pprof.NewProfileBuilder(w, &d.options, stc)
+	b := pprof.NewProfileBuilder(w, zw, &d.options, stc)
 
 	return d.impl.PrintCountCycleProfile(b, d.scaleProfile, p)
 }

godeltaprof/compat/go.mod

@@ -5,6 +5,7 @@ go 1.18
 require (
 	github.com/google/pprof v0.0.0-20231127191134-f3a68a39ae15
 	github.com/grafana/pyroscope-go/godeltaprof v0.1.5
+	github.com/klauspost/compress v1.17.8
 	github.com/stretchr/testify v1.10.0
 	golang.org/x/tools v0.16.0
 )

godeltaprof/compat/go.sum

@@ -5,6 +5,8 @@ github.com/google/pprof v0.0.0-20231127191134-f3a68a39ae15 h1:t2sLhFuGXwoomaKLTu
 github.com/google/pprof v0.0.0-20231127191134-f3a68a39ae15/go.mod h1:czg5+yv1E0ZGTi6S6vVK1mke0fV+FaUhNGcd6VRS9Ik=
 github.com/grafana/pyroscope-go/godeltaprof v0.1.5 h1:gkFVqihFRL1Nro2FCC0u6mW47jclef96Zu8I/ykq+4E=
 github.com/grafana/pyroscope-go/godeltaprof v0.1.5/go.mod h1:1HSPtjU8vLG0jE9JrTdzjgFqdJ/VgN7fvxBNq3luJko=
+github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
+github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=

godeltaprof/compat/testdata.go

@@ -7,6 +7,7 @@ import (
 	"reflect"
 	"runtime"
 
+	"github.com/klauspost/compress/gzip"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/grafana/pyroscope-go/godeltaprof/internal/pprof"
@@ -320,15 +321,17 @@ func (h *heapTestHelper) mutate(nmutations int, fs []runtime.MemProfileRecord) {
 func WriteHeapProto(dp *pprof.DeltaHeapProfiler, opt *pprof.ProfileBuilderOptions, w io.Writer,
 	p []runtime.MemProfileRecord, rate int64) error {
 	stc := pprof.HeapProfileConfig(rate)
-	b := pprof.NewProfileBuilder(w, opt, stc)
+	zw, _ := gzip.NewWriterLevel(w, gzip.BestSpeed)
+	b := pprof.NewProfileBuilder(w, zw, opt, stc)
 
 	return dp.WriteHeapProto(b, p, rate)
 }
 
 func PrintCountCycleProfile(d *pprof.DeltaMutexProfiler, opt *pprof.ProfileBuilderOptions, w io.Writer,
 	scaler pprof.MutexProfileScaler, records []runtime.BlockProfileRecord) error {
 	stc := pprof.MutexProfileConfig()
-	b := pprof.NewProfileBuilder(w, opt, stc)
+	zw, _ := gzip.NewWriterLevel(w, gzip.BestSpeed)
+	b := pprof.NewProfileBuilder(w, zw, opt, stc)
 
 	return d.PrintCountCycleProfile(b, scaler, records)
 }

godeltaprof/gzip.go

@@ -0,0 +1,21 @@
+package godeltaprof
+
+import (
+	"io"
+
+	"github.com/klauspost/compress/gzip"
+)
+
+type gz struct {
+	w *gzip.Writer
+}
+
+func (g *gz) get(w io.Writer) *gzip.Writer {
+	if g.w == nil {
+		zw, _ := gzip.NewWriterLevel(w, gzip.BestSpeed)
+		g.w = zw
+	}
+	g.w.Reset(w)
+
+	return g.w
+}

godeltaprof/gzip_test.go

@@ -0,0 +1,52 @@
+package godeltaprof
+
+import (
+	"bytes"
+	"compress/gzip"
+	"io"
+	"testing"
+)
+
+func TestGz(t *testing.T) {
+	blobs := [][]byte{
+		[]byte("Hello, World! This is the first test blob with some data to compress."),
+		[]byte("This is a second blob with different content for compression testing."),
+	}
+
+	var g gz
+	var bufs []bytes.Buffer
+	
+	for i, blob := range blobs {
+		var buf bytes.Buffer
+		gzw := g.get(&buf)
+		if _, err := gzw.Write(blob); err != nil {
+			t.Fatalf("Failed to write blob %d: %v", i, err)
+		}
+		if err := gzw.Close(); err != nil {
+			t.Fatalf("Failed to close gzip writer %d: %v", i, err)
+		}
+		bufs = append(bufs, buf)
+	}
+	
+	for i, blob := range blobs {
+		gzr, err := gzip.NewReader(&bufs[i])
+		if err != nil {
+			t.Fatalf("Failed to create gzip reader for blob %d: %v", i, err)
+		}
+		
+		decompressed, err := io.ReadAll(gzr)
+		if err != nil {
+			gzr.Close()
+			t.Fatalf("Failed to decompress blob %d: %v", i, err)
+		}
+		gzr.Close()
+		
+		if !bytes.Equal(blob, decompressed) {
+			t.Errorf("Blob %d mismatch:\nOriginal:     %q\nDecompressed: %q", i, blob, decompressed)
+		}
+		
+		if bytes.Equal(blob, bufs[i].Bytes()) {
+			t.Errorf("Buffer %d should contain compressed data, not original", i)
+		}
+	}
+}

godeltaprof/heap.go

@@ -31,6 +31,7 @@ type HeapProfiler struct {
 	impl    pprof.DeltaHeapProfiler
 	mutex   sync.Mutex
 	options pprof.ProfileBuilderOptions
+	gz      gz
 }
 
 func NewHeapProfiler() *HeapProfiler {
@@ -79,7 +80,9 @@ func (d *HeapProfiler) Profile(w io.Writer) error {
 		// Profile grew; try again.
 	}
 	rate := int64(runtime.MemProfileRate)
-	b := pprof.NewProfileBuilder(w, &d.options, pprof.HeapProfileConfig(rate))
+
+	zw := d.gz.get(w)
+	b := pprof.NewProfileBuilder(w, zw, &d.options, pprof.HeapProfileConfig(rate))
 
 	return d.impl.WriteHeapProto(b, p, rate)
 }

godeltaprof/heap_test.go

@@ -0,0 +1,17 @@
+package godeltaprof
+
+import (
+	"bytes"
+	"testing"
+)
+
+func BenchmarkHeap(b *testing.B) {
+	p := NewHeapProfiler()
+	buf := bytes.NewBuffer(nil)
+	for i := 0; i < b.N; i++ {
+		err := p.Profile(buf)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}