|
1 | 1 | // Licensed to the .NET Foundation under one or more agreements.
|
2 | 2 | // The .NET Foundation licenses this file to you under the MIT license.
|
3 | 3 |
|
4 |
| -using System.Buffers; |
5 | 4 | using System.Diagnostics;
|
6 |
| -using System.Diagnostics.CodeAnalysis; |
7 | 5 | using System.IO;
|
8 | 6 | using System.IO.Pipelines;
|
9 |
| -using System.Runtime.CompilerServices; |
| 7 | +using System.Net; |
10 | 8 | using System.Threading;
|
11 | 9 | using System.Threading.Tasks;
|
12 | 10 |
|
13 | 11 | namespace System.Text.Json
|
14 | 12 | {
|
15 | 13 | internal sealed class PooledByteBufferWriter : PipeWriter, IDisposable
|
16 | 14 | {
|
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 |
| - |
25 | 15 | private const int MinimumBufferSize = 256;
|
26 | 16 |
|
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; |
37 | 19 |
|
38 |
| - public PooledByteBufferWriter(int initialCapacity) : this() |
| 20 | + public PooledByteBufferWriter(int initialCapacity) |
39 | 21 | {
|
40 |
| - Debug.Assert(initialCapacity > 0); |
41 |
| - |
42 |
| - _rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity); |
43 |
| - _index = 0; |
| 22 | + _buffer = new ArrayBuffer(initialCapacity, usePool: true); |
44 | 23 | }
|
45 | 24 |
|
46 | 25 | public PooledByteBufferWriter(int initialCapacity, Stream stream) : this(initialCapacity)
|
47 | 26 | {
|
48 | 27 | _stream = stream;
|
49 | 28 | }
|
50 | 29 |
|
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; |
60 | 31 |
|
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; |
69 | 33 |
|
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; |
78 | 35 |
|
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); |
92 | 37 |
|
93 |
| - public void ClearAndReturnBuffers() |
94 |
| - { |
95 |
| - Debug.Assert(_rentedBuffer != null); |
| 38 | + public void ClearAndReturnBuffers() => _buffer.ClearAndReturnBuffer(); |
96 | 39 |
|
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(); |
125 | 41 |
|
126 | 42 | public void InitializeEmptyInstance(int initialCapacity)
|
127 | 43 | {
|
128 | 44 | Debug.Assert(initialCapacity > 0);
|
129 |
| - Debug.Assert(_rentedBuffer is null); |
| 45 | + Debug.Assert(_buffer.ActiveLength == 0); |
130 | 46 |
|
131 |
| - _rentedBuffer = ArrayPool<byte>.Shared.Rent(initialCapacity); |
132 |
| - _index = 0; |
| 47 | + _buffer.EnsureAvailableSpace(initialCapacity); |
133 | 48 | }
|
134 | 49 |
|
135 |
| - public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter(); |
| 50 | + public static PooledByteBufferWriter CreateEmptyInstanceForCaching() => new PooledByteBufferWriter(initialCapacity: 0); |
136 | 51 |
|
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); |
144 | 53 |
|
145 | 54 | public override Memory<byte> GetMemory(int sizeHint = MinimumBufferSize)
|
146 | 55 | {
|
147 |
| - CheckAndResizeBuffer(sizeHint); |
148 |
| - return _rentedBuffer.AsMemory(_index); |
| 56 | + Debug.Assert(sizeHint > 0); |
| 57 | + |
| 58 | + _buffer.EnsureAvailableSpace(sizeHint); |
| 59 | + return _buffer.AvailableMemory; |
149 | 60 | }
|
150 | 61 |
|
151 | 62 | public override Span<byte> GetSpan(int sizeHint = MinimumBufferSize)
|
152 | 63 | {
|
153 |
| - CheckAndResizeBuffer(sizeHint); |
154 |
| - return _rentedBuffer.AsSpan(_index); |
| 64 | + Debug.Assert(sizeHint > 0); |
| 65 | + |
| 66 | + _buffer.EnsureAvailableSpace(sizeHint); |
| 67 | + return _buffer.AvailableSpan; |
155 | 68 | }
|
156 | 69 |
|
157 | 70 | #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); |
162 | 72 | #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); |
168 | 74 | #endif
|
169 | 75 |
|
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 |
| - |
217 | 76 | public override async ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default)
|
218 | 77 | {
|
219 | 78 | Debug.Assert(_stream is not null);
|
220 |
| -#if NET |
221 | 79 | 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 |
226 | 80 | Clear();
|
227 | 81 |
|
228 | 82 | return new FlushResult(isCanceled: false, isCompleted: false);
|
229 | 83 | }
|
230 | 84 |
|
231 | 85 | public override bool CanGetUnflushedBytes => true;
|
232 |
| - public override long UnflushedBytes => _index; |
| 86 | + public override long UnflushedBytes => _buffer.ActiveLength; |
233 | 87 |
|
234 | 88 | // This type is used internally in JsonSerializer to help buffer and flush bytes to the underlying Stream.
|
235 | 89 | // It's only pretending to be a PipeWriter and doesn't need Complete or CancelPendingFlush for the internal usage.
|
236 | 90 | public override void CancelPendingFlush() => throw new NotImplementedException();
|
237 | 91 | public override void Complete(Exception? exception = null) => throw new NotImplementedException();
|
238 | 92 | }
|
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 |
| - } |
249 | 93 | }
|
0 commit comments