Consume Faster

Author: benaadams

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

@@ -6,6 +6,7 @@
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNet.Server.Kestrel.Infrastructure;
 using Microsoft.Extensions.Primitives;
 
 namespace Microsoft.AspNet.Server.Kestrel.Http
@@ -38,35 +39,40 @@ protected MessageBody(FrameContext context)
             return result;
         }
 
-        public async Task Consume(CancellationToken cancellationToken = default(CancellationToken))
+        public Task Consume(CancellationToken cancellationToken = default(CancellationToken))
         {
-            Task<int> result;
-            var send100checked = false;
-            do
+            var result = ReadAsyncImplementation(default(ArraySegment<byte>), cancellationToken);
+            if (!result.IsCompleted)
             {
-                result = ReadAsyncImplementation(default(ArraySegment<byte>), cancellationToken);
-                if (!result.IsCompleted)
-                {
-                    if (!send100checked)
-                    {
-                        if (Interlocked.Exchange(ref _send100Continue, 0) == 1)
-                        {
-                            _context.FrameControl.ProduceContinue();
-                        }
-                        send100checked = true;
-                    }
-                }
-                else if (result.GetAwaiter().GetResult() == 0) 
-                {
-                    // Completed Task, end of stream
-                    return;
-                }
-                else
+                if (Interlocked.Exchange(ref _send100Continue, 0) == 1)
                 {
-                    // Completed Task, get next Task rather than await
-                    continue;
+                    _context.FrameControl.ProduceContinue();
                 }
-            } while (await result != 0);
+
+                return ConsumeAwaited(result, cancellationToken);
+            }
+            else if (result.GetAwaiter().GetResult() == 0)
+            {
+                // Completed Task, end of stream
+                return TaskUtilities.CompletedTask;
+            }
+            else
+            {
+                // Completed Task, but non-zero get next Task and await
+                return ConsumeAwaited(ReadAsyncImplementation(default(ArraySegment<byte>), cancellationToken), cancellationToken);
+            }
+        }
+
+        private async Task ConsumeAwaited(Task<int> currentTask, CancellationToken cancellationToken)
+        {
+            var count = await currentTask;
+
+            if (count == 0) return;
+
+            while (await ReadAsyncImplementation(default(ArraySegment<byte>), cancellationToken) != 0)
+            {
+                // Consume until complete
+            }
         }
 
         public abstract Task<int> ReadAsyncImplementation(ArraySegment<byte> buffer, CancellationToken cancellationToken);
@@ -156,17 +162,37 @@ public ForContentLength(bool keepAlive, int contentLength, FrameContext context)
                 _inputLength = _contentLength;
             }
 
-            public override async Task<int> ReadAsyncImplementation(ArraySegment<byte> buffer, CancellationToken cancellationToken)
+            public override Task<int> ReadAsyncImplementation(ArraySegment<byte> buffer, CancellationToken cancellationToken)
             {
                 var input = _context.SocketInput;
 
                 var limit = buffer.Array == null ? _inputLength : Math.Min(buffer.Count, _inputLength);
                 if (limit == 0)
                 {
-                    return 0;
+                    return TaskUtilities.ZeroTask;
+                }
+
+                var task = _context.SocketInput.ReadAsync(buffer.Array, buffer.Offset, limit);
+
+                if (task.IsCompleted)
+                {
+                    var actual = task.GetAwaiter().GetResult();
+                    if (actual == 0)
+                    {
+                        throw new InvalidDataException("Unexpected end of request content");
+                    }
+                    _inputLength -= actual;
+                    return task;
                 }
+                else
+                {
+                    return ReadAsyncImplementationAwaited(task);
+                }
+            }
 
-                var actual = await _context.SocketInput.ReadAsync(buffer.Array, buffer.Offset, limit);
+            private async Task<int> ReadAsyncImplementationAwaited(Task<int> currentTask)
+            {
+                var actual = await currentTask;
                 _inputLength -= actual;
 
                 if (actual == 0)

src/Microsoft.AspNet.Server.Kestrel/Infrastructure/TaskUtilities.cs

@@ -12,5 +12,6 @@ public static class TaskUtilities
 #else
         public static Task CompletedTask = Task.FromResult<object>(null);
 #endif
+        public static Task<int> ZeroTask = Task.FromResult(0);
     }
 }