Allow RabbitMQQuery to access the transport

Travis Nickels · bording · commit e24f2b193962 · 2025-02-26T11:26:46.000-05:00
diff --git a/src/ServiceControl.Transports.RabbitMQ/IRabbitMQTransportExtentions.cs b/src/ServiceControl.Transports.RabbitMQ/IRabbitMQTransportExtentions.cs
@@ -0,0 +1,9 @@
+﻿namespace ServiceControl.Transports.RabbitMQ
+{
+    using NServiceBus;
+
+    public interface IRabbitMQTransportExtensions
+    {
+        RabbitMQTransport GetTransport();
+    }
+}
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQConventionalRoutingTransportCustomization.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQConventionalRoutingTransportCustomization.cs
@@ -9,8 +9,11 @@
     using NServiceBus;
 
     public abstract class RabbitMQConventionalRoutingTransportCustomization(QueueType queueType)
-        : TransportCustomization<RabbitMQTransport>
+        : TransportCustomization<RabbitMQTransport>, IRabbitMQTransportExtensions
     {
+
+        RabbitMQTransport rabbitMQTransport;
+
         protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
         protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
@@ -36,14 +39,13 @@ protected override RabbitMQTransport CreateTransport(TransportSettings transport
             transport.ManagementApiUrl = GetValue(connectionStringDictionary, "ManagementApiUrl", string.Empty);
             transport.UseManagementApi = disableManagementApi.Equals("false", StringComparison.OrdinalIgnoreCase);
 
+            rabbitMQTransport = transport;
+
             return transport;
         }
 
-        protected override void AddTransportForPrimaryCore(IServiceCollection services,
-            TransportSettings transportSettings)
-        {
-            services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();
-        }
+        protected override void AddTransportForPrimaryCore(IServiceCollection services, TransportSettings transportSettings)
+            => services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();
 
         protected sealed override void AddTransportForMonitoringCore(IServiceCollection services, TransportSettings transportSettings)
         {
@@ -52,8 +54,15 @@ protected sealed override void AddTransportForMonitoringCore(IServiceCollection
         }
 
         static string GetValue(Dictionary<string, string> dictionary, string key, string defaultValue)
+            => dictionary.TryGetValue(key, out var value) ? value : defaultValue;
+
+        RabbitMQTransport IRabbitMQTransportExtensions.GetTransport()
         {
-            return dictionary.TryGetValue(key, out var value) ? value : defaultValue;
+            if (rabbitMQTransport == null)
+            {
+                throw new InvalidOperationException("Transport instance has not been created yet. Make sure CreateTransport() is called before accessing the transport.");
+            };
+            return rabbitMQTransport;
         }
     }
 }
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQDirectRoutingTransportCustomization.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQDirectRoutingTransportCustomization.cs
@@ -8,11 +8,10 @@
     using Microsoft.Extensions.DependencyInjection;
     using NServiceBus;
 
-    public abstract class RabbitMQDirectRoutingTransportCustomization : TransportCustomization<RabbitMQTransport>
+    public abstract class RabbitMQDirectRoutingTransportCustomization(QueueType queueType)
+        : TransportCustomization<RabbitMQTransport>, IRabbitMQTransportExtensions
     {
-        readonly QueueType queueType;
-
-        protected RabbitMQDirectRoutingTransportCustomization(QueueType queueType) => this.queueType = queueType;
+        RabbitMQTransport rabbitMQTransport;
 
         protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
@@ -38,14 +37,13 @@ protected override RabbitMQTransport CreateTransport(TransportSettings transport
             transport.ManagementApiUrl = GetValue(connectionStringDictionary, "ManagementApiUrl", string.Empty);
             transport.UseManagementApi = disableManagementApi.Equals("false", StringComparison.OrdinalIgnoreCase);
 
+            rabbitMQTransport = transport;
+
             return transport;
         }
 
-        protected override void AddTransportForPrimaryCore(IServiceCollection services,
-            TransportSettings transportSettings)
-        {
-            services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();
-        }
+        protected override void AddTransportForPrimaryCore(IServiceCollection services, TransportSettings transportSettings)
+            => services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();
 
         protected sealed override void AddTransportForMonitoringCore(IServiceCollection services, TransportSettings transportSettings)
         {
@@ -54,8 +52,15 @@ protected sealed override void AddTransportForMonitoringCore(IServiceCollection
         }
 
         static string GetValue(Dictionary<string, string> dictionary, string key, string defaultValue)
+            => dictionary.TryGetValue(key, out var value) ? value : defaultValue;
+
+        RabbitMQTransport IRabbitMQTransportExtensions.GetTransport()
         {
-            return dictionary.TryGetValue(key, out var value) ? value : defaultValue;
+            if (rabbitMQTransport == null)
+            {
+                throw new InvalidOperationException("Transport instance has not been created yet. Make sure CreateTransport() is called before accessing the transport.");
+            };
+            return rabbitMQTransport;
         }
     }
 }
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQQuery.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQQuery.cs
@@ -16,6 +16,7 @@ namespace ServiceControl.Transports.RabbitMQ;
 using System.Threading.Tasks;
 using System.Web;
 using Microsoft.Extensions.Logging;
