SocketOutput to Tasks

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

@@ -31,7 +31,7 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
         private int _numBytesPreCompleted = 0;
         private Exception _lastWriteError;
         private WriteContext _nextWriteContext;
-        private readonly Queue<CallbackContext> _callbacksPending;
+        private readonly Queue<TaskCompletionSource<object>> _tasksPending;
 
         public SocketOutput(KestrelThread thread, UvStreamHandle socket, long connectionId, IKestrelTrace log)
         {
@@ -39,13 +39,11 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
             _socket = socket;
             _connectionId = connectionId;
             _log = log;
-            _callbacksPending = new Queue<CallbackContext>();
+            _tasksPending = new Queue<TaskCompletionSource<object>>();
         }
 
-        public void Write(
+        public Task WriteAsync(
             ArraySegment<byte> buffer,
-            Action<Exception, object, bool> callback,
-            object state,
             bool immediate = true,
             bool socketShutdownSend = false,
             bool socketDisconnect = false)
@@ -58,8 +56,8 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                 buffer = new ArraySegment<byte>(copy);
                 _log.ConnectionWrite(_connectionId, buffer.Count);
             }
-
+            
-            bool triggerCallbackNow = false;
+            TaskCompletionSource<object> tcs = null;
 
             lock (_lockObj)
             {
@@ -80,23 +78,25 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                 {
                     _nextWriteContext.SocketDisconnect = true;
                 }
+
                 // Complete the write task immediately if all previous write tasks have been completed,
                 // the buffers haven't grown too large, and the last write to the socket succeeded.
-                triggerCallbackNow = _lastWriteError == null &&
+                if (_lastWriteError == null &&
-                                     _callbacksPending.Count == 0 &&
+                    _tasksPending.Count == 0 &&
-                                     _numBytesPreCompleted + buffer.Count <= _maxBytesPreCompleted;
+                    _numBytesPreCompleted + buffer.Count <= _maxBytesPreCompleted)
-                if (triggerCallbackNow)
                 {
                     _numBytesPreCompleted += buffer.Count;
                 }
+                else if (immediate)
+                {
+                    // immediate write, which is not eligable for instant completion above
+                    tcs = new TaskCompletionSource<object>(buffer.Count);
+                    _tasksPending.Enqueue(tcs);
+                }
                 else
                 {
-                    _callbacksPending.Enqueue(new CallbackContext
+                    // immediate==false calls always return complete tasks, because there is guaranteed
-                    {
+                    // to be a subsequent immediate==true call which will go down one of the previous code-paths
-                        Callback = callback,
-                        State = state,
-                        BytesToWrite = buffer.Count
-                    });
                 }
 
                 if (_writesPending < _maxPendingWrites && immediate)
@@ -106,12 +106,8 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                 }
             }
 
-            // Make sure we call user code outside of the lock.
+            // Return TaskCompletionSource's Task if set, otherwise completed Task 
-            if (triggerCallbackNow)
+            return tcs?.Task ?? TaskUtilities.CompletedTask;
-            {
-                // callback(error, state, calledInline)
-                callback(null, state, true);
-            }
         }
 
         public void End(ProduceEndType endType)
