From 97648a2e6bd2eb852f6a84d58e86c3cf520eaf0c Mon Sep 17 00:00:00 2001 From: Pontus Leitzler Date: Fri, 25 Oct 2024 10:09:00 +0200 Subject: [PATCH] THRIFT-5828: reduce over-allocation in Go binary protocol Due to a subtle caveat in buffer.Bytes the Go binary protocol method ReadBinary() always over-allocated. We now allocate the asked size directly up to 10 MiB, and then use io.CopyN to maintain the malformed message protection. The existing benchmarks show that we reduce allocations for small payloads. Ran the existing benchmarks in lib/go/thrift with and without this change: % go test -run X -bench . -benchmem > % benchcmp -changed old.txt new.txt [...] benchmark old allocs new allocs delta BenchmarkSafeReadBytes/normal-20 5 2 -60.00% BenchmarkBinaryBinary_0-20 4 1 -75.00% BenchmarkBinaryBinary_1-20 4 1 -75.00% BenchmarkBinaryBinary_2-20 5 2 -60.00% BenchmarkCompactBinary0-20 4 1 -75.00% BenchmarkCompactBinary1-20 4 1 -75.00% BenchmarkCompactBinary2-20 5 2 -60.00% BenchmarkSerializer/baseline-20 8 5 -37.50% BenchmarkSerializer/plain-20 20 17 -15.00% BenchmarkSerializer/pool-20 8 5 -37.50% benchmark old bytes new bytes delta BenchmarkSafeReadBytes/normal-20 1656 160 -90.34% BenchmarkBinaryBinary_0-20 1608 160 -90.05% BenchmarkBinaryBinary_1-20 1608 160 -90.05% BenchmarkBinaryBinary_2-20 1634 184 -88.74% BenchmarkCompactBinary0-20 1608 160 -90.05% BenchmarkCompactBinary1-20 1608 160 -90.05% BenchmarkCompactBinary2-20 1634 184 -88.74% BenchmarkSerializer/baseline-20 1000 416 -58.40% BenchmarkSerializer/plain-20 3640 3056 -16.04% BenchmarkSerializer/pool-20 1002 417 -58.38% --- lib/go/thrift/binary_protocol.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/lib/go/thrift/binary_protocol.go b/lib/go/thrift/binary_protocol.go index eded931346e..e8156ea0176 100644 --- a/lib/go/thrift/binary_protocol.go +++ b/lib/go/thrift/binary_protocol.go @@ -555,6 +555,15 @@ func safeReadBytes(size int32, trans io.Reader) ([]byte, error) { return nil, nil } + // Fast path for reads smaller than 10 MiB that only allocates exactly + // what is asked for. + const readLimit = 10 * 1024 * 1024 + if size < readLimit { + b := make([]byte, size) + n, err := io.ReadFull(trans, b) + return b[:n], err + } + buf := new(bytes.Buffer) _, err := io.CopyN(buf, trans, int64(size)) return buf.Bytes(), err