+using NServiceBus;
 using Polly;
 using Polly.Retry;
 using ServiceControl.Transports.BrokerThroughput;
@@ -30,23 +31,35 @@ public class RabbitMQQuery : BrokerThroughputQuery
     readonly ILogger<RabbitMQQuery> logger;
     readonly TimeProvider timeProvider;
     readonly ConnectionConfiguration connectionConfiguration;
-    readonly string connectionString;
+    readonly TransportSettings transportSettings;
+    readonly RabbitMQTransport rabbitMQTransport;
 
     public RabbitMQQuery(ILogger<RabbitMQQuery> logger,
         TimeProvider timeProvider,
-        TransportSettings transportSettings) : base(logger, "RabbitMQ")
+        TransportSettings transportSettings,
+        ITransportCustomization transportCustomization) : base(logger, "RabbitMQ")
     {
         this.logger = logger;
         this.timeProvider = timeProvider;
-        connectionString = transportSettings.ConnectionString;
+        this.transportSettings = transportSettings;
+        if (transportCustomization is IRabbitMQTransportExtensions rabbitMQTransportCustomization)
+        {
+            rabbitMQTransport = rabbitMQTransportCustomization.GetTransport();
+        }
+        else
+        {
+            throw new InvalidOperationException($"Expected a RabbitMQTransport but received {transportCustomization.GetType().Name}.");
+        }
 
-        connectionConfiguration = ConnectionConfiguration.Create(connectionString, string.Empty);
+        connectionConfiguration = ConnectionConfiguration.Create(transportSettings.ConnectionString, string.Empty);
     }
 
     protected override void InitializeCore(ReadOnlyDictionary<string, string> settings)
     {
         var mangementApiUrl = GetManagementApiUrl();
 
+        // The licensing component configurations take precedence over the management API connection string configuration options
+        // https://docs.particular.net/servicecontrol/servicecontrol-instances/configuration#usage-reporting-when-using-the-rabbitmq-transport
         var userName = GetSettingsValue(settings, RabbitMQSettings.UserName, mangementApiUrl.UserName);
         var password = GetSettingsValue(settings, RabbitMQSettings.Password, mangementApiUrl.Password);
         var apiUrl = GetSettingsValue(settings, RabbitMQSettings.API, mangementApiUrl.Uri.AbsoluteUri);
@@ -107,7 +120,7 @@ string GetSettingsValue(ReadOnlyDictionary<string, string> settings, string key,
 
     UriBuilder GetManagementApiUrl()
     {
-        var dictionary = ConnectionConfiguration.ParseNServiceBusConnectionString(connectionString, new StringBuilder());
+        var dictionary = ConnectionConfiguration.ParseNServiceBusConnectionString(transportSettings.ConnectionString, new StringBuilder());
         UriBuilder uriBuilder;
 
         var managementApiUrl = GetValue(dictionary, "ManagementApiUrl", "");
diff --git a/src/ServiceControl.Transports.RabbitMQClassicConventionalRouting.Tests/RabbitMQQueryTests.cs b/src/ServiceControl.Transports.RabbitMQClassicConventionalRouting.Tests/RabbitMQQueryTests.cs
@@ -28,7 +28,7 @@ public async Task GetQueueNames_FindsQueues()
             MaxConcurrency = 1,
             EndpointName = Guid.NewGuid().ToString("N")
         };
-        var query = new RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, provider, transportSettings);
+        var query = new RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, provider, transportSettings, configuration.TransportCustomization);
         string[] additionalQueues = Enumerable.Range(1, 10).Select(i => $"myqueue{i}").ToArray();
         await configuration.TransportCustomization.ProvisionQueues(transportSettings, additionalQueues);
 
