Remove dependency on pipe events in HttpConnection

src/Servers/Kestrel/Core/src/Adapter/Internal/AdaptedPipeline.cs

@@ -5,6 +5,7 @@
 using System.IO;
 using System.IO.Pipelines;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.Extensions.Logging;
 
@@ -14,21 +15,24 @@ internal class AdaptedPipeline : IDuplexPipe
     {
         private readonly int _minAllocBufferSize;
 
-        private readonly IDuplexPipe _transport;
+        private Task _inputTask;
+        private Task _outputTask;
 
         public AdaptedPipeline(IDuplexPipe transport,
                                Pipe inputPipe,
                                Pipe outputPipe,
                                IKestrelTrace log,
                                int minAllocBufferSize)
         {
-            _transport = transport;
+            TransportStream = new RawStream(transport.Input, transport.Output, throwOnCancelled: true);
             Input = inputPipe;
             Output = outputPipe;
             Log = log;
             _minAllocBufferSize = minAllocBufferSize;
         }
 
+        public RawStream TransportStream { get; }
+
         public Pipe Input { get; }
 
         public Pipe Output { get; }
@@ -39,13 +43,31 @@ public AdaptedPipeline(IDuplexPipe transport,
 
         PipeWriter IDuplexPipe.Output => Output.Writer;
 
-        public async Task RunAsync(Stream stream)
+        public void RunAsync(Stream stream)
+        {
+            _inputTask = ReadInputAsync(stream);
+            _outputTask = WriteOutputAsync(stream);
+        }
+
+        public async Task CompleteAsync()
         {
-            var inputTask = ReadInputAsync(stream);
-            var outputTask = WriteOutputAsync(stream);
+            Output.Writer.Complete();
+            Input.Reader.Complete();
+
+            if (_outputTask == null)
+            {
+                return;
+            }
+
+            // Wait for the output task to complete, this ensures that we've copied
+            // the application data to the underlying stream
+            await _outputTask;
+
+            // Cancel the underlying stream so that the input task yields
+            TransportStream.CancelPendingRead();
 
-            await inputTask;
-            await outputTask;
+            // The input task should yield now that we've cancelled it
+            await _inputTask;
         }
 
         private async Task WriteOutputAsync(Stream stream)
@@ -97,7 +119,6 @@ private async Task WriteOutputAsync(Stream stream)
             finally
             {
                 Output.Reader.Complete();
-                _transport.Output.Complete();
             }
         }
 
@@ -115,7 +136,6 @@ private async Task ReadInputAsync(Stream stream)
 
                 while (true)
                 {
-
                     var outputBuffer = Input.Writer.GetMemory(_minAllocBufferSize);
                     var bytesRead = await stream.ReadAsync(outputBuffer);
                     Input.Writer.Advance(bytesRead);
@@ -134,6 +154,11 @@ private async Task ReadInputAsync(Stream stream)
                     }
                 }
             }
+            catch (OperationCanceledException ex)
+            {
+                // Propagate the exception if it's ConnectionAbortedException
+                error = ex as ConnectionAbortedException;
+            }
             catch (Exception ex)
             {
                 // Don't rethrow the exception. It should be handled by the Pipeline consumer.
@@ -142,9 +167,6 @@ private async Task ReadInputAsync(Stream stream)
             finally
             {
                 Input.Writer.Complete(error);
-                // The application could have ended the input pipe so complete
-                // the transport pipe as well
-                _transport.Input.Complete();
             }
         }
     }

src/Servers/Kestrel/Core/src/Adapter/Internal/RawStream.cs

