proto: amortize cost of growing a Buffer (#584)
Fixes #566
diff --git a/proto/all_test.go b/proto/all_test.go
index e7c1ff9..361f72f 100644
--- a/proto/all_test.go
+++ b/proto/all_test.go
@@ -363,6 +363,48 @@
}
}
+// Ensure that Buffer.Marshal uses O(N) memory for N messages
+func TestBufferMarshalAllocs(t *testing.T) {
+ value := &OtherMessage{Key: Int64(1)}
+ msg := &MyMessage{Count: Int32(1), Others: []*OtherMessage{value}}
+
+ reallocSize := func(t *testing.T, items int, prealloc int) (int64, int64) {
+ var b Buffer
+ b.SetBuf(make([]byte, 0, prealloc))
+
+ var allocSpace int64
+ prevCap := cap(b.Bytes())
+ for i := 0; i < items; i++ {
+ err := b.Marshal(msg)
+ if err != nil {
+ t.Errorf("Marshal err = %q", err)
+ break
+ }
+ if c := cap(b.Bytes()); prevCap != c {
+ allocSpace += int64(c)
+ prevCap = c
+ }
+ }
+ needSpace := int64(len(b.Bytes()))
+ return allocSpace, needSpace
+ }
+
+ for _, prealloc := range []int{0, 100, 10000} {
+ for _, items := range []int{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000} {
+ runtimeSpace, need := reallocSize(t, items, prealloc)
+ totalSpace := int64(prealloc) + runtimeSpace
+
+ runtimeRatio := float64(runtimeSpace) / float64(need)
+ totalRatio := float64(totalSpace) / float64(need)
+
+ if totalRatio < 1 || runtimeRatio > 4 {
+ t.Errorf("needed %dB, allocated %dB total (ratio %.1f), allocated %dB at runtime (ratio %.1f)",
+ need, totalSpace, totalRatio, runtimeSpace, runtimeRatio)
+ }
+ }
+ }
+}
+
// Simple tests for bytes
func TestBytesPrimitives(t *testing.T) {
o := old()
diff --git a/proto/table_marshal.go b/proto/table_marshal.go
index 095c13c..0f212b3 100644
--- a/proto/table_marshal.go
+++ b/proto/table_marshal.go
@@ -2643,10 +2643,7 @@
var err error
if m, ok := pb.(newMarshaler); ok {
siz := m.XXX_Size()
- // make sure buf has enough capacity
- if newCap := len(p.buf) + siz; newCap > cap(p.buf) {
- p.buf = append(make([]byte, 0, newCap), p.buf...)
- }
+ p.grow(siz) // make sure buf has enough capacity
p.buf, err = m.XXX_Marshal(p.buf, p.deterministic)
return err
}
@@ -2663,10 +2660,22 @@
}
var info InternalMessageInfo
siz := info.Size(pb)
- // make sure buf has enough capacity
- if newCap := len(p.buf) + siz; newCap > cap(p.buf) {
- p.buf = append(make([]byte, 0, newCap), p.buf...)
- }
+ p.grow(siz) // make sure buf has enough capacity
p.buf, err = info.Marshal(p.buf, pb, p.deterministic)
return err
}
+
+// grow grows the buffer's capacity, if necessary, to guarantee space for
+// another n bytes. After grow(n), at least n bytes can be written to the
+// buffer without another allocation.
+func (p *Buffer) grow(n int) {
+ need := len(p.buf) + n
+ if need <= cap(p.buf) {
+ return
+ }
+ newCap := len(p.buf) * 2
+ if newCap < need {
+ newCap = need
+ }
+ p.buf = append(make([]byte, 0, newCap), p.buf...)
+}
