Pool byte[] in RemoteRenderer (#11976)

* Pool byte[] in RemoteRenderer

src/Components/Components/src/Microsoft.AspNetCore.Components.csproj

@@ -11,6 +11,7 @@
 
   <ItemGroup>
     <Compile Include="..\..\Shared\src\JsonSerializerOptionsProvider.cs" />
+    <Compile Include="$(ComponentsSharedSourceRoot)src\ArrayBuilder.cs" LinkBase="RenderTree" />
   </ItemGroup>
 
   <ItemGroup>

src/Components/Components/src/RenderTree/ArrayBuilderExtensions.cs

@@ -0,0 +1,27 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+
+namespace Microsoft.AspNetCore.Components.RenderTree
+{
+    internal static class ArrayBuilderExtensions
+    {
+        /// <summary>
+        /// Produces an <see cref="ArrayRange{T}"/> structure describing the current contents.
+        /// </summary>
+        /// <returns>The <see cref="ArrayRange{T}"/>.</returns>
+        public static ArrayRange<T> ToRange<T>(this ArrayBuilder<T> builder)
+            => new ArrayRange<T>(builder.Buffer, builder.Count);
+
+        /// <summary>
+        /// Produces an <see cref="ArrayBuilderSegment{T}"/> structure describing the selected contents.
+        /// </summary>
+        /// <param name="builder">The <see cref="ArrayBuilder{T}"/></param>
+        /// <param name="fromIndexInclusive">The index of the first item in the segment.</param>
+        /// <param name="toIndexExclusive">One plus the index of the last item in the segment.</param>
+        /// <returns>The <see cref="ArraySegment{T}"/>.</returns>
+        public static ArrayBuilderSegment<T> ToSegment<T>(this ArrayBuilder<T> builder, int fromIndexInclusive, int toIndexExclusive)
+            => new ArrayBuilderSegment<T>(builder, fromIndexInclusive, toIndexExclusive - fromIndexInclusive);
+    }
+}

src/Components/Server/src/BlazorPack/BlazorPackHubProtocol.cs

@@ -424,8 +424,8 @@ namespace Microsoft.AspNetCore.Components.Server.BlazorPack
                     writer.Write(floatValue);
                     break;
 
-                case byte[] byteArray:
+                case ArraySegment<byte> bytes:
-                    writer.Write(byteArray);
+                    writer.Write(bytes);
                     break;
 
                 case Exception exception:

src/Components/Server/src/Circuits/ArrayBuilderMemoryStream.cs

@@ -0,0 +1,125 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.IO;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Components.Web.Rendering
+{
+    /// <summary>
+    /// Writeable memory stream backed by a an <see cref="ArrayBuilder{T}"/>.
+    /// </summary>
+    internal sealed class ArrayBuilderMemoryStream : Stream
+    {
+        public ArrayBuilderMemoryStream(ArrayBuilder<byte> arrayBuilder)
+        {
+            ArrayBuilder = arrayBuilder;
+        }
+
+        /// <inheritdoc />
+        public override bool CanRead => false;
+
+        /// <inheritdoc />
+        public override bool CanSeek => false;
+
+        /// <inheritdoc />
+        public override bool CanWrite => true;
+
+        /// <inheritdoc />
+        public override long Length => ArrayBuilder.Count;
+
+        /// <inheritdoc />
+        public override long Position
+        {
+            get => ArrayBuilder.Count;
+            set => throw new NotSupportedException();
+        }
+
+        public ArrayBuilder<byte> ArrayBuilder { get; }
+
+        /// <inheritdoc />
+        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
+
+        /// <inheritdoc />
+        public override int Read(byte[] buffer, int offset, int count)
+            => throw new NotSupportedException();
+
+        /// <inheritdoc />
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+            => throw new NotSupportedException();
+
+        /// <inheritdoc />
+        public override void Write(byte[] buffer, int offset, int count)
+        {
+            ValidateArguments(buffer, offset, count);
+
+            ArrayBuilder.Append(buffer, offset, count);
+        }
+
+        /// <inheritdoc />
+        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            ValidateArguments(buffer, offset, count);
+
+            ArrayBuilder.Append(buffer, offset, count);
+            return Task.CompletedTask;
+        }
+
+        /// <inheritdoc />
+        public override void Flush()
+        {
+            // Do nothing.
+        }
+
+        /// <inheritdoc />
+        public override Task FlushAsync(CancellationToken cancellationToken) => Task.CompletedTask;
+
+        /// <inheritdoc />
+        public override void SetLength(long value) => throw new NotSupportedException();
+
+        /// <inheritdoc />
+        protected override void Dispose(bool disposing)
+        {
+            // Do nothing.
+        }
+
+        /// <inheritdoc />
+        public override ValueTask DisposeAsync() => default;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static void ValidateArguments(byte[] buffer, int offset, int count)
+        {
+            if (buffer == null)
+            {
+                ThrowHelper.ThrowArgumentNullException(nameof(buffer));
+            }
+
+            if (offset < 0)
+            {
+                ThrowHelper.ThrowArgumentOutOfRangeException(nameof(offset));
+            }
+
+            if (count < 0)
+            {
+                ThrowHelper.ThrowArgumentOutOfRangeException(nameof(count));
+            }
+
+            if (buffer.Length - offset < count)
+            {
+                ThrowHelper.ThrowArgumentOutOfRangeException(nameof(offset));
+            }
+        }
+
+        private static class ThrowHelper
+        {
+            public static void ThrowArgumentNullException(string name)
+                => throw new ArgumentNullException(name);
+
+            public static void ThrowArgumentOutOfRangeException(string name)
+                => throw new ArgumentOutOfRangeException(name);
+        }
+    }
+}

