Complete WriteAsync Tasks early when there are less than 64KB buffered

Author: halter73

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

@@ -12,23 +12,28 @@ namespace Microsoft.AspNet.Server.Kestrel.Http
     public class SocketOutput : ISocketOutput
     {
         private const int _maxPendingWrites = 3;
+        private const int _maxBytesBufferedBeforeThrottling = 65536 / 8;
 
         private readonly KestrelThread _thread;
         private readonly UvStreamHandle _socket;
 
-        private WriteContext _nextWriteContext;
+        // This locks all access to to all the below
+        private readonly object _lockObj = new object();
 
         // The number of write operations that have been scheduled so far
         // but have not completed.
-        private int _writesSending = 0;
+        private int _writesPending = 0;
 
-        // This locks all access to _nextWriteContext and _writesSending
-        private readonly object _lockObj = new object();
+        private int _numBytesBuffered = 0;
+        private Exception _lastWriteError;
+        private WriteContext _nextWriteContext;
+        private readonly Queue<CallbackContext> _callbacksPending;
 
         public SocketOutput(KestrelThread thread, UvStreamHandle socket)
         {
             _thread = thread;
             _socket = socket;
+            _callbacksPending = new Queue<CallbackContext>();
         }
 
         public void Write(ArraySegment<byte> buffer, Action<Exception, object> callback, object state)
@@ -40,11 +45,16 @@ public void Write(ArraySegment<byte> buffer, Action<Exception, object> callback,
 
             KestrelTrace.Log.ConnectionWrite(0, buffer.Count);
 
-            var context = new WriteOperation
+            var writeOp = new WriteOperation
+            {
+                Buffer = buffer
+            };
+
+            var callbackContext = new CallbackContext
             {
-                Buffer = buffer,
                 Callback = callback,
-                State = state
+                State = state,
+                BytesToWrite = buffer.Count
             };
 
             lock (_lockObj)
@@ -54,12 +64,26 @@ public void Write(ArraySegment<byte> buffer, Action<Exception, object> callback,
                     _nextWriteContext = new WriteContext(this);
                 }
 
-                _nextWriteContext.Operations.Add(context);
+                _nextWriteContext.Operations.Enqueue(writeOp);
+                _numBytesBuffered += buffer.Count;
+
+                // 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.
+                if (_lastWriteError == null &&
+                    _callbacksPending.Count == 0 &&
+                    _numBytesBuffered < _maxBytesBufferedBeforeThrottling)
+                {
+                    TriggerCallback(callbackContext);
+                }
+                else
+                {
+                    _callbacksPending.Enqueue(callbackContext);
+                }
 
-                if (_writesSending < _maxPendingWrites)
+                if (_writesPending < _maxPendingWrites)
                 {
                     ScheduleWrite();
-                    _writesSending++;
+                    _writesPending++;
                 }
             }
         }
@@ -87,7 +111,7 @@ private void WriteAllPending()
                 }
                 else
                 {
-                    _writesSending--;
+                    _writesPending--;
                     return;
                 }
             }
@@ -115,51 +139,66 @@ private void WriteAllPending()
                 {
                     // Lock instead of using Interlocked.Decrement so _writesSending
                     // doesn't change in the middle of executing other synchronized code.
-                    _writesSending--;
+                    _writesPending--;
                 }
 
                 throw;
             }
         }
 
         // This is called on the libuv event loop
-        private void OnWriteCompleted(List<WriteOperation> completedWrites, UvWriteReq req, int status, Exception error)
+        private void OnWriteCompleted(Queue<WriteOperation> completedWrites, UvWriteReq req, int status, Exception error)
         {
             lock (_lockObj)
             {
+                _lastWriteError = error;
+
                 if (_nextWriteContext != null)
                 {
                     ScheduleWrite();
                 }
                 else
                 {
-                    _writesSending--;
+                    _writesPending--;
+                }
+
+                foreach (var writeOp in completedWrites)
+                {
+                    _numBytesBuffered -= writeOp.Buffer.Count;
+                }
+
+                var bytesLeftToBuffer = _maxBytesBufferedBeforeThrottling - _numBytesBuffered;
+                while (_callbacksPending.Count > 0 && _callbacksPending.Peek().BytesToWrite < bytesLeftToBuffer)
+                {
+                    var context = _callbacksPending.Dequeue();
+                    TriggerCallback(context);
                 }
             }
 
             req.Dispose();
+        }
 
-            foreach (var writeOp in completedWrites)
+        private void TriggerCallback(CallbackContext context)
+        {
+            context.Error = _lastWriteError;
+            ThreadPool.QueueUserWorkItem(obj =>
             {
-                KestrelTrace.Log.ConnectionWriteCallback(0, status);
-                //NOTE: pool this?
-
-                // Get off the event loop before calling user code!
-                writeOp.Error = error;
-                ThreadPool.QueueUserWorkItem(obj =>
-                {
-                    var op = (WriteOperation)obj;
-                    op.Callback(op.Error, op.State);
-                }, writeOp);
-            }
+                var c = (CallbackContext)obj;
+                c.Callback(c.Error, c.State);
+            }, context);
         }
 
-        private class WriteOperation
+        private class CallbackContext
         {
-            public ArraySegment<byte> Buffer;
             public Exception Error;
             public Action<Exception, object> Callback;
             public object State;
+            public int BytesToWrite;
+        }
+
+        private class WriteOperation
+        {
+            public ArraySegment<byte> Buffer;
         }
 
         private class WriteContext
@@ -171,13 +210,13 @@ public WriteContext(SocketOutput self)
                 WriteReq = new UvWriteReq();
                 WriteReq.Init(self._thread.Loop);
 
-                Operations = new List<WriteOperation>();
+                Operations = new Queue<WriteOperation>();
             }
 
             public SocketOutput Self;
 
             public UvWriteReq WriteReq;
-            public List<WriteOperation> Operations;
+            public Queue<WriteOperation> Operations;
         }
     }
 }