Improving Pipe/Socket cleanup. Improving IO Read loop.

Author: Stefán J. Sigurðarson

projects/RabbitMQ.Client/client/impl/SocketFrameHandler.cs

@@ -35,7 +35,6 @@
 using System.IO.Pipelines;
 using System.Net;
 using System.Net.Sockets;
-using System.Text;
 using System.Threading;
 using System.Threading.Channels;
 using System.Threading.Tasks;
@@ -226,37 +225,32 @@ public void Close()
             {
                 if (!_closed)
                 {
-                    try
-                    {
-                        _channelWriter.Complete();
-                        _writerTask.GetAwaiter().GetResult();
-                        _pipeReader.Complete();
-                        _pipeWriter.Complete();
-                    }
-                    catch (Exception)
-                    {
-                    }
+                    _channelWriter.TryComplete();
+                    _frameWriter.TryComplete();
+                    try { _writerTask.GetAwaiter().GetResult(); } catch { }
 
-                    try
-                    {
-                        _socket.Close();
-                    }
-                    catch (Exception)
-                    {
-                        // ignore, we are closing anyway
-                    }
-                    finally
+                    if (_pipe != null)
                     {
-                        _closed = true;
+                        try { _pipe.Input?.CancelPendingRead(); } catch { }
+                        try { _pipe.Input?.Complete(); } catch { }
+                        try { _pipe.Output?.CancelPendingFlush(); } catch { }
+                        try { _pipe.Output?.Complete(); } catch { }
+                        try { using (_pipe as IDisposable) { } } catch { }
                     }
+
+                    try { _socket.Close(); } catch { }
+                    _closed = true;
                 }
             }
         }
 
         public void SendHeader()
         {
-            byte[] headerBytes = ArrayPool<byte>.Shared.Rent(8);
-            Encoding.ASCII.GetBytes("AMQP", 0, 4, headerBytes, 0);
+            Span<byte> headerBytes = stackalloc byte[8];
+            headerBytes[0] = (byte)'A';
+            headerBytes[1] = (byte)'M';
+            headerBytes[2] = (byte)'Q';
+            headerBytes[3] = (byte)'P';
             if (Endpoint.Protocol.Revision != 0)
             {
                 headerBytes[4] = 0;
@@ -272,8 +266,10 @@ public void SendHeader()
                 headerBytes[7] = (byte)Endpoint.Protocol.MinorVersion;
             }
 
-            _pipeWriter.WriteAsync(headerBytes.AsMemory(0, 8)).AsTask().GetAwaiter().GetResult();
-            ArrayPool<byte>.Shared.Return(headerBytes);
+            Memory<byte> pipeMemory = _pipeWriter.GetMemory(8);
+            headerBytes.CopyTo(pipeMemory.Span);
+            _pipeWriter.Advance(headerBytes.Length);
+            _pipeWriter.FlushAsync().AsTask().GetAwaiter().GetResult();
         }
 
         public void Write(OutgoingFrame frame)
@@ -283,28 +279,51 @@ public void Write(OutgoingFrame frame)
 
         private async Task ReadLoop()
         {
-            while (!FrameReader.Completion.IsCompleted)
+            bool allowSync = true;
+            try
             {
-                // Let's read some bytes
-                if (!_pipeReader.TryRead(out ReadResult readResult))
+                while (true)
                 {
-                    readResult = await _pipeReader.ReadAsync().ConfigureAwait(false);
-                }
+                    // Let's read some bytes
+                    if (!(allowSync && _pipeReader.TryRead(out ReadResult readResult)))
+                    {
+                        readResult = await _pipeReader.ReadAsync().ConfigureAwait(false);
+                    }
 
-                ReadOnlySequence<byte> buffer = readResult.Buffer;
-                if (buffer.IsEmpty)
-                {
-                    _frameWriter.Complete(new EndOfStreamException());
-                    return;
-                }
+                    int handled = 0;
+                    ReadOnlySequence<byte> buffer = readResult.Buffer;
+                    if (!buffer.IsEmpty)
+                    {
+                        handled = ProcessBuffer(ref buffer);
+                    }
 
-                while (InboundFrame.TryReadFrame(ref buffer, out InboundFrame frame))
-                {
-                    _frameWriter.TryWrite(frame);
+                    allowSync = handled != 0;
+                    _pipeReader.AdvanceTo(buffer.Start, buffer.End);
+
+                    if (handled == 0 && readResult.IsCompleted)
+                    {
+                        break;
+                    }
                 }
+            }
+            catch (Exception e)
+            {
+                _frameWriter.TryComplete(e);
+            }
+
+            _frameWriter.TryComplete(new EndOfStreamException());
+        }
 
-                _pipeReader.AdvanceTo(buffer.Start, buffer.End);
+        private int ProcessBuffer(ref ReadOnlySequence<byte> buffer)
+        {
+            int handled = 0;
+            while (InboundFrame.TryReadFrame(ref buffer, out InboundFrame frame))
+            {
+                _frameWriter.TryWrite(frame);
+                handled++;
             }
+
+            return handled;
         }
 
         private async Task WriteLoop()

projects/RabbitMQ.Client/client/impl/TcpClientAdapter.cs

@@ -42,8 +42,9 @@ public virtual Task ConnectAsync(IPAddress ep, int port)
 
         public virtual void Close()
         {
-            _sock?.Dispose();
-            _sock = null;
+            try { _sock.Shutdown(SocketShutdown.Both); } catch { }
+            try { _sock.Close(); } catch { }
+            try { _sock.Dispose(); } catch { }
         }
 
         [Obsolete("Override Dispose(bool) instead.")]