LibUv Networking (#65)

Projects/Server/LibUv/LibuvConnection.cs

@@ -0,0 +1,309 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Buffers;
+using System.IO;
+using System.IO.Pipelines;
+using System.Net;
+using System.Threading;
+using System.Threading.Tasks;
+using Libuv.Internal;
+using Microsoft.AspNetCore.Connections;
+using Microsoft.Extensions.Logging;
+
+namespace Libuv
+{
+  public partial class LibuvConnection : TransportConnection
+  {
+    private static readonly int MinAllocBufferSize = SlabMemoryPool.BlockSize / 2;
+
+    private static readonly Action<UvStreamHandle, int, object> _readCallback = ReadCallback;
+
+    private static readonly Func<UvStreamHandle, int, object, LibuvFunctions.uv_buf_t> _allocCallback = AllocCallback;
+
+    private readonly UvStreamHandle _socket;
+    private readonly CancellationTokenSource _connectionClosedTokenSource = new CancellationTokenSource();
+
+    private volatile ConnectionAbortedException _abortReason;
+
+    private MemoryHandle _bufferHandle;
+    private Task _processingTask;
+    private readonly TaskCompletionSource<object> _waitForConnectionClosedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+    private bool _connectionClosed;
+
+    public LibuvConnection(UvStreamHandle socket,
+      ILibuvTrace log,
+      LibuvThread thread,
+      IPEndPoint remoteEndPoint,
+      IPEndPoint localEndPoint,
+      PipeOptions inputOptions = null,
+      PipeOptions outputOptions = null,
+      long? maxReadBufferSize = null,
+      long? maxWriteBufferSize = null)
+    {
+      _socket = socket;
+
+      LocalEndPoint = localEndPoint;
+      RemoteEndPoint = remoteEndPoint;
+
+      ConnectionClosed = _connectionClosedTokenSource.Token;
+      Log = log;
+      Thread = thread;
+
+      maxReadBufferSize ??= 0;
+      maxWriteBufferSize ??= 0;
+
+      inputOptions ??= new PipeOptions(MemoryPool, PipeScheduler.ThreadPool, Thread, maxReadBufferSize.Value, maxReadBufferSize.Value / 2, useSynchronizationContext: false);
+      outputOptions ??= new PipeOptions(MemoryPool, Thread, PipeScheduler.ThreadPool, maxWriteBufferSize.Value, maxWriteBufferSize.Value / 2, useSynchronizationContext: false);
+
+      var pair = DuplexPipe.CreateConnectionPair(inputOptions, outputOptions);
+
+      // Set the transport and connection id
+      Transport = pair.Transport;
+      Application = pair.Application;
+    }
+
+    public PipeWriter Input => Application.Output;
+
+    public PipeReader Output => Application.Input;
+
+    public LibuvOutputConsumer OutputConsumer { get; set; }
+    private ILibuvTrace Log { get; }
+    private LibuvThread Thread { get; }
+    public override MemoryPool<byte> MemoryPool => Thread.MemoryPool;
+
+    public void Start()
+    {
+      _processingTask = StartCore();
+    }
+
+    private async Task StartCore()
+    {
+      try
+      {
+        OutputConsumer = new LibuvOutputConsumer(Output, Thread, _socket, ConnectionId, Log);
+
+        StartReading();
+
+        Exception inputError = null;
+        Exception outputError = null;
+
+        try
+        {
+          // This *must* happen after socket.ReadStart
+          // The socket output consumer is the only thing that can close the connection. If the
+          // output pipe is already closed by the time we start then it's fine since, it'll close gracefully afterwards.
+          await OutputConsumer.WriteOutputAsync();
+        }
+        catch (UvException ex)
+        {
+          // The connection reset/error has already been logged by LibuvOutputConsumer
+          if (ex.StatusCode == LibuvConstants.ECANCELED)
+          {
+            // Connection was aborted.
+          }
+          else if (LibuvConstants.IsConnectionReset(ex.StatusCode))
+          {
+            // Don't cause writes to throw for connection resets.
+            inputError = new ConnectionResetException(ex.Message, ex);
+          }
+          else
+          {
+            // This is unexpected.
+            Log.ConnectionError(ConnectionId, ex);
+
+            inputError = ex;
+            outputError = ex;
+          }
+        }
+        finally
+        {
+          inputError ??= _abortReason ?? new ConnectionAbortedException("The libuv transport's send loop completed gracefully.");
+
+          // Now, complete the input so that no more reads can happen
+          Input.Complete(inputError);
+          Output.Complete(outputError);
+
+          // Make sure it isn't possible for a paused read to resume reading after calling uv_close
+          // on the stream handle
+          Input.CancelPendingFlush();
+
+          // Send a FIN
+          Log.ConnectionWriteFin(ConnectionId, inputError.Message);
+
+          // We're done with the socket now
+          _socket.Dispose();
+
+          // Ensure this always fires
+          FireConnectionClosed();
+
+          await _waitForConnectionClosedTcs.Task;
+        }
+      }
+      catch (Exception e)
+      {
+        Log.LogCritical(0, e, $"{nameof(LibuvConnection)}.{nameof(Start)}() {ConnectionId}");
+      }
+    }
+
+    public override void Abort(ConnectionAbortedException abortReason)
+    {
+      _abortReason = abortReason;
+
+      // Cancel WriteOutputAsync loop after setting _abortReason.
+      Output.CancelPendingRead();
+
+      // This cancels any pending I/O.
+      Thread.Post(s => s.Dispose(), _socket);
+    }
+
+    public override async ValueTask DisposeAsync()
+    {
+      Transport.Input.Complete();
+      Transport.Output.Complete();
+
+      if (_processingTask != null) await _processingTask;
+
+      _connectionClosedTokenSource.Dispose();
+    }
+
+    // Called on Libuv thread
+    private static LibuvFunctions.uv_buf_t AllocCallback(UvStreamHandle handle, int suggestedSize, object state) => ((LibuvConnection)state).OnAlloc(handle, suggestedSize);
+
+    private unsafe LibuvFunctions.uv_buf_t OnAlloc(UvStreamHandle handle, int suggestedSize)
+    {
+      var currentWritableBuffer = Input.GetMemory(MinAllocBufferSize);
+      _bufferHandle = currentWritableBuffer.Pin();
+
+      return handle.Libuv.buf_init((IntPtr)_bufferHandle.Pointer, currentWritableBuffer.Length);
+    }
+
+    private static void ReadCallback(UvStreamHandle handle, int status, object state)
+    {
+      ((LibuvConnection)state).OnRead(handle, status);
+    }
+
+    private void OnRead(UvStreamHandle handle, int status)
+    {
+      // Cleanup state from last OnAlloc. This is safe even if OnAlloc wasn't called.
+      _bufferHandle.Dispose();
+      if (status == 0)
+      {
+        // EAGAIN/EWOULDBLOCK so just return the buffer.
+        // http://docs.libuv.org/en/v1.x/stream.html#c.uv_read_cb
+      }
+      else if (status > 0)
+      {
+        Log.ConnectionRead(ConnectionId, status);
+
+        Input.Advance(status);
+        var flushTask = Input.FlushAsync();
+
+        if (!flushTask.IsCompleted)
+          // We wrote too many bytes to the reader, so pause reading and resume when
+          // we hit the low water mark.
+          _ = ApplyBackpressureAsync(flushTask);
+      }
+      else
+      {
+        // Given a negative status, it's possible that OnAlloc wasn't called.
+        _socket.ReadStop();
+
+        Exception error = null;
+
+        if (status == LibuvConstants.EOF)
+        {
+          Log.ConnectionReadFin(ConnectionId);
+        }
+        else
+        {
+          handle.Libuv.Check(status, out var uvError);
+          error = LogAndWrapReadError(uvError);
+        }
+
+        FireConnectionClosed();
+
+        // Complete after aborting the connection
+        Input.Complete(error);
+      }
+    }
+
+    private void FireConnectionClosed()
+    {
+      // Guard against scheduling this multiple times
+      if (_connectionClosed) return;
+
+      _connectionClosed = true;
+
+      ThreadPool.UnsafeQueueUserWorkItem(state =>
+        {
+          state.CancelConnectionClosedToken();
+
+          state._waitForConnectionClosedTcs.TrySetResult(null);
+        },
+        this,
+        false);
+    }
+
+    private async Task ApplyBackpressureAsync(ValueTask<FlushResult> flushTask)
+    {
+      Log.ConnectionPause(ConnectionId);
+      _socket.ReadStop();
+
+      var result = await flushTask;
+
+      // If the reader isn't complete or cancelled then resume reading
+      if (!result.IsCompleted && !result.IsCanceled)
+      {
+        Log.ConnectionResume(ConnectionId);
+        StartReading();
+      }
+    }
+
+    private void StartReading()
+    {
+      try
+      {
+        _socket.ReadStart(_allocCallback, _readCallback, this);
+      }
+      catch (UvException ex)
+      {
+        // ReadStart() can throw a UvException in some cases (e.g. socket is no longer connected).
+        // This should be treated the same as OnRead() seeing a negative status.
+        Input.Complete(LogAndWrapReadError(ex));
+      }
+    }
+
+    private Exception LogAndWrapReadError(UvException uvError)
+    {
+      if (uvError.StatusCode == LibuvConstants.ECANCELED)
+      {
+        // The operation was canceled by the server not the client. No need for additional logs.
+        return new ConnectionAbortedException(uvError.Message, uvError);
+      }
+
+      if (LibuvConstants.IsConnectionReset(uvError.StatusCode))
+      {
+        // Log connection resets at a lower (Debug) level.
+        Log.ConnectionReset(ConnectionId);
+        return new ConnectionResetException(uvError.Message, uvError);
+      }
+      // This is unexpected.
+      Log.ConnectionError(ConnectionId, uvError);
+      return new IOException(uvError.Message, uvError);
+    }
+
+    private void CancelConnectionClosedToken()
+    {
+      try
+      {
+        _connectionClosedTokenSource.Cancel();
+      }
+      catch (Exception ex)
+      {
+        Log.LogError(0, ex, $"Unexpected exception in {nameof(LibuvConnection)}.{nameof(CancelConnectionClosedToken)}.");
+      }
+    }
+  }
+}