Decompress: reuse the caller-supplied buffer before sizing from the frame

Decompress sized its destination from decompressSizeHint(src), which reads
the frame's content-size field. Frames that do not carry that field — legacy
zstd v0.5 frames, and streaming (unknown-pledged-size) frames — make the hint
fall back to its upper bound (>= decompressSizeBufferLimit, i.e. 1 MB). When
the caller passed a smaller-but-adequate buffer, Decompress discarded it and
allocated that bound, so every such decode allocated at least 1 MB regardless
of the real payload size.

Try the caller-supplied buffer first via DecompressInto (which reports a
too-small buffer without writing past it), and only fall back to the
hint-sized allocation, then the stream API, when it does not fit. For frames
that do advertise their size the behavior is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: KrisKennawayDD

zstd.go

@@ -136,13 +136,24 @@ func Decompress(dst, src []byte) ([]byte, error) {
 		return []byte{}, ErrEmptySlice
 	}
 
-	bound := decompressSizeHint(src)
-	if cap(dst) >= bound {
-		dst = dst[0:cap(dst)]
-	} else {
-		dst = make([]byte, bound)
+	// If the caller supplied a buffer, try it as-is before sizing from the
+	// frame header. Frames that do not carry their decompressed size (e.g.
+	// legacy zstd v0.5 frames) make decompressSizeHint fall back to a large
+	// upper bound, which would otherwise discard an adequate caller buffer and
+	// allocate that bound instead. DecompressInto reports a too-small buffer
+	// without writing past it, so this attempt is safe.
+	if cap(dst) > 0 {
+		written, err := DecompressInto(dst[:cap(dst)], src)
+		if err == nil {
+			return dst[:written], nil
+		}
+		if !IsDstSizeTooSmallError(err) {
+			return nil, err
+		}
 	}
 
+	// No usable caller buffer; size from the frame hint.
+	dst = make([]byte, decompressSizeHint(src))
 	written, err := DecompressInto(dst, src)
 	if err == nil {
 		return dst[:written], nil

zstd_reuse_buffer_test.go

@@ -0,0 +1,79 @@
+package zstd
+
+import (
+	"bytes"
+	"testing"
+)
+
+// unknownSizeFrame returns a zstd frame that does not advertise its decompressed
+// size in the header. The streaming writer produces such frames, as do legacy
+// zstd v0.5 frames; ZSTD_getFrameContentSize reports "unknown" for both, so
+// decompressSizeHint falls back to its (large) upper bound.
+func unknownSizeFrame(t *testing.T, payload []byte) []byte {
+	t.Helper()
+	var b bytes.Buffer
+	w := NewWriter(&b)
+	if _, err := w.Write(payload); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("close: %v", err)
+	}
+	return b.Bytes()
+}
+
+// TestDecompressReusesCallerBufferUnknownSize verifies that when a frame does
+// not advertise its content size, Decompress uses an adequately sized caller
+// buffer instead of discarding it and allocating decompressSizeBufferLimit.
+func TestDecompressReusesCallerBufferUnknownSize(t *testing.T) {
+	payload := bytes.Repeat([]byte("datadog-"), 525) // 4200 bytes
+	frame := unknownSizeFrame(t, payload)
+
+	// Sanity: the frame has unknown content size, so the hint is the upper
+	// bound (>= decompressSizeBufferLimit). Without reusing the caller buffer,
+	// Decompress would allocate that many bytes.
+	if got := decompressSizeHint(frame); got < decompressSizeBufferLimit {
+		t.Fatalf("expected unknown-size frame to hint the upper bound, got %d", got)
+	}
+
+	buf := make([]byte, 8192) // adequate for the payload, far below the bound
+	out, err := Decompress(buf, frame)
+	if err != nil {
+		t.Fatalf("Decompress: %v", err)
+	}
+	if !bytes.Equal(out, payload) {
+		t.Fatalf("round-trip mismatch")
+	}
+	if cap(out) != cap(buf) {
+		t.Fatalf("caller buffer (cap %d) should be reused, got cap %d", cap(buf), cap(out))
+	}
+}
+
+// TestDecompressUnknownSizeTooSmallBuffer verifies the fallback still works when
+// the caller buffer is too small for an unknown-size frame.
+func TestDecompressUnknownSizeTooSmallBuffer(t *testing.T) {
+	payload := bytes.Repeat([]byte("datadog-"), 525)
+	frame := unknownSizeFrame(t, payload)
+
+	out, err := Decompress(make([]byte, 8), frame) // too small; must fall back
+	if err != nil {
+		t.Fatalf("Decompress: %v", err)
+	}
+	if !bytes.Equal(out, payload) {
+		t.Fatalf("round-trip mismatch on fallback")
+	}
+}
+
+// TestDecompressUnknownSizeNilBuffer verifies nil dst still decompresses.
+func TestDecompressUnknownSizeNilBuffer(t *testing.T) {
+	payload := bytes.Repeat([]byte("datadog-"), 525)
+	frame := unknownSizeFrame(t, payload)
+
+	out, err := Decompress(nil, frame)
+	if err != nil {
+		t.Fatalf("Decompress: %v", err)
+	}
+	if !bytes.Equal(out, payload) {
+		t.Fatalf("round-trip mismatch on nil dst")
+	}
+}