Implement Http/2 CompleteAsync

src/Http/Http.Features/ref/Microsoft.AspNetCore.Http.Features.netstandard2.0.cs

@@ -200,6 +200,10 @@ public partial interface IHttpRequestTrailersFeature
         bool Available { get; }
         Microsoft.AspNetCore.Http.IHeaderDictionary Trailers { get; }
     }
+    public partial interface IHttpResponseCompletionFeature
+    {
+        System.Threading.Tasks.Task CompleteAsync();
+    }
     public partial interface IHttpResponseFeature
     {
         System.IO.Stream Body { get; set; }

src/Http/Http.Features/src/IHttpResponseCompletionFeature.cs

@@ -0,0 +1,20 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Http.Features
+{
+    /// <summary>
+    /// A feature to gracefully end a response.
+    /// </summary>
+    public interface IHttpResponseCompletionFeature
+    {
+        /// <summary>
+        /// Flush any remaining response headers, data, or trailers.
+        /// This may throw if the response is in an invalid state such as a Content-Length mismatch.
+        /// </summary>
+        /// <returns></returns>
+        Task CompleteAsync();
+    }
+}

src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.FeatureCollection.cs

@@ -277,6 +277,7 @@ protected void ResetHttp1Features()
         protected void ResetHttp2Features()
         {
             _currentIHttp2StreamIdFeature = this;
+            _currentIHttpResponseCompletionFeature = this;
             _currentIHttpResponseTrailersFeature = this;
         }
 

src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.Generated.cs

@@ -29,6 +29,7 @@ internal partial class HttpProtocol : IFeatureCollection
         private static readonly Type IFormFeatureType = typeof(IFormFeature);
         private static readonly Type IHttpUpgradeFeatureType = typeof(IHttpUpgradeFeature);
         private static readonly Type IHttp2StreamIdFeatureType = typeof(IHttp2StreamIdFeature);
+        private static readonly Type IHttpResponseCompletionFeatureType = typeof(IHttpResponseCompletionFeature);
         private static readonly Type IHttpResponseTrailersFeatureType = typeof(IHttpResponseTrailersFeature);
         private static readonly Type IResponseCookiesFeatureType = typeof(IResponseCookiesFeature);
         private static readonly Type IItemsFeatureType = typeof(IItemsFeature);
@@ -58,6 +59,7 @@ internal partial class HttpProtocol : IFeatureCollection
         private object _currentIFormFeature;
         private object _currentIHttpUpgradeFeature;
         private object _currentIHttp2StreamIdFeature;
+        private object _currentIHttpResponseCompletionFeature;
         private object _currentIHttpResponseTrailersFeature;
         private object _currentIResponseCookiesFeature;
         private object _currentIItemsFeature;
@@ -98,6 +100,7 @@ private void FastReset()
             _currentIQueryFeature = null;
             _currentIFormFeature = null;
             _currentIHttp2StreamIdFeature = null;
+            _currentIHttpResponseCompletionFeature = null;
             _currentIHttpResponseTrailersFeature = null;
             _currentIResponseCookiesFeature = null;
             _currentIItemsFeature = null;
@@ -224,6 +227,10 @@ object IFeatureCollection.this[Type key]
                 {
                     feature = _currentIHttp2StreamIdFeature;
                 }
+                else if (key == IHttpResponseCompletionFeatureType)
+                {
+                    feature = _currentIHttpResponseCompletionFeature;
+                }
                 else if (key == IHttpResponseTrailersFeatureType)
                 {
                     feature = _currentIHttpResponseTrailersFeature;
@@ -348,6 +355,10 @@ object IFeatureCollection.this[Type key]
                 {
                     _currentIHttp2StreamIdFeature = value;
                 }
+                else if (key == IHttpResponseCompletionFeatureType)
+                {
+                    _currentIHttpResponseCompletionFeature = value;
+                }
                 else if (key == IHttpResponseTrailersFeatureType)
                 {
                     _currentIHttpResponseTrailersFeature = value;
@@ -470,6 +481,10 @@ TFeature IFeatureCollection.Get<TFeature>()
             {
                 feature = (TFeature)_currentIHttp2StreamIdFeature;
             }
+            else if (typeof(TFeature) == typeof(IHttpResponseCompletionFeature))
+            {
+                feature = (TFeature)_currentIHttpResponseCompletionFeature;
+            }
             else if (typeof(TFeature) == typeof(IHttpResponseTrailersFeature))
             {
                 feature = (TFeature)_currentIHttpResponseTrailersFeature;
@@ -598,6 +613,10 @@ void IFeatureCollection.Set<TFeature>(TFeature feature)
             {
                 _currentIHttp2StreamIdFeature = feature;
             }
+            else if (typeof(TFeature) == typeof(IHttpResponseCompletionFeature))
+            {
+                _currentIHttpResponseCompletionFeature = feature;
+            }
             else if (typeof(TFeature) == typeof(IHttpResponseTrailersFeature))
             {
                 _currentIHttpResponseTrailersFeature = feature;
@@ -718,6 +737,10 @@ private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()
             {
                 yield return new KeyValuePair<Type, object>(IHttp2StreamIdFeatureType, _currentIHttp2StreamIdFeature);
             }
+            if (_currentIHttpResponseCompletionFeature != null)
+            {
+                yield return new KeyValuePair<Type, object>(IHttpResponseCompletionFeatureType, _currentIHttpResponseCompletionFeature);
+            }
             if (_currentIHttpResponseTrailersFeature != null)
             {
                 yield return new KeyValuePair<Type, object>(IHttpResponseTrailersFeatureType, _currentIHttpResponseTrailersFeature);

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

@@ -210,6 +210,7 @@ private void HttpVersionSetSlow(string value)
         public bool RequestTrailersAvailable { get; set; }
         public Stream RequestBody { get; set; }
         public PipeReader RequestBodyPipeReader { get; set; }
+        public HttpResponseTrailers ResponseTrailers { get; set; }
 
         private int _statusCode;
         public int StatusCode
@@ -287,7 +288,9 @@ public CancellationToken RequestAborted
 
         public bool HasResponseStarted => _requestProcessingStatus >= RequestProcessingStatus.HeadersCommitted;
 
-        public bool HasFlushedHeaders => _requestProcessingStatus == RequestProcessingStatus.HeadersFlushed;
+        public bool HasFlushedHeaders => _requestProcessingStatus >= RequestProcessingStatus.HeadersFlushed;
+
+        public bool HasResponseCompleted => _requestProcessingStatus == RequestProcessingStatus.ResponseCompleted;
 
         protected HttpRequestHeaders HttpRequestHeaders { get; }
 
@@ -632,9 +635,18 @@ private async Task ProcessRequests<TContext>(IHttpApplication<TContext> applicat
                     // Run the application code for this request
                     await application.ProcessRequestAsync(context);
 
-                    if (!_connectionAborted)
+                    // Trigger OnStarting if it hasn't been called yet and the app hasn't
+                    // already failed. If an OnStarting callback throws we can go through
+                    // our normal error handling in ProduceEnd.
+                    // https://github.com/aspnet/KestrelHttpServer/issues/43
+                    if (!HasResponseStarted && _applicationException == null && _onStarting?.Count > 0)
                     {
-                        VerifyResponseContentLength();
+                        await FireOnStarting();
+                    }
+
+                    if (!_connectionAborted && !VerifyResponseContentLength(out var lengthException))
+                    {
+                        ReportApplicationError(lengthException);
                     }
                 }
                 catch (BadHttpRequestException ex)
@@ -652,15 +664,6 @@ private async Task ProcessRequests<TContext>(IHttpApplication<TContext> applicat
 
                 KestrelEventSource.Log.RequestStop(this);
 
-                // Trigger OnStarting if it hasn't been called yet and the app hasn't
-                // already failed. If an OnStarting callback throws we can go through
-                // our normal error handling in ProduceEnd.
-                // https://github.com/aspnet/KestrelHttpServer/issues/43
-                if (!HasResponseStarted && _applicationException == null && _onStarting?.Count > 0)
-                {
-                    await FireOnStarting();
-                }
-
                 // At this point all user code that needs use to the request or response streams has completed.
                 // Using these streams in the OnCompleted callback is not allowed.
                 StopBodies();
@@ -898,7 +901,7 @@ private void CheckLastWrite()
             }
         }
 
-        protected void VerifyResponseContentLength()
+        protected bool VerifyResponseContentLength(out Exception ex)
         {
             var responseHeaders = HttpResponseHeaders;
 
@@ -915,9 +918,13 @@ protected void VerifyResponseContentLength()
                     _keepAlive = false;
                 }
 
-                ReportApplicationError(new InvalidOperationException(
-                    CoreStrings.FormatTooFewBytesWritten(_responseBytesWritten, responseHeaders.ContentLength.Value)));
+                ex = new InvalidOperationException(
+                    CoreStrings.FormatTooFewBytesWritten(_responseBytesWritten, responseHeaders.ContentLength.Value));
+                return false;
             }
+
+            ex = null;
+            return true;
         }
 
         public void ProduceContinue()
@@ -1045,6 +1052,11 @@ protected Task ProduceEnd()
 
         private Task WriteSuffix()
         {
+            if (HasResponseCompleted)
+            {
+                return Task.CompletedTask;
+            }
+
             // _autoChunk should be checked after we are sure ProduceStart() has been called
             // since ProduceStart() may set _autoChunk to true.
             if (_autoChunk || _httpVersion == Http.HttpVersion.Http2)
@@ -1064,7 +1076,7 @@ private Task WriteSuffix()
 
             if (!HasFlushedHeaders)
             {
-                _requestProcessingStatus = RequestProcessingStatus.HeadersFlushed;
+                _requestProcessingStatus = RequestProcessingStatus.ResponseCompleted;
                 return FlushAsyncInternal();
             }
 
@@ -1080,6 +1092,8 @@ private async Task WriteSuffixAwaited()
 
             await Output.WriteStreamSuffixAsync();
 
+            _requestProcessingStatus = RequestProcessingStatus.ResponseCompleted;
+
             if (_keepAlive)
             {
                 Log.ConnectionKeepAlive(ConnectionId);
@@ -1244,6 +1258,7 @@ private void SetErrorResponseHeaders(int statusCode)
 
             var responseHeaders = HttpResponseHeaders;
             responseHeaders.Reset();
+            ResponseTrailers?.Reset();
             var dateHeaderValues = DateHeaderValueManager.GetDateHeaderValues();
 
             responseHeaders.SetRawDate(dateHeaderValues.String, dateHeaderValues.Bytes);

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

@@ -10,6 +10,7 @@ internal enum RequestProcessingStatus
         ParsingHeaders,
         AppStarted,
         HeadersCommitted,
-        HeadersFlushed
+        HeadersFlushed,
+        ResponseCompleted
     }
 }

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

@@ -151,7 +151,7 @@ public void WriteResponseHeaders(int statusCode, string ReasonPhrase, HttpRespon
                 // 2. There is no trailing HEADERS frame.
                 Http2HeadersFrameFlags http2HeadersFrame;
 
-                if (appCompleted && !_startedWritingDataFrames && (_stream.Trailers == null || _stream.Trailers.Count == 0))
+                if (appCompleted && !_startedWritingDataFrames && (_stream.ResponseTrailers == null || _stream.ResponseTrailers.Count == 0))
                 {
                     _streamEnded = true;
                     http2HeadersFrame = Http2HeadersFrameFlags.END_STREAM;
@@ -313,7 +313,7 @@ private async ValueTask<FlushResult> ProcessDataWrites()
                 {
                     readResult = await _dataPipe.Reader.ReadAsync();
 
-                    if (readResult.IsCompleted && _stream.Trailers?.Count > 0)
+                    if (readResult.IsCompleted && _stream.ResponseTrailers?.Count > 0)
                     {
                         // Output is ending and there are trailers to write
                         // Write any remaining content then write trailers
@@ -322,7 +322,8 @@ private async ValueTask<FlushResult> ProcessDataWrites()
                             flushResult = await _frameWriter.WriteDataAsync(_streamId, _flowControl, readResult.Buffer, endStream: false);
                         }
 
-                        flushResult = await _frameWriter.WriteResponseTrailers(_streamId, _stream.Trailers);
+                        _stream.ResponseTrailers.SetReadOnly();
+                        flushResult = await _frameWriter.WriteResponseTrailers(_streamId, _stream.ResponseTrailers);
                     }
                     else if (readResult.IsCompleted && _streamEnded)
                     {

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

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
@@ -11,21 +12,25 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
 {
     internal partial class Http2Stream : IHttp2StreamIdFeature,
                                          IHttpMinRequestBodyDataRateFeature,
+                                         IHttpResponseCompletionFeature,
                                          IHttpResponseTrailersFeature
 
     {
-        internal HttpResponseTrailers Trailers { get; set; }
         private IHeaderDictionary _userTrailers;
 
         IHeaderDictionary IHttpResponseTrailersFeature.Trailers
         {
             get
             {
-                if (Trailers == null)
+                if (ResponseTrailers == null)
                 {
-                    Trailers = new HttpResponseTrailers();
+                    ResponseTrailers = new HttpResponseTrailers();
+                    if (HasResponseCompleted)
+                    {
+                        ResponseTrailers.SetReadOnly();
+                    }
                 }
-                return _userTrailers ?? Trailers;
+                return _userTrailers ?? ResponseTrailers;
             }
             set
             {
@@ -48,5 +53,25 @@ MinDataRate IHttpMinRequestBodyDataRateFeature.MinDataRate
                 MinRequestBodyDataRate = value;
             }
         }
+
+        async Task IHttpResponseCompletionFeature.CompleteAsync()
+        {
+            // Finalize headers
+            if (!HasResponseStarted)
+            {
+                await FireOnStarting();
+            }
+
+            // Flush headers, body, trailers...
+            if (!HasResponseCompleted)
+            {
+                if (!VerifyResponseContentLength(out var lengthException))
+                {
+                    throw lengthException;
+                }
+
+                await ProduceEnd();
+            }
+        }
     }
 }