bufpool: Improve perf by using buckets of different size classes

2 files changed, 28 insertions, 2 deletions

diff --git a/internal/flatex/decompress_bytes.go b/internal/flatex/decompress_bytes.go
index 5a660b0a..5e29f82d 100644
--- a/internal/flatex/decompress_bytes.go
+++ b/internal/flatex/decompress_bytes.go
@@ -26,13 +26,20 @@ var bufferDecompressorPool = sync.Pool{
 // Decompress inflates the provided DEFLATE stream and returns the full output
 // in a pooled bufpool.Buffer along with the number of consumed bytes from src.
 func Decompress(src []byte) (bufpool.Buffer, int, error) {
-	return DecompressDict(src, nil)
+	return DecompressDictSized(src, nil, 0)
 }
 
 // DecompressDict inflates the provided DEFLATE stream using dict as the preset
 // dictionary and returns the full output in a pooled bufpool.Buffer. The second
 // returned value reports how many bytes of src were consumed.
 func DecompressDict(src []byte, dict []byte) (bufpool.Buffer, int, error) {
+	return DecompressDictSized(src, dict, 0)
+}
+
+// DecompressDictSized is like DecompressDict but allows providing an expected
+// output size to pre-size the destination buffer and avoid repeated growth.
+// A non-positive sizeHint falls back to the default buffer capacity.
+func DecompressDictSized(src []byte, dict []byte, sizeHint int) (bufpool.Buffer, int, error) {
 	d := bufferDecompressorPool.Get().(*bufferDecompressor)
 	defer bufferDecompressorPool.Put(d)
 
@@ -40,7 +47,7 @@ func DecompressDict(src []byte, dict []byte) (bufpool.Buffer, int, error) {
 		return bufpool.Buffer{}, 0, err
 	}
 
-	out := bufpool.Borrow(bufpool.DefaultBufferCap)
+	out := bufpool.Borrow(sizeHint)
 	out.Resize(0)
 
 	for {
diff --git a/internal/flatex/decompress_test.go b/internal/flatex/decompress_test.go
index afb84292..7c290555 100644
--- a/internal/flatex/decompress_test.go
+++ b/internal/flatex/decompress_test.go
@@ -74,3 +74,22 @@ func TestDecompressDictMissing(t *testing.T) {
 		t.Fatalf("expected error when dictionary missing")
 	}
 }
+
+func TestDecompressDictSizedUsesHint(t *testing.T) {
+	payload := []byte("short")
+	compressed := compressDeflate(t, payload, nil)
+
+	const hint = 1 << 19
+	out, _, err := DecompressDictSized(compressed, nil, hint)
+	if err != nil {
+		t.Fatalf("DecompressDictSized: %v", err)
+	}
+	defer out.Release()
+
+	if !bytes.Equal(out.Bytes(), payload) {
+		t.Fatalf("unexpected payload: got %q", out.Bytes())
+	}
+	if cap(out.Bytes()) < hint {
+		t.Fatalf("expected capacity >= %d, got %d", hint, cap(out.Bytes()))
+	}
+}