Add GetMessageBytes to IHubProtocol (#1915)

Author: benaadams

benchmarks/Microsoft.AspNetCore.SignalR.Microbenchmarks/DefaultHubDispatcherBenchmark.cs

@@ -69,6 +69,11 @@ public bool TryParseMessage(ref ReadOnlySequence<byte> input, IInvocationBinder
             public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
             {
             }
+
+            public byte[] GetMessageBytes(HubMessage message)
+            {
+                return HubProtocolExtensions.GetMessageBytes(this, message);
+            }
         }
 
         public class NoErrorHubConnectionContext : HubConnectionContext

benchmarks/Microsoft.AspNetCore.SignalR.Microbenchmarks/HubProtocolBenchmark.cs

@@ -50,7 +50,7 @@ public void GlobalSetup()
                     break;
             }
 
-            _binaryInput = _hubProtocol.WriteToArray(_hubMessage);
+            _binaryInput = _hubProtocol.GetMessageBytes(_hubMessage);
             _binder = new TestBinder(_hubMessage);
         }
 
@@ -67,7 +67,7 @@ public void ReadSingleMessage()
         [Benchmark]
         public void WriteSingleMessage()
         {
-            var bytes = _hubProtocol.WriteToArray(_hubMessage);
+            var bytes = _hubProtocol.GetMessageBytes(_hubMessage);
             if (bytes.Length != _binaryInput.Length)
             {
                 throw new InvalidOperationException("Failed to write message");

benchmarks/Microsoft.AspNetCore.SignalR.Microbenchmarks/RedisHubLifetimeManagerBenchmark.cs

@@ -194,6 +194,11 @@ public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
                 _innerProtocol.WriteMessage(message, output);
             }
 
+            public byte[] GetMessageBytes(HubMessage message)
+            {
+                return HubProtocolExtensions.GetMessageBytes(this, message);
+            }
+
             public bool IsVersionSupported(int version)
             {
                 return _innerProtocol.IsVersionSupported(version);

benchmarks/Microsoft.AspNetCore.SignalR.Microbenchmarks/RedisProtocolBenchmark.cs

@@ -143,6 +143,11 @@ public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
             {
                 output.Write(_fixedOutput);
             }
+
+            public byte[] GetMessageBytes(HubMessage message)
+            {
+                return HubProtocolExtensions.GetMessageBytes(this, message);
+            }
         }
     }
 }

benchmarks/Microsoft.AspNetCore.SignalR.Microbenchmarks/ServerSentEventsBenchmark.cs

@@ -39,7 +39,7 @@ public void GlobalSetup()
             }
 
             _parser = new ServerSentEventsMessageParser();
-            _rawData = hubProtocol.WriteToArray(hubMessage);
+            _rawData = hubProtocol.GetMessageBytes(hubMessage);
             var ms = new MemoryStream();
             ServerSentEventsMessageFormatter.WriteMessage(_rawData, ms);
             _sseFormattedData = ms.ToArray();

src/Common/MemoryBufferWriter.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Buffers;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Runtime.CompilerServices;
 using System.Threading;
@@ -217,6 +218,35 @@ public byte[] ToArray()
             return result;
         }
 
+        public void CopyTo(Span<byte> span)
+        {
+            Debug.Assert(span.Length >= _bytesWritten);
+
+            if (_currentSegment == null)
+            {
+                return;
+            }
+
+            var totalWritten = 0;
+
+            if (_fullSegments != null)
+            {
+                // Copy full segments
+                var count = _fullSegments.Count;
+                for (var i = 0; i < count; i++)
+                {
+                    var segment = _fullSegments[i];
+                    segment.AsSpan().CopyTo(span.Slice(totalWritten));
+                    totalWritten += segment.Length;
+                }
+            }
+
+            // Copy current incomplete segment
+            _currentSegment.AsSpan(0, _position).CopyTo(span.Slice(totalWritten));
+
+            Debug.Assert(_bytesWritten == totalWritten + _position);
+        }
+
         public override void Flush() { }
         public override Task FlushAsync(CancellationToken cancellationToken) => Task.CompletedTask;
         public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();

src/Microsoft.AspNetCore.SignalR.Common/Internal/Formatters/BinaryMessageFormatter.cs

@@ -10,15 +10,21 @@ public static class BinaryMessageFormatter
     {
         public static void WriteLengthPrefix(long length, IBufferWriter<byte> output)
         {
-            // This code writes length prefix of the message as a VarInt. Read the comment in
-            // the BinaryMessageParser.TryParseMessage for details.
-
             Span<byte> lenBuffer = stackalloc byte[5];
 
+            var lenNumBytes = WriteLengthPrefix(length, lenBuffer);
+
+            output.Write(lenBuffer.Slice(0, lenNumBytes));
+        }
+
+        public static int WriteLengthPrefix(long length, Span<byte> output)
+        {
+            // This code writes length prefix of the message as a VarInt. Read the comment in
+            // the BinaryMessageParser.TryParseMessage for details.
             var lenNumBytes = 0;
             do
             {
-                ref var current = ref lenBuffer[lenNumBytes];
+                ref var current = ref output[lenNumBytes];
                 current = (byte)(length & 0x7f);
                 length >>= 7;
                 if (length > 0)
@@ -29,7 +35,20 @@ public static void WriteLengthPrefix(long length, IBufferWriter<byte> output)
             }
             while (length > 0);
 
-            output.Write(lenBuffer.Slice(0, lenNumBytes));
+            return lenNumBytes;
+        }
+
+        public static int LengthPrefixLength(long length)
+        {
+            var lenNumBytes = 0;
+            do
+            {
+                length >>= 7;
+                lenNumBytes++;
+            }
+            while (length > 0);
+
+            return lenNumBytes;
         }
     }
 }

