Change how ManagementClient gets passed in

Author: bording

src/ServiceControl.Transports.RabbitMQ/IManagementClientProvider.cs

@@ -0,0 +1,8 @@
+namespace ServiceControl.Transports.RabbitMQ;
+
+using NServiceBus.Transport.RabbitMQ.ManagementApi;
+
+interface IManagementClientProvider
+{
+    ManagementClient ManagementClient { get; }
+}

src/ServiceControl.Transports.RabbitMQ/IRabbitMQTransportExtentions.cs

@@ -5,12 +5,15 @@
     using BrokerThroughput;
     using Microsoft.Extensions.DependencyInjection;
     using NServiceBus;
+    using NServiceBus.Transport.RabbitMQ.ManagementApi;
 
-    public abstract class RabbitMQConventionalRoutingTransportCustomization(QueueType queueType) : TransportCustomization<RabbitMQTransport>, IRabbitMQTransportExtensions
+    public abstract class RabbitMQConventionalRoutingTransportCustomization(NServiceBus.QueueType queueType) : TransportCustomization<RabbitMQTransport>, IManagementClientProvider
     {
         RabbitMQTransport rabbitMQTransport;
 
-        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
+        ManagementClient IManagementClientProvider.ManagementClient => rabbitMQTransport?.ManagementClient ?? throw new InvalidOperationException("Transport instance has not been created yet.");
+
+        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
 
         protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
@@ -27,7 +30,6 @@ protected override RabbitMQTransport CreateTransport(TransportSettings transport
             transport.TransportTransactionMode = transport.GetSupportedTransactionModes().Contains(preferredTransactionMode) ? preferredTransactionMode : TransportTransactionMode.ReceiveOnly;
             transport.SetCustomSettingsFromConnectionString(transportSettings.ConnectionString);
 
-            rabbitMQTransport = transport;
             return transport;
         }
 
@@ -39,15 +41,5 @@ protected sealed override void AddTransportForMonitoringCore(IServiceCollection
             services.AddSingleton<IProvideQueueLength, QueueLengthProvider>();
             services.AddHostedService(provider => provider.GetRequiredService<IProvideQueueLength>());
         }
-
-        RabbitMQTransport IRabbitMQTransportExtensions.GetTransport()
-        {
-            if (rabbitMQTransport == null)
-            {
-                throw new InvalidOperationException("Transport instance has not been created yet. Make sure CreateTransport() is called before accessing the transport.");
-            }
-
-            return rabbitMQTransport;
-        }
     }
 }

src/ServiceControl.Transports.RabbitMQ/RabbitMQConventionalRoutingTransportCustomization.cs

@@ -5,13 +5,15 @@
     using BrokerThroughput;
     using Microsoft.Extensions.DependencyInjection;
     using NServiceBus;
+    using NServiceBus.Transport.RabbitMQ.ManagementApi;
 
-    public abstract class RabbitMQDirectRoutingTransportCustomization(QueueType queueType)
-        : TransportCustomization<RabbitMQTransport>, IRabbitMQTransportExtensions
+    public abstract class RabbitMQDirectRoutingTransportCustomization(NServiceBus.QueueType queueType) : TransportCustomization<RabbitMQTransport>, IManagementClientProvider
     {
         RabbitMQTransport rabbitMQTransport;
 
-        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
+        ManagementClient IManagementClientProvider.ManagementClient => rabbitMQTransport?.ManagementClient ?? throw new InvalidOperationException("Transport instance has not been created yet.");
+
+        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
 
         protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
@@ -28,7 +30,6 @@ protected override RabbitMQTransport CreateTransport(TransportSettings transport
             transport.TransportTransactionMode = transport.GetSupportedTransactionModes().Contains(preferredTransactionMode) ? preferredTransactionMode : TransportTransactionMode.ReceiveOnly;
             transport.SetCustomSettingsFromConnectionString(transportSettings.ConnectionString);
 
-            rabbitMQTransport = transport;
             return transport;
         }
 
@@ -40,15 +41,5 @@ protected sealed override void AddTransportForMonitoringCore(IServiceCollection
             services.AddSingleton<IProvideQueueLength, QueueLengthProvider>();
             services.AddHostedService(provider => provider.GetRequiredService<IProvideQueueLength>());
         }
-
-        RabbitMQTransport IRabbitMQTransportExtensions.GetTransport()
-        {
-            if (rabbitMQTransport == null)
-            {
-                throw new InvalidOperationException("Transport instance has not been created yet. Make sure CreateTransport() is called before accessing the transport.");
-            }
-
-            return rabbitMQTransport;
-        }
     }
 }

src/ServiceControl.Transports.RabbitMQ/RabbitMQDirectRoutingTransportCustomization.cs

@@ -11,7 +11,6 @@ namespace ServiceControl.Transports.RabbitMQ;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
-using NServiceBus;
 using NServiceBus.Transport.RabbitMQ.ManagementApi;
 using Polly;
 using Polly.Retry;
