Make chunked writes atomic

src/Microsoft.AspNet.Server.Kestrel/Filter/StreamSocketOutput.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.IO;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNet.Server.Kestrel.Http;
@@ -12,27 +13,44 @@ namespace Microsoft.AspNet.Server.Kestrel.Filter
 {
     public class StreamSocketOutput : ISocketOutput
     {
+        private static readonly byte[] _endChunkBytes = Encoding.ASCII.GetBytes("\r\n");
         private static readonly byte[] _nullBuffer = new byte[0];
 
         private readonly Stream _outputStream;
         private readonly MemoryPool2 _memory;
         private MemoryPoolBlock2 _producingBlock;
 
+        private object _writeLock = new object();
+
         public StreamSocketOutput(Stream outputStream, MemoryPool2 memory)
         {
             _outputStream = outputStream;
             _memory = memory;
         }
 
-        void ISocketOutput.Write(ArraySegment<byte> buffer, bool immediate)
+        public void Write(ArraySegment<byte> buffer, bool immediate, bool chunk)
         {
-            _outputStream.Write(buffer.Array ?? _nullBuffer, buffer.Offset, buffer.Count);
+            lock (_writeLock)
+            {
+                if (chunk && buffer.Array != null)
+                {
+                    var beginChunkBytes = ChunkWriter.BeginChunkBytes(buffer.Count);
+                    _outputStream.Write(beginChunkBytes.Array, beginChunkBytes.Offset, beginChunkBytes.Count);
+                }
+
+                _outputStream.Write(buffer.Array ?? _nullBuffer, buffer.Offset, buffer.Count);
+
+                if (chunk && buffer.Array != null)
+                {
+                    _outputStream.Write(_endChunkBytes, 0, _endChunkBytes.Length);
+                }
+            }
         }
 
-        Task ISocketOutput.WriteAsync(ArraySegment<byte> buffer, bool immediate, CancellationToken cancellationToken)
+        public Task WriteAsync(ArraySegment<byte> buffer, bool immediate, bool chunk, CancellationToken cancellationToken)
         {
             // TODO: Use _outputStream.WriteAsync
-            _outputStream.Write(buffer.Array ?? _nullBuffer, buffer.Offset, buffer.Count);
+            Write(buffer, immediate, chunk);
             return TaskUtilities.CompletedTask;
         }
 

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

@@ -0,0 +1,62 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Text;
+using Microsoft.AspNet.Server.Kestrel.Infrastructure;
+
+namespace Microsoft.AspNet.Server.Kestrel.Http
+{
+    public static class ChunkWriter
+    {
+        private static readonly ArraySegment<byte> _endChunkBytes = CreateAsciiByteArraySegment("\r\n");
+        private static readonly byte[] _hex = Encoding.ASCII.GetBytes("0123456789abcdef");
+
+        private static ArraySegment<byte> CreateAsciiByteArraySegment(string text)
+        {
+            var bytes = Encoding.ASCII.GetBytes(text);
+            return new ArraySegment<byte>(bytes);
+        }
+
+        public static ArraySegment<byte> BeginChunkBytes(int dataCount)
+        {
+            var bytes = new byte[10]
+            {
+                _hex[((dataCount >> 0x1c) & 0x0f)],
+                _hex[((dataCount >> 0x18) & 0x0f)],
+                _hex[((dataCount >> 0x14) & 0x0f)],
+                _hex[((dataCount >> 0x10) & 0x0f)],
+                _hex[((dataCount >> 0x0c) & 0x0f)],
+                _hex[((dataCount >> 0x08) & 0x0f)],
+                _hex[((dataCount >> 0x04) & 0x0f)],
+                _hex[((dataCount >> 0x00) & 0x0f)],
+                (byte)'\r',
+                (byte)'\n',
+            };
+
+            // Determine the most-significant non-zero nibble
+            int total, shift;
+            total = (dataCount > 0xffff) ? 0x10 : 0x00;
+            dataCount >>= total;
+            shift = (dataCount > 0x00ff) ? 0x08 : 0x00;
+            dataCount >>= shift;
+            total |= shift;
+            total |= (dataCount > 0x000f) ? 0x04 : 0x00;
+
+            var offset = 7 - (total >> 2);
+            return new ArraySegment<byte>(bytes, offset, 10 - offset);
+        }
+
+        public static int WriteBeginChunkBytes(ref MemoryPoolIterator2 start, int dataCount)
+        {
+            var chunkSegment = BeginChunkBytes(dataCount);
+            start.CopyFrom(chunkSegment);
+            return chunkSegment.Count;
+        }
+
+        public static void WriteEndChunkBytes(ref MemoryPoolIterator2 start)
+        {
+            start.CopyFrom(_endChunkBytes);
+        }
+    }
+}

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

@@ -23,11 +23,9 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
     public abstract partial class Frame : FrameContext, IFrameControl
     {
         private static readonly Encoding _ascii = Encoding.ASCII;
-        private static readonly ArraySegment<byte> _endChunkBytes = CreateAsciiByteArraySegment("\r\n");
         private static readonly ArraySegment<byte> _endChunkedResponseBytes = CreateAsciiByteArraySegment("0\r\n\r\n");
         private static readonly ArraySegment<byte> _continueBytes = CreateAsciiByteArraySegment("HTTP/1.1 100 Continue\r\n\r\n");
         private static readonly ArraySegment<byte> _emptyData = new ArraySegment<byte>(new byte[0]);
-        private static readonly byte[] _hex = Encoding.ASCII.GetBytes("0123456789abcdef");
 
         private static readonly byte[] _bytesConnectionClose = Encoding.ASCII.GetBytes("\r\nConnection: close");
         private static readonly byte[] _bytesConnectionKeepAlive = Encoding.ASCII.GetBytes("\r\nConnection: keep-alive");
@@ -472,45 +470,12 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
 
         private void WriteChunked(ArraySegment<byte> data)
         {
-            SocketOutput.Write(BeginChunkBytes(data.Count), immediate: false);
+            SocketOutput.Write(data, immediate: false, chunk: true);
-            SocketOutput.Write(data, immediate: false);
-            SocketOutput.Write(_endChunkBytes, immediate: true);
         }
 
         private async Task WriteChunkedAsync(ArraySegment<byte> data, CancellationToken cancellationToken)
         {
-            await SocketOutput.WriteAsync(BeginChunkBytes(data.Count), immediate: false, cancellationToken: cancellationToken);
+            await SocketOutput.WriteAsync(data, immediate: false, chunk: true, cancellationToken: cancellationToken);
-            await SocketOutput.WriteAsync(data, immediate: false, cancellationToken: cancellationToken);
-            await SocketOutput.WriteAsync(_endChunkBytes, immediate: true, cancellationToken: cancellationToken);
-        }
-
-        public static ArraySegment<byte> BeginChunkBytes(int dataCount)
-        {
-            var bytes = new byte[10]
-            {
-                _hex[((dataCount >> 0x1c) & 0x0f)],
-                _hex[((dataCount >> 0x18) & 0x0f)],
-                _hex[((dataCount >> 0x14) & 0x0f)],
-                _hex[((dataCount >> 0x10) & 0x0f)],
-                _hex[((dataCount >> 0x0c) & 0x0f)],
-                _hex[((dataCount >> 0x08) & 0x0f)],
-                _hex[((dataCount >> 0x04) & 0x0f)],
-                _hex[((dataCount >> 0x00) & 0x0f)],
-                (byte)'\r',
-                (byte)'\n',
-            };
-
-            // Determine the most-significant non-zero nibble
-            int total, shift;
-            total = (dataCount > 0xffff) ? 0x10 : 0x00;
-            dataCount >>= total;
-            shift = (dataCount > 0x00ff) ? 0x08 : 0x00;
-            dataCount >>= shift;
-            total |= shift;
-            total |= (dataCount > 0x000f) ? 0x04 : 0x00;
-
-            var offset = 7 - (total >> 2);
-            return new ArraySegment<byte>(bytes, offset, 10 - offset);
         }
 
         private void WriteChunkedResponseSuffix()

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

@@ -13,8 +13,8 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
     /// </summary>
     public interface ISocketOutput
     {
-        void Write(ArraySegment<byte> buffer, bool immediate = true);
+        void Write(ArraySegment<byte> buffer, bool immediate = true, bool chunk = false);
-        Task WriteAsync(ArraySegment<byte> buffer, bool immediate = true, CancellationToken cancellationToken = default(CancellationToken));
+        Task WriteAsync(ArraySegment<byte> buffer, bool immediate = true, bool chunk = false, CancellationToken cancellationToken = default(CancellationToken));
 
         /// <summary>
         /// Returns an iterator pointing to the tail of the response buffer. Response data can be appended

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

@@ -30,7 +30,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
         private readonly IKestrelTrace _log;
         private readonly IThreadPool _threadPool;
 
-        // This locks all access to _tail, _isProducing and _returnFromOnProducingComplete.
+        // This locks all access to _tail and _lastStart.
         // _head does not require a lock, since it is only used in the ctor and uv thread.
         private readonly object _returnLock = new object();
 
@@ -79,6 +79,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
         public Task WriteAsync(
             ArraySegment<byte> buffer,
             bool immediate = true,
+            bool chunk = false,
             bool socketShutdownSend = false,
             bool socketDisconnect = false)
         {
@@ -90,7 +91,19 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                 if (buffer.Count > 0)
                 {
                     var tail = ProducingStart();
+                    if (chunk)
+                    {
+                        _numBytesPreCompleted += ChunkWriter.WriteBeginChunkBytes(ref tail, buffer.Count);
+                    }
+
                     tail.CopyFrom(buffer);
+
+                    if (chunk)
+                    {
+                        ChunkWriter.WriteEndChunkBytes(ref tail);
+                        _numBytesPreCompleted += 2;
+                    }
+
                     // We do our own accounting below
                     ProducingCompleteNoPreComplete(tail);
                 }
@@ -359,9 +372,9 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
             }
         }
 
-        void ISocketOutput.Write(ArraySegment<byte> buffer, bool immediate)
+        void ISocketOutput.Write(ArraySegment<byte> buffer, bool immediate, bool chunk)
         {
-            var task = WriteAsync(buffer, immediate);
+            var task = WriteAsync(buffer, immediate, chunk);
 
             if (task.Status == TaskStatus.RanToCompletion)
             {
@@ -373,9 +386,9 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
             }
         }
 
-        Task ISocketOutput.WriteAsync(ArraySegment<byte> buffer, bool immediate, CancellationToken cancellationToken)
+        Task ISocketOutput.WriteAsync(ArraySegment<byte> buffer, bool immediate, bool chunk, CancellationToken cancellationToken)
         {
-            return WriteAsync(buffer, immediate);
+            return WriteAsync(buffer, immediate, chunk);
         }
 
         private static void BytesBetween(MemoryPoolIterator2 start, MemoryPoolIterator2 end, out int bytes, out int buffers)

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

@@ -0,0 +1,38 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Linq;
+using System.Text;
+using Microsoft.AspNet.Server.Kestrel.Http;
+using Xunit;
+
+namespace Microsoft.AspNet.Server.KestrelTests
+{
+    public class ChunkWriterTests
+    {
+        [Theory]
+        [InlineData(1, "1\r\n")]
+        [InlineData(10, "a\r\n")]
+        [InlineData(0x08, "8\r\n")]
+        [InlineData(0x10, "10\r\n")]
+        [InlineData(0x080, "80\r\n")]
+        [InlineData(0x100, "100\r\n")]
+        [InlineData(0x0800, "800\r\n")]
+        [InlineData(0x1000, "1000\r\n")]
+        [InlineData(0x08000, "8000\r\n")]
+        [InlineData(0x10000, "10000\r\n")]
+        [InlineData(0x080000, "80000\r\n")]
+        [InlineData(0x100000, "100000\r\n")]
+        [InlineData(0x0800000, "800000\r\n")]
+        [InlineData(0x1000000, "1000000\r\n")]
+        [InlineData(0x08000000, "8000000\r\n")]
+        [InlineData(0x10000000, "10000000\r\n")]
+        [InlineData(0x7fffffffL, "7fffffff\r\n")]
+        public void ChunkedPrefixMustBeHexCrLfWithoutLeadingZeros(int dataCount, string expected)
+        {
+            var beginChunkBytes = ChunkWriter.BeginChunkBytes(dataCount);
+
+            Assert.Equal(Encoding.ASCII.GetBytes(expected), beginChunkBytes.ToArray());
+        }
+    }
+}

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

@@ -13,31 +13,6 @@ namespace Microsoft.AspNet.Server.KestrelTests
 {
     public class FrameTests
     {
-        [Theory]
-        [InlineData(1, "1\r\n")]
-        [InlineData(10, "a\r\n")]
-        [InlineData(0x08, "8\r\n")]
-        [InlineData(0x10, "10\r\n")]
-        [InlineData(0x080, "80\r\n")]
-        [InlineData(0x100, "100\r\n")]
-        [InlineData(0x0800, "800\r\n")]
-        [InlineData(0x1000, "1000\r\n")]
-        [InlineData(0x08000, "8000\r\n")]
-        [InlineData(0x10000, "10000\r\n")]
-        [InlineData(0x080000, "80000\r\n")]
-        [InlineData(0x100000, "100000\r\n")]
-        [InlineData(0x0800000, "800000\r\n")]
-        [InlineData(0x1000000, "1000000\r\n")]
-        [InlineData(0x08000000, "8000000\r\n")]
-        [InlineData(0x10000000, "10000000\r\n")]
-        [InlineData(0x7fffffffL, "7fffffff\r\n")]
-        public void ChunkedPrefixMustBeHexCrLfWithoutLeadingZeros(int dataCount, string expected)
-        {
-            var beginChunkBytes = Frame.BeginChunkBytes(dataCount);
-
-            Assert.Equal(Encoding.ASCII.GetBytes(expected), beginChunkBytes.ToArray());
-        }
-        
         [Theory]
         [InlineData("Cookie: \r\n\r\n", 1)]
         [InlineData("Cookie:\r\n\r\n", 1)]