P/Invoke cleanup

src/Kestrel/ServerRequest.cs

@@ -230,7 +230,7 @@ async Task<Stream> IHttpUpgradeFeature.UpgradeAsync()
                     _frame.ResponseHeaders["Upgrade"] = values;
                 }
             }
-            _frame.ProduceStart();
+            await _frame.ProduceStartAsync();
             return _frame.DuplexStream;
         }
     }

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

@@ -35,78 +35,80 @@ public interface IConnectionControl
     {
         void Pause();
         void Resume();
-        void End(ProduceEndType endType);
+        Task EndAsync(ProduceEndType endType);
+        bool IsInKeepAlive { get; }
     }
 
     public class Connection : ConnectionContext, IConnectionControl
     {
-        private static readonly Action<UvStreamHandle, int, Exception, object> _readCallback = ReadCallback;
-        private static readonly Func<UvStreamHandle, int, object, Libuv.uv_buf_t> _allocCallback = AllocCallback;
+        private readonly Action<int, Exception> _readCallback;
+        private readonly Func<int, UvBuffer> _allocCallback;
+        private readonly UvTcpStreamHandle _socket;
+        private readonly Listener _listener;
 
-        private static Libuv.uv_buf_t AllocCallback(UvStreamHandle handle, int suggestedSize, object state)
-        {
-            return ((Connection)state).OnAlloc(handle, suggestedSize);
-        }
-
-        private static void ReadCallback(UvStreamHandle handle, int nread, Exception error, object state)
-        {
-            ((Connection)state).OnRead(handle, nread, error);
-        }
-
-        private readonly UvStreamHandle _socket;
+        private UvReadHandle _read;
         private Frame _frame;
         long _connectionId;
+        private bool _isInKeepAlive;
 
-        public Connection(ListenerContext context, UvStreamHandle socket) : base(context)
+        public Connection(Listener listener, UvTcpStreamHandle socket) : base(listener)
         {
+            _readCallback = OnRead;
+            _allocCallback = OnAlloc;
             _socket = socket;
             ConnectionControl = this;
-        }
+            _listener = listener;
 
-        public void Start()
-        {
             KestrelTrace.Log.ConnectionStart(_connectionId);
 
             SocketInput = new SocketInput(Memory);
             SocketOutput = new SocketOutput(Thread, _socket);
             _frame = new Frame(this);
-            _socket.ReadStart(_allocCallback, _readCallback, this);
+            _read = new UvReadHandle(_socket, _allocCallback, _readCallback);
+            listener.AddConnection(this);
         }
 
-        private Libuv.uv_buf_t OnAlloc(UvStreamHandle handle, int suggestedSize)
+        private UvBuffer OnAlloc(int suggestedSize)
         {
-            return handle.Libuv.buf_init(
-                SocketInput.Pin(2048),
-                2048);
+            const int bufferSize = 2048;
+            return new UvBuffer(SocketInput.Pin(bufferSize), bufferSize);
         }
 
-        private void OnRead(UvStreamHandle handle, int status, Exception error)
+        private void OnRead(int status, Exception error)
         {
             SocketInput.Unpin(status);
 
             var normalRead = error == null && status > 0;
             var normalDone = status == 0 || status == -4077 || status == -4095;
-            var errorDone = !(normalDone || normalRead);
 
-            if (normalRead)
-            {
-                KestrelTrace.Log.ConnectionRead(_connectionId, status);
-            }
-            else if (normalDone || errorDone)
+            if (!normalRead)
             {
                 KestrelTrace.Log.ConnectionReadFin(_connectionId);
                 SocketInput.RemoteIntakeFin = true;
-                _socket.ReadStop();
+                if (status != -4095)
+                {
+                    _read.Dispose();
+                    _read = null;
+                    _listener.RemoveConnection(this);
+                    _socket.Dispose();
+
+                    // Not sure if this is right
+                    // It should be, but there are some interesting code paths
+                    // while reading the message body regarding status == 0 && RemoteIntakeFin
+                    return;
+                }
 
-                if (errorDone && error != null)
+                if (!normalDone && error != null)
                 {
                     Trace.WriteLine("Connection.OnRead " + error.ToString());
                 }
             }
 
+            KestrelTrace.Log.ConnectionRead(_connectionId, status);
 
             try
             {
+                _isInKeepAlive = false;
                 _frame.Consume();
             }
             catch (Exception ex)
@@ -115,56 +117,66 @@ private void OnRead(UvStreamHandle handle, int status, Exception error)
             }
         }
 