@@ -119,13 +115,13 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
             switch (endType)
             {
                 case ProduceEndType.SocketShutdownSend:
-                    Write(default(ArraySegment<byte>), (error, state, calledInline) => { }, null,
+                    WriteAsync(default(ArraySegment<byte>),
                         immediate: true,
                         socketShutdownSend: true,
                         socketDisconnect: false);
                     break;
                 case ProduceEndType.SocketDisconnect:
-                    Write(default(ArraySegment<byte>), (error, state, calledInline) => { }, null,
+                    WriteAsync(default(ArraySegment<byte>),
                         immediate: true,
                         socketShutdownSend: false,
                         socketDisconnect: true);
@@ -198,20 +194,30 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                     // completed writes that we haven't triggered callbacks for yet.
                     _numBytesPreCompleted -= writeBuffer.Count;
                 }
-
+                
-
                 // bytesLeftToBuffer can be greater than _maxBytesPreCompleted
                 // This allows large writes to complete once they've actually finished.
                 var bytesLeftToBuffer = _maxBytesPreCompleted - _numBytesPreCompleted;
-                while (_callbacksPending.Count > 0 &&
+                while (_tasksPending.Count > 0 &&
-                       _callbacksPending.Peek().BytesToWrite <= bytesLeftToBuffer)
+                       (int)(_tasksPending.Peek().Task.AsyncState) <= bytesLeftToBuffer)
                 {
-                    var callbackContext = _callbacksPending.Dequeue();
+                    var tcs = _tasksPending.Dequeue();
 
-                    _numBytesPreCompleted += callbackContext.BytesToWrite;
+                    _numBytesPreCompleted += (int)(tcs.Task.AsyncState);
 
-                    // callback(error, state, calledInline)
+                    if (error == null)
-                    callbackContext.Callback(_lastWriteError, callbackContext.State, false);
+                    {
+                        ThreadPool.QueueUserWorkItem(
+                            (o) => ((TaskCompletionSource<object>)o).SetResult(null), 
+                            tcs);
+                    }
+                    else
+                    {
+                        // error is closure captured 
+                        ThreadPool.QueueUserWorkItem(
+                            (o) => ((TaskCompletionSource<object>)o).SetException(error), 
+                            tcs);
+                    }
                 }
 
                 // Now that the while loop has completed the following invariants should hold true:
@@ -222,104 +228,21 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
 
         void ISocketOutput.Write(ArraySegment<byte> buffer, bool immediate)
         {
-            if (!immediate)
+            var task = WriteAsync(buffer, immediate);
+
+            if (task.Status == TaskStatus.RanToCompletion)
             {
-                // immediate==false calls always return complete tasks, because there is guaranteed
-                // to be a subsequent immediate==true call which will go down the following code-path
-                Write(
-                    buffer,
-                    (error, state, calledInline) => { },
-                    null,
-                    immediate: false);
                 return;
             }
-
+            else
-            // TODO: Optimize task being used, and remove callback model from the underlying Write
-            var tcs = new TaskCompletionSource<int>();
-
-            Write(
-                buffer,
-                (error, state, calledInline) =>
-                {
-                    if (error != null)
-                    {
-                        tcs.SetException(error);
-                    }
-                    else
-                    {
-                        tcs.SetResult(0);
-                    }
-                },
-                tcs,
-                immediate: true);
-
-            if (tcs.Task.Status != TaskStatus.RanToCompletion)
             {
-                tcs.Task.GetAwaiter().GetResult();
+                task.GetAwaiter().GetResult();
             }
         }
 
         Task ISocketOutput.WriteAsync(ArraySegment<byte> buffer, bool immediate, CancellationToken cancellationToken)
         {
-            if (!immediate)
+            return WriteAsync(buffer, immediate);
-            {
-                // immediate==false calls always return complete tasks, because there is guaranteed
-                // to be a subsequent immediate==true call which will go down the following code-path
-                Write(
-                    buffer,
-                    (error, state, calledInline) => { },
-                    null,
-                    immediate: false);
-                return TaskUtilities.CompletedTask;
-            }
-
-            // TODO: Optimize task being used, and remove callback model from the underlying Write
-            var tcs = new TaskCompletionSource<int>();
-
-            Write(
-                buffer,
-                (error, state, calledInline) =>
-                {
-                    if (!calledInline)
-                    {
-                        ThreadPool.QueueUserWorkItem(state2 =>
-                        {
-                            var tcs2 = (TaskCompletionSource<int>)state2;
-                            if (error != null)
-                            {
-                                tcs2.SetException(error);
-                            }
-                            else
-                            {
-                                tcs2.SetResult(0);
-                            }
-                        }, state);
-                    }
-                    else
-                    {
-                        var tcs2 = (TaskCompletionSource<int>)state;
-                        if (error != null)
-                        {
-                            tcs2.SetException(error);
-                        }
-                        else
-                        {
-                            tcs2.SetResult(0);
-                        }
-                    }
-                },
-                tcs,
-                immediate: true);
-
-            return tcs.Task;
-        }
-
-        private class CallbackContext
-        {
-            // callback(error, state, calledInline)
-            public Action<Exception, object, bool> Callback;
-            public object State;
-            public int BytesToWrite;
         }
 
         private class WriteContext
@@ -341,15 +264,6 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
                 Buffers = new Queue<ArraySegment<byte>>();
             }
 
-            /// <summary>
-            /// Perform any actions needed by this work item. The individual tasks are non-blocking and
-            /// will continue through to each other in order.
-            /// </summary>
-            public void Execute()
-            {
-                DoWriteIfNeeded();
-            }
-
             /// <summary>
             /// First step: initiate async write if needed, otherwise go to next step
             /// </summary>

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

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNet.Server.Kestrel;
 using Microsoft.AspNet.Server.Kestrel.Http;
 using Microsoft.AspNet.Server.Kestrel.Infrastructure;
@@ -45,15 +46,15 @@ namespace Microsoft.AspNet.Server.KestrelTests
                 var bufferSize = 1048576;
                 var buffer = new ArraySegment<byte>(new byte[bufferSize], 0, bufferSize);
                 var completedWh = new ManualResetEventSlim();
-                Action<Exception, object, bool> onCompleted = (ex, state, calledInline) =>
-                {
-                    Assert.Null(ex);
-                    Assert.Null(state);
-                    completedWh.Set();
-                };
 
                 // Act
-                socketOutput.Write(buffer, onCompleted, null);
+                socketOutput.WriteAsync(buffer).ContinueWith(
+                    (t) =>
+                    {
+                        Assert.Null(t.Exception);
+                        completedWh.Set();
+                    }
+                );
 
                 // Assert
                 Assert.True(completedWh.Wait(1000));
@@ -89,22 +90,21 @@ namespace Microsoft.AspNet.Server.KestrelTests
                 var bufferSize = maxBytesPreCompleted;
                 var buffer = new ArraySegment<byte>(new byte[bufferSize], 0, bufferSize);
                 var completedWh = new ManualResetEventSlim();
-                Action<Exception, object, bool> onCompleted = (ex, state, calledInline) =>
+                Action<Task> onCompleted = (Task t) =>
                 {
-                    Assert.Null(ex);
+                    Assert.Null(t.Exception);
-                    Assert.Null(state);
                     completedWh.Set();
                 };
 
                 // Act 
-                socketOutput.Write(buffer, onCompleted, null);
+                socketOutput.WriteAsync(buffer).ContinueWith(onCompleted);
                 // Assert
                 // The first write should pre-complete since it is <= _maxBytesPreCompleted.
                 Assert.True(completedWh.Wait(1000));
                 // Arrange
                 completedWh.Reset();
                 // Act
-                socketOutput.Write(buffer, onCompleted, null);
+                socketOutput.WriteAsync(buffer).ContinueWith(onCompleted);
                 // Assert 
                 // Too many bytes are already pre-completed for the second write to pre-complete.
                 Assert.False(completedWh.Wait(1000));