Gracefully handle connection close in SocketOutput.ProducingComplete

src/Microsoft.AspNetCore.Server.Kestrel/Http/SocketOutput.cs

@@ -248,7 +248,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Http
 
         public void ProducingComplete(MemoryPoolIterator end)
         {
-            Debug.Assert(!_lastStart.IsDefault);
+            if(_lastStart.IsDefault)
+            {
+                return;
+            }
 
             int bytesProduced, buffersIncluded;
             BytesBetween(_lastStart, end, out bytesProduced, out buffersIncluded);
@@ -265,9 +268,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Http
         {
             MemoryPoolBlock blockToReturn = null;
 
-
             lock (_returnLock)
             {
+                // Both ProducingComplete and WriteAsync should not call this method
+                // if _lastStart was not set.
                 Debug.Assert(!_lastStart.IsDefault);
 
                 // If the socket has been closed, return the produced blocks

test/Microsoft.AspNetCore.Server.KestrelTests/SocketOutputTests.cs

@@ -24,14 +24,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             // _maxBytesPreCompleted even after the write actually completed.
 
             // Arrange
-            var mockLibuv = new MockLibuv
-            {
-                OnWrite = (socket, buffers, triggerCompleted) =>
-                {
-                    triggerCompleted(0);
-                    return 0;
-                }
-            };
+            var mockLibuv = new MockLibuv();
 
             using (var memory = new MemoryPool())
             using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
@@ -661,5 +654,37 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                     default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
             }
         }
+
+        [Fact]
+        public void ProducingStartAndProducingCompleteCanBeCalledAfterConnectionClose()
+        {
+            var mockLibuv = new MockLibuv();
+
+            using (var memory = new MemoryPool())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
+            {
+                kestrelEngine.Start(count: 1);
+
+                var kestrelThread = kestrelEngine.Threads[0];
+                var socket = new MockSocket(mockLibuv, kestrelThread.Loop.ThreadId, new TestKestrelTrace());
+                var trace = new KestrelTrace(new TestKestrelTrace());
+                var ltp = new LoggingThreadPool(trace);
+                var connection = new MockConnection();
+                var socketOutput = new SocketOutput(kestrelThread, socket, memory, connection, "0", trace, ltp, new Queue<UvWriteReq>());
+
+                // Close SocketOutput
+                var cleanupTask = socketOutput.WriteAsync(
+                    default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
+
+                Assert.True(connection.SocketClosed.Wait(1000));
+
+                var start = socketOutput.ProducingStart();
+
+                Assert.True(start.IsDefault);
+                // ProducingComplete should not throw given a default iterator
+                socketOutput.ProducingComplete(start);
+
+            }
+        }
     }
 }

test/Microsoft.AspNetCore.Server.KestrelTests/TestHelpers/MockConnection.cs

@@ -3,12 +3,15 @@
 
 using System;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Server.Kestrel.Http;
 
 namespace Microsoft.AspNetCore.Server.KestrelTests.TestHelpers
 {
     public class MockConnection : Connection, IDisposable
     {
+        private TaskCompletionSource<object> _socketClosedTcs = new TaskCompletionSource<object>();
+
         public MockConnection()
         {
             RequestAbortedSource = new CancellationTokenSource();
@@ -24,10 +27,19 @@ namespace Microsoft.AspNetCore.Server.KestrelTests.TestHelpers
 
         public override void OnSocketClosed()
         {
+            _socketClosedTcs.SetResult(null);
         }
 
         public CancellationTokenSource RequestAbortedSource { get; }
 
+        public Task SocketClosed
+        {
+            get
+            {
+                return _socketClosedTcs.Task;
+            }
+        }
+
         public void Dispose()
         {
             RequestAbortedSource.Dispose();

test/Microsoft.AspNetCore.Server.KestrelTests/TestHelpers/MockLibuv.cs

@@ -12,12 +12,10 @@ namespace Microsoft.AspNetCore.Server.KestrelTests.TestHelpers
         private bool _stopLoop;
         private readonly ManualResetEventSlim _loopWh = new ManualResetEventSlim();
 
-        private Func<UvStreamHandle, int, Action<int>, int> _onWrite;
-
         unsafe public MockLibuv()
             : base(onlyForTesting: true)
         {
-            _onWrite = (socket, buffers, triggerCompleted) =>
+            OnWrite = (socket, buffers, triggerCompleted) =>
             {
                 triggerCompleted(0);
                 return 0;
@@ -80,17 +78,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests.TestHelpers
             _uv_read_stop = handle => 0;
         }
 
-        public Func<UvStreamHandle, int, Action<int>, int> OnWrite
-        {
-            get
-            {
-                return _onWrite;
-            }
-            set
-            {
-                _onWrite = value;
-            }
-        }
+        public Func<UvStreamHandle, int, Action<int>, int> OnWrite { get; set; }
 
         public uv_alloc_cb AllocCallback { get; set; }
 
@@ -107,7 +95,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests.TestHelpers
 
         unsafe private int UvWrite(UvRequest req, UvStreamHandle handle, uv_buf_t* bufs, int nbufs, uv_write_cb cb)
         {
-            return _onWrite(handle, nbufs, status => cb(req.InternalGetHandle(), status));
+            return OnWrite(handle, nbufs, status => cb(req.InternalGetHandle(), status));
         }
     }
 }