Skip to content

Commit d454419

Browse files
authored
Switch Json's PooledByteBufferWriter to shared ArrayBuffer helper (#111348)
* 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
1 parent 5a395ed commit d454419

25 files changed

+58
-361
lines changed
Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
11
// Licensed to the .NET Foundation under one or more agreements.
22
// The .NET Foundation licenses this file to you under the MIT license.
33

4-
using System;
54
using System.Buffers;
6-
using System.Collections.Generic;
75
using System.Runtime.InteropServices;
8-
using System.Text;
96
using System.Threading;
107
using System.Threading.Tasks;
118

12-
namespace System.IO.Pipelines
9+
namespace System.IO
1310
{
1411
// Helpers to write Memory<byte> to Stream on netstandard 2.0
1512
internal static class StreamExtensions
Lines changed: 26 additions & 182 deletions
Original file line numberDiff line numberDiff line change
@@ -1,249 +1,93 @@
11
// Licensed to the .NET Foundation under one or more agreements.
22
// The .NET Foundation licenses this file to you under the MIT license.
33

4-
using System.Buffers;
54
using System.Diagnostics;
6-
using System.Diagnostics.CodeAnalysis;
75
using System.IO;
86
using System.IO.Pipelines;
9-
using System.Runtime.CompilerServices;
7+
using System.Net;
108
using System.Threading;
119
using System.Threading.Tasks;
1210

1311
namespace System.Text.Json
1412
{
1513
internal sealed class PooledByteBufferWriter : PipeWriter, IDisposable
1614
{
17-
// This class allows two possible configurations: if rentedBuffer is not null then
18-
// it can be used as an IBufferWriter and holds a buffer that should eventually be
19-
// returned to the shared pool. If rentedBuffer is null, then the instance is in a
20-
// cleared/disposed state and it must re-rent a buffer before it can be used again.
21-
private byte[]? _rentedBuffer;
22-
private int _index;
23-
private readonly Stream? _stream;
24-
2515
private const int MinimumBufferSize = 256;
2616

27-
// Value copied from Array.MaxLength in System.Private.CoreLib/src/libraries/System.Private.CoreLib/src/System/Array.cs.
28-
public const int MaximumBufferSize = 0X7FFFFFC7;
29-
30-
private PooledByteBufferWriter()
31-
{
32-
#if NET
33-
// Ensure we are in sync with the Array.MaxLength implementation.
34-
Debug.Assert(MaximumBufferSize == Array.MaxLength);
35-
#endif
36-
}
17+
private ArrayBuffer _buffer;
18+
private readonly Stream? _stream;
3719

38-
public PooledByteBufferWriter(int initialCapacity) : this()
20+
public PooledByteBufferWriter(int initialCapacity)
3921
{
40-
Debug.Assert(initialCapacity > 0);
41-
42-
_rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity);
43-
_index = 0;
22+
_buffer = new ArrayBuffer(initialCapacity, usePool: true);
4423
}
4524

4625
public PooledByteBufferWriter(int initialCapacity, Stream stream) : this(initialCapacity)
4726
{
4827
_stream = stream;
4928
}
5029

51-
public ReadOnlyMemory<byte> WrittenMemory
52-
{
53-
get
54-
{
55-
Debug.Assert(_rentedBuffer != null);
56-
Debug.Assert(_index <= _rentedBuffer.Length);
57-
return _rentedBuffer.AsMemory(0, _index);
58-
}
59-
}
30+
public ReadOnlySpan<byte> WrittenSpan => _buffer.ActiveSpan;
6031

61-
public int WrittenCount
62-
{
63-
get
64-
{
65-
Debug.Assert(_rentedBuffer != null);
66-
return _index;
67-
}
68-
}
32+
public ReadOnlyMemory<byte> WrittenMemory => _buffer.ActiveMemory;
6933

70-
public int Capacity
71-
{
72-
get
73-
{
74-
Debug.Assert(_rentedBuffer != null);
75-
return _rentedBuffer.Length;
76-
}
77-
}
34+
public int Capacity => _buffer.Capacity;
7835

79-
public int FreeCapacity
80-
{
81-
get
82-
{
83-
Debug.Assert(_rentedBuffer != null);
84-
return _rentedBuffer.Length - _index;
85-
}
86-
}
87-
88-
public void Clear()
89-
{
90-
ClearHelper();
91-
}
36+
public void Clear() => _buffer.Discard(_buffer.ActiveLength);
9237

93-
public void ClearAndReturnBuffers()
94-
{
95-
Debug.Assert(_rentedBuffer != null);
38+
public void ClearAndReturnBuffers() => _buffer.ClearAndReturnBuffer();
9639

97-
ClearHelper();
98-
byte[] toReturn = _rentedBuffer;
99-
_rentedBuffer = null;
100-
ArrayPool<byte>.Shared.Return(toReturn);
101-
}
102-
103-
private void ClearHelper()
104-
{
105-
Debug.Assert(_rentedBuffer != null);
106-
Debug.Assert(_index <= _rentedBuffer.Length);
107-
108-
_rentedBuffer.AsSpan(0, _index).Clear();
109-
_index = 0;
110-
}
111-
112-
// Returns the rented buffer back to the pool
113-
public void Dispose()
114-
{
115-
if (_rentedBuffer == null)
116-
{
117-
return;
118-
}
119-
120-
ClearHelper();
121-
byte[] toReturn = _rentedBuffer;
122-
_rentedBuffer = null;
123-
ArrayPool<byte>.Shared.Return(toReturn);
124-
}
40+
public void Dispose() => _buffer.Dispose();
12541

12642
public void InitializeEmptyInstance(int initialCapacity)
12743
{
12844
Debug.Assert(initialCapacity > 0);
129-
Debug.Assert(_rentedBuffer is null);
45+
Debug.Assert(_buffer.ActiveLength == 0);
13046

131-
_rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity);
132-
_index = 0;
47+
_buffer.EnsureAvailableSpace(initialCapacity);
13348
}
13449

135-
public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter();
50+
public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter(initialCapacity: 0);
13651

137-
public override void Advance(int count)
138-
{
139-
Debug.Assert(_rentedBuffer != null);
140-
Debug.Assert(count >= 0);
141-
Debug.Assert(_index <= _rentedBuffer.Length - count);
142-
_index += count;
143-
}
52+
public override void Advance(int count) => _buffer.Commit(count);
14453

14554
public override Memory<byte> GetMemory(int sizeHint = MinimumBufferSize)
14655
{
147-
CheckAndResizeBuffer(sizeHint);
148-
return _rentedBuffer.AsMemory(_index);
56+
Debug.Assert(sizeHint > 0);
57+
58+
_buffer.EnsureAvailableSpace(sizeHint);
59+
return _buffer.AvailableMemory;
14960
}
15061

15162
public override Span<byte> GetSpan(int sizeHint = MinimumBufferSize)
15263
{
153-
CheckAndResizeBuffer(sizeHint);
154-
return _rentedBuffer.AsSpan(_index);
64+
Debug.Assert(sizeHint > 0);
65+
66+
_buffer.EnsureAvailableSpace(sizeHint);
67+
return _buffer.AvailableSpan;
15568
}
15669

15770
#if NET
158-
internal void WriteToStream(Stream destination)
159-
{
160-
destination.Write(WrittenMemory.Span);
161-
}
71+
internal void WriteToStream(Stream destination) => destination.Write(_buffer.ActiveSpan);
16272
#else
163-
internal void WriteToStream(Stream destination)
164-
{
165-
Debug.Assert(_rentedBuffer != null);
166-
destination.Write(_rentedBuffer, 0, _index);
167-
}
73+
internal void WriteToStream(Stream destination) => destination.Write(_buffer.ActiveMemory);
16874
#endif
16975

170-
private void CheckAndResizeBuffer(int sizeHint)
171-
{
172-
Debug.Assert(_rentedBuffer != null);
173-
Debug.Assert(sizeHint > 0);
174-
175-
int currentLength = _rentedBuffer.Length;
176-
int availableSpace = currentLength - _index;
177-
178-
// If we've reached ~1GB written, grow to the maximum buffer
179-
// length to avoid incessant minimal growths causing perf issues.
180-
if (_index >= MaximumBufferSize / 2)
181-
{
182-
sizeHint = Math.Max(sizeHint, MaximumBufferSize - currentLength);
183-
}
184-
185-
if (sizeHint > availableSpace)
186-
{
187-
int growBy = Math.Max(sizeHint, currentLength);
188-
189-
int newSize = currentLength + growBy;
190-
191-
if ((uint)newSize > MaximumBufferSize)
192-
{
193-
newSize = currentLength + sizeHint;
194-
if ((uint)newSize > MaximumBufferSize)
195-
{
196-
ThrowHelper.ThrowOutOfMemoryException_BufferMaximumSizeExceeded((uint)newSize);
197-
}
198-
}
199-
200-
byte[] oldBuffer = _rentedBuffer;
201-
202-
_rentedBuffer = ArrayPool<byte>.Shared.Rent(newSize);
203-
204-
Debug.Assert(oldBuffer.Length >= _index);
205-
Debug.Assert(_rentedBuffer.Length >= _index);
206-
207-
Span<byte> oldBufferAsSpan = oldBuffer.AsSpan(0, _index);
208-
oldBufferAsSpan.CopyTo(_rentedBuffer);
209-
oldBufferAsSpan.Clear();
210-
ArrayPool<byte>.Shared.Return(oldBuffer);
211-
}
212-
213-
Debug.Assert(_rentedBuffer.Length - _index > 0);
214-
Debug.Assert(_rentedBuffer.Length - _index >= sizeHint);
215-
}
216-
21776
public override async ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default)
21877
{
21978
Debug.Assert(_stream is not null);
220-
#if NET
22179
await _stream.WriteAsync(WrittenMemory, cancellationToken).ConfigureAwait(false);
222-
#else
223-
Debug.Assert(_rentedBuffer != null);
224-
await _stream.WriteAsync(_rentedBuffer, 0, _index, cancellationToken).ConfigureAwait(false);
225-
#endif
22680
Clear();
22781

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

23185
public override bool CanGetUnflushedBytes => true;
232-
public override long UnflushedBytes => _index;
86+
public override long UnflushedBytes => _buffer.ActiveLength;
23387

23488
// This type is used internally in JsonSerializer to help buffer and flush bytes to the underlying Stream.
23589
// It's only pretending to be a PipeWriter and doesn't need Complete or CancelPendingFlush for the internal usage.
23690
public override void CancelPendingFlush() => throw new NotImplementedException();
23791
public override void Complete(Exception? exception = null) => throw new NotImplementedException();
23892
}
239-
240-
internal static partial class ThrowHelper
241-
{
242-
[DoesNotReturn]
243-
[MethodImpl(MethodImplOptions.NoInlining)]
244-
public static void ThrowOutOfMemoryException_BufferMaximumSizeExceeded(uint capacity)
245-
{
246-
throw new OutOfMemoryException(SR.Format(SR.BufferMaximumSizeExceeded, capacity));
247-
}
248-
}
24993
}

src/libraries/Common/tests/System/IO/StreamSpanExtensions.netstandard.cs

Lines changed: 0 additions & 38 deletions
This file was deleted.

src/libraries/Common/tests/System/Net/Http/DefaultCredentialsTest.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// The .NET Foundation licenses this file to you under the MIT license.
33

44
using System.Collections.Generic;
5+
using System.IO;
56
using System.Security.Principal;
67
using System.Threading.Tasks;
78

src/libraries/Common/tests/System/Net/StreamArrayExtensions.cs

Lines changed: 0 additions & 35 deletions
This file was deleted.

src/libraries/Microsoft.Extensions.Logging.Console/src/JsonConsoleFormatter.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ private void WriteInternal(IExternalScopeProvider? scopeProvider, TextWriter tex
9696
writer.Flush();
9797
}
9898

99-
var messageBytes = output.WrittenMemory.Span;
99+
var messageBytes = output.WrittenSpan;
100100
var logMessageBuffer = ArrayPool<char>.Shared.Rent(Encoding.UTF8.GetMaxCharCount(messageBytes.Length));
101101
try
102102
{

0 commit comments

Comments
 (0)