Await writes so natural back pressure is applied (#1195)

src/System.IO.Pipelines.Networking.Libuv/Internal/LibuvAwaitable.cs

@@ -0,0 +1,72 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Net;
+using System.Threading.Tasks;
+using System.IO.Pipelines.Networking.Libuv.Interop;
+using System.Runtime.CompilerServices;
+using System.Threading;
+
+namespace System.IO.Pipelines.Networking.Libuv
+{
+    public class LibuvAwaitable<TRequest> : ICriticalNotifyCompletion where TRequest : UvRequest
+    {
+        private readonly static Action CALLBACK_RAN = () => { };
+
+        private Action _callback;
+
+        private Exception _exception;
+
+        private int _status;
+
+        public static Action<TRequest, int, object> Callback = (req, status, state) =>
+        {
+            var awaitable = (LibuvAwaitable<TRequest>)state;
+
+            Exception exception;
+            req.Libuv.Check(status, out exception);
+            awaitable._exception = exception;
+            awaitable._status = status;
+
+            var continuation = Interlocked.Exchange(ref awaitable._callback, CALLBACK_RAN);
+
+            continuation?.Invoke();
+        };
+
+        public LibuvAwaitable<TRequest> GetAwaiter() => this;
+        public bool IsCompleted => _callback == CALLBACK_RAN;
+
+        public int GetResult()
+        {
+            var exception = _exception;
+            var status = _status;
+
+            // Reset the awaitable state
+            _exception = null;
+            _status = 0;
+            _callback = null;
+
+            if (exception != null)
+            {
+                throw exception;
+            }
+
+            return status;
+        }
+
+        public void OnCompleted(Action continuation)
+        {
+            if (_callback == CALLBACK_RAN ||
+                Interlocked.CompareExchange(ref _callback, continuation, null) == CALLBACK_RAN)
+            {
+                Task.Run(continuation);
+            }
+        }
+
+        public void UnsafeOnCompleted(Action continuation)
+        {
+            OnCompleted(continuation);
+        }
+    }
+}

src/System.IO.Pipelines.Networking.Libuv/Interop/UvWriteReq.cs

@@ -17,12 +17,13 @@ public class UvWriteReq : UvRequest
         private IntPtr _bufs;
 
         private Action<UvWriteReq, int, object> _callback;
-        private PreservedBuffer _buffer;
         private object _state;
         private const int BUFFER_COUNT = 4;
 
         private List<GCHandle> _pins = new List<GCHandle>(BUFFER_COUNT + 1);
 
+        private LibuvAwaitable<UvWriteReq> _awaitable = new LibuvAwaitable<UvWriteReq>();
+
         public UvWriteReq() : base()
         {
         }
@@ -38,18 +39,22 @@ public void Init(UvLoopHandle loop)
             _bufs = handle + requestSize;
         }
 
-        public unsafe void Write(
+        public unsafe LibuvAwaitable<UvWriteReq> WriteAsync(
+            UvStreamHandle handle,
+            ReadableBuffer buffer)
+        {
+            Write(handle, buffer, LibuvAwaitable<UvWriteReq>.Callback, _awaitable);
+            return _awaitable;
+        }
+
+        private unsafe void Write(
             UvStreamHandle handle,
             ReadableBuffer buffer,
             Action<UvWriteReq, int, object> callback,
             object state)
         {
             try
             {
-                // Preserve the buffer for the async call
-                _buffer = buffer.Preserve();
-                buffer = _buffer.Buffer;
-
                 int nBuffers = 0;
                 if (buffer.IsSingleSpan)
                 {
@@ -114,7 +119,6 @@ public unsafe void Write(
             {
                 _callback = null;
                 _state = null;
-                _buffer.Dispose();
                 Unpin(this);
                 throw;
             }
@@ -134,8 +138,6 @@ private static void UvWriteCb(IntPtr ptr, int status)
             var req = FromIntPtr<UvWriteReq>(ptr);
             Unpin(req);
 
-            req._buffer.Dispose();
-
             var callback = req._callback;
             req._callback = null;
 

src/System.IO.Pipelines.Networking.Libuv/UvTcpConnection.cs

@@ -15,17 +15,13 @@ public class UvTcpConnection : IPipeConnection
 
         private static readonly Action<UvStreamHandle, int, object> _readCallback = ReadCallback;
         private static readonly Func<UvStreamHandle, int, object, Uv.uv_buf_t> _allocCallback = AllocCallback;
-        private static readonly Action<UvWriteReq, int, object> _writeCallback = WriteCallback;
 
         protected readonly IPipe _input;
         protected readonly IPipe _output;
         private readonly UvThread _thread;
         private readonly UvTcpHandle _handle;
         private volatile bool _stopping;
 
-        private int _pendingWrites;
-
-        private TaskCompletionSource<object> _drainWrites;
         private Task _sendingTask;
         private WritableBuffer? _inputBuffer;
 
@@ -104,7 +100,7 @@ private async Task ProcessWrites()
 
                         if (!buffer.IsEmpty)
                         {
-                            BeginWrite(buffer);
+                            await WriteAsync(buffer);
                         }
                     }
                     finally
@@ -121,47 +117,25 @@ private async Task ProcessWrites()
             {
                 _output.Reader.Complete();
 
-                // Drain the pending writes
-                if (_pendingWrites > 0)
-                {
-                    _drainWrites = new TaskCompletionSource<object>();
-
-                    await _drainWrites.Task;
-                }
-
                 _handle.Dispose();
 
                 // We'll never call the callback after disposing the handle
                 _input.Writer.Complete();
             }
         }
 
-        private void BeginWrite(ReadableBuffer buffer)
+        private async Task WriteAsync(ReadableBuffer buffer)
         {
             var writeReq = _thread.WriteReqPool.Allocate();
 
-            _pendingWrites++;
-
-            writeReq.Write(_handle, buffer, _writeCallback, this);
-        }
-
-        private static void WriteCallback(UvWriteReq writeReq, int status, object state)
-        {
-            ((UvTcpConnection)state).EndWrite(writeReq);
-        }
-
-        private void EndWrite(UvWriteReq writeReq)
-        {
-            _pendingWrites--;
-
-            _thread.WriteReqPool.Return(writeReq);
+            try
+            {
+                await writeReq.WriteAsync(_handle, buffer);
 
-            if (_drainWrites != null)
+            }
+            finally
             {
-                if (_pendingWrites == 0)
-                {
-                    _drainWrites.TrySetResult(null);
-                }
+                _thread.WriteReqPool.Return(writeReq);
             }
         }