Skip to content

Commit d5ae64d

Browse files
sebastienrossebastienros
authored
[release/5.0] Use ArrayPool as default pool fallback (#36559)
1 parent cf7a1d4 commit d5ae64d

File tree

6 files changed

+294
-14
lines changed

6 files changed

+294
-14
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
@@ -419,6 +420,12 @@ public void Dispose()
419420
_fakeMemoryOwner = null;
420421
}
421422

423+
if (_fakeMemory != null)
424+
{
425+
ArrayPool<byte>.Shared.Return(_fakeMemory);
426+
_fakeMemory = null;
427+
}
428+
422429
// Call dispose on any memory that wasn't written.
423430
if (_completedSegments != null)
424431
{
@@ -656,13 +663,48 @@ private void WriteCurrentChunkMemoryToPipeWriter(ref BufferWriter<PipeWriter> wr
656663
_advancedBytesForChunk = 0;
657664
}
658665

659-
private Memory<byte> GetFakeMemory(int sizeHint)
666+
internal Memory<byte> GetFakeMemory(int minSize)
660667
{
661-
if (_fakeMemoryOwner == null)
668+
// Try to reuse _fakeMemoryOwner
669+
if (_fakeMemoryOwner != null)
670+
{
671+
if (_fakeMemoryOwner.Memory.Length < minSize)
672+
{
673+
_fakeMemoryOwner.Dispose();
674+
_fakeMemoryOwner = null;
675+
}
676+
else
677+
{
678+
return _fakeMemoryOwner.Memory;
679+
}
680+
}
681+
682+
// Try to reuse _fakeMemory
683+
if (_fakeMemory != null)
684+
{
685+
if (_fakeMemory.Length < minSize)
686+
{
687+
ArrayPool<byte>.Shared.Return(_fakeMemory);
688+
_fakeMemory = null;
689+
}
690+
else
691+
{
692+
return _fakeMemory;
693+
}
694+
}
695+
696+
// Requesting a bigger buffer could throw.
697+
if (minSize <= _memoryPool.MaxBufferSize)
698+
{
699+
// Use the specified pool as it fits.
700+
_fakeMemoryOwner = _memoryPool.Rent(minSize);
701+
return _fakeMemoryOwner.Memory;
702+
}
703+
else
662704
{
663-
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
705+
// Use the array pool. Its MaxBufferSize is int.MaxValue.
706+
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
664707
}
665-
return _fakeMemoryOwner.Memory;
666708
}
667709

668710
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;
@@ -119,6 +120,12 @@ public void Complete()
119120
_fakeMemoryOwner.Dispose();
120121
_fakeMemoryOwner = null;
121122
}
123+
124+
if (_fakeMemory != null)
125+
{
126+
ArrayPool<byte>.Shared.Return(_fakeMemory);
127+
_fakeMemory = null;
128+
}
122129
}
123130
}
124131

@@ -485,14 +492,48 @@ static void ThrowUnexpectedState()
485492
}
486493
}
487494

488-
private Memory<byte> GetFakeMemory(int sizeHint)
495+
internal Memory<byte> GetFakeMemory(int minSize)
489496
{
490-
if (_fakeMemoryOwner == null)
497+
// Try to reuse _fakeMemoryOwner
498+
if (_fakeMemoryOwner != null)
491499
{
492-
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
500+
if (_fakeMemoryOwner.Memory.Length < minSize)
501+
{
502+
_fakeMemoryOwner.Dispose();
503+
_fakeMemoryOwner = null;
504+
}
505+
else
506+
{
507+
return _fakeMemoryOwner.Memory;
508+
}
493509
}
494510

495-
return _fakeMemoryOwner.Memory;
511+
// Try to reuse _fakeMemory
512+
if (_fakeMemory != null)
513+
{
514+
if (_fakeMemory.Length < minSize)
515+
{
516+
ArrayPool<byte>.Shared.Return(_fakeMemory);
517+
_fakeMemory = null;
518+
}
519+
else
520+
{
521+
return _fakeMemory;
522+
}
523+
}
524+
525+
// Requesting a bigger buffer could throw.
526+
if (minSize <= _memoryPool.MaxBufferSize)
527+
{
528+
// Use the specified pool as it fits.
529+
_fakeMemoryOwner = _memoryPool.Rent(minSize);
530+
return _fakeMemoryOwner.Memory;
531+
}
532+
else
533+
{
534+
// Use the array pool. Its MaxBufferSize is int.MaxValue.
535+
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
536+
}
496537
}
497538

498539
[StackTraceHidden]

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

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

