Write async chunks async

aspnet · Jan 17, 2016 · b5b7158 · b5b7158
1 parent 129a5ad
commit b5b7158
Showing 1 changed file with 21 additions and 16 deletions.
diff --git 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(