From cd563b81ec334abe62cda34a33b8640ea67ef5ba Mon Sep 17 00:00:00 2001
From: Muir Manders <muir@retailnext.net>
Date: Sat, 2 Dec 2017 08:32:05 +0900
Subject: [PATCH] protoCodec: avoid buffer allocations if
 proto.Marshaler/Unmarshaler (#1689)

* protoCodec: return early if proto.Marshaler

If the object to marshal implements proto.Marshaler, delegate to that
immediately instead of pre-allocating a buffer. (*proto.Buffer).Marshal
has the same logic, so the []byte buffer we pre-allocate in codec.go
would never be used.

This is mainly for users of gogoproto. If you turn on the "marshaler"
and "sizer" gogoproto options, the generated Marshal() method already
pre-allocates a buffer of the appropriate size for marshaling the
entire object.

* protoCodec: return early if proto.Unmarshaler

If the object to unmarshal implements proto.Unmarshaler, delegate to
that immediately. This perhaps saves a bit of work preparing the
cached the proto.Buffer object which would not end up being used for
the proto.Unmarshaler case.

Note that I moved the obj.Reset() call above the delegation to
obj.Unmarshal(). This maintains the grpc behavior of
proto.Unmarshalers always being Reset() before being delegated to,
which is consistent to how proto.Unmarshal() behaves (proto.Buffer
does not call Reset() in Unmarshal).
---
 codec.go | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/codec.go b/codec.go
index b452a4ae..43d81ed2 100644
--- a/codec.go
+++ b/codec.go
@@ -69,6 +69,11 @@ func (p protoCodec) marshal(v interface{}, cb *cachedProtoBuffer) ([]byte, error
 }
 
 func (p protoCodec) Marshal(v interface{}) ([]byte, error) {
+	if pm, ok := v.(proto.Marshaler); ok {
+		// object can marshal itself, no need for buffer
+		return pm.Marshal()
+	}
+
 	cb := protoBufferPool.Get().(*cachedProtoBuffer)
 	out, err := p.marshal(v, cb)
 
@@ -79,10 +84,17 @@ func (p protoCodec) Marshal(v interface{}) ([]byte, error) {
 }
 
 func (p protoCodec) Unmarshal(data []byte, v interface{}) error {
+	protoMsg := v.(proto.Message)
+	protoMsg.Reset()
+
+	if pu, ok := protoMsg.(proto.Unmarshaler); ok {
+		// object can unmarshal itself, no need for buffer
+		return pu.Unmarshal(data)
+	}
+
 	cb := protoBufferPool.Get().(*cachedProtoBuffer)
 	cb.SetBuf(data)
-	v.(proto.Message).Reset()
-	err := cb.Unmarshal(v.(proto.Message))
+	err := cb.Unmarshal(protoMsg)
 	cb.SetBuf(nil)
 	protoBufferPool.Put(cb)
 	return err