From e3736564a3e9f035f015852bab03ba4f165f9706 Mon Sep 17 00:00:00 2001
From: MHSanaei <ho3ein.sanaei@gmail.com>
Date: Sat, 31 Oct 2026 09:09:22 +0000
Subject: [PATCH] perf(wirecodec): decode a config envelope into one buffer

The shared zstd decoder grew its output per frame, copying everything
decoded so far for each extra frame, so 30 KB of tiny frames cost about
8 GB of allocation and seconds of CPU per request. The output buffer is
now sized once from the first frame's declared size (Compress writes a
single frame) with decoding capped at that capacity, and callers waiting
for a decoder worker no longer each hold a buffer.
---
 internal/util/wirecodec/wirecodec.go      | 34 +++++++++++++++++++----
 internal/util/wirecodec/wirecodec_test.go | 15 ++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/internal/util/wirecodec/wirecodec.go b/internal/util/wirecodec/wirecodec.go
index 0b3a130a1..7a69bfd1f 100644
--- a/internal/util/wirecodec/wirecodec.go
+++ b/internal/util/wirecodec/wirecodec.go
@@ -7,6 +7,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"errors"
+	"runtime"
 
 	"github.com/klauspost/compress/zstd"
 )
@@ -35,8 +36,13 @@ const (
 
 // EncodeAll/DecodeAll on these shared instances are safe for concurrent use.
 var (
+	decoderWorkers = runtime.GOMAXPROCS(0)
 	zstdEncoder, _ = zstd.NewWriter(nil)
-	zstdDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderMaxMemory(maxDecodeBytes))
+	zstdDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderMaxMemory(maxDecodeBytes),
+		zstd.WithDecodeAllCapLimit(true), zstd.WithDecoderConcurrency(decoderWorkers))
+	// decodeSlots keeps callers waiting for a decoder worker from each holding an output buffer.
+	decodeSlots        = make(chan struct{}, decoderWorkers)
+	errDecodedTooLarge = errors.New("wirecodec: decompressed body exceeds limit")
 )
 
 // Compress zstd-compresses b.
@@ -47,12 +53,30 @@ func Compress(b []byte) []byte {
 // Decompress zstd-decompresses src, rejecting output larger than maxOut (and any
 // input that would blow the in-memory bomb ceiling).
 func Decompress(src []byte, maxOut int) ([]byte, error) {
-	out, err := zstdDecoder.DecodeAll(src, nil)
-	if err != nil {
+	if maxOut <= 0 || maxOut > maxDecodeBytes {
+		maxOut = maxDecodeBytes
+	}
+	var header zstd.Header
+	if err := header.Decode(src); err != nil {
 		return nil, err
 	}
-	if maxOut > 0 && len(out) > maxOut {
-		return nil, errors.New("wirecodec: decompressed body exceeds limit")
+	// The output buffer is sized once, from the first frame (Compress writes a single
+	// frame): letting it grow per frame copied all output so far for every extra frame.
+	size := maxOut
+	if header.HasFCS {
+		if header.FrameContentSize > uint64(maxOut) {
+			return nil, errDecodedTooLarge
+		}
+		size = int(header.FrameContentSize)
+	}
+	decodeSlots <- struct{}{}
+	defer func() { <-decodeSlots }()
+	out, err := zstdDecoder.DecodeAll(src, make([]byte, 0, size))
+	if errors.Is(err, zstd.ErrDecoderSizeExceeded) {
+		return nil, errDecodedTooLarge
+	}
+	if err != nil {
+		return nil, err
 	}
 	return out, nil
 }
diff --git a/internal/util/wirecodec/wirecodec_test.go b/internal/util/wirecodec/wirecodec_test.go
index 8bd9137df..08ad22b30 100644
--- a/internal/util/wirecodec/wirecodec_test.go
+++ b/internal/util/wirecodec/wirecodec_test.go
@@ -2,6 +2,7 @@ package wirecodec
 
 import (
 	"bytes"
+	"runtime"
 	"strings"
 	"testing"
 )
@@ -49,3 +50,17 @@ func TestSha256HexStableAndSensitive(t *testing.T) {
 		t.Fatalf("expected 64 hex chars, got %d", len(a))
 	}
 }
+
+// Every extra frame made the decoder copy all output so far into a larger buffer, so a
+// few KiB of tiny frames cost gigabytes of allocation for one authenticated request.
+func TestDecompressOfManyFramesAllocatesTheOutputOnce(t *testing.T) {
+	frame := Compress(bytes.Repeat([]byte("A"), 4096))
+	src := bytes.Repeat(frame, 2000)
+	var before, after runtime.MemStats
+	runtime.ReadMemStats(&before)
+	_, _ = Decompress(src, 8<<20)
+	runtime.ReadMemStats(&after)
+	if allocated := after.TotalAlloc - before.TotalAlloc; allocated > 64<<20 {
+		t.Fatalf("decoding %d bytes of frames allocated %d MiB", len(src), allocated>>20)
+	}
+}