src/Components/Server/src/Circuits/RemoteRenderer.cs

@@ -3,7 +3,6 @@
 
 using System;
 using System.Collections.Concurrent;
-using System.IO;
 using System.Linq;
 using System.Text.Encodings.Web;
 using System.Threading;
@@ -102,12 +101,13 @@ namespace Microsoft.AspNetCore.Components.Web.Rendering
         protected override void Dispose(bool disposing)
         {
             _disposing = true;
-            base.Dispose(true);
+            _rendererRegistry.TryRemove(Id);
             while (PendingRenderBatches.TryDequeue(out var entry))
             {
                 entry.CompletionSource.TrySetCanceled();
+                entry.Data.Dispose();
             }
-            _rendererRegistry.TryRemove(Id);
+            base.Dispose(true);
         }
 
         /// <inheritdoc />
@@ -123,30 +123,34 @@ namespace Microsoft.AspNetCore.Components.Web.Rendering
             // SignalR's SendAsync can wait an arbitrary duration before serializing the params.
             // The RenderBatch buffer will get reused by subsequent renders, so we need to
             // snapshot its contents now.
-            // TODO: Consider using some kind of array pool instead of allocating a new
+            var arrayBuilder = new ArrayBuilder<byte>(2048);
-            //       buffer on every render.
+            using var memoryStream = new ArrayBuilderMemoryStream(arrayBuilder);
-            byte[] batchBytes;
+            PendingRender pendingRender;
-            using (var memoryStream = new MemoryStream())
+            try
             {
                 using (var renderBatchWriter = new RenderBatchWriter(memoryStream, false))
                 {
                     renderBatchWriter.Write(in batch);
                 }
 
-                batchBytes = memoryStream.ToArray();
+                var renderId = Interlocked.Increment(ref _nextRenderId);
+
+                pendingRender = new PendingRender(
+                    renderId,
+                    arrayBuilder,
+                    new TaskCompletionSource<object>());
+
+                // Buffer the rendered batches no matter what. We'll send it down immediately when the client
+                // is connected or right after the client reconnects.
+
+                PendingRenderBatches.Enqueue(pendingRender);
+            }
+            catch
+            {
+                // if we throw prior to queueing the write, dispose the builder.
+                arrayBuilder.Dispose();
+                throw;
             }
