internal/compress/zlib: Pool bytes.Reader with a NewReaderBytes

2 files changed, 21 insertions, 1 deletions

diff --git a/internal/compress/zlib/reader.go b/internal/compress/zlib/reader.go
index 2de37af4..74357525 100644
--- a/internal/compress/zlib/reader.go
+++ b/internal/compress/zlib/reader.go
@@ -34,6 +34,7 @@ and to read that data back:
 package zlib
 
 import (
+	"bytes"
 	"encoding/binary"
 	"errors"
 	"hash"
@@ -75,6 +76,7 @@ type Reader struct {
 	trailerRead  uint64
 	err          error
 	scratch      [4]byte
+	br           bytes.Reader
 }
 
 // NewReader creates a new ReadCloser.
@@ -100,6 +102,23 @@ func NewReaderDict(r io.Reader, dict []byte) (*Reader, error) {
 	return z, nil
 }
 
+// NewReaderBytes is like [NewReader] but reads directly from payload,
+// reusing a [bytes.Reader] pooled with the returned Reader
+// instead of allocating a fresh one per call.
+// It is the caller's responsibility to call Close on the ReadCloser when done.
+func NewReaderBytes(payload []byte) (*Reader, error) {
+	z := readerPool.Get()
+
+	z.br.Reset(payload)
+
+	err := z.reset(&z.br, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return z, nil
+}
+
 // Read decompresses bytes from receiver into p.
 func (z *Reader) Read(p []byte) (int, error) {
 	if z.err != nil {
@@ -168,6 +187,7 @@ func (z *Reader) Close() error {
 		return z.err
 	}
 
+	z.br.Reset(nil)
 	readerPool.Put(z)
 
 	return nil
diff --git a/internal/compress/zlib/reader_reset.go b/internal/compress/zlib/reader_reset.go
index 8d641678..6a9340c2 100644
--- a/internal/compress/zlib/reader_reset.go
+++ b/internal/compress/zlib/reader_reset.go
@@ -18,7 +18,7 @@ import (
 
 // reset resets receiver to read a new zlib stream.
 func (z *Reader) reset(r io.Reader, dict []byte) error {
-	*z = Reader{decompressor: z.decompressor, digest: z.digest}
+	*z = Reader{decompressor: z.decompressor, digest: z.digest, br: z.br}
 
 	var input flate.Reader
 	if fr, ok := r.(flate.Reader); ok {