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