Skip to content

Commit a57a046

Browse files
github-actions[bot]sebastienroshalter73
authored
[release/6.0] Use ArrayPool as default pool fallback (#36253)
* Use ArrayPool as SlabMemoryPool fallback Fixes #30364 * Fix comments * Typos [ci skip] * Support increase buffer size * Add function tests, rename argument * Update src/Servers/Kestrel/Core/src/Internal/Http3/Http3OutputProducer.cs Co-authored-by: Stephen Halter <halter73@gmail.com> * Update src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProducer.cs Co-authored-by: Stephen Halter <halter73@gmail.com> * Reduce allocations Co-authored-by: Sebastien Ros <sebastienros@gmail.com> Co-authored-by: Stephen Halter <halter73@gmail.com>
1 parent 66f9857 commit a57a046

File tree

7 files changed

+298
-18
lines changed

7 files changed

+298
-18
lines changed

‎src/Servers/Kestrel/Core/src/Internal/Http/Http1OutputProducer.cs

Lines changed: 46 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ internal class Http1OutputProducer : IHttpOutputProducer, IDisposable
4949

5050
private readonly ConcurrentPipeWriter _pipeWriter;
5151
private IMemoryOwner<byte>? _fakeMemoryOwner;
52+
private byte[]? _fakeMemory;
5253

5354
// Chunked responses need to be treated uniquely when using GetMemory + Advance.
5455
// We need to know the size of the data written to the chunk before calling Advance on the
@@ -414,6 +415,12 @@ public void Dispose()
414415
_fakeMemoryOwner = null;
415416
}
416417

418+
if (_fakeMemory != null)
419+
{
420+
ArrayPool<byte>.Shared.Return(_fakeMemory);
421+
_fakeMemory = null;
422+
}
423+
417424
// Call dispose on any memory that wasn't written.
418425
if (_completedSegments != null)
419426
{
@@ -650,13 +657,48 @@ private void WriteCurrentChunkMemoryToPipeWriter(ref BufferWriter<PipeWriter> wr
650657
_advancedBytesForChunk = 0;
651658
}
652659

653-
private Memory<byte> GetFakeMemory(int sizeHint)
660+
internal Memory<byte> GetFakeMemory(int minSize)
654661
{
655-
if (_fakeMemoryOwner == null)
662+
// Try to reuse _fakeMemoryOwner
663+
if (_fakeMemoryOwner != null)
664+
{
665+
if (_fakeMemoryOwner.Memory.Length < minSize)
666+
{
667+
_fakeMemoryOwner.Dispose();
668+
_fakeMemoryOwner = null;
669+
}
670+
else
671+
{
672+
return _fakeMemoryOwner.Memory;
673+
}
674+
}
675+
676+
// Try to reuse _fakeMemory
677+
if (_fakeMemory != null)
678+
{
679+
if (_fakeMemory.Length < minSize)
680+
{
681+
ArrayPool<byte>.Shared.Return(_fakeMemory);
682+
_fakeMemory = null;
683+
}
684+
else
685+
{
686+
return _fakeMemory;
687+
}
688+
}
689+
690+
// Requesting a bigger buffer could throw.
691+
if (minSize <= _memoryPool.MaxBufferSize)
692+
{
693+
// Use the specified pool as it fits.
694+
_fakeMemoryOwner = _memoryPool.Rent(minSize);
695+
return _fakeMemoryOwner.Memory;
696+
}
697+
else
656698
{
657-
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
699+
// Use the array pool. Its MaxBufferSize is int.MaxValue.
700+
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
658701
}
659-
return _fakeMemoryOwner.Memory;
660702
}
661703

662704
private Memory<byte> LeasedMemory(int sizeHint)

‎src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProducer.cs

Lines changed: 45 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ internal class Http2OutputProducer : IHttpOutputProducer, IHttpOutputAborter, IV
3636
private readonly PipeReader _pipeReader;
3737
private readonly ManualResetValueTaskSource<object?> _resetAwaitable = new ManualResetValueTaskSource<object?>();
3838
private IMemoryOwner<byte>? _fakeMemoryOwner;
39+
private byte[]? _fakeMemory;
3940
private bool _startedWritingDataFrames;
4041
private bool _streamCompleted;
4142
private bool _suffixSent;
@@ -120,6 +121,12 @@ public void Complete()
120121
_fakeMemoryOwner.Dispose();
121122
_fakeMemoryOwner = null;
122123
}
124+
125+
if (_fakeMemory != null)
126+
{
127+
ArrayPool<byte>.Shared.Return(_fakeMemory);
128+
_fakeMemory = null;
129+
}
123130
}
124131
}
125132

