HTTP/3 response drain timeout

Author: JamesNK

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

@@ -176,38 +176,64 @@ public void Tick(DateTimeOffset now)
                 return;
             }
 
-            UpdateStartingStreams(now);
+            UpdateStreamTimeouts(now);
         }
 
-        private void UpdateStartingStreams(DateTimeOffset now)
+        private void UpdateStreamTimeouts(DateTimeOffset now)
         {
+            // This method checks for timeouts:
+            // 1. When a stream first starts and waits to receive headers.
+            //    Uses RequestHeadersTimeout.
+            // 2. When a stream finished and is waiting for underlying transport to drain.
+            //    Uses MinResponseDataRate.
+
             var ticks = now.Ticks;
 
             lock (_streams)
             {
                 foreach (var stream in _streams.Values)
                 {
-                    if (stream.ReceivedHeader)
+                    if (stream.IsReceivingHeader)
                     {
-                        continue;
-                    }
+                        if (stream.StreamTimeoutTicks == default)
+                        {
+                            // On expiration overflow, use max value.
+                            var expirationTicks = ticks + _context.ServiceContext.ServerOptions.Limits.RequestHeadersTimeout.Ticks;
+                            stream.StreamTimeoutTicks = expirationTicks >= 0 ? expirationTicks : long.MaxValue;
+                        }
 
-                    if (stream.HeaderTimeoutTicks == default)
-                    {
-                        // On expiration overflow, use max value.
-                        var expirationTicks = ticks + _context.ServiceContext.ServerOptions.Limits.RequestHeadersTimeout.Ticks;
-                        stream.HeaderTimeoutTicks = expirationTicks >= 0 ? expirationTicks : long.MaxValue;
+                        if (stream.StreamTimeoutTicks < ticks)
+                        {
+                            if (stream.IsRequestStream)
+                            {
+                                stream.Abort(new ConnectionAbortedException(CoreStrings.BadRequest_RequestHeadersTimeout), Http3ErrorCode.RequestRejected);
+                            }
+                            else
+                            {
+                                stream.Abort(new ConnectionAbortedException(CoreStrings.Http3ControlStreamHeaderTimeout), Http3ErrorCode.StreamCreationError);
+                            }
+                        }
                     }
-
-                    if (stream.HeaderTimeoutTicks < ticks)
+                    else if (stream.IsDraining)
                     {
-                        if (stream.IsRequestStream)
+                        var minDataRate = _context.ServiceContext.ServerOptions.Limits.MinResponseDataRate;
+                        if (minDataRate == null)
                         {
-                            stream.Abort(new ConnectionAbortedException(CoreStrings.BadRequest_RequestHeadersTimeout), Http3ErrorCode.RequestRejected);
+                            continue;
                         }
-                        else
+
+                        if (stream.StreamTimeoutTicks == default)
+                        {
+                            var expirationTicks = Math.Max(
+                                _context.TimeoutControl.GetWriteTimingTimeoutTimestamp(),
+                                ticks + minDataRate.GracePeriod.Ticks);
+
+                            stream.StreamTimeoutTicks = expirationTicks >= 0 ? expirationTicks : long.MaxValue;
+                        }
+
+                        if (stream.StreamTimeoutTicks < ticks)
                         {
-                            stream.Abort(new ConnectionAbortedException(CoreStrings.Http3ControlStreamHeaderTimeout), Http3ErrorCode.StreamCreationError);
+                            stream.Abort(new ConnectionAbortedException(CoreStrings.ConnectionTimedBecauseResponseMininumDataRateNotSatisfied), Http3ErrorCode.RequestCancelled);
                         }
                     }
                 }
@@ -396,7 +422,6 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> appl
                     }
 
                     _context.TimeoutControl.CancelTimeout();
-                    _context.TimeoutControl.StartDrainTimeout(Limits.MinResponseDataRate, Limits.MaxResponseBufferSize);
                 }
                 catch
                 {
@@ -648,7 +673,7 @@ void IHttp3StreamLifetimeHandler.OnInboundControlStreamSetting(Http3SettingType
 
         void IHttp3StreamLifetimeHandler.OnStreamHeaderReceived(IHttp3Stream stream)
         {
-            Debug.Assert(stream.ReceivedHeader);
+            Debug.Assert(!stream.IsReceivingHeader);
         }
 
         public void HandleRequestHeadersTimeout()

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

@@ -67,8 +67,9 @@ private void OnStreamClosed()
         public PipeReader Input => _context.Transport.Input;
         public IKestrelTrace Log => _context.ServiceContext.Log;
 
-        public long HeaderTimeoutTicks { get; set; }
-        public bool ReceivedHeader => _headerType >= 0;
+        public long StreamTimeoutTicks { get; set; }
+        public bool IsReceivingHeader => _headerType == -1;
+        public bool IsDraining => false;
 
         public bool IsRequestStream => false;
 

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

@@ -1,20 +1,15 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
 using System.Buffers;
 using System.Diagnostics;
 using System.IO.Pipelines;
 using System.Net.Http;
 using System.Net.Http.QPack;
-using System.Net.Quic;
 using System.Runtime.CompilerServices;
-using System.Threading;
-using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Hosting.Server;
-using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.Extensions.Logging;
@@ -73,8 +68,9 @@ internal abstract partial class Http3Stream : HttpProtocol, IHttp3Stream, IHttpH
         public KestrelServerLimits Limits => _context.ServiceContext.ServerOptions.Limits;
         public long StreamId => _streamIdFeature.StreamId;
 
-        public long HeaderTimeoutTicks { get; set; }
-        public bool ReceivedHeader => _appCompleted != null; // TCS is assigned once headers are received
+        public long StreamTimeoutTicks { get; set; }
+        public bool IsReceivingHeader => _appCompleted == null; // TCS is assigned once headers are received
+        public bool IsDraining => _appCompleted?.Task.IsCompleted ?? false; // Draining starts once app is complete
 
         public bool IsRequestStream => true;
 
@@ -97,7 +93,7 @@ public void Initialize(Http3StreamContext context)
             _totalParsedHeaderSize = 0;
             _isMethodConnect = false;
             _completionState = default;
-            HeaderTimeoutTicks = 0;
+            StreamTimeoutTicks = 0;
 
             if (_frameWriter == null)
             {
@@ -409,7 +405,6 @@ private bool TryClose()
                 return true;
             }
 
-            // TODO make this actually close the Http3Stream by telling quic to close the stream.
             return false;
         }
 
@@ -501,12 +496,21 @@ public async Task ProcessRequestAsync<TContext>(IHttpApplication<TContext> appli
                 }
                 finally
                 {
-                    ApplyCompletionFlag(StreamCompletionFlags.Completed);
+                    // Completing the stream needs to happen in the middle of DisposeAsync.
+                    // After the stream has drained but before the stream is returned to the pool.
+                    // OnCompleted callback executes in the right place.
+                    _context.StreamContext.Features.Get<IConnectionCompleteFeature>()!.OnCompleted(static (state) =>
+                    {
+                        var s = (Http3Stream)state;
+
+                        s.ApplyCompletionFlag(StreamCompletionFlags.Completed);
+
+                        // Tells the connection to remove the stream from its active collection.
+                        s._context.StreamLifetimeHandler.OnStreamCompleted(s);
 
-                    // Tells the connection to remove the stream from its active collection.
-                    _context.StreamLifetimeHandler.OnStreamCompleted(this);
+                        return Task.CompletedTask;
+                    }, this);
 
-                    // Dispose must happen after stream is no longer active.
                     await _context.StreamContext.DisposeAsync();
                 }
             }