3738
public Http3OutputProducer(
3839
Http3FrameWriter frameWriter,
@@ -72,6 +73,12 @@ public void Dispose()
7273
_fakeMemoryOwner.Dispose();
7374
_fakeMemoryOwner = null;
7475
}
76+
77+
if (_fakeMemory != null)
78+
{
79+
ArrayPool<byte>.Shared.Return(_fakeMemory);
80+
_fakeMemory = null;
81+
}
7582
}
7683
}
7784

@@ -192,14 +199,48 @@ public Span<byte> GetSpan(int sizeHint = 0)
192199
}
193200
}
194201

195-
private Memory<byte> GetFakeMemory(int sizeHint)
202+
internal Memory<byte> GetFakeMemory(int minSize)
196203
{
197-
if (_fakeMemoryOwner == null)
204+
// Try to reuse _fakeMemoryOwner
205+
if (_fakeMemoryOwner != null)
206+
{
207+
if (_fakeMemoryOwner.Memory.Length < minSize)
208+
{
209+
_fakeMemoryOwner.Dispose();
210+
_fakeMemoryOwner = null;
211+
}
212+
else
213+
{
214+
return _fakeMemoryOwner.Memory;
215+
}
216+
}
217+
218+
// Try to reuse _fakeMemory
219+
if (_fakeMemory != null)
198220
{
199-
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
221+
if (_fakeMemory.Length < minSize)
222+
{
223+
ArrayPool<byte>.Shared.Return(_fakeMemory);
224+
_fakeMemory = null;
225+
}
226+
else
227+
{
228+
return _fakeMemory;
229+
}
200230
}
201231

202-
return _fakeMemoryOwner.Memory;
232+
// Requesting a bigger buffer could throw.
233+
if (minSize <= _memoryPool.MaxBufferSize)
234+
{
235+
// Use the specified pool as it fits.
236+
_fakeMemoryOwner = _memoryPool.Rent(minSize);
237+
return _fakeMemoryOwner.Memory;
238+
}
239+
else
240+
{
241+
// Use the array pool. Its MaxBufferSize is int.MaxValue.
242+
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
243+
}
203244
}
204245

205246
[StackTraceHidden]

‎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 = SlabMemoryPoolFactory.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)

‎src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2StreamTests.cs

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4785,5 +4785,53 @@ await WaitForConnectionErrorAsync<Http2ConnectionErrorException>(
47854785
expectedErrorCode: Http2ErrorCode.PROTOCOL_ERROR,
47864786
expectedErrorMessage: CoreStrings.BadRequest_MalformedRequestInvalidHeaders);
47874787
}
4788+
4789+
[Theory]
4790+
[InlineData(1000)]
4791+
[InlineData(4096)]
4792+
[InlineData(8000)] // Greater than the default max pool size (4096)
4793+
public async Task GetMemory_AfterAbort_GetsFakeMemory(int sizeHint)
4794+
{
4795+
var headers = new[]
4796+
{
4797+
new KeyValuePair<string, string>(HeaderNames.Method, "GET"),
4798+
new KeyValuePair<string, string>(HeaderNames.Path, "/"),
4799+
new KeyValuePair<string, string>(HeaderNames.Scheme, "http"),
4800+
};
4801+
await InitializeConnectionAsync(async httpContext =>
4802+
{
4803+
var response = httpContext.Response;
4804+
4805+
await response.BodyWriter.FlushAsync();
4806+
4807+
httpContext.Abort();
4808+
4809+
var memory = response.BodyWriter.GetMemory(sizeHint);
4810+
Assert.True(memory.Length >= sizeHint);
4811+
4812+
var fisrtPartOfResponse = Encoding.ASCII.GetBytes(new String('a', sizeHint));
4813+
fisrtPartOfResponse.CopyTo(memory);
4814+
response.BodyWriter.Advance(sizeHint);
4815+
});
4816+
4817+
await StartStreamAsync(1, headers, endStream: true);
4818+
4819+
var headersFrame = await ExpectAsync(Http2FrameType.HEADERS,
4820+
withLength: 32,
4821+
withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
4822+
withStreamId: 1);
4823+
await ExpectAsync(Http2FrameType.RST_STREAM,
4824+
withLength: 4,
4825+
withFlags: (byte)Http2DataFrameFlags.NONE,
4826+
withStreamId: 1);
4827+
4828+
await StopConnectionAsync(expectedLastStreamId: 1, ignoreNonGoAwayFrames: false);
4829+
4830+
_hpackDecoder.Decode(headersFrame.PayloadSequence, endHeaders: false, handler: this);
4831+
4832+
Assert.Equal(2, _decodedHeaders.Count);
4833+
Assert.Contains("date", _decodedHeaders.Keys, StringComparer.OrdinalIgnoreCase);
4834+
Assert.Equal("200", _decodedHeaders[HeaderNames.Status]);
4835+
}
47884836
}
47894837
}

0 commit comments

Comments
 (0)