refactor(profiler): alternative zstd encoder reuse approach

Author: felixge

profiler/compression.go

@@ -32,7 +32,6 @@ import (
 	"io"
 	"strconv"
 	"strings"
-	"sync"
 
 	kgzip "github.com/klauspost/compress/gzip"
 	"github.com/klauspost/compress/zstd"
@@ -138,42 +137,29 @@ func getZstdLevelOrDefault(level int) zstd.EncoderLevel {
 	return zstd.SpeedDefault
 }
 
-type sema struct {
-	c chan struct{}
+type compressionPipelineBuilder struct {
+	zstdEncoders map[zstd.EncoderLevel]*sharedZstdEncoder
 }
 
-func (s *sema) Lock()   { s.c <- struct{}{} }
-func (s *sema) Unlock() { <-s.c }
-
-var (
-	// compressionMux protects zstdEncoder. It must be locked
-	// when doing compression that _might_ use zstdEncoder.
-	//
-	// It's a channel-based semaphore rather than a mutex
-	// so that contention on it doesn't appear in the mutex
-	// profile (we expect it to be contended).
-	// This is a kludge. We only really want it for zstdEncoder,
-	// but the places where we actually do compression just
-	// take a compressor interface. It's easier for now to just
-	// have this global semaphore than to plumb the locking
-	// through the existing abstractions.
-	compressionMux = &sema{c: make(chan struct{}, 1)}
-
-	zstdEncoderOnce sync.Once
-	zstdEncoder     *zstd.Encoder
-	zstdEncoderErr  error
-)
-
-func getZstdEncoder(opts ...zstd.EOption) (*zstd.Encoder, error) {
-	zstdEncoderOnce.Do(func() {
-		zstdEncoder, zstdEncoderErr = zstd.NewWriter(nil, opts...)
-	})
-	return zstdEncoder, zstdEncoderErr
+func (b *compressionPipelineBuilder) getZstdEncoder(level zstd.EncoderLevel) (*sharedZstdEncoder, error) {
+	if b.zstdEncoders == nil {
+		b.zstdEncoders = make(map[zstd.EncoderLevel]*sharedZstdEncoder)
+	}
+	encoder, ok := b.zstdEncoders[level]
+	if !ok {
+		var err error
+		encoder, err = newSharedZstdEncoder(level)
+		if err != nil {
+			return nil, err
+		}
+		b.zstdEncoders[level] = encoder
+	}
+	return encoder, nil
 }
 
-// newCompressionPipeline returns a compressor that converts the data written to
-// it from the expected input compression to the given output compression.
-func newCompressionPipeline(in compression, out compression) (compressor, error) {
+// Build returns a compressor that converts the data written to it from the
+// expected input compression to the given output compression.
+func (b *compressionPipelineBuilder) Build(in compression, out compression) (compressor, error) {
 	if in == out {
 		return newPassthroughCompressor(), nil
 	}
@@ -183,11 +169,15 @@ func newCompressionPipeline(in compression, out compression) (compressor, error)
 	}
 
 	if in == noCompression && out.algorithm == compressionAlgorithmZstd {
-		return getZstdEncoder(zstd.WithEncoderLevel(getZstdLevelOrDefault(out.level)))
+		return b.getZstdEncoder(getZstdLevelOrDefault(out.level))
 	}
 
 	if in.algorithm == compressionAlgorithmGzip && out.algorithm == compressionAlgorithmZstd {
-		return newZstdRecompressor(getZstdLevelOrDefault(out.level))
+		encoder, err := b.getZstdEncoder(getZstdLevelOrDefault(out.level))
+		if err != nil {
+			return nil, err
+		}
+		return newZstdRecompressor(encoder), nil
 	}
 
 	return nil, fmt.Errorf("unsupported recompression: %s -> %s", in, out)
@@ -198,8 +188,11 @@ func newCompressionPipeline(in compression, out compression) (compressor, error)
 // the data from one format and then re-compresses it into another format.
 type compressor interface {
 	io.Writer
-	io.Closer
+	// Reset reset the compressor to the given writer. It may also acquire a
+	// shared underlaying resource, so callers must always call Close().
 	Reset(w io.Writer)
+	// Close closes the compressor and releases any shared underlaying resource.
+	Close() error
 }
 
 // newPassthroughCompressor returns a compressor that simply passes all data
@@ -220,20 +213,16 @@ func (r *passthroughCompressor) Close() error {
 	return nil
 }
 
-func newZstdRecompressor(level zstd.EncoderLevel) (*zstdRecompressor, error) {
-	zstdOut, err := getZstdEncoder(zstd.WithEncoderLevel(level))
-	if err != nil {
-		return nil, err
-	}
-	return &zstdRecompressor{zstdOut: zstdOut, err: make(chan error)}, nil
+func newZstdRecompressor(encoder *sharedZstdEncoder) *zstdRecompressor {
+	return &zstdRecompressor{zstdOut: encoder, err: make(chan error)}
 }
 
 type zstdRecompressor struct {
 	// err synchronizes finishing writes after closing pw and reports any
 	// error during recompression
 	err     chan error
 	pw      io.WriteCloser
-	zstdOut *zstd.Encoder
+	zstdOut *sharedZstdEncoder
 }
 
 func (r *zstdRecompressor) Reset(w io.Writer) {
@@ -260,3 +249,36 @@ func (r *zstdRecompressor) Close() error {
 	err := <-r.err
 	return cmp.Or(err, r.zstdOut.Close())
 }
+
+// newSharedZstdEncoder creates a new shared Zstd encoder with the given level.
+// It expects the Reset and Close method to be used in an acquire and release
+// fashion.
+func newSharedZstdEncoder(level zstd.EncoderLevel) (*sharedZstdEncoder, error) {
+	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(level))
+	if err != nil {
+		return nil, err
+	}
+	return &sharedZstdEncoder{encoder: encoder, sema: make(chan struct{}, 1)}, nil
+}
+
+type sharedZstdEncoder struct {
+	encoder *zstd.Encoder
+	sema    chan struct{}
+}
+
+// Reset acquires the semaphore and resets the encoder to the given writer.
+func (s *sharedZstdEncoder) Reset(w io.Writer) {
+	s.sema <- struct{}{}
+	s.encoder.Reset(w)
+}
+
+func (s *sharedZstdEncoder) Write(p []byte) (int, error) {
+	return s.encoder.Write(p)
+}
+
+// Close releases the semaphore and closes the encoder.
+func (s *sharedZstdEncoder) Close() error {
+	err := s.encoder.Close()
+	<-s.sema
+	return err
+}

profiler/compression_test.go

@@ -43,7 +43,8 @@ func TestNewCompressionPipeline(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(fmt.Sprintf("%s->%s", test.in, test.out), func(t *testing.T) {
-			pipeline, err := newCompressionPipeline(test.in, test.out)
+			var pipelineBuilder compressionPipelineBuilder
+			pipeline, err := pipelineBuilder.Build(test.in, test.out)
 			require.NoError(t, err)
 			buf := &bytes.Buffer{}
 			pipeline.Reset(buf)
@@ -172,11 +173,13 @@ func BenchmarkRecompression(b *testing.B) {
 		b.Run(fmt.Sprintf("%s-%s", in.inAlg.String(), in.outLevel), func(b *testing.B) {
 			data := compressData(b, inputdata, in.inAlg)
 			b.ResetTimer()
+			var pipelineBuilder compressionPipelineBuilder
 			for i := 0; i < b.N; i++ {
-				z, err := newZstdRecompressor(in.outLevel)
+				encoder, err := pipelineBuilder.getZstdEncoder(in.outLevel)
 				if err != nil {
 					b.Fatal(err)
 				}
+				z := newZstdRecompressor(encoder)
 				z.Reset(io.Discard)
 				if _, err := z.Write(data); err != nil {
 					b.Fatal(err)

profiler/profile.go

@@ -114,9 +114,8 @@ var profileTypes = map[ProfileType]profileType{
 			p.stopCPUProfile()
 
 			c := p.compressors[CPUProfile]
-			compressionMux.Lock()
-			defer compressionMux.Unlock()
 			c.Reset(&buf)
+			defer c.Close()
 			if _, err := outBuf.WriteTo(c); err != nil {
 				return nil, err
 			}
@@ -183,9 +182,8 @@ var profileTypes = map[ProfileType]profileType{
 			}
 
 			c := p.compressors[expGoroutineWaitProfile]
-			compressionMux.Lock()
-			defer compressionMux.Unlock()
 			c.Reset(pprof)
+			defer c.Close()
 			err := goroutineDebug2ToPprof(text, c, now)
 			err = cmp.Or(err, c.Close())
 			return pprof.Bytes(), err
@@ -197,9 +195,8 @@ var profileTypes = map[ProfileType]profileType{
 		Collect: func(p *profiler) ([]byte, error) {
 			var buf bytes.Buffer
 			c := p.compressors[MetricsProfile]
-			compressionMux.Lock()
-			defer compressionMux.Unlock()
 			c.Reset(&buf)
+			defer c.Close()
 			interrupted := p.interruptibleSleep(p.cfg.period)
 			err := p.met.report(now(), c)
 			err = cmp.Or(err, c.Close())
@@ -229,9 +226,8 @@ var profileTypes = map[ProfileType]profileType{
 			trace.Stop()
 
 			c := p.compressors[executionTrace]
-			compressionMux.Lock()
-			defer compressionMux.Unlock()
 			c.Reset(buf)
+			defer c.Close()
 			if _, err := outBuf.WriteTo(c); err != nil {
 				return nil, err
 			}
@@ -303,9 +299,8 @@ func collectGenericProfile(name string, pt ProfileType) func(p *profiler) ([]byt
 		dp, ok := p.deltas[pt]
 		if !ok || !p.cfg.deltaProfiles {
 			c := p.compressors[pt]
-			compressionMux.Lock()
-			defer compressionMux.Unlock()
 			c.Reset(&buf)
+			defer c.Close()
 			err := p.lookupProfile(name, c, 0)
 			err = cmp.Or(err, c.Close())
 			return buf.Bytes(), err
@@ -456,9 +451,8 @@ func (fdp *fastDeltaProfiler) Delta(data []byte) (b []byte, err error) {
 
 	fdp.buf.Reset()
 	c := fdp.compressor
-	compressionMux.Lock()
-	defer compressionMux.Unlock()
 	c.Reset(&fdp.buf)
+	defer c.Close()
 
 	if err = fdp.dc.Delta(data, c); err != nil {
 		return nil, fmt.Errorf("error computing delta: %s", err.Error())

profiler/profiler.go

@@ -259,10 +259,11 @@ func newProfiler(opts ...Option) (*profiler, error) {
 	if p.cfg.traceConfig.Enabled {
 		types = append(types, executionTrace)
 	}
+	var pipelineBuilder compressionPipelineBuilder
 	for _, pt := range types {
 		isDelta := p.cfg.deltaProfiles && len(profileTypes[pt].DeltaValues) > 0
 		in, out := compressionStrategy(pt, isDelta, p.cfg.compressionConfig)
-		compressor, err := newCompressionPipeline(in, out)
+		compressor, err := pipelineBuilder.Build(in, out)
 		if err != nil {
 			return nil, err
 		}