Improve ArrayBufferWriter re-alloc perf when size is > int.MaxSize / 2 (#42950)

Author: steveharter

src/libraries/Common/src/System/Buffers/ArrayBufferWriter.cs

@@ -15,10 +15,14 @@ namespace System.Buffers
 #endif
     sealed class ArrayBufferWriter<T> : IBufferWriter<T>
     {
+        // Copy of Array.MaxArrayLength. For byte arrays the limit is slightly larger
+        private const int MaxArrayLength = 0X7FEFFFFF;
+
+        private const int DefaultInitialBufferSize = 256;
+
         private T[] _buffer;
         private int _index;
 
-        private const int DefaultInitialBufferSize = 256;
 
         /// <summary>
         /// Creates an instance of an <see cref="ArrayBufferWriter{T}"/>, in which data can be written to,
@@ -167,6 +171,8 @@ private void CheckAndResizeBuffer(int sizeHint)
             if (sizeHint > FreeCapacity)
             {
                 int currentLength = _buffer.Length;
+
+                // Attempt to grow by the larger of the sizeHint and double the current size.
                 int growBy = Math.Max(sizeHint, currentLength);
 
                 if (currentLength == 0)
@@ -178,11 +184,16 @@ private void CheckAndResizeBuffer(int sizeHint)
 
                 if ((uint)newSize > int.MaxValue)
                 {
-                    newSize = currentLength + sizeHint;
-                    if ((uint)newSize > int.MaxValue)
+                    // Attempt to grow to MaxArrayLength.
+                    uint needed = (uint)(currentLength - FreeCapacity + sizeHint);
+                    Debug.Assert(needed > currentLength);
+
+                    if (needed > MaxArrayLength)
                     {
-                        ThrowOutOfMemoryException((uint)newSize);
+                        ThrowOutOfMemoryException(needed);
                     }
+
+                    newSize = MaxArrayLength;
                 }
 
                 Array.Resize(ref _buffer, newSize);

src/libraries/System.Memory/tests/ArrayBufferWriter/ArrayBufferWriterTests.Byte.cs

@@ -98,10 +98,20 @@ public async Task WriteAndCopyToStreamAsync()
         [OuterLoop]
         public void GetMemory_ExceedMaximumBufferSize()
         {
-            var output = new ArrayBufferWriter<byte>(int.MaxValue / 2 + 1);
-            output.Advance(int.MaxValue / 2 + 1);
-            Memory<byte> memory = output.GetMemory(1); // Validate we can't double the buffer size, but can grow by sizeHint
-            Assert.Equal(1, memory.Length);
+            const int MaxArrayLength = 0X7FEFFFFF;
+
+            int initialCapacity = int.MaxValue / 2 + 1;
+
+            var output = new ArrayBufferWriter<byte>(initialCapacity);
+            output.Advance(initialCapacity);
+
+            // Validate we can't double the buffer size, but can grow
+            Memory<byte> memory = output.GetMemory(1);
+
+            // The buffer should grow more than the 1 byte requested otherwise performance will not be usable
+            // between 1GB and 2GB. The current implementation maxes out the buffer size to MaxArrayLength.
+            Assert.Equal(MaxArrayLength - initialCapacity, memory.Length);
+
             Assert.Throws<OutOfMemoryException>(() => output.GetMemory(int.MaxValue));
         }
     }

src/libraries/System.Text.Json/src/System/Text/Json/JsonHelpers.cs

@@ -11,6 +11,9 @@ namespace System.Text.Json
 {
     internal static partial class JsonHelpers
     {
+        // Copy of Array.MaxArrayLength. For byte arrays the limit is slightly larger
+        private const int MaxArrayLength = 0X7FEFFFFF;
+
         /// <summary>
         /// Returns the span for the given reader.
         /// </summary>
@@ -143,5 +146,14 @@ public static bool IsValidNumberHandlingValue(JsonNumberHandling handling) =>
                 JsonNumberHandling.AllowReadingFromString |
                 JsonNumberHandling.WriteAsString |
                 JsonNumberHandling.AllowNamedFloatingPointLiterals));
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void ValidateInt32MaxArrayLength(uint length)
+        {
+            if (length > MaxArrayLength)
+            {
+                ThrowHelper.ThrowOutOfMemoryException(length);
+            }
+        }
     }
 }

src/libraries/System.Text.Json/src/System/Text/Json/ThrowHelper.cs

@@ -506,6 +506,12 @@ public static InvalidOperationException GetInvalidOperationException(ExceptionRe
             return ex;
         }
 
+        [DoesNotReturn]
+        public static void ThrowOutOfMemoryException(uint capacity)
+        {
+            throw new OutOfMemoryException(SR.Format(SR.BufferMaximumSizeExceeded, capacity));
+        }
+
         // This function will convert an ExceptionResource enum value to the resource string.
         [MethodImpl(MethodImplOptions.NoInlining)]
         private static string GetResourceString(ExceptionResource resource, int currentDepth, byte token, JsonTokenType tokenType)

src/libraries/System.Text.Json/src/System/Text/Json/Writer/Utf8JsonWriter.cs

@@ -1004,7 +1004,10 @@ private void Grow(int requiredSize)
             {
                 Debug.Assert(_arrayBufferWriter != null);
 
-                _memory = _arrayBufferWriter.GetMemory(checked(BytesPending + sizeHint));
+                int needed = BytesPending + sizeHint;
+                JsonHelpers.ValidateInt32MaxArrayLength((uint)needed);
+
+                _memory = _arrayBufferWriter.GetMemory(needed);
 
                 Debug.Assert(_memory.Length >= sizeHint);
             }

src/libraries/System.Text.Json/tests/Utf8JsonWriterTests.cs

@@ -744,24 +744,25 @@ public void WriteLargeJsonToStreamWithoutFlushing()
                 }
                 Assert.Equal(150_097_503, writer.BytesPending);
 
-                for (int i = 0; i < 6; i++)
+                for (int i = 0; i < 13; i++)
                 {
                     writer.WriteStringValue(text3);
                 }
-                Assert.Equal(1_050_097_521, writer.BytesPending);
+                Assert.Equal(2_100_097_542, writer.BytesPending);
+
+                // Next write forces a grow beyond max array length
 
-                // Next write forces a grow beyond 2 GB
                 Assert.Throws<OutOfMemoryException>(() => writer.WriteStringValue(text3));
 
-                Assert.Equal(1_050_097_521, writer.BytesPending);
+                Assert.Equal(2_100_097_542, writer.BytesPending);
 
                 var text4 = JsonEncodedText.Encode(largeArray.AsSpan(0, 1));
                 for (int i = 0; i < 10_000_000; i++)
                 {
                     writer.WriteStringValue(text4);
                 }
 
-                Assert.Equal(1_050_097_521 + (4 * 10_000_000), writer.BytesPending);
+                Assert.Equal(2_100_097_542 + (4 * 10_000_000), writer.BytesPending);
             }
         }