@@ -21,7 +20,8 @@ public class RabbitMQQuery : BrokerThroughputQuery
 {
     readonly ILogger<RabbitMQQuery> logger;
     readonly TimeProvider timeProvider;
-    readonly RabbitMQTransport rabbitMQTransport;
+    readonly ManagementClient managementClient;
+
     readonly ResiliencePipeline pipeline = new ResiliencePipelineBuilder()
       .AddRetry(new RetryStrategyOptions()) // Add retry using the default options
       .AddTimeout(TimeSpan.FromMinutes(2)) // Add timeout if it keeps failing
@@ -31,7 +31,15 @@ public RabbitMQQuery(ILogger<RabbitMQQuery> logger, TimeProvider timeProvider, T
     {
         this.logger = logger;
         this.timeProvider = timeProvider;
-        rabbitMQTransport = GetRabbitMQTransport(transportCustomization);
+
+        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));
+        }
     }
 
     protected override void InitializeCore(ReadOnlyDictionary<string, string> settings)
@@ -43,16 +51,6 @@ protected override void InitializeCore(ReadOnlyDictionary<string, string> settin
         CheckLegacySettings(settings, RabbitMQSettings.API);
     }
 
-    static RabbitMQTransport GetRabbitMQTransport(ITransportCustomization transportCustomization)
-    {
-        if (transportCustomization is IRabbitMQTransportExtensions rabbitMQTransportCustomization)
-        {
-            return rabbitMQTransportCustomization.GetTransport();
-        }
-
-        throw new InvalidOperationException($"Expected a RabbitMQTransport but received {transportCustomization.GetType().Name}.");
-    }
-
     void CheckLegacySettings(ReadOnlyDictionary<string, string> settings, string key)
     {
         if (settings.TryGetValue(key, out _))
@@ -65,7 +63,7 @@ void CheckLegacySettings(ReadOnlyDictionary<string, string> settings, string key
     public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBrokerQueue brokerQueue, DateOnly startDate, [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
         var queue = (RabbitMQBrokerQueueDetails)brokerQueue;
-        var response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetQueue(queue.QueueName, token), cancellationToken);
+        var response = await pipeline.ExecuteAsync(async token => await managementClient.GetQueue(queue.QueueName, token), cancellationToken);
 
         if (response.Value is null)
         {
@@ -81,7 +79,7 @@ public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBro
         {
             await Task.Delay(TimeSpan.FromMinutes(15), timeProvider, cancellationToken);
 
-            response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetQueue(queue.QueueName, token), cancellationToken);
+            response = await pipeline.ExecuteAsync(async token => await managementClient.GetQueue(queue.QueueName, token), cancellationToken);
 
             if (response.Value is null)
             {
@@ -101,7 +99,7 @@ public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBro
 
     async Task GetRabbitDetails(CancellationToken cancellationToken)
     {
-        var response = await pipeline.ExecuteAsync(async async => await rabbitMQTransport.ManagementClient.GetOverview(cancellationToken), cancellationToken);
+        var response = await pipeline.ExecuteAsync(async async => await managementClient.GetOverview(cancellationToken), cancellationToken);
 
         ValidateResponse(response);
 
@@ -162,7 +160,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
     {
         try
         {
-            var response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetBindingsForQueue(brokerQueue.QueueName, token), cancellationToken);
+            var response = await pipeline.ExecuteAsync(async token => await managementClient.GetBindingsForQueue(brokerQueue.QueueName, token), cancellationToken);
 
             // Check if conventional binding is found
             if (response.Value.Any(binding => binding?.Source == brokerQueue.QueueName
@@ -181,7 +179,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
 
         try
         {
-            var response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetBindingsForExchange(brokerQueue.QueueName, token), cancellationToken);
+            var response = await pipeline.ExecuteAsync(async token => await managementClient.GetBindingsForExchange(brokerQueue.QueueName, token), cancellationToken);
 
             // Check if delayed binding is found
             if (response.Value.Any(binding => binding?.Source is "nsb.v2.delay-delivery" or "nsb.delay-delivery"
@@ -200,7 +198,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
 
     internal async Task<(List<RabbitMQBrokerQueueDetails>?, bool morePages)> GetPage(int page, CancellationToken cancellationToken)
     {
-        var (StatusCode, Reason, Value, MorePages) = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetQueues(page, 500, token), cancellationToken);
+        var (StatusCode, Reason, Value, MorePages) = await pipeline.ExecuteAsync(async token => await managementClient.GetQueues(page, 500, token), cancellationToken);
 
         ValidateResponse((StatusCode, Reason, Value));
 
@@ -229,7 +227,7 @@ static List<RabbitMQBrokerQueueDetails> MaterializeQueueDetails(List<Queue> item
     {
         try
         {
-            var (statusCode, reason, value) = await rabbitMQTransport.ManagementClient.GetOverview(cancellationToken);
+            var (statusCode, reason, value) = await managementClient.GetOverview(cancellationToken);
 
             if (value is not null)
             {