+        bool IConnectionControl.IsInKeepAlive => _isInKeepAlive;
+
         void IConnectionControl.Pause()
         {
             KestrelTrace.Log.ConnectionPause(_connectionId);
-            _socket.ReadStop();
+            Debug.Assert(_read != null);
+            _read.Dispose();
+            _read = null;
         }
 
         void IConnectionControl.Resume()
         {
             KestrelTrace.Log.ConnectionResume(_connectionId);
-            _socket.ReadStart(_allocCallback, _readCallback, this);
+            Debug.Assert(_read == null);
+            _read = new UvReadHandle(_socket, _allocCallback, _readCallback);
         }
 
-        void IConnectionControl.End(ProduceEndType endType)
+        async Task IConnectionControl.EndAsync(ProduceEndType endType)
         {
             switch (endType)
             {
                 case ProduceEndType.SocketShutdownSend:
                     KestrelTrace.Log.ConnectionWriteFin(_connectionId, 0);
-                    Thread.Post(
-                        x =>
-                        {
-                            KestrelTrace.Log.ConnectionWriteFin(_connectionId, 1);
-                            var self = (Connection)x;
-                            var shutdown = new UvShutdownReq();
-                            shutdown.Init(self.Thread.Loop);
-                            shutdown.Shutdown(self._socket, (req, status, state) =>
+                    await Thread.PostAsync(() =>
+                    {
+                        if (_read == null)
+                            return;
+
+                        KestrelTrace.Log.ConnectionWriteFin(_connectionId, 1);
+                        new UvShutdownReq(
+                            Thread.Loop,
+                            _socket,
+                            (req, status) =>
                             {
                                 KestrelTrace.Log.ConnectionWriteFin(_connectionId, 1);
-                                req.Dispose();
-                            }, null);
-                        },
-                        this);
+                                // This connection is now done
+                            });
+                    });
                     break;
                 case ProduceEndType.ConnectionKeepAlive:
                     KestrelTrace.Log.ConnectionKeepAlive(_connectionId);
                     _frame = new Frame(this);
-                    Thread.Post(
-                        x => ((Frame)x).Consume(),
-                        _frame);
+                    _isInKeepAlive = true;
+                    await Thread.PostAsync(_frame.Consume);
                     break;
                 case ProduceEndType.SocketDisconnect:
                     KestrelTrace.Log.ConnectionDisconnect(_connectionId);
-                    Thread.Post(
-                        x =>
-                        {
-                            KestrelTrace.Log.ConnectionStop(_connectionId);
-                            ((UvHandle)x).Dispose();
-                        },
-                        _socket);
+                    await Thread.PostAsync(() =>
+                    {
+                        _listener.RemoveConnection(this);
+                        if (_read == null)
+                            return;
+
+                        _read.Dispose();
+                        _socket.Dispose();
+                        KestrelTrace.Log.ConnectionStop(_connectionId);
+                    });
                     break;
+                default:
+                    throw new ArgumentException(nameof(endType));
             }
         }
     }

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

@@ -38,8 +38,8 @@ public FrameContext(ConnectionContext context) : base(context)
 
     public interface IFrameControl
     {
-        void ProduceContinue();
-        void Write(ArraySegment<byte> data, Action<Exception, object> callback, object state);
+        Task ProduceContinueAsync();
+        Task WriteAsync(ArraySegment<byte> data);
     }
 
     public class Frame : FrameContext, IFrameControl
@@ -230,53 +230,50 @@ private async Task ExecuteAsync()
             }
             finally
             {
-                ProduceEnd(error);
+                await ProduceEndAsync(error);
             }
         }
 
 
-        public void Write(ArraySegment<byte> data, Action<Exception, object> callback, object state)
+        public async Task WriteAsync(ArraySegment<byte> data)
         {
-            ProduceStart();
-            SocketOutput.Write(data, callback, state);
+            await ProduceStartAsync();
+            await SocketOutput.WriteAsync(data);
         }
 
-        public void Upgrade(IDictionary<string, object> options, Func<object, Task> callback)
+        public Task Upgrade(IDictionary<string, object> options, Func<object, Task> callback)
         {
             _keepAlive = false;
-            ProduceStart();
+            return ProduceStartAsync();
 
             // NOTE: needs changes
             //_upgradeTask = callback(_callContext);
         }
 
         byte[] _continueBytes = Encoding.ASCII.GetBytes("HTTP/1.1 100 Continue\r\n\r\n");
 