@@ -486,14 +493,48 @@ static void ThrowUnexpectedState()
486493
}
487494
}
488495

489-
private Memory<byte> GetFakeMemory(int sizeHint)
496+
internal Memory<byte> GetFakeMemory(int minSize)
490497
{
491-
if (_fakeMemoryOwner == null)
498+
// Try to reuse _fakeMemoryOwner
499+
if (_fakeMemoryOwner != null)
500+
{
501+
if (_fakeMemoryOwner.Memory.Length < minSize)
502+
{
503+
_fakeMemoryOwner.Dispose();
504+
_fakeMemoryOwner = null;
505+
}
506+
else
507+
{
508+
return _fakeMemoryOwner.Memory;
509+
}
510+
}
511+
512+
// Try to reuse _fakeMemory
513+
if (_fakeMemory != null)
492514
{
493-
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
515+
if (_fakeMemory.Length < minSize)
516+
{
517+
ArrayPool<byte>.Shared.Return(_fakeMemory);
518+
_fakeMemory = null;
519+
}
520+
else
521+
{
522+
return _fakeMemory;
523+
}
494524
}
495525

496-
return _fakeMemoryOwner.Memory;
526+
// Requesting a bigger buffer could throw.
527+
if (minSize <= _memoryPool.MaxBufferSize)
528+
{
529+
// Use the specified pool as it fits.
530+
_fakeMemoryOwner = _memoryPool.Rent(minSize);
531+
return _fakeMemoryOwner.Memory;
532+
}
533+
else
534+
{
535+
// Use the array pool. Its MaxBufferSize is int.MaxValue.
536+
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
537+
}
497538
}
498539

499540
[StackTraceHidden]

‎src/Servers/Kestrel/Core/src/Internal/Http3/Http3OutputProducer.cs

Lines changed: 45 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ internal class Http3OutputProducer : IHttpOutputProducer, IHttpOutputAborter
3434
private bool _disposed;
3535
private bool _suffixSent;
3636
private IMemoryOwner<byte>? _fakeMemoryOwner;
37+
private byte[]? _fakeMemory;
3738

3839
public Http3OutputProducer(
3940
Http3FrameWriter frameWriter,
@@ -88,6 +89,12 @@ public void Dispose()
8889
_fakeMemoryOwner.Dispose();
8990
_fakeMemoryOwner = null;
9091
}
92+
93+
if (_fakeMemory != null)
94+
{
95+
ArrayPool<byte>.Shared.Return(_fakeMemory);
96+
_fakeMemory = null;
97+
}
9198
}
9299
}
93100

@@ -208,14 +215,48 @@ public Span<byte> GetSpan(int sizeHint = 0)
208215
}
209216
}
210217

211-
private Memory<byte> GetFakeMemory(int sizeHint)
218+
internal Memory<byte> GetFakeMemory(int minSize)
212219
{
213-
if (_fakeMemoryOwner == null)
220+
// Try to reuse _fakeMemoryOwner
221+
if (_fakeMemoryOwner != null)
222+
{
223+
if (_fakeMemoryOwner.Memory.Length < minSize)
224+
{
225+
_fakeMemoryOwner.Dispose();
226+
_fakeMemoryOwner = null;
227+
}
228+
else
229+
{
230+
return _fakeMemoryOwner.Memory;
231+
}
232+
}
233+
234+
// Try to reuse _fakeMemory
235+
if (_fakeMemory != null)
214236
{
215-
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
237+
if (_fakeMemory.Length < minSize)
238+
{
239+
ArrayPool<byte>.Shared.Return(_fakeMemory);
240+
_fakeMemory = null;
241+
}
242+
else
243+
{
244+
return _fakeMemory;
245+
}
216246
}
217247

218-
return _fakeMemoryOwner.Memory;
248+
// Requesting a bigger buffer could throw.
249+
if (minSize <= _memoryPool.MaxBufferSize)
250+
{
251+
// Use the specified pool as it fits.
252+
_fakeMemoryOwner = _memoryPool.Rent(minSize);
253+
return _fakeMemoryOwner.Memory;
254+
}
255+
else
256+
{
257+
// Use the array pool. Its MaxBufferSize is int.MaxValue.
258+
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
259+
}
219260
}
220261

221262
[StackTraceHidden]

‎src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/ConcurrentPipeWriter.cs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -346,19 +346,19 @@ private void AllocateMemoryUnsynchronized(int sizeHint)
346346
}
347347
}
348348

