Reuse Http2MessageBody

Author: JamesNK

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

@@ -17,8 +17,8 @@ internal abstract class Http1MessageBody : MessageBody
         protected bool _completed;
 
         protected Http1MessageBody(Http1Connection context)
-            : base(context)
         {
+            Reset(context);
             _context = context;
         }
 

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

@@ -15,7 +15,7 @@ internal abstract class MessageBody
         private static readonly MessageBody _zeroContentLengthClose = new ZeroContentLengthMessageBody(keepAlive: false);
         private static readonly MessageBody _zeroContentLengthKeepAlive = new ZeroContentLengthMessageBody(keepAlive: true);
 
-        private readonly HttpProtocol _context;
+        private HttpProtocol _context;
 
         private bool _send100Continue = true;
         private long _consumedBytes;
@@ -26,11 +26,6 @@ internal abstract class MessageBody
         protected long _alreadyTimedBytes;
         protected long _examinedUnconsumedBytes;
 
-        protected MessageBody(HttpProtocol context)
-        {
-            _context = context;
-        }
-
         public static MessageBody ZeroContentLengthClose => _zeroContentLengthClose;
 
         public static MessageBody ZeroContentLengthKeepAlive => _zeroContentLengthKeepAlive;
@@ -73,6 +68,18 @@ public virtual Task StopAsync()
 
         protected virtual Task OnStopAsync() => Task.CompletedTask;
 
+        protected void Reset(HttpProtocol context)
+        {
+            _context = context;
+            _send100Continue = true;
+            _consumedBytes = 0;
+            _stopped = false;
+            _timingEnabled = false;
+            _backpressure = false;
+            _alreadyTimedBytes = 0;
+            _examinedUnconsumedBytes = 0;
+        }
+
         protected void TryProduceContinue()
         {
             if (_send100Continue)

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

@@ -11,8 +11,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
     internal sealed class ZeroContentLengthMessageBody : MessageBody
     {
         public ZeroContentLengthMessageBody(bool keepAlive)
-            : base(null)
         {
+            Reset(null);
             RequestKeepAlive = keepAlive;
         }
 

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

@@ -12,13 +12,12 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
 {
     internal sealed class Http2MessageBody : MessageBody
     {
-        private readonly Http2Stream _context;
+        private Http2Stream _context;
         private ReadResult _readResult;
 
-        private Http2MessageBody(Http2Stream context)
-            : base(context)
+        public Http2MessageBody(Http2Stream context)
         {
-            _context = context;
+            Reset(context);
         }
 
         protected override void OnReadStarting()
@@ -46,16 +45,29 @@ protected override void OnDataRead(long bytesRead)
             AddAndCheckConsumedBytes(bytesRead);
         }
 
-        public static MessageBody For(Http2Stream context)
+        public static MessageBody For(Http2Stream context, Http2MessageBody previousMessageBody)
         {
             if (context.ReceivedEmptyRequestBody)
             {
                 return ZeroContentLengthClose;
             }
 
+            if (previousMessageBody is Http2MessageBody http2MessageBody)
+            {
+                http2MessageBody.Reset(context);
+                return http2MessageBody;
+            }
+
             return new Http2MessageBody(context);
         }
 
+        public void Reset(Http2Stream context)
+        {
+            base.Reset(context);
+            _context = context;
+            _readResult = default;
+        }
+
         public override void AdvanceTo(SequencePosition consumed)
         {
             AdvanceTo(consumed, consumed);

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

@@ -23,6 +23,7 @@ internal abstract partial class Http2Stream : HttpProtocol, IThreadPoolWorkItem
         private Http2OutputProducer _http2Output;
         private StreamInputFlowControl _inputFlowControl;
         private StreamOutputFlowControl _outputFlowControl;
+        private Http2MessageBody _messageBody;
 
         private bool _decrementCalled;
 
@@ -170,7 +171,23 @@ protected override string CreateRequestId()
             => StringUtilities.ConcatAsHexSuffix(ConnectionId, ':', (uint)StreamId);
 
         protected override MessageBody CreateMessageBody()
-            => Http2MessageBody.For(this);
+        {
+            if (ReceivedEmptyRequestBody)
+            {
+                return MessageBody.ZeroContentLengthClose;
+            }
+
+            if (_messageBody != null)
+            {
+                _messageBody.Reset(this);
+            }
+            else
+            {
+                _messageBody = new Http2MessageBody(this);
+            }
+
+            return _messageBody;
+        }
 
         // Compare to Http1Connection.OnStartLine
         protected override bool TryParseRequest(ReadResult result, out bool endConnection)

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

@@ -16,8 +16,8 @@ internal sealed class Http3MessageBody : MessageBody
         private ReadResult _readResult;
 
         private Http3MessageBody(Http3Stream context)
-            : base(context)
         {
+            Reset(context);
             _context = context;
         }
         protected override void OnReadStarting()

src/Servers/Kestrel/Core/test/BodyControlTests.cs

@@ -141,8 +141,8 @@ public async Task ResponsePipeThrowsObjectDisposedExceptionAfterStop()
         private class MockMessageBody : MessageBody
         {
             public MockMessageBody(bool upgradeable = false)
-                : base(null)
             {
+                Reset(null);
                 RequestUpgrade = upgradeable;
             }
 

src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs

@@ -886,7 +886,7 @@ public async Task ConsumesRequestWhenApplicationDoesNotConsumeIt()
                 var buffer = new byte[10];
                 await context.Response.Body.WriteAsync(buffer, 0, 10);
             });
-            var mockMessageBody = new Mock<MessageBody>(null);
+            var mockMessageBody = new Mock<MessageBody>();
             _http1Connection.NextMessageBody = mockMessageBody.Object;
 
             var requestProcessingTask = _http1Connection.ProcessRequestsAsync(httpApplication);

src/Servers/Kestrel/Core/test/HttpRequestStreamTests.cs

@@ -108,7 +108,7 @@ public async Task SynchronousReadsThrowIfDisallowedByIHttpBodyControlFeature()
 
             var mockBodyControl = new Mock<IHttpBodyControlFeature>();
             mockBodyControl.Setup(m => m.AllowSynchronousIO).Returns(() => allowSynchronousIO);
-            var mockMessageBody = new Mock<MessageBody>(null);
+            var mockMessageBody = new Mock<MessageBody>();
             mockMessageBody.Setup(m => m.ReadAsync(CancellationToken.None)).Returns(new ValueTask<ReadResult>(new ReadResult(default, isCanceled: false, isCompleted: true)));
 
             var pipeReader = new HttpRequestPipeReader();