HTTP/3: Support canceling requests that aren't reading a body (dotnet#36017)

Author: JamesNK

Diff for: src/Servers/Connections.Abstractions/src/Features/IProtocolErrorCodeFeature.cs

@@ -9,7 +9,7 @@ namespace Microsoft.AspNetCore.Connections.Features
     public interface IProtocolErrorCodeFeature
     {
         /// <summary>
-        /// Gets or sets the error code.
+        /// Gets or sets the error code. The property returns -1 if the error code hasn't been set.
         /// </summary>
         long Error { get; set; }
     }

Diff for: src/Servers/Kestrel/Core/src/Internal/Http3/Http3Stream.cs

@@ -133,6 +133,11 @@ public void InitializeWithExistingContext(IDuplexPipe transport)
         }
 
         public void Abort(ConnectionAbortedException abortReason, Http3ErrorCode errorCode)
+        {
+            AbortCore(abortReason, errorCode);
+        }
+
+        private void AbortCore(Exception exception, Http3ErrorCode errorCode)
         {
             lock (_completionLock)
             {
@@ -148,26 +153,27 @@ public void Abort(ConnectionAbortedException abortReason, Http3ErrorCode errorCo
                     return;
                 }
 
-                Log.Http3StreamAbort(TraceIdentifier, errorCode, abortReason);
+                if (!(exception is ConnectionAbortedException abortReason))
+                {
+                    abortReason = new ConnectionAbortedException(exception.Message, exception);
+                }
 
-                _errorCodeFeature.Error = (long)errorCode;
-                _frameWriter.Abort(abortReason);
+                Log.Http3StreamAbort(TraceIdentifier, errorCode, abortReason);
 
-                AbortCore(abortReason);
-            }
-        }
+                // Call _http3Output.Stop() prior to poisoning the request body stream or pipe to
+                // ensure that an app that completes early due to the abort doesn't result in header frames being sent.
+                _http3Output.Stop();
 
-        private void AbortCore(Exception exception)
-        {
-            // Call _http3Output.Stop() prior to poisoning the request body stream or pipe to
-            // ensure that an app that completes early due to the abort doesn't result in header frames being sent.
-            _http3Output.Stop();
+                CancelRequestAbortedToken();
 
-            CancelRequestAbortedToken();
+                // Unblock the request body.
+                PoisonBody(exception);
+                RequestBodyPipe.Writer.Complete(exception);
 
-            // Unblock the request body.
-            PoisonBody(exception);
-            RequestBodyPipe.Writer.Complete(exception);
+                // Abort framewriter and underlying transport after stopping output.
+                _errorCodeFeature.Error = (long)errorCode;
+                _frameWriter.Abort(abortReason);
+            }
         }
 
         protected override void OnErrorAfterResponseStarted()
