Implement grpc.Compressor.DecompressedSize for snappy to optimize memory allocations (#5213)

Signed-off-by: Xiaochao Dong (@damnever) <[email protected]>

Author: damnever

CHANGELOG.md

@@ -23,6 +23,7 @@
 * [ENHANCEMENT] Update Go version to 1.20.1. #5159
 * [ENHANCEMENT] Distributor: Reuse byte slices when serializing requests from distributors to ingesters. #5193
 * [ENHANCEMENT] Query Frontend: Add number of chunks and samples fetched in query stats. #5198
+* [ENHANCEMENT] Implement grpc.Compressor.DecompressedSize for snappy to optimize memory allocations. #5213
 * [FEATURE] Querier/Query Frontend: support Prometheus /api/v1/status/buildinfo API. #4978
 * [FEATURE] Ingester: Add active series to all_user_stats page. #4972
 * [FEATURE] Ingester: Added `-blocks-storage.tsdb.head-chunks-write-queue-size` allowing to configure the size of the in-memory queue used before flushing chunks to the disk . #5000

pkg/util/grpcencoding/snappy/snappy.go

@@ -51,6 +51,20 @@ func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
 	return reader{dr, &c.readersPool}, nil
 }
 
+// If a Compressor implements DecompressedSize(compressedBytes []byte) int,
+// gRPC will call it to determine the size of the buffer allocated for the
+// result of decompression.
+// Return -1 to indicate unknown size.
+//
+// This is an EXPERIMENTAL feature of grpc-go.
+func (c *compressor) DecompressedSize(compressedBytes []byte) int {
+	decompressedSize, err := snappy.DecodedLen(compressedBytes)
+	if err != nil {
+		return -1
+	}
+	return decompressedSize
+}
+
 type writeCloser struct {
 	writer *snappy.Writer
 	pool   *sync.Pool