349-
private BufferSegment AllocateSegmentUnsynchronized(int sizeHint)
349+
private BufferSegment AllocateSegmentUnsynchronized(int minSize)
350350
{
351351
BufferSegment newSegment = CreateSegmentUnsynchronized();
352352

353-
if (sizeHint <= _pool.MaxBufferSize)
353+
if (minSize <= _pool.MaxBufferSize)
354354
{
355355
// Use the specified pool if it fits
356-
newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(sizeHint, _pool.MaxBufferSize)));
356+
newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(minSize, _pool.MaxBufferSize)));
357357
}
358358
else
359359
{
360360
// We can't use the recommended pool so use the ArrayPool
361-
newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeHint));
361+
newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(minSize));
362362
}
363363

364364
_tailMemory = newSegment.AvailableMemory;

‎src/Servers/Kestrel/Core/test/OutputProducerTests.cs renamed to ‎src/Servers/Kestrel/Core/test/Http1OutputProducerTests.cs

Lines changed: 85 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,11 @@
1515

1616
namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
1717
{
18-
public class OutputProducerTests : IDisposable
18+
public class Http1OutputProducerTests : IDisposable
1919
{
2020
private readonly MemoryPool<byte> _memoryPool;
2121

22-
public OutputProducerTests()
22+
public Http1OutputProducerTests()
2323
{
2424
_memoryPool = PinnedBlockMemoryPoolFactory.Create();
2525
}
@@ -74,6 +74,89 @@ public void AbortsTransportEvenAfterDispose()
7474
mockConnectionContext.Verify(f => f.Abort(null), Times.Once());
7575
}
7676

77+
[Fact]
78+
public void AllocatesFakeMemorySmallerThanMaxBufferSize()
79+
{
80+
var pipeOptions = new PipeOptions
81+
(
82+
pool: _memoryPool,
83+
readerScheduler: Mock.Of<PipeScheduler>(),
84+
writerScheduler: PipeScheduler.Inline,
85+
useSynchronizationContext: false
86+
);
87+
88+
using (var socketOutput = CreateOutputProducer(pipeOptions))
89+
{
90+
var bufferSize = 1;
91+
var fakeMemory = socketOutput.GetFakeMemory(bufferSize);
92+
93+
Assert.True(fakeMemory.Length >= bufferSize);
94+
}
95+
}
96+
97+
[Fact]
98+
public void AllocatesFakeMemoryBiggerThanMaxBufferSize()
99+
{
100+
var pipeOptions = new PipeOptions
101+
(
102+
pool: _memoryPool,
103+
readerScheduler: Mock.Of<PipeScheduler>(),
104+
writerScheduler: PipeScheduler.Inline,
105+
useSynchronizationContext: false
106+
);
107+
108+
using (var socketOutput = CreateOutputProducer(pipeOptions))
109+
{
110+
var bufferSize = _memoryPool.MaxBufferSize * 2;
111+
var fakeMemory = socketOutput.GetFakeMemory(bufferSize);
112+
113+
Assert.True(fakeMemory.Length >= bufferSize);
114+
}
115+
}
116+
117+
[Fact]
118+
public void AllocatesIncreasingFakeMemory()
119+
{
120+
var pipeOptions = new PipeOptions
121+
(
122+
pool: _memoryPool,
123+
readerScheduler: Mock.Of<PipeScheduler>(),
124+
writerScheduler: PipeScheduler.Inline,
125+
useSynchronizationContext: false
126+
);
127+
128+
using (var socketOutput = CreateOutputProducer(pipeOptions))
129+
{
130+
var bufferSize1 = 1024;
131+
var bufferSize2 = 2048;
132+
var fakeMemory = socketOutput.GetFakeMemory(bufferSize1);
133+
fakeMemory = socketOutput.GetFakeMemory(bufferSize2);
134+
135+
Assert.True(fakeMemory.Length >= bufferSize2);
136+
}
137+
}
138+
139+
[Fact]
140+
public void ReusesFakeMemory()
141+
{
142+
var pipeOptions = new PipeOptions
143+
(
144+
pool: _memoryPool,
145+
readerScheduler: Mock.Of<PipeScheduler>(),
146+
writerScheduler: PipeScheduler.Inline,
147+
useSynchronizationContext: false
148+
);
149+
150+
using (var socketOutput = CreateOutputProducer(pipeOptions))
151+
{
152+
var bufferSize = 1024;
153+
var fakeMemory1 = socketOutput.GetFakeMemory(bufferSize);
154+
var fakeMemory2 = socketOutput.GetFakeMemory(bufferSize);
155+
156+
Assert.Equal(fakeMemory1, fakeMemory2);
157+
}
158+
}
159+
77160
private TestHttpOutputProducer CreateOutputProducer(
78161
PipeOptions pipeOptions = null,
79162
ConnectionContext connectionContext = null)

0 commit comments

Comments
 (0)