Added tests for non-copying body

Author: GerardSmit

projects/RabbitMQ.Client/client/RentedOutgoingMemory.cs

@@ -18,6 +18,8 @@ internal class RentedOutgoingMemory : IDisposable, IResettable
 
         internal int Size => (int) Data.Length;
 
+        public int RentedArraySize => _rentedArray?.Length ?? 0;
+
         internal ReadOnlySequence<byte> Data { get; private set; }
 
         /// <summary>
@@ -100,7 +102,7 @@ bool IResettable.TryReset()
             return true;
         }
 
-        public static RentedOutgoingMemory GetAndInitialize(ReadOnlySequence<byte> mem, byte[] buffer, bool waitSend = false)
+        public static RentedOutgoingMemory GetAndInitialize(ReadOnlySequence<byte> mem, byte[]? buffer = null, bool waitSend = false)
         {
             var rented = s_pool.Get();
 
@@ -115,7 +117,7 @@ public static RentedOutgoingMemory GetAndInitialize(ReadOnlySequence<byte> mem,
             return rented;
         }
 
-        public static RentedOutgoingMemory GetAndInitialize(ReadOnlyMemory<byte> mem, byte[] buffer, bool waitSend = false)
+        public static RentedOutgoingMemory GetAndInitialize(ReadOnlyMemory<byte> mem, byte[]? buffer = null, bool waitSend = false)
         {
             return GetAndInitialize(new ReadOnlySequence<byte>(mem), buffer, waitSend);
         }

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

@@ -763,7 +763,12 @@ public async ValueTask<IConnection> CreateConnectionAsync(IEndpointResolver endp
             }
         }
 
