Make chunked writes atomic

This commit is contained in:
Stephen Halter 2016-01-07 20:44:21 -08:00
parent daa2b7e383
commit ab5ef547e1
7 changed files with 144 additions and 73 deletions

View File

@ -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;
}

View File

@ -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);
}
}
}

View File

@ -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);
SocketOutput.Write(_endChunkBytes, immediate: true);
SocketOutput.Write(data, immediate: false, chunk: 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, 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);
await SocketOutput.WriteAsync(data, immediate: false, chunk: true, cancellationToken: cancellationToken);
}
private void WriteChunkedResponseSuffix()

View File

@ -13,8 +13,8 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
/// </summary>
public interface ISocketOutput
{
void Write(ArraySegment<byte> buffer, bool immediate = true);
Task WriteAsync(ArraySegment<byte> buffer, bool immediate = true, CancellationToken cancellationToken = default(CancellationToken));
void Write(ArraySegment<byte> buffer, bool immediate = true, bool chunk = false);
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

View File

@ -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)

View File

@ -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());
}
}
}

View File

@ -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)]