Skip to content

Commit

Permalink
[release/5.0] Use ArrayPool as default pool fallback (#36559)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sebastienros
sebastienros committed Sep 17, 2021
1 parent cf7a1d4 commit d5ae64d
Show file tree
Hide file tree
Showing 6 changed files with 294 additions and 14 deletions.
50 changes: 46 additions & 4 deletions src/Servers/Kestrel/Core/src/Internal/Http/Http1OutputProducer.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ internal class Http1OutputProducer : IHttpOutputProducer, IDisposable

private readonly ConcurrentPipeWriter _pipeWriter;
private IMemoryOwner<byte> _fakeMemoryOwner;
private byte[] _fakeMemory;

// Chunked responses need to be treated uniquely when using GetMemory + Advance.
// We need to know the size of the data written to the chunk before calling Advance on the
Expand Down Expand Up @@ -419,6 +420,12 @@ public void Dispose()
_fakeMemoryOwner = null;
}

if (_fakeMemory != null)
{
ArrayPool<byte>.Shared.Return(_fakeMemory);
_fakeMemory = null;
}

// Call dispose on any memory that wasn't written.
if (_completedSegments != null)
{
Expand Down Expand Up @@ -656,13 +663,48 @@ private void WriteCurrentChunkMemoryToPipeWriter(ref BufferWriter<PipeWriter> wr
_advancedBytesForChunk = 0;
}

private Memory<byte> GetFakeMemory(int sizeHint)
internal Memory<byte> GetFakeMemory(int minSize)
{
if (_fakeMemoryOwner == null)
// Try to reuse _fakeMemoryOwner
if (_fakeMemoryOwner != null)
{
if (_fakeMemoryOwner.Memory.Length < minSize)
{
_fakeMemoryOwner.Dispose();
_fakeMemoryOwner = null;
}
else
{
return _fakeMemoryOwner.Memory;
}
}

// Try to reuse _fakeMemory
if (_fakeMemory != null)
{
if (_fakeMemory.Length < minSize)
{
ArrayPool<byte>.Shared.Return(_fakeMemory);
_fakeMemory = null;
}
else
{
return _fakeMemory;
}
}

// Requesting a bigger buffer could throw.
if (minSize <= _memoryPool.MaxBufferSize)
{
// Use the specified pool as it fits.
_fakeMemoryOwner = _memoryPool.Rent(minSize);
return _fakeMemoryOwner.Memory;
}
else
{
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
// Use the array pool. Its MaxBufferSize is int.MaxValue.
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
}
return _fakeMemoryOwner.Memory;
}

private Memory<byte> LeasedMemory(int sizeHint)
Expand Down
49 changes: 45 additions & 4 deletions src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProducer.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ internal class Http2OutputProducer : IHttpOutputProducer, IHttpOutputAborter, IV
private readonly PipeReader _pipeReader;
private readonly ManualResetValueTaskSource<object> _resetAwaitable = new ManualResetValueTaskSource<object>();
private IMemoryOwner<byte> _fakeMemoryOwner;
private byte[] _fakeMemory;
private bool _startedWritingDataFrames;
private bool _streamCompleted;
private bool _suffixSent;
Expand Down Expand Up @@ -119,6 +120,12 @@ public void Complete()
_fakeMemoryOwner.Dispose();
_fakeMemoryOwner = null;
}

if (_fakeMemory != null)
{
ArrayPool<byte>.Shared.Return(_fakeMemory);
_fakeMemory = null;
}
}
}

Expand Down Expand Up @@ -485,14 +492,48 @@ static void ThrowUnexpectedState()
}
}

private Memory<byte> GetFakeMemory(int sizeHint)
internal Memory<byte> GetFakeMemory(int minSize)
{
if (_fakeMemoryOwner == null)
// Try to reuse _fakeMemoryOwner
if (_fakeMemoryOwner != null)
{
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
if (_fakeMemoryOwner.Memory.Length < minSize)
{
_fakeMemoryOwner.Dispose();
_fakeMemoryOwner = null;
}
else
{
return _fakeMemoryOwner.Memory;
}
}

return _fakeMemoryOwner.Memory;
// Try to reuse _fakeMemory
if (_fakeMemory != null)
{
if (_fakeMemory.Length < minSize)
{
ArrayPool<byte>.Shared.Return(_fakeMemory);
_fakeMemory = null;
}
else
{
return _fakeMemory;
}
}

// Requesting a bigger buffer could throw.
if (minSize <= _memoryPool.MaxBufferSize)
{
// Use the specified pool as it fits.
_fakeMemoryOwner = _memoryPool.Rent(minSize);
return _fakeMemoryOwner.Memory;
}
else
{
// Use the array pool. Its MaxBufferSize is int.MaxValue.
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
}
}