-        private ConnectionConfig CreateConfig(string clientProvidedName)
+        internal ConnectionConfig CreateConfig()
+        {
+            return CreateConfig(ClientProvidedName);
+        }
+
+        internal ConnectionConfig CreateConfig(string clientProvidedName)
         {
             return new ConnectionConfig(
                 VirtualHost,

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

@@ -140,14 +140,14 @@ public static ValueTask BasicPublishAsync<T>(this IChannel channel, PublicationA
         public static void BasicPublish(this IChannel channel, string exchange, string routingKey, ReadOnlyMemory<byte> body = default, bool mandatory = false)
             => channel.BasicPublish(exchange, routingKey, in EmptyBasicProperty.Empty, body, mandatory);
 
-        public static ValueTask BasicPublishAsync(this IChannel channel, string exchange, string routingKey, ReadOnlyMemory<byte> body = default, bool mandatory = false)
-            => channel.BasicPublishAsync(exchange, routingKey, in EmptyBasicProperty.Empty, body, mandatory);
+        public static ValueTask BasicPublishAsync(this IChannel channel, string exchange, string routingKey, ReadOnlyMemory<byte> body = default, bool mandatory = false, bool? copyBody = null)
+            => channel.BasicPublishAsync(exchange, routingKey, in EmptyBasicProperty.Empty, body, mandatory, copyBody);
 
-        public static void BasicPublish(this IChannel channel, CachedString exchange, CachedString routingKey, ReadOnlyMemory<byte> body = default, bool mandatory = false)
+        public static void BasicPublish(this IChannel channel, CachedString exchange, CachedString routingKey, ReadOnlyMemory<byte> body = default, bool mandatory = false, bool? copyBody = null)
             => channel.BasicPublish(exchange, routingKey, in EmptyBasicProperty.Empty, body, mandatory);
 
-        public static ValueTask BasicPublishAsync(this IChannel channel, CachedString exchange, CachedString routingKey, ReadOnlyMemory<byte> body = default, bool mandatory = false)
-            => channel.BasicPublishAsync(exchange, routingKey, in EmptyBasicProperty.Empty, body, mandatory);
+        public static ValueTask BasicPublishAsync(this IChannel channel, CachedString exchange, CachedString routingKey, ReadOnlyMemory<byte> body = default, bool mandatory = false, bool? copyBody = null)
+            => channel.BasicPublishAsync(exchange, routingKey, in EmptyBasicProperty.Empty, body, mandatory, copyBody);
 #nullable disable
 
         /// <summary>

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

@@ -125,6 +125,11 @@ public interface IConnection : INetworkConnection, IDisposable
         /// </summary>
         IEnumerable<ShutdownReportEntry> ShutdownReport { get; }
 
+        /// <summary>
+        /// The threshold for when to copy the body to a temporary array.
+        /// </summary>
+        int CopyBodyToMemoryThreshold { get; }
+
         /// <summary>
         /// Application-specific connection name, will be displayed in the management UI
         /// if RabbitMQ server supports it. This value doesn't have to be unique and cannot

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

@@ -50,7 +50,7 @@ internal sealed partial class AutorecoveringConnection : IConnection
         private Connection _innerConnection;
         private bool _disposed;
 
-        private Connection InnerConnection
+        internal Connection InnerConnection
         {
             get
             {
@@ -181,6 +181,8 @@ public event EventHandler<RecoveringConsumerEventArgs> RecoveringConsumer
 
         public IEnumerable<ShutdownReportEntry> ShutdownReport => InnerConnection.ShutdownReport;
 
+        public int CopyBodyToMemoryThreshold => InnerConnection.CopyBodyToMemoryThreshold;
+
         public IProtocol Protocol => Endpoint.Protocol;
 
         public RecoveryAwareChannel CreateNonRecoveringChannel()

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

@@ -58,6 +58,9 @@ internal sealed partial class Connection : IConnection
         private ShutdownEventArgs? _closeReason;
         public ShutdownEventArgs? CloseReason => Volatile.Read(ref _closeReason);
 
+        internal bool TrackRentedBytes = false;
+        internal uint RentedBytes;
+
         internal Connection(ConnectionConfig config, IFrameHandler frameHandler)
         {
             _config = config;
@@ -552,9 +555,23 @@ internal void Write(RentedOutgoingMemory frames)
 
         internal ValueTask WriteAsync(RentedOutgoingMemory frames)
         {
+            TrackRented(frames.RentedArraySize);
+
             return _frameHandler.WriteAsync(frames);
         }
 
+        private void TrackRented(int size)
+        {
+            if (TrackRentedBytes && size > 0)
+            {
+#if NET
+                Interlocked.Add(ref RentedBytes, (uint)size);
+#else
+                Interlocked.Add(ref Unsafe.As<uint, int>(ref RentedBytes), size);
+#endif
+            }
+        }
+
         public void Dispose()
         {
             if (_disposed)

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

@@ -180,7 +180,7 @@ public ValueTask TransmitAsync<TMethod, THeader>(in TMethod cmd, in THeader head
                 ThrowAlreadyClosedException();
             }
 
-            copyBody ??= body.Length >= Connection.CopyBodyToMemoryThreshold;
+            copyBody ??= body.Length > Connection.CopyBodyToMemoryThreshold;
 
             return Connection.WriteAsync(Framing.SerializeToFrames(ref Unsafe.AsRef(cmd), ref Unsafe.AsRef(header), body, ChannelNumber, Connection.MaxPayloadSize, copyBody.Value));
         }

projects/Test/AsyncIntegration/TestBasicPublishAsync.cs

@@ -67,5 +67,76 @@ public async Task TestQueuePurgeAsync()
             Assert.True(await publishSyncSource.Task);
             Assert.Equal((uint)messageCount, await _channel.QueuePurgeAsync(q));
         }
+
+        [Fact]
+        public async Task TestNonCopyingBody()
+        {
+            const int size = 1024;
+
+            QueueDeclareOk q = await _channel.QueueDeclareAsync(string.Empty, false, false, true, false, null);
+            byte[] body = GetRandomBody(size);
+
+            uint rentedBytes;
+
+            using (var result = await TrackRentedBytes())
+            {
+                await _channel.BasicPublishAsync(string.Empty, q, body, copyBody: false);
+                rentedBytes = result.RentedBytes;
+            }
+
+            Assert.Equal((uint)1, await _channel.QueuePurgeAsync(q));
+
+            // It is expected that the rented bytes is smaller than the size of the body
+            // since we're not copying the body. Only the frame headers are rented.
+            Assert.True(rentedBytes < size);
+        }
+
+        [Fact]
+        public async Task TestCopyingBody()
+        {
+            const int size = 1024;
+
+            QueueDeclareOk q = await _channel.QueueDeclareAsync(string.Empty, false, false, true, false, null);
+            byte[] body = GetRandomBody(size);
+
+            uint rentedBytes;
+
+            using (var result = await TrackRentedBytes())
+            {
+                await _channel.BasicPublishAsync(string.Empty, q, body, copyBody: true);
+                rentedBytes = result.RentedBytes;
+            }
+
+            Assert.Equal((uint)1, await _channel.QueuePurgeAsync(q));
+
+            // It is expected that the rented bytes is larger than the size of the body
+            // since the body is copied with the frame headers.
+            Assert.True(rentedBytes >= size);
+        }
+
+        [Fact]
+        public async Task TestDefaultCopyingBody()
+        {
+            Assert.Equal(int.MaxValue, _conn.CopyBodyToMemoryThreshold);
+
+            const int size = 1024;
+
+            QueueDeclareOk q = await _channel.QueueDeclareAsync(string.Empty, false, false, true, false, null);
+            byte[] body = GetRandomBody(size);
+
+            uint rentedBytes;
+
+            using (var result = await TrackRentedBytes())
+            {
+                await _channel.BasicPublishAsync(string.Empty, q, body, copyBody: true);
+                rentedBytes = result.RentedBytes;
+            }
+
+            Assert.Equal((uint)1, await _channel.QueuePurgeAsync(q));
+
+            // It is expected that the rented bytes is larger than the size of the body
+            // since the body is copied with the frame headers.
+            Assert.True(rentedBytes >= size);
+        }
     }
 }

projects/Test/AsyncIntegration/TestBasicPublishCopyBodyAsync.cs

@@ -0,0 +1,66 @@
+using System.Threading.Tasks;
+using RabbitMQ.Client;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Test.AsyncIntegration;
+
+public class TestBasicPublishCopyBodyAsync : AsyncIntegrationFixture
+{
+    public TestBasicPublishCopyBodyAsync(ITestOutputHelper output) : base(output)
+    {
+    }
+
+    protected override ConnectionFactory CreateConnectionFactory()
+    {
+        var factory = base.CreateConnectionFactory();
+        factory.CopyBodyToMemoryThreshold = 1024;
+        return factory;
+    }
+
+    [Theory(Skip = "Parallelization is disabled for this collection")]
+    [InlineData(512)]
+    [InlineData(1024)]
+    public async Task TestNonCopyingBody(ushort size)
+    {
+        QueueDeclareOk q = await _channel.QueueDeclareAsync(string.Empty, false, false, true, false, null);
+        byte[] body = GetRandomBody(size);
+
+        uint rentedBytes;
+
+        using (var result = await TrackRentedBytes())
+        {
+            await _channel.BasicPublishAsync(string.Empty, q, body);
+            rentedBytes = result.RentedBytes;
+        }
+
+        Assert.Equal((uint)1, await _channel.QueuePurgeAsync(q));
+
+        // It is expected that the rented bytes is smaller than the size of the body
+        // since we're not copying the body. Only the frame headers are rented.
+        Assert.True(rentedBytes < size);
+    }
+
+    [Theory]
+    [InlineData(1025)]
+    [InlineData(2048)]
+    public async Task TestCopyingBody(ushort size)
+    {
+        QueueDeclareOk q = await _channel.QueueDeclareAsync(string.Empty, false, false, true, false, null);
+        byte[] body = GetRandomBody(size);
+
+        uint rentedBytes;
+
+        using (var result = await TrackRentedBytes())
+        {
+            await _channel.BasicPublishAsync(string.Empty, q, body);
+            rentedBytes = result.RentedBytes;
+        }
+
+        Assert.Equal((uint)1, await _channel.QueuePurgeAsync(q));
+
+        // It is expected that the rented bytes is larger than the size of the body
+        // since the body is copied with the frame headers.
+        Assert.True(rentedBytes >= size);
+    }
+}

projects/Test/Common/IntegrationFixtureBase.cs

@@ -35,6 +35,7 @@
 using System.Reflection;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 using RabbitMQ.Client;
 using RabbitMQ.Client.Framing.Impl;
 using Xunit;
@@ -44,6 +45,8 @@ namespace Test
 {
     public abstract class IntegrationFixtureBase : IDisposable
     {
+        private readonly SemaphoreSlim _byteTrackingLock = new SemaphoreSlim(1, 1);
+
         private static bool s_isRunningInCI = false;
         private static bool s_isWindows = false;
         private static bool s_isVerbose = false;
@@ -292,7 +295,7 @@ protected void Wait(ManualResetEventSlim latch, TimeSpan timeSpan, string desc)
                 $"waiting {timeSpan.TotalSeconds} seconds on a latch for '{desc}' timed out");
         }
 
-        protected ConnectionFactory CreateConnectionFactory()
+        protected virtual ConnectionFactory CreateConnectionFactory()
         {
             string now = DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture);
             return new ConnectionFactory
@@ -339,6 +342,39 @@ protected void HandleChannelShutdown(IChannel ch, ShutdownEventArgs args, Action
             a(args);
         }
 
+        protected async Task<TrackRentedByteResult> TrackRentedBytes()
+        {
+            Connection connection;
+
+            if (_conn is AutorecoveringConnection autorecoveringConnection)
+            {
+                connection = autorecoveringConnection.InnerConnection as Connection;
+            }
+            else
+            {
+                connection = _conn as Connection;
+            }
+
+            if (connection is null)
+            {
+                throw new InvalidOperationException("Cannot track rented bytes without a connection");
+            }
+
+            await _byteTrackingLock.WaitAsync();
+
+            try
+            {
+                connection.RentedBytes = 0;
+                connection.TrackRentedBytes = true;
+                return new TrackRentedByteResult(connection, _byteTrackingLock);
+            }
+            catch
+            {
+                _byteTrackingLock.Release();
+                throw;
+            }
+        }
+
         private static void InitIsRunningInCI()
         {
             bool ci;

projects/Test/Common/TrackRentedByteResult.cs

@@ -0,0 +1,24 @@
+using System;
+using System.Threading;
+using RabbitMQ.Client.Framing.Impl;
+
+namespace Test;
+
+public sealed class TrackRentedByteResult : IDisposable
+{
+    private readonly Connection _connection;
+    private readonly SemaphoreSlim _byteTrackingLock;
+
+    internal TrackRentedByteResult(Connection connection, SemaphoreSlim byteTrackingLock)
+    {
+        _connection = connection;
+        _byteTrackingLock = byteTrackingLock;
+    }
+
+    public uint RentedBytes => _connection.RentedBytes;
+
+    public void Dispose()
+    {
+        _byteTrackingLock.Release();
+    }
+}