bufpool: Improve perf by using buckets of different size classes

2 files changed, 33 insertions, 2 deletions

diff --git a/internal/zlibx/decompress.go b/internal/zlibx/decompress.go
index 34e62c6f..c6eb65e5 100644
--- a/internal/zlibx/decompress.go
+++ b/internal/zlibx/decompress.go
@@ -12,13 +12,25 @@ import (
 // Decompress inflates the provided zlib wrapped stream and returns the
 // uncompressed data inside a pooled bufpool.Buffer.
 func Decompress(src []byte) (bufpool.Buffer, error) {
-	return DecompressDict(src, nil)
+	return DecompressSized(src, 0)
+}
+
+// DecompressSized inflates the provided zlib stream, using sizeHint to
+// preallocate the output buffer when known (e.g. packfile entries).
+func DecompressSized(src []byte, sizeHint int) (bufpool.Buffer, error) {
+	return DecompressDictSized(src, nil, sizeHint)
 }
 
 // DecompressDict is like Decompress but accepts a preset dictionary. The
 // dictionary must match the checksum embedded in the stream if the dictionary
 // flag is present.
 func DecompressDict(src []byte, dict []byte) (bufpool.Buffer, error) {
+	return DecompressDictSized(src, dict, 0)
+}
+
+// DecompressDictSized is like DecompressDict but allows providing an expected
+// uncompressed size to avoid buffer growth copies.
+func DecompressDictSized(src []byte, dict []byte, sizeHint int) (bufpool.Buffer, error) {
 	if len(src) < 6 {
 		return bufpool.Buffer{}, io.ErrUnexpectedEOF
 	}
@@ -50,7 +62,7 @@ func DecompressDict(src []byte, dict []byte) (bufpool.Buffer, error) {
 	}
 
 	deflateData := src[offset:]
-	out, consumed, err := flatex.DecompressDict(deflateData, dict)
+	out, consumed, err := flatex.DecompressDictSized(deflateData, dict, sizeHint)
 	if err != nil {
 		return bufpool.Buffer{}, err
 	}
diff --git a/internal/zlibx/decompress_test.go b/internal/zlibx/decompress_test.go
index a4e9c608..3dfc07a5 100644
--- a/internal/zlibx/decompress_test.go
+++ b/internal/zlibx/decompress_test.go
@@ -82,3 +82,22 @@ func TestDecompressChecksumError(t *testing.T) {
 		t.Fatalf("expected ErrChecksum, got %v", err)
 	}
 }
+
+func TestDecompressSizedUsesHint(t *testing.T) {
+	payload := []byte("tiny payload")
+	compressed := compressZlib(t, payload, nil)
+
+	const hint = 1 << 20
+	out, err := DecompressSized(compressed, hint)
+	if err != nil {
+		t.Fatalf("DecompressSized: %v", err)
+	}
+	defer out.Release()
+
+	if !bytes.Equal(out.Bytes(), payload) {
+		t.Fatalf("unexpected payload %q", out.Bytes())
+	}
+	if cap(out.Bytes()) < hint {
+		t.Fatalf("expected capacity >= %d, got %d", hint, cap(out.Bytes()))
+	}
+}