From e9cc0b0c238c053f2431db232d3513e751c62827 Mon Sep 17 00:00:00 2001 From: James Newton-King Date: Tue, 15 Dec 2020 21:41:39 +1300 Subject: [PATCH 1/4] Optimize MapField serialization by removing MessageAdapter --- .../Google.Protobuf/Collections/MapField.cs | 118 +++--------------- .../ParsingPrimitivesMessages.cs | 63 ++++++++++ 2 files changed, 79 insertions(+), 102 deletions(-) diff --git a/csharp/src/Google.Protobuf/Collections/MapField.cs b/csharp/src/Google.Protobuf/Collections/MapField.cs index d60ebc5a8a5..fc9b047b35b 100644 --- a/csharp/src/Google.Protobuf/Collections/MapField.cs +++ b/csharp/src/Google.Protobuf/Collections/MapField.cs @@ -448,12 +448,10 @@ public void AddEntriesFrom(CodedInputStream input, Codec codec) [SecuritySafeCritical] public void AddEntriesFrom(ref ParseContext ctx, Codec codec) { - var adapter = new Codec.MessageAdapter(codec); do { - adapter.Reset(); - ctx.ReadMessage(adapter); - this[adapter.Key] = adapter.Value; + KeyValuePair entry = ParsingPrimitivesMessages.ReadMapEntry(ref ctx, codec); + this[entry.Key] = entry.Value; } while (ParsingPrimitives.MaybeConsumeTag(ref ctx.buffer, ref ctx.state, codec.MapTag)); } @@ -485,13 +483,13 @@ public void WriteTo(CodedOutputStream output, Codec codec) [SecuritySafeCritical] public void WriteTo(ref WriteContext ctx, Codec codec) { - var message = new Codec.MessageAdapter(codec); foreach (var entry in list) { - message.Key = entry.Key; - message.Value = entry.Value; ctx.WriteTag(codec.MapTag); - ctx.WriteMessage(message); + + WritingPrimitives.WriteLength(ref ctx.buffer, ref ctx.state, CalculateEntrySize(codec, entry)); + codec.KeyCodec.WriteTagAndValue(ref ctx, entry.Key); + codec.ValueCodec.WriteTagAndValue(ref ctx, entry.Value); } } @@ -506,18 +504,22 @@ public int CalculateSize(Codec codec) { return 0; } - var message = new Codec.MessageAdapter(codec); int size = 0; foreach (var entry in list) { - message.Key = entry.Key; - message.Value = entry.Value; + int entrySize = CalculateEntrySize(codec, entry); + size += CodedOutputStream.ComputeRawVarint32Size(codec.MapTag); - size += CodedOutputStream.ComputeMessageSize(message); + size += CodedOutputStream.ComputeLengthSize(entrySize) + entrySize; } return size; } + private static int CalculateEntrySize(Codec codec, KeyValuePair entry) + { + return codec.KeyCodec.CalculateSizeWithTag(entry.Key) + codec.ValueCodec.CalculateSizeWithTag(entry.Value); + } + /// /// Returns a string representation of this repeated field, in the same /// way as it would be represented by the default JSON formatter. @@ -659,96 +661,8 @@ public Codec(FieldCodec keyCodec, FieldCodec valueCodec, uint mapT /// internal uint MapTag { get { return mapTag; } } - /// - /// A mutable message class, used for parsing and serializing. This - /// delegates the work to a codec, but implements the interface - /// for interop with and . - /// This is nested inside Codec as it's tightly coupled to the associated codec, - /// and it's simpler if it has direct access to all its fields. - /// - internal class MessageAdapter : IMessage, IBufferMessage - { - private static readonly byte[] ZeroLengthMessageStreamData = new byte[] { 0 }; - - private readonly Codec codec; - internal TKey Key { get; set; } - internal TValue Value { get; set; } - - internal MessageAdapter(Codec codec) - { - this.codec = codec; - } - - internal void Reset() - { - Key = codec.keyCodec.DefaultValue; - Value = codec.valueCodec.DefaultValue; - } - - public void MergeFrom(CodedInputStream input) - { - // Message adapter is an internal class and we know that all the parsing will happen via InternalMergeFrom. - throw new NotImplementedException(); - } - - [SecuritySafeCritical] - public void InternalMergeFrom(ref ParseContext ctx) - { - uint tag; - while ((tag = ctx.ReadTag()) != 0) - { - if (tag == codec.keyCodec.Tag) - { - Key = codec.keyCodec.Read(ref ctx); - } - else if (tag == codec.valueCodec.Tag) - { - Value = codec.valueCodec.Read(ref ctx); - } - else - { - ParsingPrimitivesMessages.SkipLastField(ref ctx.buffer, ref ctx.state); - } - } - - // Corner case: a map entry with a key but no value, where the value type is a message. - // Read it as if we'd seen input with no data (i.e. create a "default" message). - if (Value == null) - { - if (ctx.state.CodedInputStream != null) - { - // the decoded message might not support parsing from ParseContext, so - // we need to allow fallback to the legacy MergeFrom(CodedInputStream) parsing. - Value = codec.valueCodec.Read(new CodedInputStream(ZeroLengthMessageStreamData)); - } - else - { - ParseContext.Initialize(new ReadOnlySequence(ZeroLengthMessageStreamData), out ParseContext zeroLengthCtx); - Value = codec.valueCodec.Read(ref zeroLengthCtx); - } - } - } - - public void WriteTo(CodedOutputStream output) - { - // Message adapter is an internal class and we know that all the writing will happen via InternalWriteTo. - throw new NotImplementedException(); - } - - [SecuritySafeCritical] - public void InternalWriteTo(ref WriteContext ctx) - { - codec.keyCodec.WriteTagAndValue(ref ctx, Key); - codec.valueCodec.WriteTagAndValue(ref ctx, Value); - } - - public int CalculateSize() - { - return codec.keyCodec.CalculateSizeWithTag(Key) + codec.valueCodec.CalculateSizeWithTag(Value); - } - - MessageDescriptor IMessage.Descriptor { get { return null; } } - } + internal FieldCodec KeyCodec => keyCodec; + internal FieldCodec ValueCodec => valueCodec; } private class MapView : ICollection, ICollection diff --git a/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs b/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs index b7097a2ad04..44b02afea4e 100644 --- a/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs +++ b/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs @@ -32,9 +32,11 @@ using System; using System.Buffers; +using System.Collections.Generic; using System.IO; using System.Runtime.CompilerServices; using System.Security; +using Google.Protobuf.Collections; namespace Google.Protobuf { @@ -134,6 +136,67 @@ public static void ReadMessage(ref ParseContext ctx, IMessage message) SegmentedBufferHelper.PopLimit(ref ctx.state, oldLimit); } + private static readonly byte[] ZeroLengthMessageStreamData = new byte[] { 0 }; + + public static KeyValuePair ReadMapEntry(ref ParseContext ctx, MapField.Codec codec) + { + int length = ParsingPrimitives.ParseLength(ref ctx.buffer, ref ctx.state); + if (ctx.state.recursionDepth >= ctx.state.recursionLimit) + { + throw InvalidProtocolBufferException.RecursionLimitExceeded(); + } + int oldLimit = SegmentedBufferHelper.PushLimit(ref ctx.state, length); + ++ctx.state.recursionDepth; + + TKey key = default; + TValue value = default; + + uint tag; + while ((tag = ctx.ReadTag()) != 0) + { + if (tag == codec.KeyCodec.Tag) + { + key = codec.KeyCodec.Read(ref ctx); + } + else if (tag == codec.ValueCodec.Tag) + { + value = codec.ValueCodec.Read(ref ctx); + } + else + { + SkipLastField(ref ctx.buffer, ref ctx.state); + } + } + + // Corner case: a map entry with a key but no value, where the value type is a message. + // Read it as if we'd seen input with no data (i.e. create a "default" message). + if (value == null) + { + if (ctx.state.CodedInputStream != null) + { + // the decoded message might not support parsing from ParseContext, so + // we need to allow fallback to the legacy MergeFrom(CodedInputStream) parsing. + value = codec.ValueCodec.Read(new CodedInputStream(ZeroLengthMessageStreamData)); + } + else + { + ParseContext.Initialize(new ReadOnlySequence(ZeroLengthMessageStreamData), out ParseContext zeroLengthCtx); + value = codec.ValueCodec.Read(ref zeroLengthCtx); + } + } + + CheckReadEndOfStreamTag(ref ctx.state); + // Check that we've read exactly as much data as expected. + if (!SegmentedBufferHelper.IsReachedLimit(ref ctx.state)) + { + throw InvalidProtocolBufferException.TruncatedMessage(); + } + --ctx.state.recursionDepth; + SegmentedBufferHelper.PopLimit(ref ctx.state, oldLimit); + + return new KeyValuePair(key, value); + } + public static void ReadGroup(ref ParseContext ctx, IMessage message) { if (ctx.state.recursionDepth >= ctx.state.recursionLimit) From 6cf068c439ca627eb34a75cf1d850bc869e4ffd6 Mon Sep 17 00:00:00 2001 From: James Newton-King Date: Tue, 15 Dec 2020 22:09:07 +1300 Subject: [PATCH 2/4] Clean up --- csharp/src/Google.Protobuf/Collections/MapField.cs | 13 ++++++++++--- .../Google.Protobuf/ParsingPrimitivesMessages.cs | 4 ++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/csharp/src/Google.Protobuf/Collections/MapField.cs b/csharp/src/Google.Protobuf/Collections/MapField.cs index fc9b047b35b..6b7d0f101a1 100644 --- a/csharp/src/Google.Protobuf/Collections/MapField.cs +++ b/csharp/src/Google.Protobuf/Collections/MapField.cs @@ -657,12 +657,19 @@ public Codec(FieldCodec keyCodec, FieldCodec valueCodec, uint mapT } /// - /// The tag used in the enclosing message to indicate map entries. + /// The key codec. /// - internal uint MapTag { get { return mapTag; } } - internal FieldCodec KeyCodec => keyCodec; + + /// + /// The value codec. + /// internal FieldCodec ValueCodec => valueCodec; + + /// + /// The tag used in the enclosing message to indicate map entries. + /// + internal uint MapTag => mapTag; } private class MapView : ICollection, ICollection diff --git a/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs b/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs index 44b02afea4e..99d44019ee0 100644 --- a/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs +++ b/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs @@ -46,6 +46,8 @@ namespace Google.Protobuf [SecuritySafeCritical] internal static class ParsingPrimitivesMessages { + private static readonly byte[] ZeroLengthMessageStreamData = new byte[] { 0 }; + public static void SkipLastField(ref ReadOnlySpan buffer, ref ParserInternalState state) { if (state.lastTag == 0) @@ -136,8 +138,6 @@ public static void ReadMessage(ref ParseContext ctx, IMessage message) SegmentedBufferHelper.PopLimit(ref ctx.state, oldLimit); } - private static readonly byte[] ZeroLengthMessageStreamData = new byte[] { 0 }; - public static KeyValuePair ReadMapEntry(ref ParseContext ctx, MapField.Codec codec) { int length = ParsingPrimitives.ParseLength(ref ctx.buffer, ref ctx.state); From 542c0cc9d326402d3a520c4b455a5562a817d626 Mon Sep 17 00:00:00 2001 From: James Newton-King Date: Thu, 21 Jan 2021 13:29:39 +1300 Subject: [PATCH 3/4] Fix missing key and test --- .../Collections/MapFieldTest.cs | 28 +++++++++++++++++++ .../ParsingPrimitivesMessages.cs | 4 +-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs b/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs index d8cdee0f0f7..80892a32206 100644 --- a/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs +++ b/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs @@ -600,6 +600,8 @@ public void AddEntriesFrom_CodedInputStream() output.WriteString("the_value"); output.Flush(); + Console.WriteLine(BitConverter.ToString(memoryStream.ToArray())); + var field = new MapField(); var mapCodec = new MapField.Codec(FieldCodec.ForString(keyTag, ""), FieldCodec.ForString(valueTag, ""), 10); var input = new CodedInputStream(memoryStream.ToArray()); @@ -611,6 +613,32 @@ public void AddEntriesFrom_CodedInputStream() Assert.IsTrue(input.IsAtEnd); } + [Test] + public void AddEntriesFrom_CodedInputStream_MissingKey() + { + // map will have string key and string value + var keyTag = WireFormat.MakeTag(1, WireFormat.WireType.LengthDelimited); + var valueTag = WireFormat.MakeTag(2, WireFormat.WireType.LengthDelimited); + + var memoryStream = new MemoryStream(); + var output = new CodedOutputStream(memoryStream); + output.WriteLength(11); // total of valueTag + value + output.WriteTag(valueTag); + output.WriteString("the_value"); + output.Flush(); + + Console.WriteLine(BitConverter.ToString(memoryStream.ToArray())); + + var field = new MapField(); + var mapCodec = new MapField.Codec(FieldCodec.ForString(keyTag, ""), FieldCodec.ForString(valueTag, ""), 10); + var input = new CodedInputStream(memoryStream.ToArray()); + + field.AddEntriesFrom(input, mapCodec); + CollectionAssert.AreEquivalent(new[] { "" }, field.Keys); + CollectionAssert.AreEquivalent(new[] { "the_value" }, field.Values); + Assert.IsTrue(input.IsAtEnd); + } + #if !NET35 [Test] public void IDictionaryKeys_Equals_IReadOnlyDictionaryKeys() diff --git a/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs b/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs index 99d44019ee0..eabaf96d5c9 100644 --- a/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs +++ b/csharp/src/Google.Protobuf/ParsingPrimitivesMessages.cs @@ -148,8 +148,8 @@ public static void ReadMessage(ref ParseContext ctx, IMessage message) int oldLimit = SegmentedBufferHelper.PushLimit(ref ctx.state, length); ++ctx.state.recursionDepth; - TKey key = default; - TValue value = default; + TKey key = codec.KeyCodec.DefaultValue; + TValue value = codec.ValueCodec.DefaultValue; uint tag; while ((tag = ctx.ReadTag()) != 0) From 69223b8386b36f128d5d110d6f5f759c8374a423 Mon Sep 17 00:00:00 2001 From: James Newton-King Date: Thu, 21 Jan 2021 13:36:00 +1300 Subject: [PATCH 4/4] Clean up --- csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs | 2 -- 1 file changed, 2 deletions(-) diff --git a/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs b/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs index 80892a32206..d4c63dc6880 100644 --- a/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs +++ b/csharp/src/Google.Protobuf.Test/Collections/MapFieldTest.cs @@ -600,8 +600,6 @@ public void AddEntriesFrom_CodedInputStream() output.WriteString("the_value"); output.Flush(); - Console.WriteLine(BitConverter.ToString(memoryStream.ToArray())); - var field = new MapField(); var mapCodec = new MapField.Codec(FieldCodec.ForString(keyTag, ""), FieldCodec.ForString(valueTag, ""), 10); var input = new CodedInputStream(memoryStream.ToArray());