Update QueueLengthProvider to use management client

Author: bording

src/Directory.Packages.props

@@ -38,7 +38,7 @@
     <PackageVersion Include="NServiceBus.Metrics" Version="5.0.1" />
     <PackageVersion Include="NServiceBus.Metrics.ServiceControl" Version="5.0.0" />
     <PackageVersion Include="NServiceBus.Persistence.NonDurable" Version="2.0.1" />
-    <PackageVersion Include="NServiceBus.RabbitMQ" Version="10.0.0-beta.1" />
+    <PackageVersion Include="NServiceBus.RabbitMQ" Version="10.0.0-beta.2" />
     <PackageVersion Include="NServiceBus.SagaAudit" Version="5.0.2" />
     <PackageVersion Include="NServiceBus.Testing" Version="9.0.1" />
     <PackageVersion Include="NServiceBus.Transport.AzureServiceBus" Version="5.0.0" />

src/ServiceControl.Transports.RabbitMQ/QueueLengthProvider.cs

@@ -4,17 +4,21 @@
     using System.Collections.Concurrent;
     using System.Threading;
     using System.Threading.Tasks;
-    using global::RabbitMQ.Client;
     using NServiceBus.Logging;
-    using NServiceBus.Transport.RabbitMQ;
-    using ConnectionFactory = NServiceBus.Transport.RabbitMQ.ConnectionFactory;
+    using NServiceBus.Transport.RabbitMQ.ManagementApi;
 
     class QueueLengthProvider : AbstractQueueLengthProvider
     {
-        public QueueLengthProvider(TransportSettings settings, Action<QueueLengthEntry[], EndpointToQueueMapping> store) : base(settings, store)
+        public QueueLengthProvider(TransportSettings settings, Action<QueueLengthEntry[], EndpointToQueueMapping> store, ITransportCustomization transportCustomization) : base(settings, store)
         {
-            queryExecutor = new QueryExecutor(ConnectionString);
-            queryExecutor.Initialize();
+            if (transportCustomization is IManagementClientProvider provider)
+            {
+                managementClient = provider.ManagementClient;
+            }
+            else
+            {
+                throw new ArgumentException($"Transport customization does not implement {nameof(IManagementClientProvider)}. Type: {transportCustomization.GetType().Name}", nameof(transportCustomization));
+            }
         }
 
         public override void TrackEndpointInputQueue(EndpointToQueueMapping queueToTrack) =>
@@ -76,90 +80,37 @@ async Task FetchQueueLengths(CancellationToken cancellationToken)
         {
             foreach (var endpointQueuePair in endpointQueues)
             {
-                await queryExecutor.Execute(async m =>
-                {
-                    var queueName = endpointQueuePair.Value;
-
-                    try
-                    {
-                        var size = (int)await m.MessageCountAsync(queueName, cancellationToken).ConfigureAwait(false);
-
-                        sizes.AddOrUpdate(queueName, _ => size, (_, __) => size);
-                    }
-                    catch (Exception e)
-                    {
-                        Logger.Warn($"Error querying queue length for {queueName}", e);
-                    }
-                }, cancellationToken);
-            }
-        }
-
-        readonly QueryExecutor queryExecutor;
-        static readonly TimeSpan QueryDelayInterval = TimeSpan.FromMilliseconds(200);
-
-        readonly ConcurrentDictionary<string, string> endpointQueues = new ConcurrentDictionary<string, string>();
-        readonly ConcurrentDictionary<string, int> sizes = new ConcurrentDictionary<string, int>();
-
-        static readonly ILog Logger = LogManager.GetLogger<QueueLengthProvider>();
-
-        class QueryExecutor(string connectionString) : IDisposable
-        {
-
-            public void Initialize()
-            {
-                var connectionConfiguration = ConnectionConfiguration.Create(connectionString);
-
-                // TODO Fix this up
-                //var dbConnectionStringBuilder = new DbConnectionStringBuilder { ConnectionString = connectionString };
-
-                connectionFactory = new("ServiceControl.Monitoring",
-                    connectionConfiguration,
-                    null, //providing certificates is not supported yet
-                    false, // TODO Fix dbConnectionStringBuilder.GetBooleanValue("DisableRemoteCertificateValidation"),
-                    false, // TODO fix dbConnectionStringBuilder.GetBooleanValue("UseExternalAuthMechanism"),
-                    TimeSpan.FromSeconds(60), // value would come from config API in actual transport
-                    TimeSpan.FromSeconds(10), // value would come from config API in actual transport
-                    null); // value would come from config API in actual transport
-            }
+                var queueName = endpointQueuePair.Value;
 
-            public async Task Execute(Action<IChannel> action, CancellationToken cancellationToken = default)
-            {
                 try
                 {
-                    connection ??= await connectionFactory.CreateConnection("queue length monitor", cancellationToken: cancellationToken);
+                    var (statusCode, reason, queue) = await managementClient.GetQueue(queueName, cancellationToken);
 
-                    //Connection implements reconnection logic
-                    while (!connection.IsOpen)
+                    if (queue is not null)
                     {
-                        await Task.Delay(ReconnectionDelay, cancellationToken);
+                        var size = queue.MessageCount;
+                        sizes.AddOrUpdate(queueName, _ => size, (_, _) => size);
                     }
-
-                    if (channel == null || channel.IsClosed)
+                    else
                     {
-                        channel?.Dispose();
-
-                        channel = await connection.CreateChannelAsync(cancellationToken: cancellationToken).ConfigureAwait(false);
+                        Logger.Warn($"Error querying queue length for {queueName}. {statusCode}: {reason}");
                     }
 
-                    action(channel);
-                }
-                catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
-                {
-                    // no-op
                 }
                 catch (Exception e)
                 {
-                    Logger.Warn("Error querying queue length.", e);
+                    Logger.Warn($"Error querying queue length for {queueName}", e);
                 }
             }
+        }
 
-            public void Dispose() => connection?.Dispose();
+        static readonly TimeSpan QueryDelayInterval = TimeSpan.FromMilliseconds(200);
 
-            IConnection connection;
-            IChannel channel;
-            ConnectionFactory connectionFactory;
+        readonly ConcurrentDictionary<string, string> endpointQueues = new();
+        readonly ConcurrentDictionary<string, long> sizes = new();
 
-            static readonly TimeSpan ReconnectionDelay = TimeSpan.FromSeconds(5);
-        }
+        static readonly ILog Logger = LogManager.GetLogger<QueueLengthProvider>();
+
+        readonly ManagementClient managementClient;
     }
 }

src/ServiceControl.Transports.RabbitMQ/RabbitMQConventionalRoutingTransportCustomization.cs

@@ -17,7 +17,7 @@ public abstract class RabbitMQConventionalRoutingTransportCustomization(NService
 
         protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
-        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
+        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
 
         protected override RabbitMQTransport CreateTransport(TransportSettings transportSettings, TransportTransactionMode preferredTransactionMode = TransportTransactionMode.ReceiveOnly)
         {

src/ServiceControl.Transports.RabbitMQ/RabbitMQDirectRoutingTransportCustomization.cs

@@ -17,7 +17,7 @@ public abstract class RabbitMQDirectRoutingTransportCustomization(NServiceBus.Qu
 
         protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
-        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
+        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
 
         protected override RabbitMQTransport CreateTransport(TransportSettings transportSettings, TransportTransactionMode preferredTransactionMode = TransportTransactionMode.ReceiveOnly)
         {