zstd.Encoder uses much more memory than it should

[g7r]

zstd.Encoder supports two modes of operation: as an io.Writer for stream compression and as a small block compressor using EncodeAll. The client code can't really opt-out of EncodeAll because zstd.Encoder implicitly calls it in the Close method. The problem with EncodeAll is that it allocates a number of compressors that are intended to be used concurrently. The default number of compressors equals to GOMAXPROCS (16 in my case).

My code uses zstd.Encoder for stream compression only. As recommended by the documentation, zstd.Encoder is being reused. After some number of reuses zstd.Encoder allocates 16x more memory than expected.

To demonstrate the issue I'm going to use the Flush method of zstd.Encoder and the fact that it allows me to prevent calling the EncodeAll method. Therefore if I use Flush before Close, zstd.Encoder doesn't allocate 16x more memory. However, Flush isn't recommended for general use. I consider that "optimization" using Flush isn't the way Flush is intended to be used. Here's the repro code:

import (
	"io"
	"testing"

	"github.com/klauspost/compress/zstd"
)

var data = []byte{1, 2, 3}

func newZstdWriter() (*zstd.Encoder, error) {
	return zstd.NewWriter(
		io.Discard,
		zstd.WithEncoderLevel(zstd.SpeedBetterCompression),
		zstd.WithEncoderConcurrency(16), // we implicitly get this concurrency level if we run on 16 core CPU
	)
}

func BenchmarkMemWithFlush(b *testing.B) {
	for i := 0; i < b.N; i++ {
		w, err := newZstdWriter()
		if err != nil {
			b.Fatal(err)
		}

		for j := 0; j < 16; j++ {
			w.Reset(io.Discard)

			if _, err := w.Write(data); err != nil {
				b.Fatal(err)
			}

			if err := w.Flush(); err != nil {
				b.Fatal(err)
			}

			if err := w.Close(); err != nil {
				b.Fatal(err)
			}
		}
	}
}

func BenchmarkMemWithoutFlush(b *testing.B) {
	for i := 0; i < b.N; i++ {
		w, err := newZstdWriter()
		if err != nil {
			b.Fatal(err)
		}

		for j := 0; j < 16; j++ {
			w.Reset(io.Discard)

			if _, err := w.Write(data); err != nil {
				b.Fatal(err)
			}

			if err := w.Close(); err != nil {
				b.Fatal(err)
			}
		}
	}
}

Running the benchmark locally I get the following result:

goos: linux
goarch: amd64
pkg: mytestpackage
cpu: 11th Gen Intel(R) Core(TM) i7-11800H @ 2.30GHz
BenchmarkMemWithFlush
BenchmarkMemWithFlush-16       	    664	  1741676 ns/op	22057007 B/op	    109 allocs/op
BenchmarkMemWithoutFlush
BenchmarkMemWithoutFlush-16    	     82	 13833828 ns/op	346925504 B/op	    254 allocs/op
PASS
ok  	mytestpackage	3.466s

Note that BenchmarkMemWithoutFlush uses 16x more memory. The expected result is that memory usage should be more or less the same.

[klauspost]

If you just intend to use it for a stream, use zstd.WithEncoderConcurrency(2) or something in that area.

[yusufozturk]

I am sorry for my very unclear message. I have a scenario where I constantly write files with zstd compression and then read those files. I am using all cores possible, so there are many concurrent read and write operations on different files. So 5x more memory is a case where this scenario happens. I believe just the cost of encoder or decoder is too much, and when there are multiple concurrent operations, this increases the overall memory usage. Is my assumption correct?

[klauspost]

Then use zstd.WithEncoderConcurrency(1) which will only allocate one internal encoder.

[klauspost]

WithLowerEncoderMem(true) cuts memory usage in ~half.

Using a smaller window size WithWindowSize(1<<20) cuts off a lot of allocs as well.

The difference between flush/no flush is that a single block Write+Close gets translated into an EncodeAll call, so multiple encoders will end up being used - this is faster than shoving a small block through the stream encoder.

[klauspost]

#1007 should eliminate the extra allocs on small encodes.

[g7r]

Thank you very much for the fix! 🙏🏻🙏🏻🙏🏻