reduce the amount of times we rent / return from arraypool

projects/Apigen/apigen/Apigen.cs

@@ -1410,6 +1410,10 @@ string GetParameterString(ParameterInfo pi)
                 {
                     return "cmd.Body";
                 }
+                else if (Attribute(pi, typeof(AmqpContentBodyArrayMappingAttribute)) != null)
+                {
+                    return "cmd.TakeoverPayload()";
+                }
                 else
                 {
                     return $"__impl._{(!(Attribute(pi, typeof(AmqpFieldMappingAttribute)) is AmqpFieldMappingAttribute fieldMapping) ? pi.Name : fieldMapping.m_fieldName)}";

projects/RabbitMQ.Client/client/api/BasicGetResult.cs

@@ -30,15 +30,18 @@
 //---------------------------------------------------------------------------
 
 using System;
+using System.Buffers;
 
 namespace RabbitMQ.Client
 {
     /// <summary>Represents Basic.GetOk responses from the server.</summary>
     /// <remarks>
     /// Basic.Get either returns an instance of this class, or null if a Basic.GetEmpty was received.
     /// </remarks>
-    public class BasicGetResult
+    public sealed class BasicGetResult : IDisposable
     {
+        private readonly byte[] _rentedArray;
+
         /// <summary>
         /// Sets the new instance's properties from the arguments passed in.
         /// </summary>
@@ -48,9 +51,9 @@ public class BasicGetResult
         /// <param name="routingKey">Routing key with which the message was published.</param>
         /// <param name="messageCount">The number of messages pending on the queue, excluding the message being delivered.</param>
         /// <param name="basicProperties">The Basic-class content header properties for the message.</param>
-        /// <param name="body"></param>
-        public BasicGetResult(ulong deliveryTag, bool redelivered, string exchange,
-            string routingKey, uint messageCount, IBasicProperties basicProperties, ReadOnlyMemory<byte> body)
+        /// <param name="body">The body</param>
+        public BasicGetResult(ulong deliveryTag, bool redelivered, string exchange, string routingKey,
+            uint messageCount, IBasicProperties basicProperties, ReadOnlyMemory<byte> body)
         {
             DeliveryTag = deliveryTag;
             Redelivered = redelivered;
@@ -61,6 +64,30 @@ public BasicGetResult(ulong deliveryTag, bool redelivered, string exchange,
             Body = body;
         }
 
+        /// <summary>
+        /// Sets the new instance's properties from the arguments passed in.
+        /// </summary>
+        /// <param name="deliveryTag">Delivery tag for the message.</param>
+        /// <param name="redelivered">Redelivered flag for the message</param>
+        /// <param name="exchange">The exchange this message was published to.</param>
+        /// <param name="routingKey">Routing key with which the message was published.</param>
+        /// <param name="messageCount">The number of messages pending on the queue, excluding the message being delivered.</param>
+        /// <param name="basicProperties">The Basic-class content header properties for the message.</param>
+        /// <param name="body">The body</param>
+        /// <param name="rentedArray">The rented array which body is part of.</param>
+        public BasicGetResult(ulong deliveryTag, bool redelivered, string exchange, string routingKey,
+            uint messageCount, IBasicProperties basicProperties, ReadOnlyMemory<byte> body, byte[] rentedArray)
+        {
+            DeliveryTag = deliveryTag;
+            Redelivered = redelivered;
+            Exchange = exchange;
+            RoutingKey = routingKey;
+            MessageCount = messageCount;
+            BasicProperties = basicProperties;
+            Body = body;
+            _rentedArray = rentedArray;
+        }
+
         /// <summary>
         /// Retrieves the Basic-class content header properties for this message.
         /// </summary>
@@ -99,5 +126,14 @@ public BasicGetResult(ulong deliveryTag, bool redelivered, string exchange,
         /// Retrieve the routing key with which this message was published.
         /// </summary>
         public string RoutingKey { get; private set; }
+
+        /// <inheritdoc />
+        public void Dispose()
+        {
+            if (!(_rentedArray is null))
+            {
+                ArrayPool<byte>.Shared.Return(_rentedArray);
+            }
+        }
     }
 }

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

@@ -45,12 +45,10 @@ public void HandleBasicDeliver(IBasicConsumer consumer,
             string exchange,
             string routingKey,
             IBasicProperties basicProperties,
-            ReadOnlySpan<byte> body)
+            ReadOnlyMemory<byte> body,
+            byte[] rentedArray)
         {
-            byte[] bodyBytes = ArrayPool<byte>.Shared.Rent(body.Length);
-            Memory<byte> bodyCopy = new Memory<byte>(bodyBytes, 0, body.Length);
-            body.CopyTo(bodyCopy.Span);
-            ScheduleUnlessShuttingDown(new BasicDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, bodyCopy));
+            ScheduleUnlessShuttingDown(new BasicDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body, rentedArray));
         }
 
         public void HandleBasicCancelOk(IBasicConsumer consumer, string consumerTag)
