Decompress: reuse the caller-supplied buffer for unknown-size frames

Decompress sized its destination from decompressSizeHint(src), which reads the
frame's content-size field. Frames that do not carry that field make the hint
fall back to a pessimistic upper bound (>= decompressSizeBufferLimit, i.e. 1 MB):
this is the case for legacy zstd v0.5 frames and for streaming frames compressed
without a pledged source size. 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.

decompressSizeHint now also reports whether the size was read from the frame
(foundHint). When it was not, Decompress and ctx.Decompress try the
caller-supplied buffer first via DecompressInto -- which reports a too-small
buffer without writing past it -- before falling back to the hint-sized
allocation and then the stream API. When the size is known (the common case
where the same zstd version compressed and decompressed) the original path is
unchanged, so callers passing a too-small buffer do not pay for a failed attempt.

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

Author: KrisKennawayDD

zstd.go

@@ -58,30 +58,34 @@ func cCompressBound(srcSize int) int {
 }
 
 // decompressSizeHint tries to give a hint on how much of the output buffer size we should have
-// based on zstd frame descriptors. To prevent DOS from maliciously-created payloads, limit the size
-func decompressSizeHint(src []byte) int {
+// based on zstd frame descriptors. To prevent DOS from maliciously-created payloads, limit the size.
+// foundHint reports whether the size was read from the frame header. When it is false the returned
+// hint is only a pessimistic upper bound: the frame does not advertise its decompressed size (as is
+// the case for legacy zstd v0.5 frames and streaming frames compressed without a pledged size).
+func decompressSizeHint(src []byte) (hint int, foundHint bool) {
 	// 1 MB or 50x input size
 	upperBound := 50 * len(src)
 	if upperBound < decompressSizeBufferLimit {
 		upperBound = decompressSizeBufferLimit
 	}
 
-	hint := upperBound
+	hint = upperBound
 	if len(src) >= zstdFrameHeaderSizeMin {
-		hint = int(C.ZSTD_getFrameContentSize(unsafe.Pointer(&src[0]), C.size_t(len(src))))
-		if hint < 0 { // On error, just use upperBound
-			hint = upperBound
-		}
-		if hint == 0 { // When compressing the empty slice, we need an output of at least 1 to pass down to the C lib
-			hint = 1
+		contentSize := int(C.ZSTD_getFrameContentSize(unsafe.Pointer(&src[0]), C.size_t(len(src))))
+		if contentSize >= 0 { // a negative value means the size is unknown or the header is in error
+			foundHint = true
+			hint = contentSize
+			if hint == 0 { // When compressing the empty slice, we need an output of at least 1 to pass down to the C lib
+				hint = 1
+			}
 		}
 	}
 
 	// Take the minimum of both
 	if hint > upperBound {
-		return upperBound
+		return upperBound, foundHint
 	}
-	return hint
+	return hint, foundHint
 }
 
 // Compress src into dst.  If you have a buffer to use, you can pass it to
@@ -136,11 +140,31 @@ func Decompress(dst, src []byte) ([]byte, error) {
 		return []byte{}, ErrEmptySlice
 	}
 
-	bound := decompressSizeHint(src)
-	if cap(dst) >= bound {
+	hint, foundHint := decompressSizeHint(src)
+
+	// When the frame does not advertise its decompressed size, the hint is only
+	// a pessimistic upper bound (>= decompressSizeBufferLimit). This happens for
+	// legacy zstd v0.5 frames and for streaming frames compressed without a
+	// pledged source size. Rather than discard an adequate caller-supplied
+	// buffer and allocate that bound, try the caller buffer as-is first.
+	// DecompressInto reports a too-small buffer without writing past it, so this
+	// attempt is safe. When the size is known (the common case where the same
+	// zstd version compressed and decompressed) we skip this and size from the
+	// hint, so callers passing a too-small buffer do not pay for a failed attempt.
+	if !foundHint && cap(dst) > 0 {
+		written, err := DecompressInto(dst[:cap(dst)], src)
+		if err == nil {
+			return dst[:written], nil
+		}
+		if !IsDstSizeTooSmallError(err) {
+			return nil, err
+		}
+	}
+
+	if cap(dst) >= hint {
 		dst = dst[0:cap(dst)]
 	} else {
-		dst = make([]byte, bound)
+		dst = make([]byte, hint)
 	}
 
 	written, err := DecompressInto(dst, src)

zstd_bulk.go

@@ -110,7 +110,7 @@ func (p *BulkProcessor) Decompress(dst, src []byte) ([]byte, error) {
 		return nil, ErrEmptySlice
 	}
 
-	contentSize := decompressSizeHint(src)
+	contentSize, _ := decompressSizeHint(src)
 	if cap(dst) >= contentSize {
 		dst = dst[0:cap(dst)]
 	} else {

zstd_ctx.go

@@ -41,14 +41,14 @@ type ctx struct {
 }
 
 // Create a new ZStd Context.
-//  When compressing/decompressing many times, it is recommended to allocate a
-//  context just once, and re-use it for each successive compression operation.
-//  This will make workload friendlier for system's memory.
-//  Note : re-using context is just a speed / resource optimization.
-//         It doesn't change the compression ratio, which remains identical.
-//  Note 2 : In multi-threaded environments,
-//         use one different context per thread for parallel execution.
 //
+//	When compressing/decompressing many times, it is recommended to allocate a
+//	context just once, and re-use it for each successive compression operation.
+//	This will make workload friendlier for system's memory.
+//	Note : re-using context is just a speed / resource optimization.
+//	       It doesn't change the compression ratio, which remains identical.
+//	Note 2 : In multi-threaded environments,
+//	       use one different context per thread for parallel execution.
 func NewCtx() Ctx {
 	c := &ctx{
 		cctx: C.ZSTD_createCCtx(),
@@ -106,11 +106,24 @@ func (c *ctx) Decompress(dst, src []byte) ([]byte, error) {
 		return []byte{}, ErrEmptySlice
 	}
 
-	bound := decompressSizeHint(src)
-	if cap(dst) >= bound {
+	hint, foundHint := decompressSizeHint(src)
+
+	// See Decompress: when the frame does not advertise its size, prefer a
+	// usable caller buffer over allocating the pessimistic upper bound.
+	if !foundHint && cap(dst) > 0 {
+		written, err := c.DecompressInto(dst[:cap(dst)], src)
+		if err == nil {
+			return dst[:written], nil
+		}
+		if !IsDstSizeTooSmallError(err) {
+			return nil, err
+		}
+	}
+
+	if cap(dst) >= hint {
 		dst = dst[0:cap(dst)]
 	} else {
-		dst = make([]byte, bound)
+		dst = make([]byte, hint)
 	}
 
 	written, err := c.DecompressInto(dst, src)

zstd_reuse_buffer_test.go

@@ -0,0 +1,153 @@
+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()
+}
+
+type namedDecompressor struct {
+	name string
+	fn   func(dst, src []byte) ([]byte, error)
+}
+
+// decompressors are the entry points that share the caller-buffer-reuse logic.
+func decompressors() []namedDecompressor {
+	return []namedDecompressor{
+		{"Decompress", Decompress},
+		{"ctx.Decompress", NewCtx().Decompress},
+	}
+}
+
+// TestDecompressReusesCallerBufferUnknownSize verifies that when a frame does
+// not advertise its content size, decompression 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,
+	// decompression would allocate that many bytes.
+	if hint, found := decompressSizeHint(frame); found || hint < decompressSizeBufferLimit {
+		t.Fatalf("expected unknown-size frame to hint the upper bound: found=%v hint=%d", found, hint)
+	}
+
+	for _, d := range decompressors() {
+		t.Run(d.name, func(t *testing.T) {
+			buf := make([]byte, 8192) // adequate for the payload, far below the bound
+			out, err := d.fn(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)
+
+	for _, d := range decompressors() {
+		t.Run(d.name, func(t *testing.T) {
+			out, err := d.fn(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)
+
+	for _, d := range decompressors() {
+		t.Run(d.name, func(t *testing.T) {
+			out, err := d.fn(nil, frame)
+			if err != nil {
+				t.Fatalf("decompress: %v", err)
+			}
+			if !bytes.Equal(out, payload) {
+				t.Fatalf("round-trip mismatch on nil dst")
+			}
+		})
+	}
+}
+
+// TestDecompressKnownSizeReusesCallerBuffer verifies that for frames that do
+// advertise their size, an adequate caller buffer is still reused (the common
+// path is unchanged).
+func TestDecompressKnownSizeReusesCallerBuffer(t *testing.T) {
+	payload := bytes.Repeat([]byte("datadog-"), 525)
+	frame, err := Compress(nil, payload)
+	if err != nil {
+		t.Fatalf("Compress: %v", err)
+	}
+	if _, found := decompressSizeHint(frame); !found {
+		t.Fatal("Compress frame should advertise its content size")
+	}
+
+	for _, d := range decompressors() {
+		t.Run(d.name, func(t *testing.T) {
+			buf := make([]byte, 8192)
+			out, err := d.fn(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))
+			}
+		})
+	}
+}
+
+// TestDecompressSizeHintFound verifies foundHint distinguishes frames that
+// advertise their content size (produced by Compress) from those that do not
+// (streaming, as here, and legacy zstd v0.5 frames).
+func TestDecompressSizeHintFound(t *testing.T) {
+	payload := bytes.Repeat([]byte("datadog-"), 525)
+
+	known, err := Compress(nil, payload)
+	if err != nil {
+		t.Fatalf("Compress: %v", err)
+	}
+	if hint, found := decompressSizeHint(known); !found || hint != len(payload) {
+		t.Fatalf("known-size frame: found=%v hint=%d, want found=true hint=%d", found, hint, len(payload))
+	}
+
+	if _, found := decompressSizeHint(unknownSizeFrame(t, payload)); found {
+		t.Fatalf("streaming frame should not advertise its content size")
+	}
+}