Skip to content

Use ArrayPool as default pool fallback #35719

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Sep 7, 2021
Merged

Use ArrayPool as default pool fallback #35719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
1cd3e96
Use ArrayPool as SlabMemoryPool fallback
sebastienros Aug 25, 2021
f8db646
Fix comments
sebastienros Aug 25, 2021
a17efe5
Typos
sebastienros Aug 25, 2021
a1358d0
Support increase buffer size
sebastienros Aug 27, 2021
0fe3535
Add function tests, rename argument
sebastienros Aug 28, 2021
c0ee4c0
Update src/Servers/Kestrel/Core/src/Internal/Http3/Http3OutputProduce…
sebastienros Sep 1, 2021
6348dc5
Update src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProduce…
sebastienros Sep 1, 2021
a58d167
Reduce allocations
sebastienros Sep 2, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 46 additions & 4 deletions src/Servers/Kestrel/Core/src/Internal/Http/Http1OutputProducer.cs
View file
Open in desktop
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 @@ -414,6 +415,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 @@ -650,13 +657,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
Open in desktop
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 @@ -120,6 +121,12 @@ public void Complete()
_fakeMemoryOwner.Dispose();
_fakeMemoryOwner = null;
}

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

Expand Down Expand Up @@ -486,14 +493,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)
{
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
49 changes: 45 additions & 4 deletions src/Servers/Kestrel/Core/src/Internal/Http3/Http3OutputProducer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,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 @@ -88,6 +89,12 @@ public void Dispose()
_fakeMemoryOwner.Dispose();
_fakeMemoryOwner = null;
}

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

Expand Down Expand Up @@ -208,14 +215,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
8 changes: 4 additions & 4 deletions ...ervers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/ConcurrentPipeWriter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -346,19 +346,19 @@ private void AllocateMemoryUnsynchronized(int sizeHint)
}
}

private BufferSegment AllocateSegmentUnsynchronized(int sizeHint)
private BufferSegment AllocateSegmentUnsynchronized(int minSize)
{
BufferSegment newSegment = CreateSegmentUnsynchronized();

if (sizeHint <= _pool.MaxBufferSize)
if (minSize <= _pool.MaxBufferSize)
{
// Use the specified pool if it fits
newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(sizeHint, _pool.MaxBufferSize)));
newSegment.SetOwnedMemory(_pool.Rent(GetSegmentSize(minSize, _pool.MaxBufferSize)));
}
else
{
// We can't use the recommended pool so use the ArrayPool
newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(sizeHint));
newSegment.SetOwnedMemory(ArrayPool<byte>.Shared.Rent(minSize));
}

_tailMemory = newSegment.AvailableMemory;
Expand Down
87 changes: 85 additions & 2 deletions .../Kestrel/Core/test/OutputProducerTests.cs → ...rel/Core/test/Http1OutputProducerTests.cs
View file
Open in desktop
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 = PinnedBlockMemoryPoolFactory.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
Loading