Skip to content

Commit

Permalink
Switch Json's PooledByteBufferWriter to shared ArrayBuffer helper (#1…
Browse files Browse the repository at this point in the history 
…11348)

* Switch Json's PooledByteBufferWriter to use ArrayBuffer

* Revert useless change

* Delete unused error message

* Merge a couple Stream polyfills

* Commit the moved file 🤦

* Use the polyfills in more places
  • Loading branch information
@MihaZupan
MihaZupan authored Jan 15, 2025
1 parent 5a395ed commit d454419
Show file tree
Hide file tree
Showing 25 changed files with 58 additions and 361 deletions.
5 changes: 1 addition & 4 deletions ...Pipelines/StreamExtensions.netstandard.cs → ...System/IO/StreamExtensions.netstandard.cs
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,12 @@
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace System.IO.Pipelines
namespace System.IO
{
// Helpers to write Memory<byte> to Stream on netstandard 2.0
internal static class StreamExtensions
Expand Down
208 changes: 26 additions & 182 deletions src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs
View file
Original file line number Diff line number Diff line change
@@ -1,249 +1,93 @@
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System.Buffers;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.IO.Pipelines;
using System.Runtime.CompilerServices;
using System.Net;
using System.Threading;
using System.Threading.Tasks;

namespace System.Text.Json
{
internal sealed class PooledByteBufferWriter : PipeWriter, IDisposable
{
// This class allows two possible configurations: if rentedBuffer is not null then
// it can be used as an IBufferWriter and holds a buffer that should eventually be
// returned to the shared pool. If rentedBuffer is null, then the instance is in a
// cleared/disposed state and it must re-rent a buffer before it can be used again.
private byte[]? _rentedBuffer;
private int _index;
private readonly Stream? _stream;

private const int MinimumBufferSize = 256;

// Value copied from Array.MaxLength in System.Private.CoreLib/src/libraries/System.Private.CoreLib/src/System/Array.cs.
public const int MaximumBufferSize = 0X7FFFFFC7;

private PooledByteBufferWriter()
{
#if NET
// Ensure we are in sync with the Array.MaxLength implementation.
Debug.Assert(MaximumBufferSize == Array.MaxLength);
#endif
}
private ArrayBuffer _buffer;
private readonly Stream? _stream;

public PooledByteBufferWriter(int initialCapacity) : this()
public PooledByteBufferWriter(int initialCapacity)
{
Debug.Assert(initialCapacity > 0);

_rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity);
_index = 0;
_buffer = new ArrayBuffer(initialCapacity, usePool: true);
}

public PooledByteBufferWriter(int initialCapacity, Stream stream) : this(initialCapacity)
{
_stream = stream;
}

public ReadOnlyMemory<byte> WrittenMemory
{
get
{
Debug.Assert(_rentedBuffer != null);
Debug.Assert(_index <= _rentedBuffer.Length);
return _rentedBuffer.AsMemory(0, _index);
}
}
public ReadOnlySpan<byte> WrittenSpan => _buffer.ActiveSpan;

public int WrittenCount
{
get
{
Debug.Assert(_rentedBuffer != null);
return _index;
}
}
public ReadOnlyMemory<byte> WrittenMemory => _buffer.ActiveMemory;

public int Capacity
{
get
{
Debug.Assert(_rentedBuffer != null);
return _rentedBuffer.Length;
}
}
public int Capacity => _buffer.Capacity;

public int FreeCapacity
{
get
{
Debug.Assert(_rentedBuffer != null);
return _rentedBuffer.Length - _index;
}
}

public void Clear()
{
ClearHelper();
}
public void Clear() => _buffer.Discard(_buffer.ActiveLength);

public void ClearAndReturnBuffers()
{
Debug.Assert(_rentedBuffer != null);
public void ClearAndReturnBuffers() => _buffer.ClearAndReturnBuffer();

ClearHelper();
byte[] toReturn = _rentedBuffer;
_rentedBuffer = null;
ArrayPool<byte>.Shared.Return(toReturn);
}

private void ClearHelper()
{
Debug.Assert(_rentedBuffer != null);
Debug.Assert(_index <= _rentedBuffer.Length);

_rentedBuffer.AsSpan(0, _index).Clear();
_index = 0;
}

// Returns the rented buffer back to the pool
public void Dispose()
{
if (_rentedBuffer == null)
{
return;
}

ClearHelper();
byte[] toReturn = _rentedBuffer;
_rentedBuffer = null;
ArrayPool<byte>.Shared.Return(toReturn);
}
public void Dispose() => _buffer.Dispose();

public void InitializeEmptyInstance(int initialCapacity)
{
Debug.Assert(initialCapacity > 0);
Debug.Assert(_rentedBuffer is null);
Debug.Assert(_buffer.ActiveLength == 0);

_rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity);
_index = 0;
_buffer.EnsureAvailableSpace(initialCapacity);
}

public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter();
public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter(initialCapacity: 0);

public override void Advance(int count)
{
Debug.Assert(_rentedBuffer != null);
Debug.Assert(count >= 0);
Debug.Assert(_index <= _rentedBuffer.Length - count);
_index += count;
}
public override void Advance(int count) => _buffer.Commit(count);

public override Memory<byte> GetMemory(int sizeHint = MinimumBufferSize)
{
CheckAndResizeBuffer(sizeHint);
return _rentedBuffer.AsMemory(_index);
Debug.Assert(sizeHint > 0);

_buffer.EnsureAvailableSpace(sizeHint);
return _buffer.AvailableMemory;
}

public override Span<byte> GetSpan(int sizeHint = MinimumBufferSize)
{
CheckAndResizeBuffer(sizeHint);
return _rentedBuffer.AsSpan(_index);
Debug.Assert(sizeHint > 0);

_buffer.EnsureAvailableSpace(sizeHint);
return _buffer.AvailableSpan;
}

#if NET
internal void WriteToStream(Stream destination)
{
destination.Write(WrittenMemory.Span);
}
internal void WriteToStream(Stream destination) => destination.Write(_buffer.ActiveSpan);
#else
internal void WriteToStream(Stream destination)
{
Debug.Assert(_rentedBuffer != null);
destination.Write(_rentedBuffer, 0, _index);
}
internal void WriteToStream(Stream destination) => destination.Write(_buffer.ActiveMemory);
#endif

private void CheckAndResizeBuffer(int sizeHint)
{
Debug.Assert(_rentedBuffer != null);
Debug.Assert(sizeHint > 0);

int currentLength = _rentedBuffer.Length;
int availableSpace = currentLength - _index;

// If we've reached ~1GB written, grow to the maximum buffer
// length to avoid incessant minimal growths causing perf issues.
if (_index >= MaximumBufferSize / 2)
{
sizeHint = Math.Max(sizeHint, MaximumBufferSize - currentLength);
}

if (sizeHint > availableSpace)
{
int growBy = Math.Max(sizeHint, currentLength);

int newSize = currentLength + growBy;

if ((uint)newSize > MaximumBufferSize)
{
newSize = currentLength + sizeHint;
if ((uint)newSize > MaximumBufferSize)
{
ThrowHelper.ThrowOutOfMemoryException_BufferMaximumSizeExceeded((uint)newSize);
}
}

byte[] oldBuffer = _rentedBuffer;

_rentedBuffer = ArrayPool<byte>.Shared.Rent(newSize);

Debug.Assert(oldBuffer.Length >= _index);
Debug.Assert(_rentedBuffer.Length >= _index);

Span<byte> oldBufferAsSpan = oldBuffer.AsSpan(0, _index);
oldBufferAsSpan.CopyTo(_rentedBuffer);
oldBufferAsSpan.Clear();
ArrayPool<byte>.Shared.Return(oldBuffer);
}

Debug.Assert(_rentedBuffer.Length - _index > 0);
Debug.Assert(_rentedBuffer.Length - _index >= sizeHint);
}

public override async ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default)
{
Debug.Assert(_stream is not null);
#if NET
await _stream.WriteAsync(WrittenMemory, cancellationToken).ConfigureAwait(false);
#else
Debug.Assert(_rentedBuffer != null);
await _stream.WriteAsync(_rentedBuffer, 0, _index, cancellationToken).ConfigureAwait(false);
#endif
Clear();

return new FlushResult(isCanceled: false, isCompleted: false);
}

