Optimize client buffers during serialization (#775)

JamesNK · web-flow · commit 2cb89bf44666 · 2020-02-20T21:23:38.000+13:00
diff --git a/perf/Grpc.AspNetCore.Microbenchmarks/Program.cs b/perf/Grpc.AspNetCore.Microbenchmarks/Program.cs
@@ -31,18 +31,18 @@ static void Main(string[] args)
         // Profiling option. This will call methods explicitly, in-process
         static async System.Threading.Tasks.Task Main(string[] args)
         {
-            var benchmark = new Client.CompressedUnaryServerCallHandlerBenchmark();
+            var benchmark = new Client.UnaryClientBenchmark();
             benchmark.GlobalSetup();
             for (var i = 0; i < 10000; i++)
             {
-                await benchmark.CompressedSayHelloAsync();
+                await benchmark.SayHelloAsync();
             }
 
             System.Console.WriteLine("Press any key to start.");
             System.Console.ReadKey();
             for (var i = 0; i < 1; i++)
             {
-                await benchmark.CompressedSayHelloAsync();
+                await benchmark.SayHelloAsync();
             }
 
             System.Console.WriteLine("Done. Press any key to exit.");
diff --git a/src/Grpc.Net.Client/Internal/DefaultSerializationContext.cs b/src/Grpc.Net.Client/Internal/DefaultSerializationContext.cs
@@ -19,6 +19,7 @@
 using System;
 using System.Buffers;
 using System.Buffers.Binary;
+using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using Grpc.Core;
 
@@ -112,6 +113,9 @@ public override IBufferWriter<byte> GetBufferWriter()
 
         private IBufferWriter<byte> ResolveBufferWriter()
         {
+            // TODO(JamesNK): I believe length should be known by the context before the buffer writer is
+            // fetched for the first time. Should be able to initialize a custom buffer writer with pooled
+            // array of the required size.
             return _bufferWriter ??= new ArrayBufferWriter<byte>();
         }
 
@@ -131,5 +135,28 @@ public override void Complete()
                     break;
             }
         }
+
+        public Memory<byte> GetHeader(bool isCompressed, int length)
+        {
+            // TODO(JamesNK): We can optimize header allocation when IBufferWriter is being used.
+            // IBufferWriter can be used to provide a buffer, either before or after message content.
+            // https://github.com/grpc/grpc-dotnet/issues/784
+            var buffer = new byte[GrpcProtocolConstants.HeaderSize];
+
+            // Compression flag
+            buffer[0] = isCompressed ? (byte)1 : (byte)0;
+
+            // Message length
+            EncodeMessageLength(length, buffer.AsSpan(1, 4));
+
+            return buffer;
+        }
+
+        private static void EncodeMessageLength(int messageLength, Span<byte> destination)
+        {
+            Debug.Assert(destination.Length >= GrpcProtocolConstants.MessageDelimiterSize, "Buffer too small to encode message length.");
+
+            BinaryPrimitives.WriteUInt32BigEndian(destination, (uint)messageLength);
+        }
     }
 }
diff --git a/src/Grpc.Net.Client/Internal/GrpcProtocolConstants.cs b/src/Grpc.Net.Client/Internal/GrpcProtocolConstants.cs
@@ -49,6 +49,9 @@ internal static class GrpcProtocolConstants
             // deflate is not supported. .NET's DeflateStream does not support RFC1950 - https://github.com/dotnet/corefx/issues/7570
         };
 
+        internal const int MessageDelimiterSize = 4; // how many bytes it takes to encode "Message-Length"
+        internal const int HeaderSize = MessageDelimiterSize + 1; // message length + compression flag
+
         internal static readonly string DefaultMessageAcceptEncodingValue;
 
         internal static readonly string UserAgentHeader;
