From b5b715872034fbc1836ad3e944621d3222c00e2a Mon Sep 17 00:00:00 2001
From: Ben Adams <thundercat@illyriad.co.uk>
Date: Sat, 16 Jan 2016 23:57:59 +0000
Subject: [PATCH] Write async chunks async

---
 .../Http/Frame.cs                             | 37 +++++++++++--------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs
index f315bb702..54cc69855 100644
--- a/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs
+++ b/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs
@@ -421,7 +421,7 @@ public void Write(ArraySegment<byte> data)
                 {
                     return;
                 }
-                WriteChunked(data);
+                WriteChunkedAsync(data, RequestAborted).GetAwaiter().GetResult();
             }
             else
             {
@@ -468,19 +468,14 @@ public async Task WriteAsyncAwaited(ArraySegment<byte> data, CancellationToken c
             }
         }
 
-        private void WriteChunked(ArraySegment<byte> data)
+        private Task WriteChunkedAsync(ArraySegment<byte> data, CancellationToken cancellationToken)
         {
-            SocketOutput.Write(data, immediate: false, chunk: true);
+            return SocketOutput.WriteAsync(data, immediate: false, chunk: true, cancellationToken: cancellationToken);
         }
 
-        private async Task WriteChunkedAsync(ArraySegment<byte> data, CancellationToken cancellationToken)
+        private Task WriteChunkedResponseSuffix()
         {
-            await SocketOutput.WriteAsync(data, immediate: false, chunk: true, cancellationToken: cancellationToken);
-        }
-
-        private void WriteChunkedResponseSuffix()
-        {
-            SocketOutput.Write(_endChunkedResponseBytes, immediate: true);
+            return SocketOutput.WriteAsync(_endChunkedResponseBytes, immediate: true);
         }
 
         private static ArraySegment<byte> CreateAsciiByteArraySegment(string text)
@@ -571,31 +566,41 @@ protected Task ProduceEnd()
                 return ProduceEndAwaited();
             }
 
-            WriteSuffix();
-
-            return TaskUtilities.CompletedTask;
+            return WriteSuffix();
         }
 
         private async Task ProduceEndAwaited()
         {
             await ProduceStart(immediate: true, appCompleted: true);
 
-            WriteSuffix();
+            await WriteSuffix();
         }
 
-        private void WriteSuffix()
+        private Task WriteSuffix()
         {
             // _autoChunk should be checked after we are sure ProduceStart() has been called
             // since ProduceStart() may set _autoChunk to true.
             if (_autoChunk)
             {
-                WriteChunkedResponseSuffix();
+                return WriteAutoChunkSuffixAwaited();
             }
 
             if (_keepAlive)
             {
                 ConnectionControl.End(ProduceEndType.ConnectionKeepAlive);
             }
+
+            return TaskUtilities.CompletedTask;
+        }
+
+        private async Task WriteAutoChunkSuffixAwaited()
+        {
+            await WriteChunkedResponseSuffix();
+
+            if (_keepAlive)
+            {
+                ConnectionControl.End(ProduceEndType.ConnectionKeepAlive);
+            }
         }
 
         private Task CreateResponseHeader(