-        public void ProduceContinue()
+        public Task ProduceContinueAsync()
         {
-            if (_resultStarted) return;
+            if (_resultStarted)
+                return Task.FromResult(0);
 
             string[] expect;
             if (HttpVersion.Equals("HTTP/1.1") &&
                 RequestHeaders.TryGetValue("Expect", out expect) &&
                 (expect.FirstOrDefault() ?? "").Equals("100-continue", StringComparison.OrdinalIgnoreCase))
             {
-                SocketOutput.Write(
-                    new ArraySegment<byte>(_continueBytes, 0, _continueBytes.Length),
-                    (error, _) =>
-                    {
-                        if (error != null)
-                        {
-                            Trace.WriteLine("ProduceContinue " + error.ToString());
-                        }
-                    },
-                    null);
+                return SocketOutput.WriteAsync(
+                    new ArraySegment<byte>(_continueBytes, 0, _continueBytes.Length));
             }
+
+            return Task.FromResult(0);
         }
 
-        public void ProduceStart()
+        public async Task ProduceStartAsync()
         {
-            if (_resultStarted) return;
+            if (_resultStarted)
+                return;
+
             _resultStarted = true;
 
             FireOnSendingHeaders();
@@ -286,30 +283,31 @@ public void ProduceStart()
             var status = ReasonPhrases.ToStatus(StatusCode, ReasonPhrase);
 
             var responseHeader = CreateResponseHeader(status, ResponseHeaders);
-            SocketOutput.Write(
-                responseHeader.Item1,
-                (error, x) =>
-                {
-                    if (error != null)
-                    {
-                        Trace.WriteLine("ProduceStart " + error.ToString());
-                    }
-                    ((IDisposable)x).Dispose();
-                },
-                responseHeader.Item2);
+            try
+            {
+                await SocketOutput.WriteAsync(responseHeader.Item1);
+            }
+            finally
+            {
+                responseHeader.Item2.Dispose();
+            }
         }
 
-        public void ProduceEnd(Exception ex)
+        public Task ProduceEndAsync(Exception ex)
         {
-            ProduceStart();
+            var tasks = new List<Task>(3);
+            tasks.Add(ProduceStartAsync());
 
             if (!_keepAlive)
             {
-                ConnectionControl.End(ProduceEndType.SocketShutdownSend);
+                tasks.Add(ConnectionControl.EndAsync(ProduceEndType.SocketShutdownSend));
             }
 
             //NOTE: must finish reading request body
-            ConnectionControl.End(_keepAlive ? ProduceEndType.ConnectionKeepAlive : ProduceEndType.SocketDisconnect);
+            tasks.Add(ConnectionControl.EndAsync(
+                _keepAlive ? ProduceEndType.ConnectionKeepAlive : ProduceEndType.SocketDisconnect));
+
+            return Task.WhenAll(tasks);
         }
 
         private Tuple<ArraySegment<byte>, IDisposable> CreateResponseHeader(

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

@@ -24,23 +24,7 @@ public override void Flush()
 
         public override Task FlushAsync(CancellationToken cancellationToken)
         {
-            var tcs = new TaskCompletionSource<int>();
-            _context.FrameControl.Write(
-                new ArraySegment<byte>(new byte[0]),
-                (error, arg) =>
-                {
-                    var tcsArg = (TaskCompletionSource<int>)arg;
-                    if (error != null)
-                    {
-                        tcsArg.SetException(error);
-                    }
-                    else
-                    {
-                        tcsArg.SetResult(0);
-                    }
-                },
-                tcs);
-            return tcs.Task;
+            return _context.FrameControl.WriteAsync(new ArraySegment<byte>(new byte[0]));
         }
 
         public override long Seek(long offset, SeekOrigin origin)
@@ -65,23 +49,7 @@ public override void Write(byte[] buffer, int offset, int count)
 
         public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
-            var tcs = new TaskCompletionSource<int>();
-            _context.FrameControl.Write(
-                new ArraySegment<byte>(buffer, offset, count),
-                (error, arg) =>
-                {
-                    var tcsArg = (TaskCompletionSource<int>)arg;
-                    if (error != null)
-                    {
-                        tcsArg.SetException(error);
-                    }
-                    else
-                    {
-                        tcsArg.SetResult(0);
-                    }
-                },
-                tcs);
-            return tcs.Task;
+            return _context.FrameControl.WriteAsync(new ArraySegment<byte>(buffer, offset, count));
         }
 
         public override bool CanRead