diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs index 3ea812fd14..efbd159c7f 100644 --- a/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs +++ b/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs @@ -464,14 +464,14 @@ namespace Microsoft.AspNet.Server.Kestrel.Http SocketOutput.Write(data, immediate: false, chunk: true); } - private async Task WriteChunkedAsync(ArraySegment data, CancellationToken cancellationToken) + private Task WriteChunkedAsync(ArraySegment data, CancellationToken cancellationToken) { - await SocketOutput.WriteAsync(data, immediate: false, chunk: true, cancellationToken: cancellationToken); + return SocketOutput.WriteAsync(data, immediate: false, chunk: true, cancellationToken: cancellationToken); } - private void WriteChunkedResponseSuffix() + private Task WriteChunkedResponseSuffix() { - SocketOutput.Write(_endChunkedResponseBytes, immediate: true); + return SocketOutput.WriteAsync(_endChunkedResponseBytes, immediate: true); } private static ArraySegment CreateAsciiByteArraySegment(string text) @@ -562,31 +562,41 @@ namespace Microsoft.AspNet.Server.Kestrel.Http 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( diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/SocketOutput.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/SocketOutput.cs index cf3f8dc528..3c39954494 100644 --- a/src/Microsoft.AspNet.Server.Kestrel/Http/SocketOutput.cs +++ b/src/Microsoft.AspNet.Server.Kestrel/Http/SocketOutput.cs @@ -48,7 +48,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http private int _numBytesPreCompleted = 0; private Exception _lastWriteError; private WriteContext _nextWriteContext; - private readonly Queue> _tasksPending; + private readonly Queue _tasksPending; private readonly Queue _writeContextPool; private readonly Queue _writeReqPool; @@ -68,7 +68,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http _connectionId = connectionId; _log = log; _threadPool = threadPool; - _tasksPending = new Queue>(_initialTaskQueues); + _tasksPending = new Queue(_initialTaskQueues); _writeContextPool = new Queue(_maxPooledWriteContexts); _writeReqPool = writeReqPool; @@ -81,7 +81,8 @@ namespace Microsoft.AspNet.Server.Kestrel.Http bool immediate = true, bool chunk = false, bool socketShutdownSend = false, - bool socketDisconnect = false) + bool socketDisconnect = false, + bool isSync = false) { TaskCompletionSource tcs = null; var scheduleWrite = false; @@ -147,7 +148,11 @@ namespace Microsoft.AspNet.Server.Kestrel.Http { // immediate write, which is not eligable for instant completion above tcs = new TaskCompletionSource(buffer.Count); - _tasksPending.Enqueue(tcs); + _tasksPending.Enqueue(new WaitingTask() { + CompletionSource = tcs, + BytesToWrite = buffer.Count, + IsSync = isSync + }); } if (!_writePending && immediate) @@ -316,21 +321,35 @@ namespace Microsoft.AspNet.Server.Kestrel.Http // This allows large writes to complete once they've actually finished. var bytesLeftToBuffer = _maxBytesPreCompleted - _numBytesPreCompleted; while (_tasksPending.Count > 0 && - (int)(_tasksPending.Peek().Task.AsyncState) <= bytesLeftToBuffer) + (_tasksPending.Peek().BytesToWrite) <= bytesLeftToBuffer) { - var tcs = _tasksPending.Dequeue(); - var bytesToWrite = (int)tcs.Task.AsyncState; + var waitingTask = _tasksPending.Dequeue(); + var bytesToWrite = waitingTask.BytesToWrite; _numBytesPreCompleted += bytesToWrite; bytesLeftToBuffer -= bytesToWrite; if (_lastWriteError == null) { - _threadPool.Complete(tcs); + if (waitingTask.IsSync) + { + waitingTask.CompletionSource.TrySetResult(null); + } + else + { + _threadPool.Complete(waitingTask.CompletionSource); + } } else { - _threadPool.Error(tcs, _lastWriteError); + if (waitingTask.IsSync) + { + waitingTask.CompletionSource.TrySetException(_lastWriteError); + } + else + { + _threadPool.Error(waitingTask.CompletionSource, _lastWriteError); + } } } @@ -374,16 +393,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http void ISocketOutput.Write(ArraySegment buffer, bool immediate, bool chunk) { - var task = WriteAsync(buffer, immediate, chunk); - - if (task.Status == TaskStatus.RanToCompletion) - { - return; - } - else - { - task.GetAwaiter().GetResult(); - } + WriteAsync(buffer, immediate, chunk, isSync: true).GetAwaiter().GetResult(); } Task ISocketOutput.WriteAsync(ArraySegment buffer, bool immediate, bool chunk, CancellationToken cancellationToken) @@ -634,5 +644,12 @@ namespace Microsoft.AspNet.Server.Kestrel.Http ShutdownSendStatus = 0; } } + + private struct WaitingTask + { + public bool IsSync; + public int BytesToWrite; + public TaskCompletionSource CompletionSource; + } } }