Zstd decoder (#948)

* blerg

Signed-off-by: Joe Elliott <number101010@gmail.com>

sadness

Signed-off-by: Joe Elliott <number101010@gmail.com>

removed encoding

Signed-off-by: Joe Elliott <number101010@gmail.com>

duh

Signed-off-by: Joe Elliott <number101010@gmail.com>

* updated compression library

Signed-off-by: Joe Elliott <number101010@gmail.com>

* different horribleness

Signed-off-by: Joe Elliott <number101010@gmail.com>

* klauspost

Signed-off-by: Joe Elliott <number101010@gmail.com>

* cgo zstd

* Revert "cgo zstd"

This reverts commit 04bae4a.

* http conn mucking

Signed-off-by: Joe Elliott <number101010@gmail.com>

* encoder

Signed-off-by: Joe Elliott <number101010@gmail.com>

* Revert "http conn mucking"

This reverts commit 4b1e3cd.

* Revert "encoder"

This reverts commit b081970.

* comments

Signed-off-by: Joe Elliott <number101010@gmail.com>

* changelog

Signed-off-by: Joe Elliott <number101010@gmail.com>

* remove test code

Signed-off-by: Joe Elliott <number101010@gmail.com>

* todo

Signed-off-by: Joe Elliott <number101010@gmail.com>

CHANGELOG.md

@@ -33,6 +33,7 @@
 * [ENHANCEMENT] Include additional detail when searching for traces [#916](https://github.com/grafana/tempo/pull/916) (@zalegrala)
 * [ENHANCEMENT] Add `gen index` and `gen bloom` commands to tempo-cli. [#903](https://github.com/grafana/tempo/pull/903) (@annanay25)
 * [ENHANCEMENT] Implement trace comparison in Vulture [#904](https://github.com/grafana/tempo/pull/904) (@zalegrala)
+* [ENHANCEMENT] Improve zstd read throughput using zstd.Decoder [#948](https://github.com/grafana/tempo/pull/948) (@joe-elliott)
 * [ENHANCEMENT] Dedupe search records while replaying WAL [#940](https://github.com/grafana/tempo/pull/940) (@annanay25)
 * [ENHANCEMENT] Add status endpoint to list the available endpoints [#938](https://github.com/grafana/tempo/pull/938) (@zalegrala)
 * [ENHANCEMENT] Add search block headers [#943](https://github.com/grafana/tempo/pull/943) (@mdisibio)

go.mod

@@ -32,7 +32,7 @@ require (
 	github.com/jaegertracing/jaeger v1.21.0
 	github.com/jedib0t/go-pretty/v6 v6.2.4
 	github.com/jsternberg/zap-logfmt v1.2.0
-	github.com/klauspost/compress v1.13.1
+	github.com/klauspost/compress v1.13.5
 	github.com/minio/minio-go/v7 v7.0.10
 	github.com/olekukonko/tablewriter v0.0.2
 	github.com/opentracing-contrib/go-grpc v0.0.0-20210225150812-73cb765af46e

go.sum

@@ -1260,8 +1260,9 @@ github.com/klauspost/compress v1.11.3/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYs
 github.com/klauspost/compress v1.11.7/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 github.com/klauspost/compress v1.11.12/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 github.com/klauspost/compress v1.11.13/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
-github.com/klauspost/compress v1.13.1 h1:wXr2uRxZTJXHLly6qhJabee5JqIhTRoLBhDOA74hDEQ=
 github.com/klauspost/compress v1.13.1/go.mod h1:8dP1Hq4DHOhN9w426knH3Rhby4rFm6D8eO+e+Dq5Gzg=
+github.com/klauspost/compress v1.13.5 h1:9O69jUPDcsT9fEm74W92rZL9FQY7rCdaXVneq+yyzl4=
+github.com/klauspost/compress v1.13.5/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
 github.com/klauspost/cpuid v0.0.0-20170728055534-ae7887de9fa5/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
 github.com/klauspost/cpuid v1.2.3/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
 github.com/klauspost/cpuid v1.3.1 h1:5JNjFYYQrZeKRJ0734q51WCEEn2huer72Dc7K+R/b6s=

tempodb/encoding/v2/data_reader.go

@@ -9,13 +9,15 @@ import (
 	tempo_io "github.com/grafana/tempo/pkg/io"
 	"github.com/grafana/tempo/tempodb/backend"
 	"github.com/grafana/tempo/tempodb/encoding/common"
+	"github.com/klauspost/compress/zstd"
 )
 
 type dataReader struct {
 	contextReader backend.ContextReader
 
 	pageBuffer []byte
 
+	encoding         backend.Encoding
 	pool             ReaderPool
 	compressedReader io.Reader
 }
@@ -33,6 +35,7 @@ func NewDataReader(r backend.ContextReader, encoding backend.Encoding) (common.D
 	}
 
 	return &dataReader{
+		encoding:      encoding,
 		contextReader: r,
 		pool:          pool,
 	}, nil
@@ -105,7 +108,14 @@ func (r *dataReader) Read(ctx context.Context, records []common.Record, pagesBuf
 			return nil, nil, err
 		}
 
-		pagesBuffer[i], err = tempo_io.ReadAllWithBuffer(reader, len(page), pagesBuffer[i])
+		// TODO: leaky abstraction. can a real programmer fix this in the future?
+		// zstd decoder is ~10-20% faster then the streaming io.Reader interface so prefer that
+		decoder, ok := reader.(*zstd.Decoder)
+		if ok {
+			pagesBuffer[i], err = decoder.DecodeAll(page, pagesBuffer[i][:0])
+		} else {
+			pagesBuffer[i], err = tempo_io.ReadAllWithBuffer(reader, len(page), pagesBuffer[i])
+		}
 		if err != nil {
 			return nil, nil, err
 		}
@@ -138,7 +148,15 @@ func (r *dataReader) NextPage(buffer []byte) ([]byte, uint32, error) {
 		return nil, 0, err
 	}
 
-	buffer, err = tempo_io.ReadAllWithBuffer(compressedReader, len(page.data), buffer)
+	// TODO: leaky abstraction. can a real programmer fix this in the future?
+	// zstd decoder is ~10-20% faster then the streaming io.Reader interface so prefer that
+	decoder, ok := compressedReader.(*zstd.Decoder)
+	if ok {
+		buffer, err = decoder.DecodeAll(page.data, buffer[:0])
+	} else {
+		buffer, err = tempo_io.ReadAllWithBuffer(compressedReader, len(page.data), buffer)
+	}
+
 	if err != nil {
 		return nil, 0, err
 	}
@@ -148,10 +166,18 @@ func (r *dataReader) NextPage(buffer []byte) ([]byte, uint32, error) {
 
 func (r *dataReader) getCompressedReader(page []byte) (io.Reader, error) {
 	var err error
+	var reader io.Reader
+	// we are going to use the stateless zstd decoding functionality. if you pass
+	// a non-nil reader to .GetReader() and then use .DecodeAll() the process hangs
+	// for unknown reasons. so don't do that.
+	if r.encoding != backend.EncZstd {
+		reader = bytes.NewReader(page)
+	}
+
 	if r.compressedReader == nil {
-		r.compressedReader, err = r.pool.GetReader(bytes.NewReader(page))
+		r.compressedReader, err = r.pool.GetReader(reader)
 	} else {
-		r.compressedReader, err = r.pool.ResetReader(bytes.NewReader(page), r.compressedReader)
+		r.compressedReader, err = r.pool.ResetReader(reader, r.compressedReader)
 	}
 	return r.compressedReader, err
 }