diff --git a/src/ServiceControl.Transports.RabbitMQQuorumDirectRouting.Tests/RabbitMQQueryTests.cs b/src/ServiceControl.Transports.RabbitMQQuorumDirectRouting.Tests/RabbitMQQueryTests.cs
@@ -32,7 +32,7 @@ public void Initialise()
             MaxConcurrency = 1,
             EndpointName = Guid.NewGuid().ToString("N")
         };
-        query = new RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, provider, transportSettings);
+        query = new RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, provider, transportSettings, configuration.TransportCustomization);
     }
 
     [Test]
diff --git a/src/ServiceControl.Transports.RabbitMQQuorumDirectRouting.Tests/RabbitMQQuery_ResponseParsing_Tests.cs b/src/ServiceControl.Transports.RabbitMQQuorumDirectRouting.Tests/RabbitMQQuery_ResponseParsing_Tests.cs
@@ -14,7 +14,7 @@ namespace ServiceControl.Transport.Tests;
 using System.Net;
 
 [TestFixture]
-class RabbitMQQuery_ResponseParsing_Tests
+class RabbitMQQuery_ResponseParsing_Tests : TransportTestFixture
 {
     FakeTimeProvider provider;
     TransportSettings transportSettings;
@@ -35,7 +35,7 @@ public void SetUp()
         httpHandler = new FakeHttpHandler();
         var httpClient = new HttpClient(httpHandler) { BaseAddress = new Uri("http://localhost:15672") };
 
-        rabbitMQQuery = new TestableRabbitMQQuery(provider, transportSettings, httpClient);
+        rabbitMQQuery = new TestableRabbitMQQuery(provider, transportSettings, httpClient, configuration.TransportCustomization);
         rabbitMQQuery.Initialize(ReadOnlyDictionary<string, string>.Empty);
     }
 
@@ -131,8 +131,9 @@ public async Task Should_fetch_queue_details_in_old_format()
     sealed class TestableRabbitMQQuery(
         TimeProvider timeProvider,
         TransportSettings transportSettings,
-        HttpClient customHttpClient)
-        : RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, timeProvider, transportSettings)
+        HttpClient customHttpClient,
+        ITransportCustomization transportCustomization)
+        : RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, timeProvider, transportSettings, transportCustomization)
     {
         protected override HttpClient CreateHttpClient(NetworkCredential defaultCredential, string apiUrl) => customHttpClient;
     }

Original file line number	Diff line number	Diff line change
`@@ -9,8 +9,11 @@`
`9`	`9`	`using NServiceBus;`
`10`	`10`
`11`	`11`	`public abstract class RabbitMQConventionalRoutingTransportCustomization(QueueType queueType)`
`12`		`- : TransportCustomization<RabbitMQTransport>`
	`12`	`+ : TransportCustomization<RabbitMQTransport>, IRabbitMQTransportExtensions`
`13`	`13`	`{`
	`14`	`+`
	`15`	`+ RabbitMQTransport rabbitMQTransport;`
	`16`	`+`
`14`	`17`	`protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }`
`15`	`18`
`16`	`19`	`protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }`
`@@ -36,14 +39,13 @@ protected override RabbitMQTransport CreateTransport(TransportSettings transport`
`36`	`39`	`transport.ManagementApiUrl = GetValue(connectionStringDictionary, "ManagementApiUrl", string.Empty);`
`37`	`40`	`transport.UseManagementApi = disableManagementApi.Equals("false", StringComparison.OrdinalIgnoreCase);`
`38`	`41`
	`42`	`+ rabbitMQTransport = transport;`
	`43`	`+`
`39`	`44`	`return transport;`
`40`	`45`	`}`
`41`	`46`
`42`		`- protected override void AddTransportForPrimaryCore(IServiceCollection services,`
`43`		`- TransportSettings transportSettings)`
`44`		`- {`
`45`		`- services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();`
`46`		`- }`
	`47`	`+ protected override void AddTransportForPrimaryCore(IServiceCollection services, TransportSettings transportSettings)`
	`48`	`+ => services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();`
`47`	`49`
`48`	`50`	`protected sealed override void AddTransportForMonitoringCore(IServiceCollection services, TransportSettings transportSettings)`
`49`	`51`	`{`
`@@ -52,8 +54,15 @@ protected sealed override void AddTransportForMonitoringCore(IServiceCollection`
`52`	`54`	`}`
`53`	`55`
`54`	`56`	`static string GetValue(Dictionary<string, string> dictionary, string key, string defaultValue)`
	`57`	`+ => dictionary.TryGetValue(key, out var value) ? value : defaultValue;`
	`58`	`+`
	`59`	`+ RabbitMQTransport IRabbitMQTransportExtensions.GetTransport()`
