TLS implementation - high performance

src/TurboMqtt/Client/ClientStreamInstance.cs

@@ -107,6 +107,13 @@ private void PrepareStream(MqttClientConnectOptions clientConnectOptions,
 
         // begin inbound stream
         inboundStream
+            .WatchTermination(async (_, task) =>
+            {
+                await task;
+
+                // guarantee that we kill the client in the event that the stream terminates
+                disconnectPromise.TrySetResult(DisconnectPacket.Instance);
+            })
             // setting IsOwner to false here is crucial - otherwise, we can't reboot the client after broker disconnects
             .To(ChannelSink.FromWriter(inboundPackets.Writer, false))
             .Run(_materializer);
@@ -145,7 +152,6 @@ private static (Source<MqttMessage, NotUsed> inbound, Sink<MqttPacket, NotUsed>
                 var outboundMessages = MqttClientStreams.Mqtt311OutboundPacketSink(
                     clientConnectOptions.ClientId,
                     transport,
-                    MemoryPool<byte>.Shared,
                     maxFrameSize, (int)clientConnectOptions.MaximumPacketSize,
                     clientConnectOptions.EnableOpenTelemetry);
 

src/TurboMqtt/IO/DuplexChannel.cs

@@ -0,0 +1,28 @@
+// -----------------------------------------------------------------------
+// <copyright file="DuplexChannel.cs" company="Petabridge, LLC">
+//      Copyright (C) 2024 - 2024 Petabridge, LLC <https://petabridge.com>
+// </copyright>
+// -----------------------------------------------------------------------
+
+using System.Buffers;
+using System.Threading.Channels;
+
+namespace TurboMqtt.IO;
+
+internal interface IDuplexChannel
+{
+    public ChannelWriter<(IMemoryOwner<byte> buffer, int readableBytes)> Writer { get; }
+    public ChannelReader<(IMemoryOwner<byte> buffer, int readableBytes)> Reader { get; }
+}
+
+internal sealed class DuplexChannel : IDuplexChannel
+{
+    public DuplexChannel(ChannelWriter<(IMemoryOwner<byte> buffer, int readableBytes)> writer, ChannelReader<(IMemoryOwner<byte> buffer, int readableBytes)> reader)
+    {
+        Writer = writer;
+        Reader = reader;
+    }
+
+    public ChannelWriter<(IMemoryOwner<byte> buffer, int readableBytes)> Writer { get; }
+    public ChannelReader<(IMemoryOwner<byte> buffer, int readableBytes)> Reader { get; }
+}

src/TurboMqtt/IO/FakeServerHandle.cs

@@ -91,7 +91,7 @@ public virtual void FlushPackets()
         }
         else
         {
-            Log.Debug("Successfully wrote N packets {0} [{1} bytes] to transport.", _pendingPackets.Count, totalSize);
+            Log.Debug("Successfully wrote {0} packets [{1} bytes] to transport.", _pendingPackets.Count, totalSize);
         }
 
         bufferPooled.Dispose();

src/TurboMqtt/IO/InMem/InMemoryMqttTransport.cs

@@ -5,8 +5,10 @@
 // -----------------------------------------------------------------------
 
 using System.Buffers;
+using System.IO.Pipelines;
 using System.Threading.Channels;
 using Akka.Event;
+using TurboMqtt.IO.Internal;
 using TurboMqtt.IO.Tcp;
 using TurboMqtt.PacketTypes;
 using TurboMqtt.Protocol;
@@ -42,12 +44,10 @@ internal sealed class InMemoryMqttTransport : IMqttTransport
 {
     private readonly TaskCompletionSource<DisconnectReasonCode> _terminationSource = new();
 
-    private readonly Channel<(IMemoryOwner<byte> buffer, int readableBytes)> _writesToTransport =
-        Channel.CreateUnbounded<(IMemoryOwner<byte> buffer, int readableBytes)>();
+    private readonly IDuplexPipe _transport;
+
+    private readonly IDuplexPipe _application;
 
-    private readonly Channel<(IMemoryOwner<byte> buffer, int readableBytes)> _readsFromTransport =
-        Channel.CreateUnbounded<(IMemoryOwner<byte> buffer, int readableBytes)>();
-
     private readonly CancellationTokenSource _shutdownTokenSource = new();
     private readonly IFakeServerHandle _serverHandle;
 
@@ -56,8 +56,13 @@ public InMemoryMqttTransport(int maxFrameSize, ILoggingAdapter log, MqttProtocol
         MaxFrameSize = maxFrameSize;
         Log = log;
         ProtocolVersion = protocolVersion;
-        Reader = _readsFromTransport;
-        Writer = _writesToTransport;
+        var pipeOptions = new PipeOptions(pauseWriterThreshold: MaxFrameSize,
+            resumeWriterThreshold: MaxFrameSize / 2,
+            useSynchronizationContext: false);
+        var pipes = DuplexPipe.CreateConnectionPair(pipeOptions, pipeOptions);
+        _transport = pipes.Transport;
+        _application = pipes.Application;
+        Channel = _application;
 
         _serverHandle = protocolVersion switch
         {
@@ -71,7 +76,20 @@ async Task CloseFn()
             await CloseAsync();
         }
 
-        bool PushFn((IMemoryOwner<byte> buffer, int estimatedSize) msg) => _readsFromTransport.Writer.TryWrite(msg);
+        bool PushFn((IMemoryOwner<byte> buffer, int estimatedSize) msg)
+        {
+            try
+            {
+                _transport.Output.Write(msg.buffer.Memory.Span);
+                _ = _transport.Output.FlushAsync();
+                return true;
+            }
+            finally
+            {
+                _waitForPendingWrites.TrySetResult(true);
+                msg.buffer.Dispose();
+            }
+        }
     }
 
     public MqttProtocolVersion ProtocolVersion { get; }
@@ -80,38 +98,36 @@ async Task CloseFn()
     public ConnectionStatus Status { get; private set; } = ConnectionStatus.NotStarted;
 
     public Task<DisconnectReasonCode> WhenTerminated => _terminationSource.Task;
-    
+
     private readonly TaskCompletionSource<bool> _waitForPendingWrites = new();
     public Task WaitForPendingWrites => _waitForPendingWrites.Task;
-
 
 
     public async Task<bool> CloseAsync(CancellationToken ct = default)
     {
         Status = ConnectionStatus.Disconnected;
-        _writesToTransport.Writer.TryComplete();
+        await _application.Output.CompleteAsync();
+        await _transport.Output.CompleteAsync();
         await _waitForPendingWrites.Task;
         await _shutdownTokenSource.CancelAsync();
-        _readsFromTransport.Writer.TryComplete();
         _terminationSource.TrySetResult(DisconnectReasonCode.NormalDisconnection);
         return true;
     }
 
-    public Task AbortAsync(CancellationToken ct = default)
+    public async Task AbortAsync(CancellationToken ct = default)
     {
         Status = ConnectionStatus.Disconnected;
-        _writesToTransport.Writer.TryComplete();
-        _readsFromTransport.Writer.TryComplete();
+        await _application.Output.CompleteAsync();
+        await _transport.Output.CompleteAsync();
         _terminationSource.TrySetResult(DisconnectReasonCode.UnspecifiedError);
-        return Task.CompletedTask;
     }
 
     public Task<bool> ConnectAsync(CancellationToken ct = default)
     {
         if (Status == ConnectionStatus.NotStarted)
         {
             Status = ConnectionStatus.Connected;
-            
+
 #pragma warning disable CS4014 // Because this call is not awaited, execution of the current method continues before the call is completed
             DoByteWritesAsync(_shutdownTokenSource.Token);
 #pragma warning restore CS4014 // Because this call is not awaited, execution of the current method continues before the call is completed
@@ -128,31 +144,38 @@ public Task<bool> ConnectAsync(CancellationToken ct = default)
     private async Task DoByteWritesAsync(CancellationToken ct)
     {
         Log.Debug("Starting to read from transport.");
-        while(!_writesToTransport.Reader.Completion.IsCompleted)
+        while (ct.IsCancellationRequested == false)
         {
-            if (!_writesToTransport.Reader.TryRead(out var msg))
+            try
             {
-                try
+                var msg = await _transport.Input.ReadAsync(ct);
+
+                var buffer = msg.Buffer;
+                Log.Debug("Received {0} bytes from transport.", buffer.Length);
+
+                if (!buffer.IsEmpty)
                 {
-                    await _writesToTransport.Reader.WaitToReadAsync(ct);
+                    var seqPosition = buffer.Start;
+                    var copyBuffer = new ReadOnlyMemory<byte>(buffer.ToArray());
+                    _serverHandle.HandleBytes(copyBuffer);
+
+                    if (!msg.IsCompleted)
+                    {
+                        // we will throw if we try to advance past the end of the buffer
+                        _transport.Input.AdvanceTo(buffer.End);
+                    }
                 }
-                catch(OperationCanceledException)
+
+                if (msg.IsCompleted)
                 {
-
+                    break;
                 }
-                continue;
             }
-
-            try
+            catch (Exception exception)
             {
-                ReadOnlyMemory<byte> buffer = msg.buffer.Memory.Slice(0, msg.readableBytes);
-                Log.Debug("Received {0} bytes from transport.", buffer.Length);
-                _serverHandle.HandleBytes(buffer);
-            }
-            finally
-            {
-                // have to free the shared buffer
-                msg.buffer.Dispose();
+                Log.Error(exception, "Error reading from transport.");
+                _transport.Input.CancelPendingRead();
+                break;
             }
         }
 
@@ -161,7 +184,5 @@ private async Task DoByteWritesAsync(CancellationToken ct)
     }
 
     public int MaxFrameSize { get; }
-    public ChannelWriter<(IMemoryOwner<byte> buffer, int readableBytes)> Writer { get; }
-    public ChannelReader<(IMemoryOwner<byte> buffer, int readableBytes)> Reader { get; }
-
+    public IDuplexPipe Channel { get; }
 }

src/TurboMqtt/IO/Internal/DuplexPipe.cs

@@ -0,0 +1,46 @@
+// -----------------------------------------------------------------------
+// <copyright file="DuplexPipe.cs" company="Petabridge, LLC">
+//      Copyright (C) 2024 - 2024 Petabridge, LLC <https://petabridge.com>
+// </copyright>
+// -----------------------------------------------------------------------
+
+using System.IO.Pipelines;
+
+namespace TurboMqtt.IO.Internal;
+
+internal sealed class DuplexPipe : IDuplexPipe
+{
+    public DuplexPipe(PipeReader reader, PipeWriter writer)
+    {
+        Input = reader;
+        Output = writer;
+    }
+
+    public PipeReader Input { get; }
+
+    public PipeWriter Output { get; }
+
+    public static DuplexPipePair CreateConnectionPair(PipeOptions inputOptions, PipeOptions outputOptions)
+    {
+        var input = new Pipe(inputOptions);
+        var output = new Pipe(outputOptions);
+
+        var transportToApplication = new DuplexPipe(output.Reader, input.Writer);
+        var applicationToTransport = new DuplexPipe(input.Reader, output.Writer);
+
+        return new DuplexPipePair(applicationToTransport, transportToApplication);
+    }
+
+    // This class exists to work around issues with value tuple on .NET Framework
+    public readonly struct DuplexPipePair
+    {
+        public IDuplexPipe Transport { get; }
+        public IDuplexPipe Application { get; }
+
+        public DuplexPipePair(IDuplexPipe transport, IDuplexPipe application)
+        {
+            Transport = transport;
+            Application = application;
+        }
+    }
+}