@@ -600,8 +604,6 @@ private async Task ProcessHeadersFrameAsync<TContext>(IHttpApplication<TContext>
                 case RequestHeaderParsingState.Headers:
                     break;
                 case RequestHeaderParsingState.Trailers:
-                    // trailers
-                    // TODO figure out if there is anything else to do here.
                     return;
                 default:
                     Debug.Fail("Unexpected header parsing state.");
@@ -627,6 +629,7 @@ private async Task ProcessHeadersFrameAsync<TContext>(IHttpApplication<TContext>
             }
 
             _appCompleted = new TaskCompletionSource();
+            StreamTimeoutTicks = default;
             _context.StreamLifetimeHandler.OnStreamHeaderReceived(this);
 
             ThreadPool.UnsafeQueueUserWorkItem(this, preferLocal: false);

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

@@ -14,17 +14,25 @@ internal interface IHttp3Stream
         long StreamId { get; }
 
         /// <summary>
-        /// Used to track the timeout between when the stream was started by the client, and getting a header.
-        /// Value is driven by <see cref="KestrelServerLimits.RequestHeadersTimeout"/>.
+        /// Used to track the timeout in two situations:
+        /// 1. Between when the stream was started by the client, and getting a header.
+        ///    Value is driven by <see cref="KestrelServerLimits.RequestHeadersTimeout"/>.
+        /// 2. Between when the request delegate is complete and the transport draining.
+        ///    Value is driven by <see cref="KestrelServerLimits.MinResponseDataRate"/>.
         /// </summary>
-        long HeaderTimeoutTicks { get; set; }
+        long StreamTimeoutTicks { get; set; }
 
         /// <summary>
-        /// The stream has received and parsed the header frame.
+        /// The stream is receiving the header frame.
         /// - Request streams = HEADERS frame.
         /// - Control streams = unidirectional stream header.
         /// </summary>
-        bool ReceivedHeader { get; }
+        bool IsReceivingHeader { get; }
+
+        /// <summary>
+        /// The stream request delegate is complete and the transport is draining.
+        /// </summary>
+        bool IsDraining { get; }
 
         bool IsRequestStream { get; }
 

src/Servers/Kestrel/Core/src/Internal/Infrastructure/ITimeoutControl.cs

@@ -26,5 +26,6 @@ internal interface ITimeoutControl
         void StartTimingWrite();
         void StopTimingWrite();
         void BytesWrittenToBuffer(MinDataRate minRate, long count);
+        long GetWriteTimingTimeoutTimestamp();
     }
 }