@@ -470,7 +476,8 @@ public async Task ProcessRequestAsync<TContext>(IHttpApplication<TContext> appli
             catch (ConnectionResetException ex)
             {
                 error = ex;
-                AbortCore(new IOException(CoreStrings.HttpStreamResetByClient, ex));
+                var resolvedErrorCode = _errorCodeFeature.Error >= 0 ? _errorCodeFeature.Error : 0;
+                AbortCore(new IOException(CoreStrings.HttpStreamResetByClient, ex), (Http3ErrorCode)resolvedErrorCode);
             }
             catch (Exception ex)
             {
@@ -484,10 +491,44 @@ public async Task ProcessRequestAsync<TContext>(IHttpApplication<TContext> appli
 
                 await Input.CompleteAsync();
 
-                // Make sure application func is completed before completing writer.
-                if (_appCompleted != null)
+                var appCompleted = _appCompleted?.Task ?? Task.CompletedTask;
+                if (!appCompleted.IsCompletedSuccessfully)
                 {
-                    await _appCompleted.Task;
+                    // At this point in the stream's read-side is complete. However, with HTTP/3
+                    // the write-side of the stream can still be aborted by the client on request
+                    // aborted.
+                    //
+                    // To get notification of request aborted we register to connection closed
+                    // token. It will notify this type that the client has aborted the request
+                    // and Kestrel will complete pipes and cancel the RequestAborted token.
+                    //
+                    // Only subscribe to this event after the stream's read-side is complete to
+                    // avoid interactions between reading that is in-progress and an abort.
+                    // This means while reading, read-side abort will handle getting abort notifications.
+                    //
+                    // We don't need to hang on to the CancellationTokenRegistration from register.
+                    // The CTS is cleaned up in StreamContext.DisposeAsync.
+                    //
+                    // TODO: Consider a better way to provide this notification. For perf we want to
+                    // make the ConnectionClosed CTS pay-for-play, and change this event to use
+                    // something that is more lightweight than a CTS.
+                    _context.StreamContext.ConnectionClosed.Register(static s =>
+                    {
+                        var stream = (Http3Stream)s!;
+
+                        if (!stream.IsCompleted)
+                        {
+                            // An error code value other than -1 indicates a value was set and the request didn't gracefully complete.
+                            var errorCode = stream._errorCodeFeature.Error;
+                            if (errorCode >= 0)
+                            {
+                                stream.AbortCore(new IOException(CoreStrings.HttpStreamResetByClient), (Http3ErrorCode)errorCode);
+                            }
+                        }
+                    }, this);
+
+                    // Make sure application func is completed before completing writer.
+                    await appCompleted;
                 }
 
                 try

Diff for: src/Servers/Kestrel/Transport.Quic/src/Internal/IQuicTrace.cs

@@ -19,7 +19,8 @@ internal interface IQuicTrace : ILogger
         void StreamPause(QuicStreamContext streamContext);
         void StreamResume(QuicStreamContext streamContext);
         void StreamShutdownWrite(QuicStreamContext streamContext, string reason);
-        void StreamAborted(QuicStreamContext streamContext, long errorCode, Exception ex);
+        void StreamAbortedRead(QuicStreamContext streamContext, long errorCode);
+        void StreamAbortedWrite(QuicStreamContext streamContext, long errorCode);
         void StreamAbort(QuicStreamContext streamContext, long errorCode, string reason);
         void StreamAbortRead(QuicStreamContext streamContext, long errorCode, string reason);
         void StreamAbortWrite(QuicStreamContext streamContext, long errorCode, string reason);

Diff for: src/Servers/Kestrel/Transport.Quic/src/Internal/QuicConnectionContext.FeatureCollection.cs

@@ -16,8 +16,13 @@ internal sealed partial class QuicConnectionContext : IProtocolErrorCodeFeature,
     {
         private X509Certificate2? _clientCert;
         private Task<X509Certificate2?>? _clientCertTask;
+        private long? _error;
 
-        public long Error { get; set; }
+        public long Error
+        {
+            get => _error ?? -1;
+            set => _error = value;
+        }
 
         public X509Certificate2? ClientCertificate
         {

Diff for: src/Servers/Kestrel/Transport.Quic/src/Internal/QuicConnectionContext.cs

@@ -83,9 +83,10 @@ public override void Abort(ConnectionAbortedException abortReason)
                     return;
                 }
 
+                var resolvedErrorCode = _error ?? 0;
                 _abortReason = ExceptionDispatchInfo.Capture(abortReason);
-                _log.ConnectionAbort(this, Error, abortReason.Message);
-                _closeTask = _connection.CloseAsync(errorCode: Error).AsTask();
+                _log.ConnectionAbort(this, resolvedErrorCode, abortReason.Message);
+                _closeTask = _connection.CloseAsync(errorCode: resolvedErrorCode).AsTask();
             }
         }
 
@@ -127,7 +128,7 @@ public override void Abort(ConnectionAbortedException abortReason)
             catch (QuicConnectionAbortedException ex)
             {
                 // Shutdown initiated by peer, abortive.
-                Error = ex.ErrorCode;
+                _error = ex.ErrorCode;
                 _log.ConnectionAborted(this, ex.ErrorCode, ex);
 
                 ThreadPool.UnsafeQueueUserWorkItem(state =>

Diff for: src/Servers/Kestrel/Transport.Quic/src/Internal/QuicStreamContext.FeatureCollection.cs

@@ -9,11 +9,16 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Quic.Internal
     internal sealed partial class QuicStreamContext : IPersistentStateFeature, IStreamDirectionFeature, IProtocolErrorCodeFeature, IStreamIdFeature, IStreamAbortFeature
     {
         private IDictionary<object, object?>? _persistentState;
+        private long? _error;
 
         public bool CanRead { get; private set; }
         public bool CanWrite { get; private set; }
 
-        public long Error { get; set; }
+        public long Error
+        {
+            get => _error ?? -1;
+            set => _error = value;
+        }
 
         public long StreamId { get; private set; }