@@ -14,11 +14,20 @@ internal sealed class RawStream : Stream
     {
         private readonly PipeReader _input;
         private readonly PipeWriter _output;
+        private readonly bool _throwOnCancelled;
+        private volatile bool _cancelCalled;
 
-        public RawStream(PipeReader input, PipeWriter output)
+        public RawStream(PipeReader input, PipeWriter output, bool throwOnCancelled = false)
         {
             _input = input;
             _output = output;
+            _throwOnCancelled = throwOnCancelled;
+        }
+
+        public void CancelPendingRead()
+        {
+            _cancelCalled = true;
+            _input.CancelPendingRead();
         }
 
         public override bool CanRead => true;
@@ -61,17 +70,17 @@ public override int Read(byte[] buffer, int offset, int count)
         {
             // ValueTask uses .GetAwaiter().GetResult() if necessary
             // https://github.com/dotnet/corefx/blob/f9da3b4af08214764a51b2331f3595ffaf162abe/src/System.Threading.Tasks.Extensions/src/System/Threading/Tasks/ValueTask.cs#L156
-            return ReadAsyncInternal(new Memory<byte>(buffer, offset, count)).Result;
+            return ReadAsyncInternal(new Memory<byte>(buffer, offset, count), default).Result;
         }
 
-        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default)
         {
-            return ReadAsyncInternal(new Memory<byte>(buffer, offset, count)).AsTask();
+            return ReadAsyncInternal(new Memory<byte>(buffer, offset, count), cancellationToken).AsTask();
         }
 
         public override ValueTask<int> ReadAsync(Memory<byte> destination, CancellationToken cancellationToken = default)
         {
-            return ReadAsyncInternal(destination);
+            return ReadAsyncInternal(destination, cancellationToken);
         }
 
         public override void Write(byte[] buffer, int offset, int count)
@@ -105,14 +114,21 @@ public override Task FlushAsync(CancellationToken cancellationToken)
             return WriteAsync(null, 0, 0, cancellationToken);
         }
 
-        private async ValueTask<int> ReadAsyncInternal(Memory<byte> destination)
+        private async ValueTask<int> ReadAsyncInternal(Memory<byte> destination, CancellationToken cancellationToken)
         {
             while (true)
             {
-                var result = await _input.ReadAsync();
+                var result = await _input.ReadAsync(cancellationToken);
                 var readableBuffer = result.Buffer;
                 try
                 {
+                    if (_throwOnCancelled && result.IsCanceled && _cancelCalled)
+                    {
+                        // Reset the bool
+                        _cancelCalled = false;
+                        throw new OperationCanceledException();
+                    }
+
                     if (!readableBuffer.IsEmpty)
                     {
                         // buffer.Count is int

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

@@ -68,8 +68,6 @@ public Http1Connection(HttpConnectionContext context)
 
         protected override void OnRequestProcessingEnded()
         {
-            Input.Complete();
-
             TimeoutControl.StartDrainTimeout(MinResponseDataRate, ServerOptions.Limits.MaxResponseBufferSize);
 
             // Prevent RequestAborted from firing. Free up unneeded feature references.

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

@@ -86,6 +86,9 @@ public Http1OutputProducer(
             _memoryPool = memoryPool;
         }
 
+        // For tests
+        internal PipeWriter PipeWriter => _pipeWriter;
+
         public Task WriteDataAsync(ReadOnlySpan<byte> buffer, CancellationToken cancellationToken = default)
         {
             if (cancellationToken.IsCancellationRequested)
@@ -402,7 +405,6 @@ private void CompletePipe()
                 _log.ConnectionDisconnect(_connectionId);
                 _pipeWriterCompleted = true;
                 _completed = true;
-                _pipeWriter.Complete();
             }
         }
 

src/Servers/Kestrel/Core/src/Internal/HttpConnection.cs

@@ -82,7 +82,6 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> http
             try
             {
                 AdaptedPipeline adaptedPipeline = null;
-                var adaptedPipelineTask = Task.CompletedTask;
 
                 // _adaptedTransport must be set prior to wiring up callbacks
                 // to allow the connection to be aborted prior to protocol selection.
@@ -120,8 +119,8 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> http
                     if (adaptedPipeline != null)
                     {
                         // Stream can be null here and run async will close the connection in that case
-                        var stream = await ApplyConnectionAdaptersAsync();
-                        adaptedPipelineTask = adaptedPipeline.RunAsync(stream);
+                        var stream = await ApplyConnectionAdaptersAsync(adaptedPipeline.TransportStream);
+                        adaptedPipeline.RunAsync(stream);
                     }
 
                     IRequestProcessor requestProcessor = null;
@@ -160,20 +159,19 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> http
                         }
                     }
 
-                    _context.Transport.Input.OnWriterCompleted(
-                        (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
-                        this);
+                    var closedRegistration = _context.ConnectionContext.ConnectionClosed.Register(state => ((HttpConnection)state).OnInputOrOutputCompleted(), this);
 
-                    _context.Transport.Output.OnReaderCompleted(
-                        (_, state) => ((HttpConnection)state).OnInputOrOutputCompleted(),
-                        this);
-
-                    if (requestProcessor != null)
+                    // We don't care about callbacks once all requests are processed
+                    using (closedRegistration)
                     {
-                        await requestProcessor.ProcessRequestsAsync(httpApplication);
+                        if (requestProcessor != null)
+                        {
+                            await requestProcessor.ProcessRequestsAsync(httpApplication);
+                        }
                     }
 
-                    await adaptedPipelineTask;
+                    // Complete the pipeline after the method runs
+                    await (adaptedPipeline?.CompleteAsync() ?? Task.CompletedTask);
                 }
             }
             catch (Exception ex)
@@ -277,10 +275,9 @@ private void Abort(ConnectionAbortedException ex)
             }
         }
 
