diff --git a/rpc/message/test/message.go b/rpc/message/test/message.go index df769d2..435f20a 100644 --- a/rpc/message/test/message.go +++ b/rpc/message/test/message.go @@ -17,6 +17,7 @@ type jsonMessage interface { type binaryMessage interface { StableMarshal([]byte) []byte + StableSize() int Unmarshal([]byte) error } @@ -53,6 +54,11 @@ func TestRPCMessage(t *testing.T, msgGens ...func(empty bool) message.Message) { } if bm, ok := msg.(binaryMessage); ok { + t.Run(fmt.Sprintf("%T.StableSize() does no allocations", bm), func(t *testing.T) { + require.Zero(t, testing.AllocsPerRun(1000, func() { + _ = bm.StableSize() + })) + }) t.Run(fmt.Sprintf("Binary_%T", msg), func(t *testing.T) { data := bm.StableMarshal(nil) diff --git a/session/marshal.go b/session/marshal.go index 3c56cd4..cda9579 100644 --- a/session/marshal.go +++ b/session/marshal.go @@ -211,7 +211,7 @@ func (c *ObjectSessionContext) StableMarshal(buf []byte) []byte { } offset := proto.EnumMarshal(objectCtxVerbField, buf, int32(c.verb)) - proto.NestedStructureMarshal(objectCtxTargetField, buf[offset:], &objectSessionContextTarget{ + proto.NestedStructureMarshal(objectCtxTargetField, buf[offset:], objectSessionContextTarget{ cnr: c.cnr, objs: c.objs, }) @@ -225,7 +225,7 @@ func (c *ObjectSessionContext) StableSize() (size int) { } size += proto.EnumSize(objectCtxVerbField, int32(c.verb)) - size += proto.NestedStructureSize(objectCtxTargetField, &objectSessionContextTarget{ + size += proto.NestedStructureSize(objectCtxTargetField, objectSessionContextTarget{ cnr: c.cnr, objs: c.objs, }) diff --git a/status/marshal.go b/status/marshal.go index 78064c1..1868a43 100644 --- a/status/marshal.go +++ b/status/marshal.go @@ -69,6 +69,10 @@ func (x *Status) StableMarshal(buf []byte) []byte { } func (x *Status) StableSize() (size int) { + if x == nil { + return 0 + } + size += protoutil.UInt32Size(statusCodeFNum, CodeToGRPC(x.code)) size += protoutil.StringSize(statusMsgFNum, x.msg) diff --git a/util/proto/marshal.go b/util/proto/marshal.go index a82478b..b1f55a3 100644 --- a/util/proto/marshal.go +++ b/util/proto/marshal.go @@ -10,7 +10,6 @@ import ( "encoding/binary" "math" "math/bits" - "reflect" ) type ( @@ -21,13 +20,21 @@ type ( ) func BytesMarshal(field int, buf, v []byte) int { - if len(v) == 0 { - return 0 - } return bytesMarshal(field, buf, v) } -func bytesMarshal(field int, buf, v []byte) int { +func BytesSize(field int, v []byte) int { + return bytesSize(field, v) +} + +func bytesMarshal[T ~[]byte | ~string](field int, buf []byte, v T) int { + if len(v) == 0 { + return 0 + } + return bytesMarshalNoCheck(field, buf, v) +} + +func bytesMarshalNoCheck[T ~[]byte | ~string](field int, buf []byte, v T) int { prefix := field<<3 | 0x2 // buf length check can prevent panic at PutUvarint, but it will make @@ -39,26 +46,25 @@ func bytesMarshal(field int, buf, v []byte) int { return i } -func BytesSize(field int, v []byte) int { - ln := len(v) - if ln == 0 { +func bytesSize[T ~[]byte | ~string](field int, v T) int { + if len(v) == 0 { return 0 } - return bytesSize(field, v) + return bytesSizeNoCheck(field, v) } -func bytesSize(field int, v []byte) int { +func bytesSizeNoCheck[T ~[]byte | ~string](field int, v T) int { prefix := field<<3 | 0x2 return VarUIntSize(uint64(prefix)) + VarUIntSize(uint64(len(v))) + len(v) } func StringMarshal(field int, buf []byte, v string) int { - return BytesMarshal(field, buf, []byte(v)) + return bytesMarshal(field, buf, v) } func StringSize(field int, v string) int { - return BytesSize(field, []byte(v)) + return bytesSize(field, v) } func BoolMarshal(field int, buf []byte, v bool) int { @@ -147,7 +153,7 @@ func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) int { var offset int for i := range v { - offset += bytesMarshal(field, buf[offset:], v[i]) + offset += bytesMarshalNoCheck(field, buf[offset:], v[i]) } return offset @@ -155,7 +161,7 @@ func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) int { func RepeatedBytesSize(field int, v [][]byte) (size int) { for i := range v { - size += bytesSize(field, v[i]) + size += bytesSizeNoCheck(field, v[i]) } return size @@ -165,7 +171,7 @@ func RepeatedStringMarshal(field int, buf []byte, v []string) int { var offset int for i := range v { - offset += bytesMarshal(field, buf[offset:], []byte(v[i])) + offset += bytesMarshalNoCheck(field, buf[offset:], v[i]) } return offset @@ -173,36 +179,19 @@ func RepeatedStringMarshal(field int, buf []byte, v []string) int { func RepeatedStringSize(field int, v []string) (size int) { for i := range v { - size += bytesSize(field, []byte(v[i])) + size += bytesSizeNoCheck(field, v[i]) } return size } -func RepeatedUInt64Marshal(field int, buf []byte, v []uint64) int { - if len(v) == 0 { - return 0 - } - - prefix := field<<3 | 0x02 - offset := binary.PutUvarint(buf, uint64(prefix)) - - _, arrSize := RepeatedUInt64Size(field, v) - offset += binary.PutUvarint(buf[offset:], uint64(arrSize)) - for i := range v { - offset += binary.PutUvarint(buf[offset:], v[i]) - } - - return offset -} - -func RepeatedUInt64Size(field int, v []uint64) (size, arraySize int) { +func repeatedUIntSize[T ~uint64 | ~int64 | ~uint32 | ~int32](field int, v []T) (size, arraySize int) { if len(v) == 0 { return 0, 0 } for i := range v { - size += VarUIntSize(v[i]) + size += VarUIntSize(uint64(v[i])) } arraySize = size @@ -214,82 +203,53 @@ func RepeatedUInt64Size(field int, v []uint64) (size, arraySize int) { return size, arraySize } -func RepeatedInt64Marshal(field int, buf []byte, v []int64) int { +func repeatedUIntMarshal[T ~uint64 | ~int64 | ~uint32 | ~int32](field int, buf []byte, v []T) int { if len(v) == 0 { return 0 } - convert := make([]uint64, len(v)) + prefix := field<<3 | 0x02 + offset := binary.PutUvarint(buf, uint64(prefix)) + + _, arrSize := repeatedUIntSize(field, v) + offset += binary.PutUvarint(buf[offset:], uint64(arrSize)) for i := range v { - convert[i] = uint64(v[i]) + offset += binary.PutUvarint(buf[offset:], uint64(v[i])) } - return RepeatedUInt64Marshal(field, buf, convert) + return offset +} + +func RepeatedUInt64Marshal(field int, buf []byte, v []uint64) int { + return repeatedUIntMarshal(field, buf, v) +} + +func RepeatedUInt64Size(field int, v []uint64) (size, arraySize int) { + return repeatedUIntSize(field, v) +} + +func RepeatedInt64Marshal(field int, buf []byte, v []int64) int { + return repeatedUIntMarshal(field, buf, v) } func RepeatedInt64Size(field int, v []int64) (size, arraySize int) { - if len(v) == 0 { - return 0, 0 - } - - convert := make([]uint64, len(v)) - for i := range v { - convert[i] = uint64(v[i]) - } - - return RepeatedUInt64Size(field, convert) + return repeatedUIntSize(field, v) } func RepeatedUInt32Marshal(field int, buf []byte, v []uint32) int { - if len(v) == 0 { - return 0 - } - - convert := make([]uint64, len(v)) - for i := range v { - convert[i] = uint64(v[i]) - } - - return RepeatedUInt64Marshal(field, buf, convert) + return repeatedUIntMarshal(field, buf, v) } func RepeatedUInt32Size(field int, v []uint32) (size, arraySize int) { - if len(v) == 0 { - return 0, 0 - } - - convert := make([]uint64, len(v)) - for i := range v { - convert[i] = uint64(v[i]) - } - - return RepeatedUInt64Size(field, convert) + return repeatedUIntSize(field, v) } func RepeatedInt32Marshal(field int, buf []byte, v []int32) int { - if len(v) == 0 { - return 0 - } - - convert := make([]uint64, len(v)) - for i := range v { - convert[i] = uint64(v[i]) - } - - return RepeatedUInt64Marshal(field, buf, convert) + return repeatedUIntMarshal(field, buf, v) } func RepeatedInt32Size(field int, v []int32) (size, arraySize int) { - if len(v) == 0 { - return 0, 0 - } - - convert := make([]uint64, len(v)) - for i := range v { - convert[i] = uint64(v[i]) - } - - return RepeatedUInt64Size(field, convert) + return repeatedUIntSize(field, v) } // VarUIntSize returns length of varint byte sequence for uint64 value 'x'. @@ -302,28 +262,27 @@ func NestedStructurePrefix(field int64) (prefix uint64, ln int) { return prefix, VarUIntSize(prefix) } -func NestedStructureMarshal(field int64, buf []byte, v stableMarshaller) int { - if v == nil || reflect.ValueOf(v).IsNil() { +func NestedStructureMarshal[T stableMarshaller](field int64, buf []byte, v T) int { + n := v.StableSize() + if n == 0 { return 0 } prefix, _ := NestedStructurePrefix(field) offset := binary.PutUvarint(buf, prefix) - - n := v.StableSize() offset += binary.PutUvarint(buf[offset:], uint64(n)) v.StableMarshal(buf[offset:]) return offset + n } -func NestedStructureSize(field int64, v stableMarshaller) (size int) { - if v == nil || reflect.ValueOf(v).IsNil() { +func NestedStructureSize[T stableMarshaller](field int64, v T) (size int) { + n := v.StableSize() + if n == 0 { return 0 } _, ln := NestedStructurePrefix(field) - n := v.StableSize() size = ln + VarUIntSize(uint64(n)) + n return size