`55`	`60`	`{`
`56`		`- return dictionary.TryGetValue(key, out var value) ? value : defaultValue;`
	`61`	`+ if (rabbitMQTransport == null)`
	`62`	`+ {`
	`63`	`+ throw new InvalidOperationException("Transport instance has not been created yet. Make sure CreateTransport() is called before accessing the transport.");`
	`64`	`+ };`
	`65`	`+ return rabbitMQTransport;`
`57`	`66`	`}`
`58`	`67`	`}`
`59`	`68`	`}`
Original file line number	Diff line number	Diff line change
`@@ -8,11 +8,10 @@`
`8`	`8`	`using Microsoft.Extensions.DependencyInjection;`
`9`	`9`	`using NServiceBus;`
`10`	`10`
`11`		`- public abstract class RabbitMQDirectRoutingTransportCustomization : TransportCustomization<RabbitMQTransport>`
	`11`	`+ public abstract class RabbitMQDirectRoutingTransportCustomization(QueueType queueType)`
	`12`	`+ : TransportCustomization<RabbitMQTransport>, IRabbitMQTransportExtensions`
`12`	`13`	`{`
`13`		`- readonly QueueType queueType;`
`14`		`-`
`15`		`- protected RabbitMQDirectRoutingTransportCustomization(QueueType queueType) => this.queueType = queueType;`
	`14`	`+ RabbitMQTransport rabbitMQTransport;`
`16`	`15`
`17`	`16`	`protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }`
`18`	`17`
`@@ -38,14 +37,13 @@ protected override RabbitMQTransport CreateTransport(TransportSettings transport`
`38`	`37`	`transport.ManagementApiUrl = GetValue(connectionStringDictionary, "ManagementApiUrl", string.Empty);`
`39`	`38`	`transport.UseManagementApi = disableManagementApi.Equals("false", StringComparison.OrdinalIgnoreCase);`
`40`	`39`
	`40`	`+ rabbitMQTransport = transport;`
	`41`	`+`
`41`	`42`	`return transport;`
`42`	`43`	`}`
`43`	`44`
`44`		`- protected override void AddTransportForPrimaryCore(IServiceCollection services,`
`45`		`- TransportSettings transportSettings)`
`46`		`- {`
`47`		`- services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();`
`48`		`- }`
	`45`	`+ protected override void AddTransportForPrimaryCore(IServiceCollection services, TransportSettings transportSettings)`
	`46`	`+ => services.AddSingleton<IBrokerThroughputQuery, RabbitMQQuery>();`
`49`	`47`
`50`	`48`	`protected sealed override void AddTransportForMonitoringCore(IServiceCollection services, TransportSettings transportSettings)`
`51`	`49`	`{`
`@@ -54,8 +52,15 @@ protected sealed override void AddTransportForMonitoringCore(IServiceCollection`
`54`	`52`	`}`
`55`	`53`
`56`	`54`	`static string GetValue(Dictionary<string, string> dictionary, string key, string defaultValue)`
	`55`	`+ => dictionary.TryGetValue(key, out var value) ? value : defaultValue;`
	`56`	`+`
	`57`	`+ RabbitMQTransport IRabbitMQTransportExtensions.GetTransport()`
`57`	`58`	`{`
`58`		`- return dictionary.TryGetValue(key, out var value) ? value : defaultValue;`
	`59`	`+ if (rabbitMQTransport == null)`
	`60`	`+ {`
	`61`	`+ throw new InvalidOperationException("Transport instance has not been created yet. Make sure CreateTransport() is called before accessing the transport.");`
	`62`	`+ };`
	`63`	`+ return rabbitMQTransport;`
`59`	`64`	`}`
`60`	`65`	`}`
`61`	`66`	`}`
Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ public void Initialise()`
`32`	`32`	`MaxConcurrency = 1,`
`33`	`33`	`EndpointName = Guid.NewGuid().ToString("N")`
`34`	`34`	`};`
`35`		`- query = new RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, provider, transportSettings);`
	`35`	`+ query = new RabbitMQQuery(NullLogger<RabbitMQQuery>.Instance, provider, transportSettings, configuration.TransportCustomization);`
`36`	`36`	`}`
`37`	`37`
`38`	`38`	`[Test]`