-        private async Task<Stream> ApplyConnectionAdaptersAsync()
+        private async Task<Stream> ApplyConnectionAdaptersAsync(RawStream stream)
         {
             var connectionAdapters = _context.ConnectionAdapters;
-            var stream = new RawStream(_context.Transport.Input, _context.Transport.Output);
             var adapterContext = new ConnectionAdapterContext(_context.ConnectionContext, stream);
             _adaptedConnections = new List<IAdaptedConnection>(connectionAdapters.Count);
 
@@ -367,12 +364,7 @@ private void Tick()
 
         private void CloseUninitializedConnection(ConnectionAbortedException abortReason)
         {
-            Debug.Assert(_adaptedTransport != null);
-
             _context.ConnectionContext.Abort(abortReason);
-
-            _adaptedTransport.Input.Complete();
-            _adaptedTransport.Output.Complete();
         }
 
         public void OnTimeout(TimeoutReason reason)

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

@@ -47,9 +47,10 @@ public async Task WritesNoopAfterConnectionCloses()
 
                 await socketOutput.WriteDataAsync(new byte[] { 1, 2, 3, 4 }, default);
 
-                Assert.True(socketOutput.Pipe.Reader.TryRead(out var result));
-                Assert.True(result.IsCompleted);
-                Assert.True(result.Buffer.IsEmpty);
+                Assert.False(socketOutput.Pipe.Reader.TryRead(out var result));
+
+                socketOutput.Pipe.Writer.Complete();
+                socketOutput.Pipe.Reader.Complete();
             }
         }
 

src/Servers/Kestrel/Transport.Libuv/src/Internal/LibuvConnection.cs

@@ -31,6 +31,8 @@ internal partial class LibuvConnection : TransportConnection
 
         private MemoryHandle _bufferHandle;
         private Task _processingTask;
+        private readonly TaskCompletionSource<object> _waitForConnectionClosedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+        private bool _connectionClosed;
 
         public LibuvConnection(UvStreamHandle socket,
                                ILibuvTrace log,
@@ -135,21 +137,10 @@ private async Task StartCore()
                     // We're done with the socket now
                     _socket.Dispose();
 
-                    // Fire the connection closed token and wait for it to complete
-                    var waitForConnectionClosedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+                    // Ensure this always fires
+                    FireConnectionClosed();
 
-                    ThreadPool.UnsafeQueueUserWorkItem(state =>
-                    {
-                        (var connection, var tcs) = state;
-
-                        connection.CancelConnectionClosedToken();
-
-                        tcs.TrySetResult(null);
-                    },
-                    (this, waitForConnectionClosedTcs),
-                    preferLocal: false);
-
-                    await waitForConnectionClosedTcs.Task;
+                    await _waitForConnectionClosedTcs.Task;
                 }
             }
             catch (Exception e)
@@ -241,11 +232,33 @@ private void OnRead(UvStreamHandle handle, int status)
                     error = LogAndWrapReadError(uvError);
                 }
 
+                FireConnectionClosed();
+
                 // Complete after aborting the connection
                 Input.Complete(error);
             }
         }
 
+        private void FireConnectionClosed()
+        {
+            // Guard against scheduling this multiple times
+            if (_connectionClosed)
+            {
+                return;
+            }
+
+            _connectionClosed = true;
+
+            ThreadPool.UnsafeQueueUserWorkItem(state =>
+            {
+                state.CancelConnectionClosedToken();
+
+                state._waitForConnectionClosedTcs.TrySetResult(null);
+            },
+            this,
+            preferLocal: false);
+        }
+
         private async Task ApplyBackpressureAsync(ValueTask<FlushResult> flushTask)
         {
             Log.ConnectionPause(ConnectionId);