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.
This commit is contained in:
MHSanaei
2026-10-31 09:09:22 +00:00
parent cea69a4117
commit e3736564a3
2 changed files with 44 additions and 5 deletions
+29 -5
View File
@@ -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
}