public override bool CanGetUnflushedBytes => true;
public override long UnflushedBytes => _index;
public override long UnflushedBytes => _buffer.ActiveLength;

// This type is used internally in JsonSerializer to help buffer and flush bytes to the underlying Stream.
// It's only pretending to be a PipeWriter and doesn't need Complete or CancelPendingFlush for the internal usage.
public override void CancelPendingFlush() => throw new NotImplementedException();
public override void Complete(Exception? exception = null) => throw new NotImplementedException();
}

internal static partial class ThrowHelper
{
[DoesNotReturn]
[MethodImpl(MethodImplOptions.NoInlining)]
public static void ThrowOutOfMemoryException_BufferMaximumSizeExceeded(uint capacity)
{
throw new OutOfMemoryException(SR.Format(SR.BufferMaximumSizeExceeded, capacity));
}
}
}
38 changes: 0 additions & 38 deletions src/libraries/Common/tests/System/IO/StreamSpanExtensions.netstandard.cs
View file

This file was deleted.

1 change: 1 addition & 0 deletions src/libraries/Common/tests/System/Net/Http/DefaultCredentialsTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// The .NET Foundation licenses this file to you under the MIT license.

using System.Collections.Generic;
using System.IO;
using System.Security.Principal;
using System.Threading.Tasks;

Expand Down
35 changes: 0 additions & 35 deletions src/libraries/Common/tests/System/Net/StreamArrayExtensions.cs
View file

This file was deleted.

2 changes: 1 addition & 1 deletion src/libraries/Microsoft.Extensions.Logging.Console/src/JsonConsoleFormatter.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ private void WriteInternal(IExternalScopeProvider? scopeProvider, TextWriter tex
writer.Flush();
}

var messageBytes = output.WrittenMemory.Span;
var messageBytes = output.WrittenSpan;
var logMessageBuffer = ArrayPool<char>.Shared.Rent(Encoding.UTF8.GetMaxCharCount(messageBytes.Length));
try
{
Expand Down
Loading

0 comments on commit d454419

Please sign in to comment.