-
-            var renderId = Interlocked.Increment(ref _nextRenderId);
-
-            var pendingRender = new PendingRender(
-                renderId,
-                batchBytes,
-                new TaskCompletionSource<object>());
-
-            // Buffer the rendered batches no matter what. We'll send it down immediately when the client
-            // is connected or right after the client reconnects.
-
-            PendingRenderBatches.Enqueue(pendingRender);
 
             // Fire and forget the initial send for this batch (if connected). Otherwise it will be sent
             // as soon as the client reconnects.
@@ -181,8 +185,9 @@ namespace Microsoft.AspNetCore.Components.Web.Rendering
                     return;
                 }
 
-                Log.BeginUpdateDisplayAsync(_logger, _client.ConnectionId, pending.BatchId, pending.Data.Length);
+                Log.BeginUpdateDisplayAsync(_logger, _client.ConnectionId, pending.BatchId, pending.Data.Count);
-                await _client.SendAsync("JS.RenderBatch", Id, pending.BatchId, pending.Data);
+                var segment = new ArraySegment<byte>(pending.Data.Buffer, 0, pending.Data.Count);
+                await _client.SendAsync("JS.RenderBatch", Id, pending.BatchId, segment);
             }
             catch (Exception e)
             {
@@ -207,22 +212,33 @@ namespace Microsoft.AspNetCore.Components.Web.Rendering
             // line (i.e., matching the order in which we received batch completion messages) based on the fact that SignalR
             // synchronizes calls to hub methods. That is, it won't issue more than one call to this method from the same hub
             // at the same time on different threads.
-            if (!PendingRenderBatches.TryDequeue(out var entry) || entry.BatchId != incomingBatchId)
+            if (!PendingRenderBatches.TryDequeue(out var entry))
             {
                 HandleException(
                     new InvalidOperationException($"Received a notification for a rendered batch when not expecting it. Batch id '{incomingBatchId}'."));
             }
             else
             {
-                if (errorMessageOrNull == null)
+                entry.Data.Dispose();
+
+                if (entry.BatchId != incomingBatchId)
                 {
-                    Log.CompletingBatchWithoutError(_logger, entry.BatchId);
+                    HandleException(
+                        new InvalidOperationException($"Received a notification for a rendered batch when not expecting it. Batch id '{incomingBatchId}'."));
                 }
                 else
                 {
-                    Log.CompletingBatchWithError(_logger, entry.BatchId, errorMessageOrNull);
+                    if (errorMessageOrNull == null)
+                    {
+                        Log.CompletingBatchWithoutError(_logger, entry.BatchId);
+                    }
+                    else
+                    {
+                        Log.CompletingBatchWithError(_logger, entry.BatchId, errorMessageOrNull);
+                    }
                 }
 
+
                 CompleteRender(entry.CompletionSource, errorMessageOrNull);
             }
         }
@@ -241,7 +257,7 @@ namespace Microsoft.AspNetCore.Components.Web.Rendering
 
         internal readonly struct PendingRender
         {
-            public PendingRender(long batchId, byte[] data, TaskCompletionSource<object> completionSource)
+            public PendingRender(long batchId, ArrayBuilder<byte> data, TaskCompletionSource<object> completionSource)
             {
                 BatchId = batchId;
                 Data = data;
@@ -249,7 +265,7 @@ namespace Microsoft.AspNetCore.Components.Web.Rendering
             }
 
             public long BatchId { get; }
-            public byte[] Data { get; }
+            public ArrayBuilder<byte> Data { get; }
             public TaskCompletionSource<object> CompletionSource { get; }
         }
 

