cd563b81ec
* 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).
115 lines
2.8 KiB
Go
115 lines
2.8 KiB
Go
/*
|
|
*
|
|
* Copyright 2014 gRPC authors.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*
|
|
*/
|
|
|
|
package grpc
|
|
|
|
import (
|
|
"math"
|
|
"sync"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
)
|
|
|
|
// Codec defines the interface gRPC uses to encode and decode messages.
|
|
// Note that implementations of this interface must be thread safe;
|
|
// a Codec's methods can be called from concurrent goroutines.
|
|
type Codec interface {
|
|
// Marshal returns the wire format of v.
|
|
Marshal(v interface{}) ([]byte, error)
|
|
// Unmarshal parses the wire format into v.
|
|
Unmarshal(data []byte, v interface{}) error
|
|
// String returns the name of the Codec implementation. The returned
|
|
// string will be used as part of content type in transmission.
|
|
String() string
|
|
}
|
|
|
|
// protoCodec is a Codec implementation with protobuf. It is the default codec for gRPC.
|
|
type protoCodec struct {
|
|
}
|
|
|
|
type cachedProtoBuffer struct {
|
|
lastMarshaledSize uint32
|
|
proto.Buffer
|
|
}
|
|
|
|
func capToMaxInt32(val int) uint32 {
|
|
if val > math.MaxInt32 {
|
|
return uint32(math.MaxInt32)
|
|
}
|
|
return uint32(val)
|
|
}
|
|
|
|
func (p protoCodec) marshal(v interface{}, cb *cachedProtoBuffer) ([]byte, error) {
|
|
protoMsg := v.(proto.Message)
|
|
newSlice := make([]byte, 0, cb.lastMarshaledSize)
|
|
|
|
cb.SetBuf(newSlice)
|
|
cb.Reset()
|
|
if err := cb.Marshal(protoMsg); err != nil {
|
|
return nil, err
|
|
}
|
|
out := cb.Bytes()
|
|
cb.lastMarshaledSize = capToMaxInt32(len(out))
|
|
return out, nil
|
|
}
|
|
|
|
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)
|
|
|
|
// put back buffer and lose the ref to the slice
|
|
cb.SetBuf(nil)
|
|
protoBufferPool.Put(cb)
|
|
return out, err
|
|
}
|
|
|
|
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)
|
|
err := cb.Unmarshal(protoMsg)
|
|
cb.SetBuf(nil)
|
|
protoBufferPool.Put(cb)
|
|
return err
|
|
}
|
|
|
|
func (protoCodec) String() string {
|
|
return "proto"
|
|
}
|
|
|
|
var protoBufferPool = &sync.Pool{
|
|
New: func() interface{} {
|
|
return &cachedProtoBuffer{
|
|
Buffer: proto.Buffer{},
|
|
lastMarshaledSize: 16,
|
|
}
|
|
},
|
|
}
|