src/Servers/Kestrel/Core/src/Internal/Infrastructure/TimeoutControl.cs

@@ -330,5 +330,10 @@ void IConnectionTimeoutFeature.ResetTimeout(TimeSpan timeSpan)
 
             ResetTimeout(timeSpan.Ticks, TimeoutReason.TimeoutFeature);
         }
+
+        public long GetWriteTimingTimeoutTimestamp()
+        {
+            return _writeTimingTimeoutTimestamp;
+        }
     }
 }

src/Servers/Kestrel/Transport.Quic/src/Internal/QuicStreamContext.FeatureCollection.cs

@@ -4,12 +4,14 @@
 using System.Net.Sockets;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Connections.Features;
+using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Quic.Internal
 {
-    internal sealed partial class QuicStreamContext : IPersistentStateFeature, IStreamDirectionFeature, IProtocolErrorCodeFeature, IStreamIdFeature, IStreamAbortFeature
+    internal sealed partial class QuicStreamContext : IPersistentStateFeature, IStreamDirectionFeature, IProtocolErrorCodeFeature, IStreamIdFeature, IStreamAbortFeature, IConnectionCompleteFeature
     {
         private IDictionary<object, object?>? _persistentState;
+        private Stack<KeyValuePair<Func<object, Task>, object>>? _onCompleted;
 
         public bool CanRead { get; private set; }
         public bool CanWrite { get; private set; }
@@ -31,15 +33,18 @@ public void AbortRead(long errorCode, ConnectionAbortedException abortReason)
         {
             lock (_shutdownLock)
             {
-                if (_stream.CanRead)
+                if (_stream != null)
                 {
-                    _shutdownReadReason = abortReason;
-                    _log.StreamAbortRead(this, errorCode, abortReason.Message);
-                    _stream.AbortRead(errorCode);
-                }
-                else
-                {
-                    throw new InvalidOperationException("Unable to abort reading from a stream that doesn't support reading.");
+                    if (_stream.CanRead)
+                    {
+                        _shutdownReadReason = abortReason;
+                        _log.StreamAbortRead(this, errorCode, abortReason.Message);
+                        _stream.AbortRead(errorCode);
+                    }
+                    else
+                    {
+                        throw new InvalidOperationException("Unable to abort reading from a stream that doesn't support reading.");
+                    }
                 }
             }
         }
@@ -48,26 +53,39 @@ public void AbortWrite(long errorCode, ConnectionAbortedException abortReason)
         {
             lock (_shutdownLock)
             {
-                if (_stream.CanWrite)
+                if (_stream != null)
                 {
-                    _shutdownWriteReason = abortReason;
-                    _log.StreamAbortWrite(this, errorCode, abortReason.Message);
-                    _stream.AbortWrite(errorCode);
-                }
-                else
-                {
-                    throw new InvalidOperationException("Unable to abort writing to a stream that doesn't support writing.");
+                    if (_stream.CanWrite)
+                    {
+                        _shutdownWriteReason = abortReason;
+                        _log.StreamAbortWrite(this, errorCode, abortReason.Message);
+                        _stream.AbortWrite(errorCode);
+                    }
+                    else
+                    {
+                        throw new InvalidOperationException("Unable to abort writing to a stream that doesn't support writing.");
+                    }
                 }
             }
         }
 
+        public void OnCompleted(Func<object, Task> callback, object state)
+        {
+            if (_onCompleted == null)
+            {
+                _onCompleted = new Stack<KeyValuePair<Func<object, Task>, object>>();
+            }
+            _onCompleted.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
+        }
+
         private void InitializeFeatures()
         {
             _currentIPersistentStateFeature = this;
             _currentIStreamDirectionFeature = this;
             _currentIProtocolErrorCodeFeature = this;
             _currentIStreamIdFeature = this;
             _currentIStreamAbortFeature = this;
+            _currentIConnectionCompleteFeature = this;
         }
     }
 }