diff --git a/src/Grpc.Net.Client/Internal/StreamExtensions.cs b/src/Grpc.Net.Client/Internal/StreamExtensions.cs
@@ -37,9 +37,6 @@ namespace Grpc.Net.Client
 {
     internal static partial class StreamExtensions
     {
-        private const int MessageDelimiterSize = 4; // how many bytes it takes to encode "Message-Length"
-        private const int HeaderSize = MessageDelimiterSize + 1; // message length + compression flag
-
         private static readonly Status SendingMessageExceedsLimitStatus = new Status(StatusCode.ResourceExhausted, "Sending message exceeds the maximum configured message size.");
         private static readonly Status ReceivedMessageExceedsLimitStatus = new Status(StatusCode.ResourceExhausted, "Received message exceeds the maximum configured message size.");
         private static readonly Status NoMessageEncodingMessageStatus = new Status(StatusCode.Internal, "Request did not include grpc-encoding value with compressed message.");
@@ -49,21 +46,21 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
             return new Status(StatusCode.Unimplemented, $"Unsupported grpc-encoding value '{unsupportedEncoding}'. Supported encodings: {string.Join(", ", supportedEncodings)}");
         }
 
-        private static async Task<(uint length, bool compressed)?> ReadHeaderAsync(Stream responseStream, Memory<byte> header, CancellationToken cancellationToken)
+        private static async Task<(int length, bool compressed)?> ReadHeaderAsync(Stream responseStream, Memory<byte> header, CancellationToken cancellationToken)
         {
             int read;
             var received = 0;
-            while ((read = await responseStream.ReadAsync(header.Slice(received, header.Length - received), cancellationToken).ConfigureAwait(false)) > 0)
+            while ((read = await responseStream.ReadAsync(header.Slice(received, GrpcProtocolConstants.HeaderSize - received), cancellationToken).ConfigureAwait(false)) > 0)
             {
                 received += read;
 
-                if (received == header.Length)
+                if (received == GrpcProtocolConstants.HeaderSize)
                 {
                     break;
                 }
             }
 
-            if (received < header.Length)
+            if (received < GrpcProtocolConstants.HeaderSize)
             {
                 if (received == 0)
                 {
@@ -73,10 +70,18 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
                 throw new InvalidDataException("Unexpected end of content while reading the message header.");
             }
 
+            // Read the header first
+            // - 1 byte flag for compression
+            // - 4 bytes for the content length
             var compressed = ReadCompressedFlag(header.Span[0]);
-            var length = BinaryPrimitives.ReadUInt32BigEndian(header.Span.Slice(1));
+            var length = BinaryPrimitives.ReadUInt32BigEndian(header.Span.Slice(1, 4));
 
-            return (length, compressed);
+            if (length > int.MaxValue)
+            {
+                throw new InvalidDataException("Message too large.");
+            }
+
+            return ((int)length, compressed);
         }
 
         public static async ValueTask<TResponse?> ReadMessageAsync<TResponse>(
@@ -90,17 +95,19 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
             CancellationToken cancellationToken)
             where TResponse : class
         {
+            byte[]? buffer = null;
+
             try
             {
                 GrpcCallLog.ReadingMessage(logger);
                 cancellationToken.ThrowIfCancellationRequested();
 
-                // Read the header first
-                // - 1 byte flag for compression
-                // - 4 bytes for the content length
-                var header = new byte[HeaderSize];
+                // Buffer is used to read header, then message content.
+                // This size was randomly chosen to hopefully be big enough for many small messages.
+                // If the message is larger then the array will be replaced when the message size is known.
+                buffer = ArrayPool<byte>.Shared.Rent(minimumLength: 4096);
 
-                var headerDetails = await ReadHeaderAsync(responseStream, header, cancellationToken).ConfigureAwait(false);
+                var headerDetails = await ReadHeaderAsync(responseStream, buffer, cancellationToken).ConfigureAwait(false);
 
                 if (headerDetails == null)
                 {
@@ -111,17 +118,22 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
                 var length = headerDetails.Value.length;
                 var compressed = headerDetails.Value.compressed;
 
-                if (length > int.MaxValue)
+                if (length > 0)
                 {
-                    throw new InvalidDataException("Message too large.");
-                }
+                    if (length > maximumMessageSize)
+                    {
+                        throw new RpcException(ReceivedMessageExceedsLimitStatus);
+                    }
 
-                if (length > maximumMessageSize)
-                {
-                    throw new RpcException(ReceivedMessageExceedsLimitStatus);
-                }
+                    // Replace buffer if the message doesn't fit
+                    if (buffer.Length < length)
+                    {
+                        ArrayPool<byte>.Shared.Return(buffer);
+                        buffer = ArrayPool<byte>.Shared.Rent(length);
+                    }
 
-                var messageData = await ReadMessageContent(responseStream, length, cancellationToken).ConfigureAwait(false);
+                    await ReadMessageContent(responseStream, buffer, length, cancellationToken).ConfigureAwait(false);
+                }
 
                 cancellationToken.ThrowIfCancellationRequested();
 
@@ -138,7 +150,7 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
                     }
 
                     // Performance improvement would be to decompress without converting to an intermediary byte array
-                    if (!TryDecompressMessage(logger, grpcEncoding, compressionProviders, messageData, out var decompressedMessage))
+                    if (!TryDecompressMessage(logger, grpcEncoding, compressionProviders, buffer, length, out var decompressedMessage))
                     {
                         var supportedEncodings = new List<string>();
                         supportedEncodings.Add(GrpcProtocolConstants.IdentityGrpcEncoding);
@@ -150,10 +162,10 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
                 }
                 else
                 {
-                    payload = new ReadOnlySequence<byte>(messageData);
+                    payload = new ReadOnlySequence<byte>(buffer, 0, length);
                 }
 
-                GrpcCallLog.DeserializingMessage(logger, messageData.Length, typeof(TResponse));
+                GrpcCallLog.DeserializingMessage(logger, length, typeof(TResponse));
 
                 var deserializationContext = new DefaultDeserializationContext();
                 deserializationContext.SetPayload(payload);
@@ -164,7 +176,7 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
                 {
                     // Check that there is no additional content in the stream for a single message
                     // There is no ReadByteAsync on stream. Reuse header array with ReadAsync, we don't need it anymore
-                    if (await responseStream.ReadAsync(header).ConfigureAwait(false) > 0)
+                    if (await responseStream.ReadAsync(buffer).ConfigureAwait(false) > 0)
                     {
                         throw new InvalidDataException("Unexpected data after finished reading message.");
                     }
@@ -179,43 +191,44 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
                 GrpcCallLog.ErrorReadingMessage(logger, ex);
                 throw;
             }
+            finally
+            {
+                if (buffer != null)
+                {
+                    ArrayPool<byte>.Shared.Return(buffer);
+                }
+            }
         }
 
-        private static async Task<byte[]> ReadMessageContent(Stream responseStream, uint length, CancellationToken cancellationToken)
+        private static async Task ReadMessageContent(Stream responseStream, Memory<byte> messageData, int length, CancellationToken cancellationToken)
         {
             // Read message content until content length is reached
-            byte[] messageData;
-            if (length > 0)
+            var received = 0;
+            int read;
+            while ((read = await responseStream.ReadAsync(messageData.Slice(received, length - received), cancellationToken).ConfigureAwait(false)) > 0)
             {
-                var received = 0;
-                var read = 0;
-                messageData = new byte[length];
-                while ((read = await responseStream.ReadAsync(messageData.AsMemory(received, messageData.Length - received), cancellationToken).ConfigureAwait(false)) > 0)
-                {
-                    received += read;
+                received += read;
 
-                    if (received == messageData.Length)
-                    {
-                        break;
-                    }
+                if (received == length)
+                {
+                    break;
                 }
             }
-            else
+
+            if (received < length)
             {
-                messageData = Array.Empty<byte>();
+                throw new InvalidDataException("Unexpected end of content while reading the message content.");
             }
-
-            return messageData;
         }
 
-        private static bool TryDecompressMessage(ILogger logger, string compressionEncoding, Dictionary<string, ICompressionProvider> compressionProviders, byte[] messageData, [NotNullWhen(true)]out ReadOnlySequence<byte>? result)
+        private static bool TryDecompressMessage(ILogger logger, string compressionEncoding, Dictionary<string, ICompressionProvider> compressionProviders, byte[] messageData, int length, [NotNullWhen(true)]out ReadOnlySequence<byte>? result)
         {
             if (compressionProviders.TryGetValue(compressionEncoding, out var compressionProvider))
             {
                 GrpcCallLog.DecompressingMessage(logger, compressionProvider.EncodingName);
 
                 var output = new MemoryStream();
-                using (var compressionStream = compressionProvider.CreateDecompressionStream(new MemoryStream(messageData)))
+                using (var compressionStream = compressionProvider.CreateDecompressionStream(new MemoryStream(messageData, 0, length, writable: true, publiclyVisible: true)))
                 {
                     compressionStream.CopyTo(output);
                 }
@@ -244,6 +257,7 @@ private static bool ReadCompressedFlag(byte flag)
             }
         }
 
+        // TODO(JamesNK): Reuse serialization content between message writes. Improve client/duplex streaming allocations.
         public static async ValueTask WriteMessageAsync<TMessage>(
             this Stream stream,
             ILogger logger,
@@ -275,8 +289,8 @@ public static async ValueTask WriteMessageAsync<TMessage>(
                 }
 
                 var isCompressed =
-                    GrpcProtocolHelpers.CanWriteCompressed(callOptions.WriteOptions) &&
-                    !string.Equals(grpcEncoding, GrpcProtocolConstants.IdentityGrpcEncoding, StringComparison.Ordinal);
+                   GrpcProtocolHelpers.CanWriteCompressed(callOptions.WriteOptions) &&
+                   !string.Equals(grpcEncoding, GrpcProtocolConstants.IdentityGrpcEncoding, StringComparison.Ordinal);
 
                 if (isCompressed)
                 {
@@ -288,7 +302,7 @@ public static async ValueTask WriteMessageAsync<TMessage>(
                         data);
                 }
 
-                await WriteHeaderAsync(stream, data.Length, isCompressed, callOptions.CancellationToken).ConfigureAwait(false);
+                await stream.WriteAsync(serializationContext.GetHeader(isCompressed, data.Length), callOptions.CancellationToken).ConfigureAwait(false);
                 await stream.WriteAsync(data, callOptions.CancellationToken).ConfigureAwait(false);
                 await stream.FlushAsync(callOptions.CancellationToken).ConfigureAwait(false);
 
@@ -322,25 +336,5 @@ private static ReadOnlyMemory<byte> CompressMessage(ILogger logger, string compr
             // Should never reach here
             throw new InvalidOperationException($"Could not find compression provider for '{compressionEncoding}'.");
         }
-
-        private static ValueTask WriteHeaderAsync(Stream stream, int length, bool compress, CancellationToken cancellationToken)
-        {
-            var headerData = new byte[HeaderSize];
-
-            // Compression flag
-            headerData[0] = compress ? (byte)1 : (byte)0;
-
-            // Message length
-            EncodeMessageLength(length, headerData.AsSpan(1));
-
-            return stream.WriteAsync(headerData.AsMemory(0, headerData.Length), cancellationToken);
-        }
-
-        private static void EncodeMessageLength(int messageLength, Span<byte> destination)
-        {
-            Debug.Assert(destination.Length >= MessageDelimiterSize, "Buffer too small to encode message length.");
-
-            BinaryPrimitives.WriteUInt32BigEndian(destination, (uint)messageLength);
-        }
     }
 }
diff --git a/src/Shared/DefaultDeserializationContext.cs b/src/Shared/DefaultDeserializationContext.cs
@@ -34,6 +34,9 @@ public void SetPayload(in ReadOnlySequence<byte>? payload)
         public override byte[] PayloadAsNewBuffer()
         {
             Debug.Assert(_payload != null, "Payload must be set.");
+
+            // The array returned by PayloadAsNewBuffer must be the exact message size.
+            // There is no opportunity here to return a pooled array.
             return _payload.GetValueOrDefault().ToArray();
         }
 

Original file line number	Diff line number	Diff line change
`@@ -31,18 +31,18 @@ static void Main(string[] args)`
`31`	`31`	`// Profiling option. This will call methods explicitly, in-process`
`32`	`32`	`static async System.Threading.Tasks.Task Main(string[] args)`
`33`	`33`	`{`
`34`		`- var benchmark = new Client.CompressedUnaryServerCallHandlerBenchmark();`
	`34`	`+ var benchmark = new Client.UnaryClientBenchmark();`
`35`	`35`	`benchmark.GlobalSetup();`
`36`	`36`	`for (var i = 0; i < 10000; i++)`
`37`	`37`	`{`
`38`		`- await benchmark.CompressedSayHelloAsync();`
	`38`	`+ await benchmark.SayHelloAsync();`
`39`	`39`	`}`
`40`	`40`
`41`	`41`	`System.Console.WriteLine("Press any key to start.");`
`42`	`42`	`System.Console.ReadKey();`
`43`	`43`	`for (var i = 0; i < 1; i++)`
`44`	`44`	`{`
`45`		`- await benchmark.CompressedSayHelloAsync();`
	`45`	`+ await benchmark.SayHelloAsync();`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`System.Console.WriteLine("Done. Press any key to exit.");`
Original file line number	Diff line number	Diff line change
`@@ -37,9 +37,6 @@ namespace Grpc.Net.Client`
`37`	`37`	`{`
`38`	`38`	`internal static partial class StreamExtensions`
`39`	`39`	`{`
`40`		`- private const int MessageDelimiterSize = 4; // how many bytes it takes to encode "Message-Length"`
`41`		`- private const int HeaderSize = MessageDelimiterSize + 1; // message length + compression flag`
`42`		`-`
`43`	`40`	`private static readonly Status SendingMessageExceedsLimitStatus = new Status(StatusCode.ResourceExhausted, "Sending message exceeds the maximum configured message size.");`
`44`	`41`	`private static readonly Status ReceivedMessageExceedsLimitStatus = new Status(StatusCode.ResourceExhausted, "Received message exceeds the maximum configured message size.");`
`45`	`42`	`private static readonly Status NoMessageEncodingMessageStatus = new Status(StatusCode.Internal, "Request did not include grpc-encoding value with compressed message.");`
`@@ -49,21 +46,21 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`49`	`46`	`return new Status(StatusCode.Unimplemented, $"Unsupported grpc-encoding value '{unsupportedEncoding}'. Supported encodings: {string.Join(", ", supportedEncodings)}");`
`50`	`47`	`}`
`51`	`48`
`52`		`- private static async Task<(uint length, bool compressed)?> ReadHeaderAsync(Stream responseStream, Memory<byte> header, CancellationToken cancellationToken)`
	`49`	`+ private static async Task<(int length, bool compressed)?> ReadHeaderAsync(Stream responseStream, Memory<byte> header, CancellationToken cancellationToken)`
`53`	`50`	`{`
`54`	`51`	`int read;`
`55`	`52`	`var received = 0;`
`56`		`- while ((read = await responseStream.ReadAsync(header.Slice(received, header.Length - received), cancellationToken).ConfigureAwait(false)) > 0)`
	`53`	`+ while ((read = await responseStream.ReadAsync(header.Slice(received, GrpcProtocolConstants.HeaderSize - received), cancellationToken).ConfigureAwait(false)) > 0)`
`57`	`54`	`{`
`58`	`55`	`received += read;`
`59`	`56`
`60`		`- if (received == header.Length)`
	`57`	`+ if (received == GrpcProtocolConstants.HeaderSize)`
`61`	`58`	`{`
`62`	`59`	`break;`
`63`	`60`	`}`
`64`	`61`	`}`
`65`	`62`
`66`		`- if (received < header.Length)`
	`63`	`+ if (received < GrpcProtocolConstants.HeaderSize)`
`67`	`64`	`{`
`68`	`65`	`if (received == 0)`
`69`	`66`	`{`
`@@ -73,10 +70,18 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`73`	`70`	`throw new InvalidDataException("Unexpected end of content while reading the message header.");`
`74`	`71`	`}`
`75`	`72`
	`73`	`+ // Read the header first`
	`74`	`+ // - 1 byte flag for compression`
	`75`	`+ // - 4 bytes for the content length`
`76`	`76`	`var compressed = ReadCompressedFlag(header.Span[0]);`
`77`		`- var length = BinaryPrimitives.ReadUInt32BigEndian(header.Span.Slice(1));`
	`77`	`+ var length = BinaryPrimitives.ReadUInt32BigEndian(header.Span.Slice(1, 4));`
`78`	`78`
`79`		`- return (length, compressed);`
	`79`	`+ if (length > int.MaxValue)`
	`80`	`+ {`
	`81`	`+ throw new InvalidDataException("Message too large.");`
	`82`	`+ }`
	`83`	`+`
	`84`	`+ return ((int)length, compressed);`
`80`	`85`	`}`
`81`	`86`
`82`	`87`	`public static async ValueTask<TResponse?> ReadMessageAsync<TResponse>(`
`@@ -90,17 +95,19 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`90`	`95`	`CancellationToken cancellationToken)`
`91`	`96`	`where TResponse : class`
`92`	`97`	`{`
	`98`	`+ byte[]? buffer = null;`
	`99`	`+`
`93`	`100`	`try`
`94`	`101`	`{`
`95`	`102`	`GrpcCallLog.ReadingMessage(logger);`
`96`	`103`	`cancellationToken.ThrowIfCancellationRequested();`
`97`	`104`
`98`		`- // Read the header first`
`99`		`- // - 1 byte flag for compression`
`100`		`- // - 4 bytes for the content length`
`101`		`- var header = new byte[HeaderSize];`
	`105`	`+ // Buffer is used to read header, then message content.`
	`106`	`+ // This size was randomly chosen to hopefully be big enough for many small messages.`
	`107`	`+ // If the message is larger then the array will be replaced when the message size is known.`
	`108`	`+ buffer = ArrayPool<byte>.Shared.Rent(minimumLength: 4096);`
`102`	`109`
`103`		`- var headerDetails = await ReadHeaderAsync(responseStream, header, cancellationToken).ConfigureAwait(false);`
	`110`	`+ var headerDetails = await ReadHeaderAsync(responseStream, buffer, cancellationToken).ConfigureAwait(false);`
`104`	`111`
`105`	`112`	`if (headerDetails == null)`
`106`	`113`	`{`
`@@ -111,17 +118,22 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`111`	`118`	`var length = headerDetails.Value.length;`
`112`	`119`	`var compressed = headerDetails.Value.compressed;`
`113`	`120`
`114`		`- if (length > int.MaxValue)`
	`121`	`+ if (length > 0)`
`115`	`122`	`{`
`116`		`- throw new InvalidDataException("Message too large.");`
`117`		`- }`
	`123`	`+ if (length > maximumMessageSize)`
	`124`	`+ {`
	`125`	`+ throw new RpcException(ReceivedMessageExceedsLimitStatus);`
	`126`	`+ }`
`118`	`127`
`119`		`- if (length > maximumMessageSize)`
`120`		`- {`
`121`		`- throw new RpcException(ReceivedMessageExceedsLimitStatus);`
`122`		`- }`
	`128`	`+ // Replace buffer if the message doesn't fit`
	`129`	`+ if (buffer.Length < length)`
	`130`	`+ {`
	`131`	`+ ArrayPool<byte>.Shared.Return(buffer);`
	`132`	`+ buffer = ArrayPool<byte>.Shared.Rent(length);`
	`133`	`+ }`
`123`	`134`
`124`		`- var messageData = await ReadMessageContent(responseStream, length, cancellationToken).ConfigureAwait(false);`
	`135`	`+ await ReadMessageContent(responseStream, buffer, length, cancellationToken).ConfigureAwait(false);`
	`136`	`+ }`
`125`	`137`
`126`	`138`	`cancellationToken.ThrowIfCancellationRequested();`
`127`	`139`
`@@ -138,7 +150,7 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`138`	`150`	`}`
`139`	`151`
`140`	`152`	`// Performance improvement would be to decompress without converting to an intermediary byte array`
`141`		`- if (!TryDecompressMessage(logger, grpcEncoding, compressionProviders, messageData, out var decompressedMessage))`
	`153`	`+ if (!TryDecompressMessage(logger, grpcEncoding, compressionProviders, buffer, length, out var decompressedMessage))`
`142`	`154`	`{`
`143`	`155`	`var supportedEncodings = new List<string>();`
`144`	`156`	`supportedEncodings.Add(GrpcProtocolConstants.IdentityGrpcEncoding);`
`@@ -150,10 +162,10 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`150`	`162`	`}`
`151`	`163`	`else`
`152`	`164`	`{`
`153`		`- payload = new ReadOnlySequence<byte>(messageData);`
	`165`	`+ payload = new ReadOnlySequence<byte>(buffer, 0, length);`
`154`	`166`	`}`
`155`	`167`
`156`		`- GrpcCallLog.DeserializingMessage(logger, messageData.Length, typeof(TResponse));`
	`168`	`+ GrpcCallLog.DeserializingMessage(logger, length, typeof(TResponse));`
`157`	`169`
`158`	`170`	`var deserializationContext = new DefaultDeserializationContext();`
`159`	`171`	`deserializationContext.SetPayload(payload);`
`@@ -164,7 +176,7 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`164`	`176`	`{`
`165`	`177`	`// Check that there is no additional content in the stream for a single message`
`166`	`178`	`// There is no ReadByteAsync on stream. Reuse header array with ReadAsync, we don't need it anymore`
`167`		`- if (await responseStream.ReadAsync(header).ConfigureAwait(false) > 0)`
	`179`	`+ if (await responseStream.ReadAsync(buffer).ConfigureAwait(false) > 0)`
`168`	`180`	`{`
`169`	`181`	`throw new InvalidDataException("Unexpected data after finished reading message.");`
`170`	`182`	`}`
`@@ -179,43 +191,44 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport`
`179`	`191`	`GrpcCallLog.ErrorReadingMessage(logger, ex);`
`180`	`192`	`throw;`
`181`	`193`	`}`
	`194`	`+ finally`
	`195`	`+ {`
	`196`	`+ if (buffer != null)`
	`197`	`+ {`
	`198`	`+ ArrayPool<byte>.Shared.Return(buffer);`
	`199`	`+ }`
	`200`	`+ }`
`182`	`201`	`}`
`183`	`202`
`184`		`- private static async Task<byte[]> ReadMessageContent(Stream responseStream, uint length, CancellationToken cancellationToken)`
	`203`	`+ private static async Task ReadMessageContent(Stream responseStream, Memory<byte> messageData, int length, CancellationToken cancellationToken)`
`185`	`204`	`{`
`186`	`205`	`// Read message content until content length is reached`
`187`		`- byte[] messageData;`
`188`		`- if (length > 0)`
	`206`	`+ var received = 0;`
	`207`	`+ int read;`
	`208`	`+ while ((read = await responseStream.ReadAsync(messageData.Slice(received, length - received), cancellationToken).ConfigureAwait(false)) > 0)`
`189`	`209`	`{`
`190`		`- var received = 0;`
`191`		`- var read = 0;`
`192`		`- messageData = new byte[length];`
`193`		`- while ((read = await responseStream.ReadAsync(messageData.AsMemory(received, messageData.Length - received), cancellationToken).ConfigureAwait(false)) > 0)`
`194`		`- {`
`195`		`- received += read;`
	`210`	`+ received += read;`
`196`	`211`
`197`		`- if (received == messageData.Length)`
`198`		`- {`
`199`		`- break;`
`200`		`- }`
	`212`	`+ if (received == length)`
	`213`	`+ {`
	`214`	`+ break;`
`201`	`215`	`}`
`202`	`216`	`}`
`203`		`- else`
	`217`	`+`
	`218`	`+ if (received < length)`
`204`	`219`	`{`
`205`		`- messageData = Array.Empty<byte>();`
	`220`	`+ throw new InvalidDataException("Unexpected end of content while reading the message content.");`
`206`	`221`	`}`
`207`		`-`
`208`		`- return messageData;`
`209`	`222`	`}`
`210`	`223`
`211`		`- private static bool TryDecompressMessage(ILogger logger, string compressionEncoding, Dictionary<string, ICompressionProvider> compressionProviders, byte[] messageData, [NotNullWhen(true)]out ReadOnlySequence<byte>? result)`
	`224`	`+ private static bool TryDecompressMessage(ILogger logger, string compressionEncoding, Dictionary<string, ICompressionProvider> compressionProviders, byte[] messageData, int length, [NotNullWhen(true)]out ReadOnlySequence<byte>? result)`
`212`	`225`	`{`
`213`	`226`	`if (compressionProviders.TryGetValue(compressionEncoding, out var compressionProvider))`
`214`	`227`	`{`
`215`	`228`	`GrpcCallLog.DecompressingMessage(logger, compressionProvider.EncodingName);`
`216`	`229`
`217`	`230`	`var output = new MemoryStream();`
`218`		`- using (var compressionStream = compressionProvider.CreateDecompressionStream(new MemoryStream(messageData)))`
	`231`	`+ using (var compressionStream = compressionProvider.CreateDecompressionStream(new MemoryStream(messageData, 0, length, writable: true, publiclyVisible: true)))`
`219`	`232`	`{`
`220`	`233`	`compressionStream.CopyTo(output);`
`221`	`234`	`}`
`@@ -244,6 +257,7 @@ private static bool ReadCompressedFlag(byte flag)`
`244`	`257`	`}`
`245`	`258`	`}`
`246`	`259`
	`260`	`+ // TODO(JamesNK): Reuse serialization content between message writes. Improve client/duplex streaming allocations.`
`247`	`261`	`public static async ValueTask WriteMessageAsync<TMessage>(`
`248`	`262`	`this Stream stream,`
`249`	`263`	`ILogger logger,`
`@@ -275,8 +289,8 @@ public static async ValueTask WriteMessageAsync<TMessage>(`
`275`	`289`	`}`
`276`	`290`
`277`	`291`	`var isCompressed =`
`278`		`- GrpcProtocolHelpers.CanWriteCompressed(callOptions.WriteOptions) &&`
`279`		`- !string.Equals(grpcEncoding, GrpcProtocolConstants.IdentityGrpcEncoding, StringComparison.Ordinal);`
	`292`	`+ GrpcProtocolHelpers.CanWriteCompressed(callOptions.WriteOptions) &&`
	`293`	`+ !string.Equals(grpcEncoding, GrpcProtocolConstants.IdentityGrpcEncoding, StringComparison.Ordinal);`
`280`	`294`
`281`	`295`	`if (isCompressed)`
`282`	`296`	`{`
`@@ -288,7 +302,7 @@ public static async ValueTask WriteMessageAsync<TMessage>(`
`288`	`302`	`data);`
`289`	`303`	`}`
`290`	`304`
`291`		`- await WriteHeaderAsync(stream, data.Length, isCompressed, callOptions.CancellationToken).ConfigureAwait(false);`
	`305`	`+ await stream.WriteAsync(serializationContext.GetHeader(isCompressed, data.Length), callOptions.CancellationToken).ConfigureAwait(false);`
`292`	`306`	`await stream.WriteAsync(data, callOptions.CancellationToken).ConfigureAwait(false);`
`293`	`307`	`await stream.FlushAsync(callOptions.CancellationToken).ConfigureAwait(false);`
`294`	`308`
`@@ -322,25 +336,5 @@ private static ReadOnlyMemory<byte> CompressMessage(ILogger logger, string compr`
`322`	`336`	`// Should never reach here`
`323`	`337`	`throw new InvalidOperationException($"Could not find compression provider for '{compressionEncoding}'.");`
`324`	`338`	`}`
`325`		`-`
`326`		`- private static ValueTask WriteHeaderAsync(Stream stream, int length, bool compress, CancellationToken cancellationToken)`
`327`		`- {`
`328`		`- var headerData = new byte[HeaderSize];`
`329`		`-`
`330`		`- // Compression flag`
`331`		`- headerData[0] = compress ? (byte)1 : (byte)0;`
`332`		`-`
`333`		`- // Message length`
`334`		`- EncodeMessageLength(length, headerData.AsSpan(1));`
`335`		`-`
`336`		`- return stream.WriteAsync(headerData.AsMemory(0, headerData.Length), cancellationToken);`
`337`		`- }`
`338`		`-`
`339`		`- private static void EncodeMessageLength(int messageLength, Span<byte> destination)`
`340`		`- {`
`341`		`- Debug.Assert(destination.Length >= MessageDelimiterSize, "Buffer too small to encode message length.");`
`342`		`-`
`343`		`- BinaryPrimitives.WriteUInt32BigEndian(destination, (uint)messageLength);`
`344`		`- }`
`345`	`339`	`}`
`346`	`340`	`}`
Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,9 @@ public void SetPayload(in ReadOnlySequence<byte>? payload)`
`34`	`34`	`public override byte[] PayloadAsNewBuffer()`
`35`	`35`	`{`
`36`	`36`	`Debug.Assert(_payload != null, "Payload must be set.");`
	`37`	`+`
	`38`	`+ // The array returned by PayloadAsNewBuffer must be the exact message size.`
	`39`	`+ // There is no opportunity here to return a pooled array.`
`37`	`40`	`return _payload.GetValueOrDefault().ToArray();`
`38`	`41`	`}`
`39`	`42`