Update custom-serialization code to go backward

This commit is contained in:
Antoine Pelisse 2019-05-02 21:16:13 -07:00
parent c4ffec336c
commit 4f1daf0cf3
6 changed files with 115 additions and 56 deletions

View File

@ -19,6 +19,7 @@ package resource
import (
"fmt"
"io"
"math/bits"
"github.com/gogo/protobuf/proto"
)
@ -28,7 +29,7 @@ var _ proto.Sizer = &Quantity{}
func (m *Quantity) Marshal() (data []byte, err error) {
size := m.Size()
data = make([]byte, size)
n, err := m.MarshalTo(data)
n, err := m.MarshalToSizedBuffer(data[:size])
if err != nil {
return nil, err
}
@ -38,30 +39,40 @@ func (m *Quantity) Marshal() (data []byte, err error) {
// MarshalTo is a customized version of the generated Protobuf unmarshaler for a struct
// with a single string field.
func (m *Quantity) MarshalTo(data []byte) (int, error) {
var i int
size := m.Size()
return m.MarshalToSizedBuffer(data[:size])
}
// MarshalToSizedBuffer is a customized version of the generated
// Protobuf unmarshaler for a struct with a single string field.
func (m *Quantity) MarshalToSizedBuffer(data []byte) (int, error) {
i := len(data)
_ = i
var l int
_ = l
data[i] = 0xa
i++
// BEGIN CUSTOM MARSHAL
out := m.String()
i -= len(out)
copy(data[i:], out)
i = encodeVarintGenerated(data, i, uint64(len(out)))
i += copy(data[i:], out)
// END CUSTOM MARSHAL
i--
data[i] = 0xa
return i, nil
return len(data) - i, nil
}
func encodeVarintGenerated(data []byte, offset int, v uint64) int {
offset -= sovGenerated(v)
base := offset
for v >= 1<<7 {
data[offset] = uint8(v&0x7f | 0x80)
v >>= 7
offset++
}
data[offset] = uint8(v)
return offset + 1
return base
}
func (m *Quantity) Size() (n int) {
@ -77,14 +88,7 @@ func (m *Quantity) Size() (n int) {
}
func sovGenerated(x uint64) (n int) {
for {
n++
x >>= 7
if x == 0 {
break
}
}
return n
return (bits.Len64(x|1) + 6) / 7
}
// Unmarshal is a customized version of the generated Protobuf unmarshaler for a struct

View File

@ -70,3 +70,11 @@ func (m *MicroTime) MarshalTo(data []byte) (int, error) {
}
return m.ProtoMicroTime().MarshalTo(data)
}
// MarshalToSizedBuffer implements the protobuf marshalling interface.
func (m *MicroTime) MarshalToSizedBuffer(data []byte) (int, error) {
if m == nil || m.Time.IsZero() {
return 0, nil
}
return m.ProtoMicroTime().MarshalToSizedBuffer(data)
}

View File

@ -75,7 +75,7 @@ func (m *Time) Unmarshal(data []byte) error {
return nil
}
// Marshal implements the protobuf marshalling interface.
// Marshal implements the protobuf marshaling interface.
func (m *Time) Marshal() (data []byte, err error) {
if m == nil || m.Time.IsZero() {
return nil, nil
@ -83,10 +83,18 @@ func (m *Time) Marshal() (data []byte, err error) {
return m.ProtoTime().Marshal()
}
// MarshalTo implements the protobuf marshalling interface.
// MarshalTo implements the protobuf marshaling interface.
func (m *Time) MarshalTo(data []byte) (int, error) {
if m == nil || m.Time.IsZero() {
return 0, nil
}
return m.ProtoTime().MarshalTo(data)
}
// MarshalToSizedBuffer implements the protobuf reverse marshaling interface.
func (m *Time) MarshalToSizedBuffer(data []byte) (int, error) {
if m == nil || m.Time.IsZero() {
return 0, nil
}
return m.ProtoTime().MarshalToSizedBuffer(data)
}

View File

@ -203,7 +203,7 @@ func (s *Serializer) Encode(obj runtime.Object, w io.Writer) error {
switch t := obj.(type) {
case bufferedMarshaller:
// this path performs a single allocation during write but requires the caller to implement
// the more efficient Size and MarshalTo methods
// the more efficient Size and MarshalToSizedBuffer methods
encodedSize := uint64(t.Size())
estimatedSize := prefixSize + estimateUnknownSize(&unk, encodedSize)
data := make([]byte, estimatedSize)
@ -283,6 +283,12 @@ type bufferedMarshaller interface {
runtime.ProtobufMarshaller
}
// Like bufferedMarshaller, but is able to marshal backwards, which is more efficient since it doesn't call Size() as frequently.
type bufferedReverseMarshaller interface {
proto.Sizer
runtime.ProtobufReverseMarshaller
}
// estimateUnknownSize returns the expected bytes consumed by a given runtime.Unknown
// object with a nil RawJSON struct and the expected size of the provided buffer. The
// returned size will not be correct if RawJSOn is set on unk.
@ -414,6 +420,19 @@ func unmarshalToObject(typer runtime.ObjectTyper, creater runtime.ObjectCreater,
// Encode serializes the provided object to the given writer. Overrides is ignored.
func (s *RawSerializer) Encode(obj runtime.Object, w io.Writer) error {
switch t := obj.(type) {
case bufferedReverseMarshaller:
// this path performs a single allocation during write but requires the caller to implement
// the more efficient Size and MarshalToSizedBuffer methods
encodedSize := uint64(t.Size())
data := make([]byte, encodedSize)
n, err := t.MarshalToSizedBuffer(data)
if err != nil {
return err
}
_, err = w.Write(data[:n])
return err
case bufferedMarshaller:
// this path performs a single allocation during write but requires the caller to implement
// the more efficient Size and MarshalTo methods

View File

@ -24,46 +24,66 @@ type ProtobufMarshaller interface {
MarshalTo(data []byte) (int, error)
}
type ProtobufReverseMarshaller interface {
MarshalToSizedBuffer(data []byte) (int, error)
}
// NestedMarshalTo allows a caller to avoid extra allocations during serialization of an Unknown
// that will contain an object that implements ProtobufMarshaller.
// that will contain an object that implements ProtobufMarshaller or ProtobufReverseMarshaller.
func (m *Unknown) NestedMarshalTo(data []byte, b ProtobufMarshaller, size uint64) (int, error) {
var i int
_ = i
var l int
_ = l
data[i] = 0xa
i++
i = encodeVarintGenerated(data, i, uint64(m.TypeMeta.Size()))
n1, err := m.TypeMeta.MarshalTo(data[i:])
// Calculate the full size of the message.
msgSize := m.Size()
if b != nil {
msgSize += int(size) + sovGenerated(size) + 1
}
// Reverse marshal the fields of m.
i := msgSize
i -= len(m.ContentType)
copy(data[i:], m.ContentType)
i = encodeVarintGenerated(data, i, uint64(len(m.ContentType)))
i--
data[i] = 0x22
i -= len(m.ContentEncoding)
copy(data[i:], m.ContentEncoding)
i = encodeVarintGenerated(data, i, uint64(len(m.ContentEncoding)))
i--
data[i] = 0x1a
if b != nil {
if r, ok := b.(ProtobufReverseMarshaller); ok {
n1, err := r.MarshalToSizedBuffer(data[:i])
if err != nil {
return 0, err
}
i -= int(size)
if uint64(n1) != size {
// programmer error: the Size() method for protobuf does not match the results of LashramOt, which means the proto
// struct returned would be wrong.
return 0, fmt.Errorf("the Size() value of %T was %d, but NestedMarshalTo wrote %d bytes to data", b, size, n1)
}
} else {
i -= int(size)
n1, err := b.MarshalTo(data[i:])
if err != nil {
return 0, err
}
if uint64(n1) != size {
// programmer error: the Size() method for protobuf does not match the results of MarshalTo, which means the proto
// struct returned would be wrong.
return 0, fmt.Errorf("the Size() value of %T was %d, but NestedMarshalTo wrote %d bytes to data", b, size, n1)
}
}
i = encodeVarintGenerated(data, i, size)
i--
data[i] = 0x12
}
n2, err := m.TypeMeta.MarshalToSizedBuffer(data[:i])
if err != nil {
return 0, err
}
i += n1
if b != nil {
data[i] = 0x12
i++
i = encodeVarintGenerated(data, i, size)
n2, err := b.MarshalTo(data[i:])
if err != nil {
return 0, err
}
if uint64(n2) != size {
// programmer error: the Size() method for protobuf does not match the results of MarshalTo, which means the proto
// struct returned would be wrong.
return 0, fmt.Errorf("the Size() value of %T was %d, but NestedMarshalTo wrote %d bytes to data", b, size, n2)
}
i += n2
}
data[i] = 0x1a
i++
i = encodeVarintGenerated(data, i, uint64(len(m.ContentEncoding)))
i += copy(data[i:], m.ContentEncoding)
data[i] = 0x22
i++
i = encodeVarintGenerated(data, i, uint64(len(m.ContentType)))
i += copy(data[i:], m.ContentType)
return i, nil
i -= n2
i = encodeVarintGenerated(data, i, uint64(n2))
i--
data[i] = 0xa
return msgSize - i, nil
}

View File

@ -132,7 +132,7 @@ func rewriteOptionalMethods(decl ast.Decl, isOptional OptionalFunc) {
switch t.Name.Name {
case "Unmarshal":
ast.Walk(&optionalItemsVisitor{}, t.Body)
case "MarshalTo", "Size", "String":
case "MarshalTo", "Size", "String", "MarshalToSizedBuffer":
ast.Walk(&optionalItemsVisitor{}, t.Body)
fallthrough
case "Marshal":