mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-11-02 19:25:05 +03:00
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:
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user