Refactor OutgoingFrameMemory for zero-copy AMQP payload streaming

This commit eliminates large contiguous buffer allocations and redundant
payload copying during AMQP frame serialization by fully leveraging
System.IO.Pipelines.

Author: inoa-pmattos

projects/RabbitMQ.Client/Impl/Frame.cs

@@ -31,10 +31,12 @@
 
 using System;
 using System.Buffers;
+using System.Buffers.Binary;
 using System.IO;
 using System.IO.Pipelines;
 using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Channels;
 using System.Threading.Tasks;
 using RabbitMQ.Client.Exceptions;
 using RabbitMQ.Client.Framing;
@@ -113,17 +115,30 @@ internal static class BodySegment
              * +--------------+
              * | x bytes      |
              * +--------------+ */
+            public const int HeaderSize = StartPayload;
+            public const int FooterSize = 1;
             public const int FrameSize = BaseFrameSize;
 
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
             public static int WriteTo(Span<byte> span, ushort channel, ReadOnlySpan<byte> body)
             {
                 const int StartBodyArgument = StartPayload;
-                NetworkOrderSerializer.WriteUInt64(ref span.GetStart(), ((ulong)Constants.FrameBody << 56) | ((ulong)channel << 40) | ((ulong)body.Length << 8));
+                WriteHeader(span, channel, body.Length);
                 body.CopyTo(span.Slice(StartBodyArgument));
                 span[StartPayload + body.Length] = Constants.FrameEnd;
+                WriteFooter(span.Slice(StartPayload + body.Length));
                 return body.Length + BaseFrameSize;
             }
+
+            public static void WriteHeader(Span<byte> span, ushort channel, int bodyLength)
+            {
+                NetworkOrderSerializer.WriteUInt64(ref span.GetStart(), ((ulong)Constants.FrameBody << 56) | ((ulong)channel << 40) | ((ulong)bodyLength << 8));
+            }
+
+            public static void WriteFooter(Span<byte> span)
+            {
+                span[0] = Constants.FrameEnd;
+            }
         }
 
         internal static class Heartbeat