@@ -69,17 +67,15 @@ public void HandleModelShutdown(IBasicConsumer consumer, ShutdownEventArgs reaso
             Schedule(new ModelShutdown(consumer, reason, _model));
         }
 
-        private void ScheduleUnlessShuttingDown<TWork>(TWork work)
-            where TWork : Work
+        private void ScheduleUnlessShuttingDown(Work work)
         {
             if (!IsShutdown)
             {
                 Schedule(work);
             }
         }
 
-        private void Schedule<TWork>(TWork work)
-            where TWork : Work
+        private void Schedule(Work work)
         {
             _workService.Schedule(_model, work);
         }

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

@@ -18,7 +18,7 @@ public AsyncConsumerWorkService(int concurrency) : base(concurrency)
             _startNewWorkPoolFunc = model => StartNewWorkPool(model);
         }
 
-        public void Schedule<TWork>(IModel model, TWork work) where TWork : Work
+        public void Schedule(IModel model, Work work)
         {
             /*
              * rabbitmq/rabbitmq-dotnet-client#841

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

@@ -564,15 +564,15 @@ public void HandleBasicDeliver(string consumerTag,
             string exchange,
             string routingKey,
             IBasicProperties basicProperties,
-            ReadOnlyMemory<byte> body)
+            ReadOnlyMemory<byte> body,
+            byte[] rentedArray)
         {
             if (_disposed)
             {
                 throw new ObjectDisposedException(GetType().FullName);
             }
 
-            _delegate.HandleBasicDeliver(consumerTag, deliveryTag, redelivered, exchange,
-                routingKey, basicProperties, body);
+            _delegate.HandleBasicDeliver(consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body, rentedArray);
         }
 
         public void HandleBasicGetEmpty()
@@ -591,15 +591,15 @@ public void HandleBasicGetOk(ulong deliveryTag,
             string routingKey,
             uint messageCount,
             IBasicProperties basicProperties,
-            ReadOnlyMemory<byte> body)
+            ReadOnlyMemory<byte> body,
+            byte[] rentedArray)
         {
             if (_disposed)
             {
                 throw new ObjectDisposedException(GetType().FullName);
             }
 
-            _delegate.HandleBasicGetOk(deliveryTag, redelivered, exchange, routingKey,
-                messageCount, basicProperties, body);
+            _delegate.HandleBasicGetOk(deliveryTag, redelivered, exchange, routingKey, messageCount, basicProperties, body, rentedArray);
         }
 
         public void HandleBasicNack(ulong deliveryTag,
@@ -629,15 +629,15 @@ public void HandleBasicReturn(ushort replyCode,
             string exchange,
             string routingKey,
             IBasicProperties basicProperties,
-            ReadOnlyMemory<byte> body)
+            ReadOnlyMemory<byte> body,
+            byte[] rentedArray)
         {
             if (_disposed)
             {
                 throw new ObjectDisposedException(GetType().FullName);
             }
 
-            _delegate.HandleBasicReturn(replyCode, replyText, exchange,
-                routingKey, basicProperties, body);
+            _delegate.HandleBasicReturn(replyCode, replyText, exchange, routingKey, basicProperties, body, rentedArray);
         }
 
         public void HandleChannelClose(ushort replyCode,

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

@@ -1,6 +1,5 @@
 using System;
 using System.Buffers;
-using System.Runtime.InteropServices;
 using System.Threading.Tasks;
 
 namespace RabbitMQ.Client.Impl
@@ -14,6 +13,7 @@ internal sealed class BasicDeliver : Work
         private readonly string _routingKey;
         private readonly IBasicProperties _basicProperties;
         private readonly ReadOnlyMemory<byte> _body;
+        private readonly byte[] _rentedBytes;
 
         public override string Context => "HandleBasicDeliver";
 
@@ -24,7 +24,8 @@ public BasicDeliver(IBasicConsumer consumer,
             string exchange,
             string routingKey,
             IBasicProperties basicProperties,
-            ReadOnlyMemory<byte> body) : base(consumer)
+            ReadOnlyMemory<byte> body,
+            byte[] rentedBytes) : base(consumer)
         {
             _consumerTag = consumerTag;
             _deliveryTag = deliveryTag;
@@ -33,6 +34,7 @@ public BasicDeliver(IBasicConsumer consumer,
             _routingKey = routingKey;
             _basicProperties = basicProperties;
             _body = body;
+            _rentedBytes = rentedBytes;
         }
 
         protected override Task Execute(IAsyncBasicConsumer consumer)
@@ -48,10 +50,7 @@ protected override Task Execute(IAsyncBasicConsumer consumer)
 
         public override void PostExecute()
         {
-            if (MemoryMarshal.TryGetArray(_body, out ArraySegment<byte> segment))
-            {
-                ArrayPool<byte>.Shared.Return(segment.Array);
-            }
+            ArrayPool<byte>.Shared.Return(_rentedBytes);
         }
     }
 }

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

@@ -46,7 +46,7 @@ internal sealed class CommandAssembler
         private MethodBase _method;
         private ContentHeaderBase _header;
         private byte[] _bodyBytes;
-        private Memory<byte> _body;
+        private ReadOnlyMemory<byte> _body;
         private int _remainingBodyBytes;
         private int _offset;
         private AssemblyState _state;
@@ -62,14 +62,15 @@ private void Reset()
             _method = null;
             _header = null;
             _bodyBytes = null;
-            _body = Memory<byte>.Empty;
+            _body = ReadOnlyMemory<byte>.Empty;
             _remainingBodyBytes = 0;
             _offset = 0;
             _state = AssemblyState.ExpectingMethod;
         }
 
-        public IncomingCommand HandleFrame(in InboundFrame frame)
+        public bool HandleFrame(in InboundFrame frame, out IncomingCommand command)
         {
+            bool shallReturn = true;
             switch (_state)
             {
                 case AssemblyState.ExpectingMethod:
@@ -79,18 +80,19 @@ public IncomingCommand HandleFrame(in InboundFrame frame)
                     ParseHeaderFrame(in frame);
                     break;
                 case AssemblyState.ExpectingContentBody:
-                    ParseBodyFrame(in frame);
+                    shallReturn = ParseBodyFrame(in frame);
                     break;
             }
 
             if (_state != AssemblyState.Complete)
             {
-                return IncomingCommand.Empty;
+                command = IncomingCommand.Empty;
+                return true;
             }
 
-            var result = new IncomingCommand(_method, _header, _body, _bodyBytes);
+            command = new IncomingCommand(_method, _header, _body, _bodyBytes);
             Reset();
-            return result;
+            return shallReturn;
         }
 
         private void ParseMethodFrame(in InboundFrame frame)
@@ -121,14 +123,10 @@ private void ParseHeaderFrame(in InboundFrame frame)
             }
 
             _remainingBodyBytes = (int) totalBodyBytes;
-
-            // Is returned by IncomingCommand.Dispose in Session.HandleFrame
-            _bodyBytes = ArrayPool<byte>.Shared.Rent(_remainingBodyBytes);
-            _body = new Memory<byte>(_bodyBytes, 0, _remainingBodyBytes);
             UpdateContentBodyState();
         }
 
-        private void ParseBodyFrame(in InboundFrame frame)
+        private bool ParseBodyFrame(in InboundFrame frame)
         {
             if (frame.Type != FrameType.FrameBody)
             {
@@ -141,10 +139,27 @@ private void ParseBodyFrame(in InboundFrame frame)
                 throw new MalformedFrameException($"Overlong content body received - {_remainingBodyBytes} bytes remaining, {payloadLength} bytes received");
             }
 
-            frame.Payload.CopyTo(_body.Slice(_offset));
+            if (_bodyBytes is null)
+            {
+                // check for single frame payload for an early exit
+                if (payloadLength == _remainingBodyBytes)
+                {
+                    _bodyBytes = frame.TakeoverPayload();
+                    _body = frame.Payload;
+                    _state = AssemblyState.Complete;
+                    return false;
+                }
+
+                // Is returned by IncomingCommand.ReturnPayload in Session.HandleFrame
+                _bodyBytes = ArrayPool<byte>.Shared.Rent(_remainingBodyBytes);
+                _body = new ReadOnlyMemory<byte>(_bodyBytes, 0, _remainingBodyBytes);
+            }
+
+            frame.Payload.Span.CopyTo(_bodyBytes.AsSpan(_offset));
             _remainingBodyBytes -= payloadLength;
             _offset += payloadLength;
             UpdateContentBodyState();
+            return true;
         }
 
         private void UpdateContentBodyState()