src/Components/Server/src/Circuits/RenderBatchWriter.cs

@@ -1,14 +1,14 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
-using Microsoft.AspNetCore.Components.Rendering;
-using Microsoft.AspNetCore.Components.RenderTree;
 using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Text;
+using Microsoft.AspNetCore.Components.Rendering;
+using Microsoft.AspNetCore.Components.RenderTree;
 
-namespace Microsoft.AspNetCore.Components.Server.Circuits
+namespace Microsoft.AspNetCore.Components.Web.Rendering
 {
     // TODO: We should consider *not* having this type of infrastructure in the .Server
     // project, but instead in some new project called .Remote or similar, since it
@@ -33,13 +33,13 @@ namespace Microsoft.AspNetCore.Components.Server.Circuits
     /// </summary>
     internal class RenderBatchWriter : IDisposable
     {
-        private readonly List<string> _strings;
+        private readonly ArrayBuilder<string> _strings;
         private readonly Dictionary<string, int> _deduplicatedStringIndices;
         private readonly BinaryWriter _binaryWriter;
 
         public RenderBatchWriter(Stream output, bool leaveOpen)
         {
-            _strings = new List<string>();
+            _strings = new ArrayBuilder<string>();
             _deduplicatedStringIndices = new Dictionary<string, int>();
             _binaryWriter = new BinaryWriter(output, Encoding.UTF8, leaveOpen);
         }
@@ -243,7 +243,7 @@ namespace Microsoft.AspNetCore.Components.Server.Circuits
                 if (!allowDeduplication || !_deduplicatedStringIndices.TryGetValue(value, out stringIndex))
                 {
                     stringIndex = _strings.Count;
-                    _strings.Add(value);
+                    _strings.Append(value);
 
                     if (allowDeduplication)
                     {
@@ -263,7 +263,7 @@ namespace Microsoft.AspNetCore.Components.Server.Circuits
 
             for (var i = 0; i < stringsCount; i++)
             {
-                var stringValue = _strings[i];
+                var stringValue = _strings.Buffer[i];
                 locations[i] = (int)_binaryWriter.BaseStream.Position;
                 _binaryWriter.Write(stringValue);
             }
@@ -295,6 +295,7 @@ namespace Microsoft.AspNetCore.Components.Server.Circuits
 
         public void Dispose()
         {
+            _strings.Dispose();
             _binaryWriter.Dispose();
         }
     }

src/Components/Server/src/Microsoft.AspNetCore.Components.Server.csproj

@@ -8,7 +8,7 @@
     <GenerateDocumentationFile>true</GenerateDocumentationFile>
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
     <NoWarn>CS0436;$(NoWarn)</NoWarn>
-    <DefineConstants>$(DefineConstants);MESSAGEPACK_INTERNAL</DefineConstants>
+    <DefineConstants>$(DefineConstants);MESSAGEPACK_INTERNAL;COMPONENTS_SERVER</DefineConstants>
   </PropertyGroup>
 
   <ItemGroup>
@@ -28,6 +28,7 @@
 
   <ItemGroup>
     <Compile Include="$(ComponentsSharedSourceRoot)src\CacheHeaderSettings.cs" Link="Shared\CacheHeaderSettings.cs" />
+    <Compile Include="$(ComponentsSharedSourceRoot)src\ArrayBuilder.cs" LinkBase="Circuits" />
 
     <Compile Include="$(RepoRoot)src\SignalR\common\Shared\BinaryMessageFormatter.cs" LinkBase="BlazorPack" />
     <Compile Include="$(RepoRoot)src\SignalR\common\Shared\BinaryMessageParser.cs" LinkBase="BlazorPack" />

src/Components/Server/test/Circuits/RenderBatchWriterTest.cs

@@ -5,6 +5,7 @@ using Microsoft.AspNetCore.Components;
 using Microsoft.AspNetCore.Components.Rendering;
 using Microsoft.AspNetCore.Components.RenderTree;
 using Microsoft.AspNetCore.Components.Server.Circuits;
+using Microsoft.AspNetCore.Components.Web.Rendering;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging.Abstractions;
 using System;
@@ -153,7 +154,7 @@ namespace Microsoft.AspNetCore.Components.Server
                 RenderTreeEdit.UpdateMarkup(108, 109),
                 RenderTreeEdit.RemoveAttribute(110, "Some removed attribute"), // To test deduplication
             };
-            var editsBuilder = new ArrayBuilder<RenderTreeEdit>();
+            var editsBuilder = new RenderTree.ArrayBuilder<RenderTreeEdit>();
             editsBuilder.Append(edits, 0, edits.Length);
             var editsSegment = editsBuilder.ToSegment(1, edits.Length); // Skip first to show offset is respected
             var bytes = Serialize(new RenderBatch(

src/Components/Shared/src/ArrayBuilder.cs

@@ -6,7 +6,11 @@ using System.Buffers;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
 
+#if COMPONENTS_SERVER
+namespace Microsoft.AspNetCore.Components.Web.Rendering
+#else
 namespace Microsoft.AspNetCore.Components.RenderTree
+#endif
 {
     /// <summary>
     /// Implements a list that uses an array of objects to store the elements.
@@ -154,22 +158,6 @@ namespace Microsoft.AspNetCore.Components.RenderTree
             _itemsInUse = 0;
         }
 
-        /// <summary>
-        /// Produces an <see cref="ArrayRange{T}"/> structure describing the current contents.
-        /// </summary>
-        /// <returns>The <see cref="ArrayRange{T}"/>.</returns>
-        public ArrayRange<T> ToRange()
-            => new ArrayRange<T>(_items, _itemsInUse);
-
-        /// <summary>
-        /// Produces an <see cref="ArrayBuilderSegment{T}"/> structure describing the selected contents.
-        /// </summary>
-        /// <param name="fromIndexInclusive">The index of the first item in the segment.</param>
-        /// <param name="toIndexExclusive">One plus the index of the last item in the segment.</param>
-        /// <returns>The <see cref="ArraySegment{T}"/>.</returns>
-        public ArrayBuilderSegment<T> ToSegment(int fromIndexInclusive, int toIndexExclusive)
-            => new ArrayBuilderSegment<T>(this, fromIndexInclusive, toIndexExclusive - fromIndexInclusive);
-
         private void GrowBuffer(int desiredCapacity)
         {
             var newCapacity = Math.Max(desiredCapacity, _minCapacity);

src/Components/test/Ignitor.Test/RenderBatchReaderTest.cs

@@ -6,11 +6,10 @@ using System.Collections.Generic;
 using System.IO;
 using System.Text;
 using System.Threading.Tasks;
-using Castle.Core.Logging;
 using Microsoft.AspNetCore.Components;
 using Microsoft.AspNetCore.Components.Rendering;
 using Microsoft.AspNetCore.Components.RenderTree;
-using Microsoft.AspNetCore.Components.Server.Circuits;
+using Microsoft.AspNetCore.Components.Web.Rendering;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging.Abstractions;
 using Xunit;
@@ -103,7 +102,7 @@ namespace Ignitor
                 RenderTreeEdit.UpdateMarkup(108, 109),
                 RenderTreeEdit.RemoveAttribute(110, "Some removed attribute"), // To test deduplication
             };
-            var editsBuilder = new ArrayBuilder<RenderTreeEdit>();
+            var editsBuilder = new Microsoft.AspNetCore.Components.RenderTree.ArrayBuilder<RenderTreeEdit>();
             editsBuilder.Append(edits, 0, edits.Length);
             var editsSegment = editsBuilder.ToSegment(1, edits.Length); // Skip first to show offset is respected
             var bytes = RoundTripSerialize(new RenderBatch(