go: reduce memory consumption of small fields (#3212)
In the context of THRIFT-5322 commit 37c2ceb7 introduced the
safeReadBytes method, which prevents allocating very large buffers on
malformed input by using a bytes.Buffer as the backing buffer to contain
data.
Due to how io.CopyN and bytes.Buffer interact, and bytes.Buffer's grow
behaviour, the smallest buffer size is 512 bytes, but typically buffer
sizes are at least 1024 bytes. In degenerate cases (reading many small
fields) this can cause significant bloat, e.g. a value of 8 bytes in
length sits in a buffer of 1024 bytes.
This change eliminates bloat by allocating an array of exactly the
desired size, when the desired size is below bytes.Buffer's minimum
buffer size.
diff --git a/lib/go/thrift/binary_protocol.go b/lib/go/thrift/binary_protocol.go
index 431e446..3f2843e 100644
--- a/lib/go/thrift/binary_protocol.go
+++ b/lib/go/thrift/binary_protocol.go
@@ -588,18 +588,23 @@
}
}
-
-
// This function is shared between TBinaryProtocol and TCompactProtocol.
//
// It tries to read size bytes from trans, in a way that prevents large
-// allocations when size is insanely large (mostly caused by malformed message).
+// allocations when size is insanely large (mostly caused by malformed message),
+// or smaller than bytes.MinRead.
func safeReadBytes(size int32, trans io.Reader) ([]byte, error) {
if size < 0 {
return nil, nil
}
-
- buf := new(bytes.Buffer)
- _, err := io.CopyN(buf, trans, int64(size))
- return buf.Bytes(), err
+ if size > bytes.MinRead {
+ // Use bytes.Buffer to prevent allocating size bytes when size is very large
+ buf := new(bytes.Buffer)
+ _, err := io.CopyN(buf, trans, int64(size))
+ return buf.Bytes(), err
+ }
+ // Allocate size bytes
+ b := make([]byte, size)
+ n, err := io.ReadFull(trans, b)
+ return b[:n], err
}
diff --git a/lib/go/thrift/binary_protocol_test.go b/lib/go/thrift/binary_protocol_test.go
index 67f9923..6657a99 100644
--- a/lib/go/thrift/binary_protocol_test.go
+++ b/lib/go/thrift/binary_protocol_test.go
@@ -61,6 +61,11 @@
dataSize int
}{
{
+ label: "tiny",
+ askedSize: 8,
+ dataSize: 8,
+ },
+ {
label: "normal",
askedSize: 100,
dataSize: 100,
@@ -81,6 +86,9 @@
len(buf),
)
}
+ if c.dataSize < bytes.MinRead && cap(buf) != c.dataSize {
+ t.Errorf("Expected to allocate %d bytes for read, allocated %d", c.dataSize, cap(buf))
+ }
if !strings.HasPrefix(safeReadBytesSource, string(buf)) {
t.Errorf("Unexpected read data: %q", buf)
}