From 0971ec6e34c75b09f7871f425e0507856e43d1f9 Mon Sep 17 00:00:00 2001 From: Paulo Morgado <470455+paulomorgado@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:22:37 +0100 Subject: [PATCH] Add segmented storage to ArrayPoolBufferWriter Retain completed pooled arrays instead of copying data on every growth, and expose written data through GetReadOnlySequence(). Grow segments geometrically and consolidate only when a contiguous view is requested. Keep capacity and write counts within supported limits, avoid consolidation during char formatting and debugger inspection, and add tests for buffer ownership, failure recovery, and boundary conditions. --- .../Buffers/ArrayPoolBufferWriter{T}.cs | 368 ++++++++++++-- .../Buffers/Views/MemoryDebugView{T}.cs | 3 +- .../Buffers/Internals/TrackingArrayPool{T}.cs | 13 + .../Buffers/Test_ArrayPoolBufferWriter{T}.cs | 481 ++++++++++++++++++ 4 files changed, 828 insertions(+), 37 deletions(-) diff --git a/src/CommunityToolkit.HighPerformance/Buffers/ArrayPoolBufferWriter{T}.cs b/src/CommunityToolkit.HighPerformance/Buffers/ArrayPoolBufferWriter{T}.cs index 1fdc0f67f..75e8c5bea 100644 --- a/src/CommunityToolkit.HighPerformance/Buffers/ArrayPoolBufferWriter{T}.cs +++ b/src/CommunityToolkit.HighPerformance/Buffers/ArrayPoolBufferWriter{T}.cs @@ -4,6 +4,7 @@ using System; using System.Buffers; +using System.Collections.Generic; using System.Diagnostics; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; @@ -36,6 +37,11 @@ public sealed class ArrayPoolBufferWriter : IBuffer, IMemoryOwner /// private const int DefaultInitialBufferSize = 256; + /// + /// The maximum size to target when growing buffers geometrically. + /// + private const int MaximumSegmentSize = 1024 * 1024; + /// /// The instance used to rent . /// @@ -46,6 +52,16 @@ public sealed class ArrayPoolBufferWriter : IBuffer, IMemoryOwner /// private T[]? array; + /// + /// The completed buffers that were previously used as active write targets. + /// + private List? buffers; + + /// + /// The total number of written items in . + /// + private int bufferedCount; + #pragma warning disable IDE0032 // Use field over auto-property (clearer and faster) /// /// The starting offset within . @@ -121,14 +137,9 @@ public ReadOnlyMemory WrittenMemory [MethodImpl(MethodImplOptions.AggressiveInlining)] get { - T[]? array = this.array; + EnsureSingleBuffer(); - if (array is null) - { - ThrowObjectDisposedException(); - } - - return array!.AsMemory(0, this.index); + return this.array.AsMemory(0, this.index); } } @@ -138,14 +149,9 @@ public ReadOnlySpan WrittenSpan [MethodImpl(MethodImplOptions.AggressiveInlining)] get { - T[]? array = this.array; + EnsureSingleBuffer(); - if (array is null) - { - ThrowObjectDisposedException(); - } - - return array!.AsSpan(0, this.index); + return this.array.AsSpan(0, this.index); } } @@ -153,7 +159,7 @@ public ReadOnlySpan WrittenSpan public int WrittenCount { [MethodImpl(MethodImplOptions.AggressiveInlining)] - get => this.index; + get => this.bufferedCount + this.index; } /// @@ -169,7 +175,7 @@ public int Capacity ThrowObjectDisposedException(); } - return array!.Length; + return this.bufferedCount + Math.Min(array!.Length, int.MaxValue - this.bufferedCount); } } @@ -186,7 +192,7 @@ public int FreeCapacity ThrowObjectDisposedException(); } - return array!.Length - this.index; + return Math.Min(array!.Length - this.index, int.MaxValue - WrittenCount); } } @@ -200,6 +206,18 @@ public void Clear() ThrowObjectDisposedException(); } + if (this.buffers is { Count: > 0 } buffers) + { + foreach (BufferInfo buffer in buffers) + { + buffer.Array.AsSpan(0, buffer.Length).Clear(); + this.pool.Return(buffer.Array); + } + + buffers.Clear(); + this.bufferedCount = 0; + } + array.AsSpan(0, this.index).Clear(); this.index = 0; @@ -225,6 +243,11 @@ public void Advance(int count) ThrowArgumentExceptionForAdvancedTooFar(); } + if (count > int.MaxValue - WrittenCount) + { + ThrowArgumentExceptionForAdvancedTooFar(); + } + this.index += count; } @@ -233,7 +256,7 @@ public Memory GetMemory(int sizeHint = 0) { CheckBufferAndEnsureCapacity(sizeHint); - return this.array.AsMemory(this.index); + return this.array.AsMemory(this.index, FreeCapacity); } /// @@ -241,7 +264,64 @@ public Span GetSpan(int sizeHint = 0) { CheckBufferAndEnsureCapacity(sizeHint); - return this.array.AsSpan(this.index); + return this.array.AsSpan(this.index, FreeCapacity); + } + + /// + /// Gets a over the written data. + /// + /// A over the written data. + /// Thrown when the current instance has been disposed. + /// + /// The returned sequence is only valid as long as the current instance is not modified, cleared, consolidated or disposed. + /// + public ReadOnlySequence GetReadOnlySequence() + { + T[]? array = this.array; + + if (array is null) + { + ThrowObjectDisposedException(); + } + + if (this.bufferedCount + this.index == 0) + { + return ReadOnlySequence.Empty; + } + + if (this.buffers is not { Count: > 0 }) + { + return new ReadOnlySequence(array!, 0, this.index); + } + + SequenceSegment? startSegment = null; + SequenceSegment? endSegment = null; + + foreach (BufferInfo buffer in this.buffers) + { + SequenceSegment segment = new(buffer.Array.AsMemory(0, buffer.Length)); + + if (startSegment is null) + { + startSegment = segment; + } + else + { + endSegment!.SetNext(segment); + } + + endSegment = segment; + } + + if (this.index > 0) + { + SequenceSegment segment = new(array!.AsMemory(0, this.index)); + + endSegment!.SetNext(segment); + endSegment = segment; + } + + return new(startSegment!, 0, endSegment!, endSegment!.Memory.Length); } /// @@ -258,14 +338,9 @@ public Span GetSpan(int sizeHint = 0) [MethodImpl(MethodImplOptions.AggressiveInlining)] public ArraySegment DangerousGetArray() { - T[]? array = this.array; + EnsureSingleBuffer(); - if (array is null) - { - ThrowObjectDisposedException(); - } - - return new(array!, 0, this.index); + return new(this.array!, 0, this.index); } /// @@ -280,6 +355,19 @@ public void Dispose() this.array = null; + if (this.buffers is { Count: > 0 } buffers) + { + foreach (BufferInfo buffer in buffers) + { + this.pool.Return(buffer.Array); + } + + buffers.Clear(); + } + + this.bufferedCount = 0; + this.index = 0; + this.pool.Return(array); } @@ -290,11 +378,47 @@ public override string ToString() if (typeof(T) == typeof(char) && this.array is char[] chars) { - return new(chars, 0, this.index); + if (this.bufferedCount == 0) + { + return new(chars, 0, this.index); + } + +#if NETSTANDARD2_0 + ArrayPool pool = (ArrayPool)(object)this.pool; + char[] buffer = pool.Rent(WrittenCount); + + try + { + CopyWrittenCharsTo(buffer); + + return new(buffer, 0, WrittenCount); + } + finally + { + pool.Return(buffer); + } +#else + return string.Create(WrittenCount, this, static (destination, writer) => writer.CopyWrittenCharsTo(destination)); +#endif } // Same representation used in Span - return $"CommunityToolkit.HighPerformance.Buffers.ArrayPoolBufferWriter<{typeof(T)}>[{this.index}]"; + return $"CommunityToolkit.HighPerformance.Buffers.ArrayPoolBufferWriter<{typeof(T)}>[{WrittenCount}]"; + } + + /// + /// Copies the written characters to a destination without modifying the writer. + /// + /// The destination for the written characters. + private void CopyWrittenCharsTo(Span destination) + { + foreach (BufferInfo buffer in this.buffers!) + { + ((char[])(object)buffer.Array).AsSpan(0, buffer.Length).CopyTo(destination); + destination = destination.Slice(buffer.Length); + } + + ((char[])(object)this.array!).AsSpan(0, this.index).CopyTo(destination); } /// @@ -321,32 +445,143 @@ private void CheckBufferAndEnsureCapacity(int sizeHint) sizeHint = 1; } + // The public counts are ints, so never expose or retain more than int.MaxValue items. + // Check before renting to preserve ownership and the current sequence on failure. + if (sizeHint > int.MaxValue - WrittenCount) + { + ThrowArgumentOutOfRangeExceptionForExcessiveSizeHint(); + } + if (sizeHint > array!.Length - this.index) { - ResizeBuffer(sizeHint); + RentNextBuffer(sizeHint); } } /// - /// Resizes to ensure it can fit the specified number of new items. + /// Rents a new active buffer with enough space for the specified number of new items. /// /// The minimum number of items to ensure space for in . [MethodImpl(MethodImplOptions.NoInlining)] - private void ResizeBuffer(int sizeHint) + private void RentNextBuffer(int sizeHint) { - uint minimumSize = (uint)this.index + (uint)sizeHint; + Debug.Assert(this.array is not null); + + T[] array = this.array!; + + // Rent before transferring ownership of the current buffer so that a failed rent + // leaves the writer unchanged and still owning its current buffer. + T[] nextArray = this.pool.Rent(GetMinimumBufferSize(array.Length, sizeHint)); + + try + { + if (this.index > 0) + { + AddBuffer(array, this.index); + } + else + { + this.pool.Return(array); + } + } + catch + { + // If transferring the current buffer fails, return the new one to avoid a leak. + this.pool.Return(nextArray); + + throw; + } + + this.array = nextArray; + this.index = 0; + } + + /// + /// Adds a completed buffer to . + /// + /// The buffer to add. + /// The number of written items in . + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private void AddBuffer(T[] array, int length) + { + (this.buffers ??= new List()).Add(new BufferInfo(array, length)); + this.bufferedCount += length; + } + + /// + /// Gets the minimum size to use for a new active buffer. + /// + /// The length of the current active buffer. + /// The minimum required number of writable items. + /// The minimum size to request to the pool. + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static int GetMinimumBufferSize(int currentLength, int sizeHint) + { + int nextSize = currentLength < MaximumSegmentSize + ? Math.Min(currentLength * 2, MaximumSegmentSize) + : currentLength; + + uint minimumSize = (uint)Math.Max(nextSize, sizeHint); // The ArrayPool class has a maximum threshold of 1024 * 1024 for the maximum length of // pooled arrays, and once this is exceeded it will just allocate a new array every time // of exactly the requested size. In that case, we manually round up the requested size to - // the nearest power of two, to ensure that repeated consecutive writes when the array in - // use is bigger than that threshold don't end up causing a resize every single time. - if (minimumSize > 1024 * 1024) + // the nearest power of two, to avoid renting a new array for every consecutive write + // once the active buffer exceeds that threshold. + if (minimumSize > MaximumSegmentSize) { minimumSize = BitOperations.RoundUpToPowerOf2(minimumSize); } - this.pool.Resize(ref this.array, (int)minimumSize); + return (int)minimumSize; + } + + /// + /// Ensures all written data is represented by a single active buffer. + /// + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private void EnsureSingleBuffer() + { + T[]? array = this.array; + + if (array is null) + { + ThrowObjectDisposedException(); + } + + if (this.bufferedCount == 0) + { + return; + } + + ConsolidateBuffers(); + } + + /// + /// Consolidates all written data into a single active buffer. + /// + [MethodImpl(MethodImplOptions.NoInlining)] + private void ConsolidateBuffers() + { + T[] currentBuffer = this.array!; + int totalCount = this.bufferedCount + this.index; + T[] targetBuffer = this.pool.Rent(totalCount); + int offset = 0; + + foreach (BufferInfo buffer in this.buffers!) + { + buffer.Array.AsSpan(0, buffer.Length).CopyTo(targetBuffer.AsSpan(offset)); + offset += buffer.Length; + this.pool.Return(buffer.Array); + } + + currentBuffer.AsSpan(0, this.index).CopyTo(targetBuffer.AsSpan(offset)); + this.pool.Return(currentBuffer); + + this.buffers!.Clear(); + this.bufferedCount = 0; + this.array = targetBuffer; + this.index = totalCount; } /// @@ -365,6 +600,14 @@ private static void ThrowArgumentOutOfRangeExceptionForNegativeSizeHint() throw new ArgumentOutOfRangeException("sizeHint", "The size hint can't be a negative value."); } + /// + /// Throws an when the requested size exceeds the supported length. + /// + private static void ThrowArgumentOutOfRangeExceptionForExcessiveSizeHint() + { + throw new ArgumentOutOfRangeException("sizeHint", "The buffer writer cannot contain more than int.MaxValue items."); + } + /// /// Throws an when the requested count is negative. /// @@ -380,4 +623,57 @@ private static void ThrowObjectDisposedException() { throw new ObjectDisposedException("The current buffer has already been disposed."); } + + /// + /// The metadata for each completed buffer. + /// + private readonly struct BufferInfo + { + /// + /// Initializes a new instance of the struct. + /// + /// The underlying array. + /// The number of written items. + public BufferInfo(T[] array, int length) + { + this.Array = array; + this.Length = length; + } + + /// + /// Gets the underlying array. + /// + public T[] Array { get; } + + /// + /// Gets the number of written items in . + /// + public int Length { get; } + } + + /// + /// A implementation for pooled buffers. + /// + private sealed class SequenceSegment : ReadOnlySequenceSegment + { + /// + /// Initializes a new instance of the class. + /// + /// The memory to wrap. + public SequenceSegment(ReadOnlyMemory memory) + { + this.Memory = memory; + } + + /// + /// Sets the next segment and updates the running index. + /// + /// The next segment. + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void SetNext(SequenceSegment next) + { + next.RunningIndex = this.RunningIndex + this.Memory.Length; + this.Next = next; + } + } } diff --git a/src/CommunityToolkit.HighPerformance/Buffers/Views/MemoryDebugView{T}.cs b/src/CommunityToolkit.HighPerformance/Buffers/Views/MemoryDebugView{T}.cs index e52b6603d..5c2cd3e22 100644 --- a/src/CommunityToolkit.HighPerformance/Buffers/Views/MemoryDebugView{T}.cs +++ b/src/CommunityToolkit.HighPerformance/Buffers/Views/MemoryDebugView{T}.cs @@ -2,6 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. // See the LICENSE file in the project root for more information. +using System.Buffers; using System.Diagnostics; namespace CommunityToolkit.HighPerformance.Buffers.Views; @@ -18,7 +19,7 @@ internal sealed class MemoryDebugView /// The input instance with the items to display. public MemoryDebugView(ArrayPoolBufferWriter? arrayPoolBufferWriter) { - this.Items = arrayPoolBufferWriter?.WrittenSpan.ToArray(); + this.Items = arrayPoolBufferWriter?.GetReadOnlySequence().ToArray(); } /// diff --git a/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Internals/TrackingArrayPool{T}.cs b/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Internals/TrackingArrayPool{T}.cs index 6743c3758..4798da25c 100644 --- a/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Internals/TrackingArrayPool{T}.cs +++ b/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Internals/TrackingArrayPool{T}.cs @@ -2,6 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. // See the LICENSE file in the project root for more information. +using System; using System.Buffers; using System.Collections.Generic; @@ -18,9 +19,21 @@ internal sealed class TrackingArrayPool : ArrayPool /// public IReadOnlyCollection RentedArrays => this.arrays; + /// + /// Gets or sets whether the next rent should fail. + /// + public bool FailNextRent { get; set; } + /// public override T[] Rent(int minimumLength) { + if (this.FailNextRent) + { + this.FailNextRent = false; + + throw new InvalidOperationException(); + } + T[] array = this.pool.Rent(minimumLength); _ = this.arrays.Add(array); diff --git a/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Test_ArrayPoolBufferWriter{T}.cs b/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Test_ArrayPoolBufferWriter{T}.cs index 8126c834f..aee3dc6db 100644 --- a/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Test_ArrayPoolBufferWriter{T}.cs +++ b/tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Test_ArrayPoolBufferWriter{T}.cs @@ -218,6 +218,478 @@ public void Test_ArrayPoolBufferWriterOfT_AsStream() _ = Assert.ThrowsExactly(() => writer.Capacity); } + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_GetReadOnlySequence_DoesNotConsolidate() + { + Random random = new(Guid.NewGuid().GetHashCode()); + int firstLength = random.Next(16, 33); + int secondLength = random.Next(34, 65); + + byte[] first = GetRandomBytes(firstLength); + byte[] second = GetRandomBytes(secondLength); + + TrackingArrayPool pool = new(); + + using ArrayPoolBufferWriter writer = new(pool, firstLength); + + first.CopyTo(writer.GetSpan(firstLength)); + writer.Advance(firstLength); + + second.CopyTo(writer.GetSpan(secondLength)); + writer.Advance(secondLength); + + byte[] expected = first.Concat(second).ToArray(); + + Assert.HasCount(2, pool.RentedArrays); + + ReadOnlySequence sequence = writer.GetReadOnlySequence(); + + Assert.IsTrue(sequence.ToArray().SequenceEqual(expected)); + Assert.HasCount(2, pool.RentedArrays); + } + + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_CharToString_EmptyAndSingleBuffer() + { + TrackingArrayPool pool = new(); + + using (ArrayPoolBufferWriter writer = new(pool)) + { + char[] array = pool.RentedArrays.Single(); + int capacity = writer.Capacity; + + Assert.AreEqual(string.Empty, writer.ToString()); + Assert.HasCount(1, pool.RentedArrays); + Assert.IsTrue(pool.RentedArrays.Contains(array)); + Assert.AreEqual(capacity, writer.Capacity); + Assert.AreEqual(capacity, writer.FreeCapacity); + Assert.AreEqual(0, writer.WrittenCount); + + string value = Guid.NewGuid().ToString("N"); + + value.AsSpan().CopyTo(writer.GetSpan(value.Length)); + writer.Advance(value.Length); + + ReadOnlySequence sequence = writer.GetReadOnlySequence(); + int freeCapacity = writer.FreeCapacity; + + Assert.AreEqual(value, writer.ToString()); + Assert.AreEqual(value, new string(sequence.ToArray())); + Assert.AreEqual(value.Length, writer.WrittenCount); + Assert.AreEqual(capacity, writer.Capacity); + Assert.AreEqual(freeCapacity, writer.FreeCapacity); + Assert.HasCount(1, pool.RentedArrays); + Assert.IsTrue(pool.RentedArrays.Contains(array)); + } + + Assert.IsEmpty(pool.RentedArrays); + } + + [TestMethod] + [DataRow(false)] + [DataRow(true)] + public void Test_ArrayPoolBufferWriterOfT_CharToString_PreservesSegments(bool emptyActiveBuffer) + { + TrackingArrayPool pool = new(); + + using (ArrayPoolBufferWriter writer = new(pool, Guid.NewGuid().ToString("N").Length)) + { + string first = Guid.NewGuid().ToString("N"); + string second = Guid.NewGuid().ToString("N"); + string third = Guid.NewGuid().ToString("N"); + string expected = first + second + third; + + foreach (string value in new[] { first, second }) + { + value.AsSpan().CopyTo(writer.GetSpan(value.Length)); + writer.Advance(value.Length); + } + + third.AsSpan().CopyTo(writer.GetSpan(writer.FreeCapacity + 1)); + writer.Advance(third.Length); + + if (emptyActiveBuffer) + { + _ = writer.GetMemory(writer.FreeCapacity + 1); + } + + ReadOnlySequence sequence = writer.GetReadOnlySequence(); + char[][] rentedArrays = pool.RentedArrays.ToArray(); + int writtenCount = writer.WrittenCount; + int capacity = writer.Capacity; + int freeCapacity = writer.FreeCapacity; + + Assert.IsGreaterThanOrEqualTo(3, rentedArrays.Length); + + Assert.AreEqual(expected, writer.ToString()); + Assert.AreEqual(expected, writer.ToString()); + Assert.AreEqual(expected, new string(sequence.ToArray())); + Assert.AreEqual(writtenCount, writer.WrittenCount); + Assert.AreEqual(capacity, writer.Capacity); + Assert.AreEqual(freeCapacity, writer.FreeCapacity); + Assert.HasCount(rentedArrays.Length, pool.RentedArrays); + + foreach (char[] array in rentedArrays) + { + Assert.IsTrue(pool.RentedArrays.Contains(array)); + } + } + + Assert.IsEmpty(pool.RentedArrays); + } + + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_CharToString_Disposed() + { + ArrayPoolBufferWriter writer = new(); + + writer.Dispose(); + + Assert.AreEqual($"CommunityToolkit.HighPerformance.Buffers.ArrayPoolBufferWriter<{typeof(char)}>[0]", writer.ToString()); + } + + [TestMethod] + [DataRow(false)] + [DataRow(true)] + public void Test_ArrayPoolBufferWriterOfT_ReadOnlyMembers_PreserveSegments(bool emptyActiveBuffer) + { + TrackingArrayPool pool = new(); + + using (ArrayPoolBufferWriter writer = new(pool, Guid.NewGuid().ToByteArray().Length)) + { + byte[] first = Guid.NewGuid().ToByteArray(); + byte[] second = Guid.NewGuid().ToByteArray(); + byte[] expected = first.Concat(second).ToArray(); + + first.CopyTo(writer.GetSpan(first.Length)); + writer.Advance(first.Length); + + second.CopyTo(writer.GetSpan(first.Length + 1)); + writer.Advance(second.Length); + + if (emptyActiveBuffer) + { + _ = writer.GetMemory(writer.FreeCapacity + 1); + } + + ReadOnlySequence sequence = writer.GetReadOnlySequence(); + byte[][] rentedArrays = pool.RentedArrays.ToArray(); + int writtenCount = writer.WrittenCount; + int capacity = writer.Capacity; + int freeCapacity = writer.FreeCapacity; + string display = $"CommunityToolkit.HighPerformance.Buffers.ArrayPoolBufferWriter<{typeof(byte)}>[{writtenCount}]"; + + Assert.IsGreaterThanOrEqualTo(2, rentedArrays.Length); + + Assert.AreEqual(display, writer.ToString()); + Assert.AreEqual(writtenCount, writer.WrittenCount); + Assert.AreEqual(capacity, writer.Capacity); + Assert.AreEqual(freeCapacity, writer.FreeCapacity); + Assert.IsTrue(writer.GetReadOnlySequence().ToArray().SequenceEqual(expected)); + Assert.AreEqual(display, writer.ToString()); + Assert.IsTrue(sequence.ToArray().SequenceEqual(expected)); + Assert.AreEqual(writtenCount, writer.WrittenCount); + Assert.AreEqual(capacity, writer.Capacity); + Assert.AreEqual(freeCapacity, writer.FreeCapacity); + Assert.HasCount(rentedArrays.Length, pool.RentedArrays); + + foreach (byte[] array in rentedArrays) + { + Assert.IsTrue(pool.RentedArrays.Contains(array)); + } + } + + Assert.IsEmpty(pool.RentedArrays); + } + + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_WrittenMemory_ConsolidatesBuffers_AndDangerousGetArrayMatches() + { + Random random = new(Guid.NewGuid().GetHashCode()); + int firstLength = random.Next(16, 33); + int secondLength = random.Next(34, 65); + + byte[] first = GetRandomBytes(firstLength); + byte[] second = GetRandomBytes(secondLength); + + TrackingArrayPool pool = new(); + + using ArrayPoolBufferWriter writer = new(pool, firstLength); + + first.CopyTo(writer.GetSpan(firstLength)); + writer.Advance(firstLength); + + second.CopyTo(writer.GetSpan(secondLength)); + writer.Advance(secondLength); + + byte[] expected = first.Concat(second).ToArray(); + + Assert.HasCount(2, pool.RentedArrays); + + ReadOnlyMemory memory = writer.WrittenMemory; + + Assert.IsTrue(memory.Span.SequenceEqual(expected)); + Assert.HasCount(1, pool.RentedArrays); + + ArraySegment segment = writer.DangerousGetArray(); + + Assert.HasCount(memory.Length, segment); + Assert.IsTrue(memory.Span.SequenceEqual(segment.AsSpan())); + + _ = MemoryMarshal.TryGetArray(memory, out ArraySegment writtenSegment); + + Assert.AreSame(writtenSegment.Array, segment.Array); + } + + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_Clear_ReturnsRetainedBuffersOnly() + { + Random random = new(Guid.NewGuid().GetHashCode()); + int firstLength = random.Next(16, 33); + int secondLength = random.Next(34, 65); + + TrackingArrayPool pool = new(); + + using ArrayPoolBufferWriter writer = new(pool, firstLength); + + writer.GetSpan(firstLength).Slice(0, firstLength).Fill(1); + writer.Advance(firstLength); + + writer.GetSpan(secondLength).Slice(0, secondLength).Fill(2); + writer.Advance(secondLength); + + Assert.HasCount(2, pool.RentedArrays); + + writer.Clear(); + + Assert.AreEqual(0, writer.WrittenCount); + Assert.IsTrue(writer.WrittenMemory.IsEmpty); + Assert.HasCount(1, pool.RentedArrays); + } + + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_Capacity_TracksOnlyWritableSpace() + { + TrackingArrayPool pool = new(); + + using (ArrayPoolBufferWriter writer = new(pool, Guid.NewGuid().ToByteArray().Length)) + { + Assert.AreEqual(writer.Capacity - writer.WrittenCount, writer.FreeCapacity); + + byte[] first = GetRandomBytes(new Random(Guid.NewGuid().GetHashCode()).Next(1, writer.Capacity)); + + first.CopyTo(writer.GetSpan(first.Length)); + writer.Advance(first.Length); + + Assert.AreEqual(writer.Capacity - writer.WrittenCount, writer.FreeCapacity); + + byte[] second = GetRandomBytes(writer.FreeCapacity + 1); + + second.CopyTo(writer.GetSpan(second.Length)); + + Assert.HasCount(2, pool.RentedArrays); + Assert.AreEqual(first.Length, writer.WrittenCount); + Assert.AreEqual(writer.Capacity - writer.WrittenCount, writer.FreeCapacity); + + writer.Advance(second.Length); + + byte[] expected = first.Concat(second).ToArray(); + + Assert.IsTrue(writer.GetReadOnlySequence().ToArray().SequenceEqual(expected)); + Assert.AreEqual(writer.Capacity - writer.WrittenCount, writer.FreeCapacity); + + ArraySegment contiguous = writer.DangerousGetArray(); + + Assert.IsTrue(contiguous.AsSpan().SequenceEqual(expected)); + Assert.HasCount(1, pool.RentedArrays); + Assert.AreEqual(writer.Capacity - writer.WrittenCount, writer.FreeCapacity); + + writer.Clear(); + + Assert.AreEqual(0, writer.WrittenCount); + Assert.AreEqual(writer.Capacity, writer.FreeCapacity); + } + + Assert.IsEmpty(pool.RentedArrays); + } + + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_SmallWrites_GrowSegmentsGeometrically() + { + TrackingArrayPool pool = new(); + + using (ArrayPoolBufferWriter writer = new(pool)) + { + byte[] chunk = GetRandomBytes(writer.Capacity); + byte[] expected = new byte[1024 * 1024]; + + for (int offset = 0; offset < expected.Length; offset += chunk.Length) + { + chunk.CopyTo(writer.GetSpan(chunk.Length)); + writer.Advance(chunk.Length); + chunk.CopyTo(expected.AsSpan(offset)); + } + + Assert.AreEqual(expected.Length, writer.WrittenCount); + Assert.AreEqual(writer.Capacity - writer.WrittenCount, writer.FreeCapacity); + Assert.IsLessThanOrEqualTo(16, pool.RentedArrays.Count); + Assert.IsTrue(pool.RentedArrays.Any(array => array.Length > chunk.Length)); + Assert.IsTrue(writer.GetReadOnlySequence().ToArray().SequenceEqual(expected)); + } + + Assert.IsEmpty(pool.RentedArrays); + } + + [TestMethod] + [DataRow(false)] + [DataRow(true)] + public void Test_ArrayPoolBufferWriterOfT_FailedRent_PreservesCurrentBuffer(bool hasWrittenData) + { + TrackingArrayPool pool = new(); + byte[] data; + + using (ArrayPoolBufferWriter writer = new(pool)) + { + byte[] originalArray = writer.DangerousGetArray().Array!; + data = hasWrittenData ? GetRandomBytes(new Random(Guid.NewGuid().GetHashCode()).Next(1, writer.Capacity)) : Array.Empty(); + + data.CopyTo(writer.GetSpan(data.Length)); + writer.Advance(data.Length); + + int originalCapacity = writer.Capacity; + int sizeHint = writer.FreeCapacity + 1; + + pool.FailNextRent = true; + + _ = Assert.ThrowsExactly(() => writer.GetMemory(sizeHint)); + + Assert.HasCount(1, pool.RentedArrays); + Assert.IsTrue(pool.RentedArrays.Contains(originalArray)); + Assert.AreEqual(originalCapacity, writer.Capacity); + Assert.AreEqual(data.Length, writer.WrittenCount); + Assert.IsTrue(writer.WrittenSpan.SequenceEqual(data)); + + _ = writer.GetMemory(sizeHint); + + Assert.IsTrue(writer.GetReadOnlySequence().ToArray().SequenceEqual(data)); + } + + Assert.IsEmpty(pool.RentedArrays); + } + + [TestMethod] + public void Test_ArrayPoolBufferWriterOfT_IntLengthBoundary_RejectsWritesWithoutChangingOwnership() + { + TrackingArrayPool pool = new(); + byte[] first = GetRandomBytes(new Random(Guid.NewGuid().GetHashCode()).Next(8, 33)); + byte[] last = GetRandomBytes(new Random(Guid.NewGuid().GetHashCode()).Next(1, 4)); + + using (ArrayPoolBufferWriter writer = new(pool, first.Length)) + { + first.CopyTo(writer.GetSpan(first.Length)); + writer.Advance(first.Length); + _ = writer.GetMemory(writer.FreeCapacity + 1); + + FieldInfo bufferedCountField = typeof(ArrayPoolBufferWriter).GetField("bufferedCount", BindingFlags.Instance | BindingFlags.NonPublic)!; + byte[][] rentedArrays = pool.RentedArrays.ToArray(); + int physicalFreeCapacity = writer.FreeCapacity; + + // Simulate the count of already completed segments without allocating gigabytes. + // Restore the actual count before reading or disposing the writer. + try + { + bufferedCountField.SetValue(writer, int.MaxValue - physicalFreeCapacity); + + // Without the length guard, this request would rent another array. + _ = Assert.ThrowsExactly(() => writer.GetMemory(physicalFreeCapacity + 1)); + Assert.HasCount(rentedArrays.Length, pool.RentedArrays); + + bufferedCountField.SetValue(writer, int.MaxValue - last.Length); + + Assert.AreEqual(int.MaxValue - last.Length, writer.WrittenCount); + Assert.AreEqual(last.Length, writer.FreeCapacity); + Assert.AreEqual(int.MaxValue, writer.Capacity); + Assert.AreEqual(last.Length, writer.GetMemory().Length); + Assert.AreEqual(last.Length, writer.GetSpan().Length); + + _ = Assert.ThrowsExactly(() => writer.GetMemory(last.Length + 1)); + _ = Assert.ThrowsExactly(() => writer.GetSpan(last.Length + 1)); + Assert.HasCount(rentedArrays.Length, pool.RentedArrays); + + last.CopyTo(writer.GetSpan(last.Length)); + writer.Advance(last.Length); + + Assert.AreEqual(int.MaxValue, writer.WrittenCount); + Assert.AreEqual(0, writer.FreeCapacity); + Assert.AreEqual(int.MaxValue, writer.Capacity); + _ = Assert.ThrowsExactly(() => writer.Advance(1)); + _ = Assert.ThrowsExactly(() => writer.GetMemory()); + _ = Assert.ThrowsExactly(() => writer.GetSpan()); + Assert.AreEqual(int.MaxValue, writer.WrittenCount); + Assert.HasCount(rentedArrays.Length, pool.RentedArrays); + + foreach (byte[] array in rentedArrays) + { + Assert.IsTrue(pool.RentedArrays.Contains(array)); + } + } + finally + { + bufferedCountField.SetValue(writer, first.Length); + } + + CollectionAssert.AreEqual(first.Concat(last).ToArray(), writer.GetReadOnlySequence().ToArray()); + } + + Assert.IsEmpty(pool.RentedArrays); + } + + [TestMethod] + [DataRow(false)] + [DataRow(true)] + public void Test_ArrayPoolBufferWriterOfT_DebugView_PreservesSequenceAndPool(bool emptyActiveBuffer) + { + TrackingArrayPool pool = new(); + byte[] first = GetRandomBytes(new Random(Guid.NewGuid().GetHashCode()).Next(8, 33)); + byte[] second = GetRandomBytes(new Random(Guid.NewGuid().GetHashCode()).Next(1, 17)); + byte[] expected = first.Concat(second).ToArray(); + + using (ArrayPoolBufferWriter writer = new(pool, first.Length)) + { + first.CopyTo(writer.GetSpan(first.Length)); + writer.Advance(first.Length); + second.CopyTo(writer.GetSpan(writer.FreeCapacity + 1)); + writer.Advance(second.Length); + + if (emptyActiveBuffer) + { + _ = writer.GetMemory(writer.FreeCapacity + 1); + } + + ReadOnlySequence sequence = writer.GetReadOnlySequence(); + byte[][] rentedArrays = pool.RentedArrays.ToArray(); + int capacity = writer.Capacity; + int freeCapacity = writer.FreeCapacity; + Type viewType = typeof(ArrayPoolBufferWriter).Assembly.GetType("CommunityToolkit.HighPerformance.Buffers.Views.MemoryDebugView`1")!.MakeGenericType(typeof(byte)); + object view = Activator.CreateInstance(viewType, writer)!; + byte[] items = (byte[])viewType.GetProperty("Items")!.GetValue(view)!; + + CollectionAssert.AreEqual(expected, items); + CollectionAssert.AreEqual(expected, sequence.ToArray()); + Assert.AreEqual(expected.Length, writer.WrittenCount); + Assert.AreEqual(capacity, writer.Capacity); + Assert.AreEqual(freeCapacity, writer.FreeCapacity); + Assert.HasCount(rentedArrays.Length, pool.RentedArrays); + + foreach (byte[] array in rentedArrays) + { + Assert.IsTrue(pool.RentedArrays.Contains(array)); + } + } + + Assert.IsEmpty(pool.RentedArrays); + } + [TestMethod] public void Test_ArrayPoolBufferWriterOfT_AllocateAndGetArray() { @@ -243,4 +715,13 @@ public void Test_ArrayPoolBufferWriterOfT_AllocateAndGetArray() _ = Assert.ThrowsExactly(() => bufferWriter.DangerousGetArray()); } + + private static byte[] GetRandomBytes(int length) + { + byte[] data = new byte[length]; + + new Random(Guid.NewGuid().GetHashCode()).NextBytes(data); + + return data; + } }