[release/5.0] Use ArrayPool as default pool fallback

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

@@ -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
@@ -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)
                 {
@@ -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)

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

@@ -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;
@@ -119,6 +120,12 @@ public void Complete()
                     _fakeMemoryOwner.Dispose();
                     _fakeMemoryOwner = null;
                 }
+
+                if (_fakeMemory != null)
+                {
+                    ArrayPool<byte>.Shared.Return(_fakeMemory);
+                    _fakeMemory = null;
+                }                
             }
         }
 
@@ -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]

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

@@ -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,
@@ -72,6 +73,12 @@ public void Dispose()
                     _fakeMemoryOwner.Dispose();
                     _fakeMemoryOwner = null;
                 }
+
+                if (_fakeMemory != null)
+                {
+                    ArrayPool<byte>.Shared.Return(_fakeMemory);
+                    _fakeMemory = null;
+                }
             }
         }
 
@@ -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]

src/Servers/Kestrel/Core/test/OutputProducerTests.cs → src/Servers/Kestrel/Core/test/Http1OutputProducerTests.cs

@@ -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();
         }
@@ -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)

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

@@ -4785,5 +4785,53 @@ await WaitForConnectionErrorAsync<Http2ConnectionErrorException>(
                 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]);
+        }
     }
 }