src/Microsoft.AspNetCore.SignalR.Common/Internal/Protocol/HubProtocolExtensions.cs

@@ -7,7 +7,8 @@ namespace Microsoft.AspNetCore.SignalR.Internal.Protocol
 {
     public static class HubProtocolExtensions
     {
-        public static byte[] WriteToArray(this IHubProtocol hubProtocol, HubMessage message)
+        // Would work as default interface impl
+        public static byte[] GetMessageBytes(this IHubProtocol hubProtocol, HubMessage message)
         {
             var writer = MemoryBufferWriter.Get();
             try

src/Microsoft.AspNetCore.SignalR.Common/Internal/Protocol/IHubProtocol.cs

@@ -2,7 +2,6 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System.Buffers;
-using System.IO;
 using Microsoft.AspNetCore.Connections;
 
 namespace Microsoft.AspNetCore.SignalR.Internal.Protocol
@@ -19,6 +18,8 @@ public interface IHubProtocol
 
         void WriteMessage(HubMessage message, IBufferWriter<byte> output);
 
+        byte[] GetMessageBytes(HubMessage message);
+
         bool IsVersionSupported(int version);
     }
 }

src/Microsoft.AspNetCore.SignalR.Core/HubConnectionContext.cs

@@ -340,7 +340,7 @@ await WriteHandshakeResponseAsync(new HandshakeResponseMessage(
                                         transferFormatFeature.ActiveFormat = Protocol.TransferFormat;
                                     }
 
-                                    _cachedPingMessage = Protocol.WriteToArray(PingMessage.Instance);
+                                    _cachedPingMessage = Protocol.GetMessageBytes(PingMessage.Instance);
 
                                     UserIdentifier = userIdProvider.GetUserId(this);
 

src/Microsoft.AspNetCore.SignalR.Core/Internal/SerializedHubMessage.cs

@@ -39,7 +39,7 @@ public ReadOnlyMemory<byte> GetSerializedMessage(IHubProtocol protocol)
                         "This message was received from another server that did not have the requested protocol available.");
                 }
 
-                serialized = protocol.WriteToArray(Message);
+                serialized = protocol.GetMessageBytes(Message);
                 SetCache(protocol.Name, serialized);
             }
 
@@ -65,7 +65,7 @@ public static void WriteAllSerializedVersions(BinaryWriter writer, HubMessage me
             {
                 writer.Write(protocol.Name);
 
-                var buffer = protocol.WriteToArray(message);
+                var buffer = protocol.GetMessageBytes(message);
                 writer.Write(buffer.Length);
                 writer.Write(buffer);
             }

src/Microsoft.AspNetCore.SignalR.Protocols.Json/Internal/Protocol/JsonHubProtocol.cs

@@ -82,6 +82,11 @@ public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
             TextMessageFormatter.WriteRecordSeparator(output);
         }
 
+        public byte[] GetMessageBytes(HubMessage message)
+        {
+            return HubProtocolExtensions.GetMessageBytes(this, message);
+        }
+
         private HubMessage ParseMessage(Utf8BufferTextReader textReader, IInvocationBinder binder)
         {
             try

src/Microsoft.AspNetCore.SignalR.Protocols.MsgPack/Internal/Protocol/MessagePackHubProtocol.cs

@@ -303,6 +303,34 @@ public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
             }
         }
 
+        public byte[] GetMessageBytes(HubMessage message)
+        {
+            var writer = MemoryBufferWriter.Get();
+
+            try
+            {
+                // Write message to a buffer so we can get its length
+                WriteMessageCore(message, writer);
+
+                var dataLength = writer.Length;
+                var prefixLength = BinaryMessageFormatter.LengthPrefixLength(writer.Length);
+
+                var array = new byte[dataLength + prefixLength];
+                var span = array.AsSpan();
+
+                // Write length then message to output
+                var written = BinaryMessageFormatter.WriteLengthPrefix(writer.Length, span);
+                Debug.Assert(written == prefixLength);
+                writer.CopyTo(span.Slice(prefixLength));
+
+                return array;
+            }
+            finally
+            {
+                MemoryBufferWriter.Return(writer);
+            }
+        }
+
         private void WriteMessageCore(HubMessage message, Stream packer)
         {
             switch (message)

test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.cs

@@ -167,6 +167,11 @@ public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
                     throw _error;
                 }
             }
+
+            public byte[] GetMessageBytes(HubMessage message)
+            {
+                return HubProtocolExtensions.GetMessageBytes(this, message);
+            }
         }
     }
 }