Add tests for non-immediate byte counting

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

@@ -56,7 +56,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                 buffer = new ArraySegment<byte>(copy);
                 _log.ConnectionWrite(_connectionId, buffer.Count);
             }
-            
+
             TaskCompletionSource<object> tcs = null;
 
             lock (_lockObj)
@@ -79,25 +79,26 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                     _nextWriteContext.SocketDisconnect = true;
                 }
 
-                // Complete the write task immediately if all previous write tasks have been completed,
-                // the buffers haven't grown too large, and the last write to the socket succeeded.
-                if (_lastWriteError == null &&
-                    _tasksPending.Count == 0 &&
-                    _numBytesPreCompleted + buffer.Count <= _maxBytesPreCompleted)
+                if (!immediate)
                 {
+                    // immediate==false calls always return complete tasks, because there is guaranteed
+                    // to be a subsequent immediate==true call which will go down one of the previous code-paths
                     _numBytesPreCompleted += buffer.Count;
                 }
-                else if (immediate)
+                else if (_lastWriteError == null &&
+                        _tasksPending.Count == 0 &&
+                        _numBytesPreCompleted + buffer.Count <= _maxBytesPreCompleted)
+                {
+                    // Complete the write task immediately if all previous write tasks have been completed,
+                    // the buffers haven't grown too large, and the last write to the socket succeeded.
+                    _numBytesPreCompleted += buffer.Count;
+                }
+                else
                 {
                     // immediate write, which is not eligable for instant completion above
                     tcs = new TaskCompletionSource<object>(buffer.Count);
                     _tasksPending.Enqueue(tcs);
                 }
-                else
-                {
-                    // immediate==false calls always return complete tasks, because there is guaranteed
-                    // to be a subsequent immediate==true call which will go down one of the previous code-paths
-                }
 
                 if (_writesPending < _maxPendingWrites && immediate)
                 {
@@ -222,7 +223,6 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
 
                 // Now that the while loop has completed the following invariants should hold true:
                 Debug.Assert(_numBytesPreCompleted >= 0);
-                Debug.Assert(_numBytesPreCompleted <= _maxBytesPreCompleted);
             }
         }
 

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

@@ -115,6 +115,82 @@ namespace Microsoft.AspNet.Server.KestrelTests
                 Assert.True(completedWh.Wait(1000));
             }
         }
+        
+        [Fact]
+        public void WritesDontCompleteImmediatelyWhenTooManyBytesIncludingNonImmediateAreAlreadyPreCompleted()
+        {
+            // This should match _maxBytesPreCompleted in SocketOutput
+            var maxBytesPreCompleted = 65536;
+            var completeQueue = new Queue<Action<int>>();
+
+            // Arrange
+            var mockLibuv = new MockLibuv
+            {
+                OnWrite = (socket, buffers, triggerCompleted) =>
+                {
+                    completeQueue.Enqueue(triggerCompleted);
+                    return 0;
+                }
+            };
+
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
+            {
+                kestrelEngine.Start(count: 1);
+
+                var kestrelThread = kestrelEngine.Threads[0];
+                var socket = new MockSocket(kestrelThread.Loop.ThreadId, new TestKestrelTrace());
+                var trace = new KestrelTrace(new TestKestrelTrace());
+                var socketOutput = new SocketOutput(kestrelThread, socket, 0, trace);
+
+                var bufferSize = maxBytesPreCompleted;
+
+                var data = new byte[bufferSize];
+                var fullBuffer = new ArraySegment<byte>(data, 0, bufferSize);
+                var halfBuffer = new ArraySegment<byte>(data, 0, bufferSize / 2);
+
+                var completedWh = new ManualResetEventSlim();
+                Action<Task> onCompleted = (Task t) =>
+                {
+                    Assert.Null(t.Exception);
+                    completedWh.Set();
+                };
+
+                // Act 
+                socketOutput.WriteAsync(halfBuffer, false).ContinueWith(onCompleted);
+                // Assert
+                // The first write should pre-complete since it is not immediate.
+                Assert.True(completedWh.Wait(1000));
+                // Arrange
+                completedWh.Reset();
+                // Act 
+                socketOutput.WriteAsync(halfBuffer).ContinueWith(onCompleted);
+                // Assert
+                // The second write should pre-complete since it is <= _maxBytesPreCompleted.
+                Assert.True(completedWh.Wait(1000));
+                // Arrange
+                completedWh.Reset();
+                // Act 
+                socketOutput.WriteAsync(halfBuffer, false).ContinueWith(onCompleted);
+                // Assert
+                // The third write should pre-complete since it is not immediate, even though too many.
+                Assert.True(completedWh.Wait(1000));
+                // Arrange
+                completedWh.Reset();
+                // Act
+                socketOutput.WriteAsync(halfBuffer).ContinueWith(onCompleted);
+                // Assert 
+                // Too many bytes are already pre-completed for the fourth write to pre-complete.
+                Assert.False(completedWh.Wait(1000));
+                // Act
+                while (completeQueue.Count > 0)
+                {
+                    completeQueue.Dequeue()(0);
+                }
+                // Assert
+                // Finishing the first write should allow the second write to pre-complete.
+                Assert.True(completedWh.Wait(1000));
+            }
+        }
 
         private class MockSocket : UvStreamHandle
         {