@@ -170,27 +185,42 @@ public static OutgoingFrameMemory SerializeToFrames<TMethod, THeader>(ref TMetho
             where TMethod : struct, IOutgoingAmqpMethod
             where THeader : IAmqpHeader
         {
-            int remainingBodyBytes = body.Length;
-            int size = Method.FrameSize + Header.FrameSize +
-                       method.GetRequiredBufferSize() + header.GetRequiredBufferSize() +
-                       BodySegment.FrameSize * GetBodyFrameCount(maxBodyPayloadBytes, remainingBodyBytes) + remainingBodyBytes;
-
             // Will be returned by SocketFrameWriter.WriteLoop
-            IMemoryOwner<byte> buffer = MemoryPool<byte>.Shared.Rent(size);
+            IMemoryOwner<byte> bodyCopy = MemoryPool<byte>.Shared.Rent(body.Length);
+            body.CopyTo(bodyCopy.Memory);
+            return SerializeToFrames(ref method, ref header, bodyCopy, body.Length, channelNumber, maxBodyPayloadBytes);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static OutgoingFrameMemory SerializeToFrames<TMethod, THeader>(ref TMethod method, ref THeader header, IMemoryOwner<byte> body, int bodyLength, ushort channelNumber, int maxBodyPayloadBytes)
+            where TMethod : struct, IOutgoingAmqpMethod
+            where THeader : IAmqpHeader
+        {
+            // Calculate ONLY the Method and Header framing size
+            int framingSize = Method.FrameSize + Header.FrameSize +
+                              method.GetRequiredBufferSize() + header.GetRequiredBufferSize();
+
+            // Pre-calculate total final sequence size
+            int bodyFramesCount = GetBodyFrameCount(maxBodyPayloadBytes, bodyLength);
+            int totalSize = framingSize + bodyLength + (BodySegment.FrameSize * bodyFramesCount);
+
+            // Rent a smaller buffer exclusively for the Method and Header
+            IMemoryOwner<byte> buffer = MemoryPool<byte>.Shared.Rent(framingSize);
             Span<byte> bufferSpan = buffer.Memory.Span;
 
             int offset = Method.WriteTo(bufferSpan, channelNumber, ref method);
-            offset += Header.WriteTo(bufferSpan.Slice(offset), channelNumber, ref header, remainingBodyBytes);
-            ReadOnlySpan<byte> bodySpan = body.Span;
-            while (remainingBodyBytes > 0)
-            {
-                int frameSize = remainingBodyBytes > maxBodyPayloadBytes ? maxBodyPayloadBytes : remainingBodyBytes;
-                offset += BodySegment.WriteTo(bufferSpan.Slice(offset), channelNumber, bodySpan.Slice(bodySpan.Length - remainingBodyBytes, frameSize));
-                remainingBodyBytes -= frameSize;
-            }
-
-            System.Diagnostics.Debug.Assert(offset == size, $"Serialized to wrong size, expect {size}, offset {offset}");
-            return new OutgoingFrameMemory(buffer, size);
+            offset += Header.WriteTo(bufferSpan.Slice(offset), channelNumber, ref header, bodyLength);
+
+            System.Diagnostics.Debug.Assert(offset == framingSize, $"Serialized to wrong size, expect {framingSize}, offset {offset}");
+
+            return new OutgoingFrameMemory(
+                buffer,
+                framingSize,
+                body,
+                bodyLength,
+                channelNumber,
+                maxBodyPayloadBytes,
+                totalSize);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]

projects/RabbitMQ.Client/Impl/SocketFrameHandler.cs

@@ -251,7 +251,8 @@ private async Task WriteLoopAsync()
                         try
                         {
                             frames.WriteTo(_pipeWriter);
-                            await _pipeWriter.FlushAsync().ConfigureAwait(false);
+                            await _pipeWriter.FlushAsync()
+                                .ConfigureAwait(false);
                             RabbitMqClientEventSource.Log.CommandSent(frames.Size);
                         }
                         finally

projects/RabbitMQ.Client/OutgoingFrameMemory.cs

@@ -31,35 +31,92 @@
 
 using System;
 using System.Buffers;
+using System.Diagnostics;
+using static RabbitMQ.Client.Impl.Framing;
 
 namespace RabbitMQ.Client
 {
     internal struct OutgoingFrameMemory : IDisposable
     {
-        internal OutgoingFrameMemory(IMemoryOwner<byte> memoryOwner, int length)
+        private IMemoryOwner<byte>? _methodAndHeader;
+        private readonly int _methodAndHeaderLength;
+        private IMemoryOwner<byte>? _body;
+        private readonly int _bodyLength;
+        private readonly int _maxBodyPayloadBytes;
+        private readonly ushort _channelNumber;
+
+        internal OutgoingFrameMemory(
+            IMemoryOwner<byte> methodAndHeader,
+            int methodAndHeaderLength)
         {
-            _memory = memoryOwner.Memory.Slice(0, length);
-            _memoryOwner = memoryOwner;
+            _methodAndHeader = methodAndHeader;
+            _methodAndHeaderLength = methodAndHeaderLength;
+            _body = null;
+            _bodyLength = 0;
+            _channelNumber = 0;
+            _maxBodyPayloadBytes = 0;
+            Size = methodAndHeaderLength;
         }
 
-        internal readonly int Size => _memory.Length;
+        internal OutgoingFrameMemory(
+            IMemoryOwner<byte> methodAndHeader,
+            int methodAndHeaderLength,
+            IMemoryOwner<byte> body,
+            int bodyLength,
+            ushort channelNumber,
+            int maxBodyPayloadBytes,
+            int totalSize)
+        {
+            _methodAndHeader = methodAndHeader;
+            _methodAndHeaderLength = methodAndHeaderLength;
+            _body = body;
+            _bodyLength = bodyLength;
+            _channelNumber = channelNumber;
+            _maxBodyPayloadBytes = maxBodyPayloadBytes;
+            Size = totalSize;
+        }
 
-        private ReadOnlyMemory<byte> _memory;
-        private IMemoryOwner<byte>? _memoryOwner;
+        internal readonly int Size { get; }
 
-        internal readonly void WriteTo(IBufferWriter<byte> pipeWriter)
+        internal readonly void WriteTo(IBufferWriter<byte> writer)
         {
-            pipeWriter.Write(_memory.Span);
+            Debug.Assert(_methodAndHeader is not null);
+            ReadOnlySpan<byte> methodAndHeader = _methodAndHeader!.Memory.Span.Slice(0, _methodAndHeaderLength);
+            writer.Write(methodAndHeader);
+            
+            if (_bodyLength == 0)
+            {
+                return;
+            }
+
+            Debug.Assert(_body is not null);
+            ReadOnlySpan<byte> bodySpan = _body!.Memory.Span.Slice(0, _bodyLength);
+            int remainingBodyBytes = bodySpan.Length;
+            int bodyOffset = 0;
+
+            while (remainingBodyBytes > 0)
+            {
+                int payloadSize = remainingBodyBytes > _maxBodyPayloadBytes ? _maxBodyPayloadBytes : remainingBodyBytes;
+
+                Span<byte> span = writer.GetSpan(BodySegment.HeaderSize + payloadSize + BodySegment.FooterSize);
+                int offset = BodySegment.WriteTo(span, _channelNumber, bodySpan.Slice(bodyOffset, payloadSize));
+                writer.Advance(offset);
+
+                remainingBodyBytes -= payloadSize;
+                bodyOffset += payloadSize;
+            }
         }
 
         public void Dispose()
         {
-            IMemoryOwner<byte>? memoryOwner = _memoryOwner;
-            _memoryOwner = null;
+            IMemoryOwner<byte>? memoryOwner = _methodAndHeader;
+            _methodAndHeader = null;
             if (memoryOwner != null)
             {
                 memoryOwner.Dispose();
-                _memory = default;
+                _methodAndHeader = default;
+                _body?.Dispose();
+                _body = null;
             }
         }
     }