[StackTraceHidden]
Expand Down
49 changes: 45 additions & 4 deletions src/Servers/Kestrel/Core/src/Internal/Http3/Http3OutputProducer.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ internal class Http3OutputProducer : IHttpOutputProducer, IHttpOutputAborter
private bool _disposed;
private bool _suffixSent;
private IMemoryOwner<byte> _fakeMemoryOwner;
private byte[] _fakeMemory;

public Http3OutputProducer(
Http3FrameWriter frameWriter,
Expand Down Expand Up @@ -72,6 +73,12 @@ public void Dispose()
_fakeMemoryOwner.Dispose();
_fakeMemoryOwner = null;
}

if (_fakeMemory != null)
{
ArrayPool<byte>.Shared.Return(_fakeMemory);
_fakeMemory = null;
}
}
}

Expand Down Expand Up @@ -192,14 +199,48 @@ public Span<byte> GetSpan(int sizeHint = 0)
}
}

private Memory<byte> GetFakeMemory(int sizeHint)
internal Memory<byte> GetFakeMemory(int minSize)
{
if (_fakeMemoryOwner == null)
// Try to reuse _fakeMemoryOwner
if (_fakeMemoryOwner != null)
{
if (_fakeMemoryOwner.Memory.Length < minSize)
{
_fakeMemoryOwner.Dispose();
_fakeMemoryOwner = null;
}
else
{
return _fakeMemoryOwner.Memory;
}
}

// Try to reuse _fakeMemory
if (_fakeMemory != null)
{
_fakeMemoryOwner = _memoryPool.Rent(sizeHint);
if (_fakeMemory.Length < minSize)
{
ArrayPool<byte>.Shared.Return(_fakeMemory);
_fakeMemory = null;
}
else
{
return _fakeMemory;
}
}

return _fakeMemoryOwner.Memory;
// Requesting a bigger buffer could throw.
if (minSize <= _memoryPool.MaxBufferSize)
{
// Use the specified pool as it fits.
_fakeMemoryOwner = _memoryPool.Rent(minSize);
return _fakeMemoryOwner.Memory;
}
else
{
// Use the array pool. Its MaxBufferSize is int.MaxValue.
return _fakeMemory = ArrayPool<byte>.Shared.Rent(minSize);
}
}

[StackTraceHidden]
Expand Down
87 changes: 85 additions & 2 deletions .../Kestrel/Core/test/OutputProducerTests.cs → ...rel/Core/test/Http1OutputProducerTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,11 @@

namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
{
public class OutputProducerTests : IDisposable
public class Http1OutputProducerTests : IDisposable
{
private readonly MemoryPool<byte> _memoryPool;

public OutputProducerTests()
public Http1OutputProducerTests()
{
_memoryPool = SlabMemoryPoolFactory.Create();
}
Expand Down Expand Up @@ -74,6 +74,89 @@ public void AbortsTransportEvenAfterDispose()
mockConnectionContext.Verify(f => f.Abort(null), Times.Once());
}

[Fact]
public void AllocatesFakeMemorySmallerThanMaxBufferSize()
{
var pipeOptions = new PipeOptions
(
pool: _memoryPool,
readerScheduler: Mock.Of<PipeScheduler>(),
writerScheduler: PipeScheduler.Inline,
useSynchronizationContext: false
);

using (var socketOutput = CreateOutputProducer(pipeOptions))
{
var bufferSize = 1;
var fakeMemory = socketOutput.GetFakeMemory(bufferSize);

Assert.True(fakeMemory.Length >= bufferSize);
}
}

[Fact]
public void AllocatesFakeMemoryBiggerThanMaxBufferSize()
{
var pipeOptions = new PipeOptions
(
pool: _memoryPool,
readerScheduler: Mock.Of<PipeScheduler>(),
writerScheduler: PipeScheduler.Inline,
useSynchronizationContext: false
);

using (var socketOutput = CreateOutputProducer(pipeOptions))
{
var bufferSize = _memoryPool.MaxBufferSize * 2;
var fakeMemory = socketOutput.GetFakeMemory(bufferSize);

Assert.True(fakeMemory.Length >= bufferSize);
}
}

[Fact]
public void AllocatesIncreasingFakeMemory()
{
var pipeOptions = new PipeOptions
(
pool: _memoryPool,
readerScheduler: Mock.Of<PipeScheduler>(),
writerScheduler: PipeScheduler.Inline,
useSynchronizationContext: false
);

using (var socketOutput = CreateOutputProducer(pipeOptions))
{
var bufferSize1 = 1024;
var bufferSize2 = 2048;
var fakeMemory = socketOutput.GetFakeMemory(bufferSize1);
fakeMemory = socketOutput.GetFakeMemory(bufferSize2);

Assert.True(fakeMemory.Length >= bufferSize2);
}
}

[Fact]
public void ReusesFakeMemory()
{
var pipeOptions = new PipeOptions
(
pool: _memoryPool,
readerScheduler: Mock.Of<PipeScheduler>(),
writerScheduler: PipeScheduler.Inline,
useSynchronizationContext: false
);

using (var socketOutput = CreateOutputProducer(pipeOptions))
{
var bufferSize = 1024;
var fakeMemory1 = socketOutput.GetFakeMemory(bufferSize);
var fakeMemory2 = socketOutput.GetFakeMemory(bufferSize);

Assert.Equal(fakeMemory1, fakeMemory2);
}
}

private TestHttpOutputProducer CreateOutputProducer(
PipeOptions pipeOptions = null,
ConnectionContext connectionContext = null)
Expand Down
48 changes: 48 additions & 0 deletions src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2StreamTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4785,5 +4785,53 @@ public async Task HEADERS_Received_Latin1_RejectedWhenLatin1OptionIsNotConfigure
expectedErrorCode: Http2ErrorCode.PROTOCOL_ERROR,
expectedErrorMessage: CoreStrings.BadRequest_MalformedRequestInvalidHeaders);
}

[Theory]
[InlineData(1000)]
[InlineData(4096)]
[InlineData(8000)] // Greater than the default max pool size (4096)
public async Task GetMemory_AfterAbort_GetsFakeMemory(int sizeHint)
{
var headers = new[]
{
new KeyValuePair<string, string>(HeaderNames.Method, "GET"),
new KeyValuePair<string, string>(HeaderNames.Path, "/"),
new KeyValuePair<string, string>(HeaderNames.Scheme, "http"),
};
await InitializeConnectionAsync(async httpContext =>
{
var response = httpContext.Response;
await response.BodyWriter.FlushAsync();
httpContext.Abort();
var memory = response.BodyWriter.GetMemory(sizeHint);
Assert.True(memory.Length >= sizeHint);
var fisrtPartOfResponse = Encoding.ASCII.GetBytes(new String('a', sizeHint));
fisrtPartOfResponse.CopyTo(memory);
response.BodyWriter.Advance(sizeHint);
});

await StartStreamAsync(1, headers, endStream: true);

var headersFrame = await ExpectAsync(Http2FrameType.HEADERS,
withLength: 32,
withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
withStreamId: 1);
await ExpectAsync(Http2FrameType.RST_STREAM,
withLength: 4,
withFlags: (byte)Http2DataFrameFlags.NONE,
withStreamId: 1);

await StopConnectionAsync(expectedLastStreamId: 1, ignoreNonGoAwayFrames: false);

_hpackDecoder.Decode(headersFrame.PayloadSequence, endHeaders: false, handler: this);

Assert.Equal(2, _decodedHeaders.Count);
Assert.Contains("date", _decodedHeaders.Keys, StringComparer.OrdinalIgnoreCase);
Assert.Equal("200", _decodedHeaders[HeaderNames.Status]);
}
}
}
Loading

0 comments on commit d5ae64d

Please sign in to comment.