From cd563b81ec334abe62cda34a33b8640ea67ef5ba Mon Sep 17 00:00:00 2001 From: Muir Manders 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