Make IManagementClientProvider lazy to ensure endpoints have started

bording · bording · commit 122649f9136f · 2025-03-07T11:31:36.000-05:00
diff --git a/src/ServiceControl.Transports.RabbitMQ/IManagementClientProvider.cs b/src/ServiceControl.Transports.RabbitMQ/IManagementClientProvider.cs
@@ -1,8 +1,9 @@
 ﻿namespace ServiceControl.Transports.RabbitMQ;
 
+using System;
 using NServiceBus.Transport.RabbitMQ.ManagementApi;
 
 interface IManagementClientProvider
 {
-    ManagementClient ManagementClient { get; }
+    Lazy<ManagementClient> GetManagementClient();
 }
diff --git a/src/ServiceControl.Transports.RabbitMQ/QueueLengthProvider.cs b/src/ServiceControl.Transports.RabbitMQ/QueueLengthProvider.cs
@@ -13,7 +13,7 @@ public QueueLengthProvider(TransportSettings settings, Action<QueueLengthEntry[]
         {
             if (transportCustomization is IManagementClientProvider provider)
             {
-                managementClient = provider.ManagementClient;
+                managementClient = provider.GetManagementClient();
             }
             else
             {
@@ -84,7 +84,7 @@ async Task FetchQueueLengths(CancellationToken cancellationToken)
 
                 try
                 {
-                    var (statusCode, reason, queue) = await managementClient.GetQueue(queueName, cancellationToken);
+                    var (statusCode, reason, queue) = await managementClient.Value.GetQueue(queueName, cancellationToken);
 
                     if (queue is not null)
                     {
@@ -111,6 +111,6 @@ async Task FetchQueueLengths(CancellationToken cancellationToken)
 
         static readonly ILog Logger = LogManager.GetLogger<QueueLengthProvider>();
 
-        readonly ManagementClient managementClient;
+        readonly Lazy<ManagementClient> managementClient;
     }
 }
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQConventionalRoutingTransportCustomization.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQConventionalRoutingTransportCustomization.cs
@@ -9,15 +9,34 @@
 
     public abstract class RabbitMQConventionalRoutingTransportCustomization(NServiceBus.QueueType queueType) : TransportCustomization<RabbitMQTransport>, IManagementClientProvider
     {
-        RabbitMQTransport rabbitMQTransport;
+        RabbitMQTransport transport;
 
-        ManagementClient IManagementClientProvider.ManagementClient => rabbitMQTransport?.ManagementClient ?? new ManagementClient(rabbitMQTransport.ConnectionConfiguration, rabbitMQTransport.ManagementApiConfiguration);
+        Lazy<ManagementClient> IManagementClientProvider.GetManagementClient()
+        {
+            return new(() => Get());
+
+            ManagementClient Get()
+            {
+                if (transport is null)
+                {
+                    throw new InvalidOperationException("Management client not available because a CustomizeTransport method has not been called first.");
+                }
+
+                // Since some tests don't actually start an endpoint, this is needed to ensure a management client is available
+                if (transport.ManagementClient is null)
+                {
+                    return new ManagementClient(transport.ConnectionConfiguration, transport.ManagementApiConfiguration);
+                }
+
+                return transport.ManagementClient;
+            }
+        }
 
-        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
+        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => transport = transportDefinition;
 
         protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
-        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
+        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => transport = transportDefinition;
 
         protected override RabbitMQTransport CreateTransport(TransportSettings transportSettings, TransportTransactionMode preferredTransactionMode = TransportTransactionMode.ReceiveOnly)
         {
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQDirectRoutingTransportCustomization.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQDirectRoutingTransportCustomization.cs
@@ -9,15 +9,34 @@
 
     public abstract class RabbitMQDirectRoutingTransportCustomization(NServiceBus.QueueType queueType) : TransportCustomization<RabbitMQTransport>, IManagementClientProvider
     {
-        RabbitMQTransport rabbitMQTransport;
+        RabbitMQTransport transport;
 
-        ManagementClient IManagementClientProvider.ManagementClient => rabbitMQTransport?.ManagementClient ?? new ManagementClient(rabbitMQTransport.ConnectionConfiguration, rabbitMQTransport.ManagementApiConfiguration);
+        Lazy<ManagementClient> IManagementClientProvider.GetManagementClient()
+        {
+            return new(() => Get());
+
+            ManagementClient Get()
+            {
+                if (transport is null)
+                {
+                    throw new InvalidOperationException("Management client not available because a CustomizeTransport method has not been called first.");
+                }
+
+                // Since some tests don't actually start an endpoint, this is needed to ensure a management client is available
+                if (transport.ManagementClient is null)
+                {
+                    return new ManagementClient(transport.ConnectionConfiguration, transport.ManagementApiConfiguration);
+                }
+
+                return transport.ManagementClient;
+            }
+        }
 
-        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
+        protected override void CustomizeTransportForPrimaryEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => transport = transportDefinition;
 
         protected override void CustomizeTransportForAuditEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) { }
 
-        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => rabbitMQTransport = transportDefinition;
+        protected override void CustomizeTransportForMonitoringEndpoint(EndpointConfiguration endpointConfiguration, RabbitMQTransport transportDefinition, TransportSettings transportSettings) => transport = transportDefinition;
 
         protected override RabbitMQTransport CreateTransport(TransportSettings transportSettings, TransportTransactionMode preferredTransactionMode = TransportTransactionMode.ReceiveOnly)
         {
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQQuery.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQQuery.cs
@@ -20,7 +20,7 @@ public class RabbitMQQuery : BrokerThroughputQuery
 {
     readonly ILogger<RabbitMQQuery> logger;
     readonly TimeProvider timeProvider;
-    readonly ManagementClient managementClient;
+    readonly Lazy<ManagementClient> managementClient;
 
     readonly ResiliencePipeline pipeline = new ResiliencePipelineBuilder()
       .AddRetry(new RetryStrategyOptions()) // Add retry using the default options
@@ -34,7 +34,7 @@ public RabbitMQQuery(ILogger<RabbitMQQuery> logger, TimeProvider timeProvider, I
 
         if (transportCustomization is IManagementClientProvider provider)
         {
-            managementClient = provider.ManagementClient;
+            managementClient = provider.GetManagementClient();
         }
         else
         {
@@ -63,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 managementClient.GetQueue(queue.QueueName, token), cancellationToken);
+        var response = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetQueue(queue.QueueName, token), cancellationToken);
 
         if (response.Value is null)
         {
@@ -79,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 managementClient.GetQueue(queue.QueueName, token), cancellationToken);
+            response = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetQueue(queue.QueueName, token), cancellationToken);
 
             if (response.Value is null)
             {
@@ -99,7 +99,7 @@ public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBro
 
     async Task GetRabbitDetails(CancellationToken cancellationToken)
     {
-        var response = await pipeline.ExecuteAsync(async async => await managementClient.GetOverview(cancellationToken), cancellationToken);
+        var response = await pipeline.ExecuteAsync(async async => await managementClient.Value.GetOverview(cancellationToken), cancellationToken);
 
         ValidateResponse(response);
 
@@ -160,7 +160,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
     {
         try
         {
-            var response = await pipeline.ExecuteAsync(async token => await managementClient.GetBindingsForQueue(brokerQueue.QueueName, token), cancellationToken);
+            var response = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetBindingsForQueue(brokerQueue.QueueName, token), cancellationToken);
 
             // Check if conventional binding is found
             if (response.Value.Any(binding => binding?.Source == brokerQueue.QueueName
@@ -179,7 +179,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
 
         try
         {
-            var response = await pipeline.ExecuteAsync(async token => await managementClient.GetBindingsForExchange(brokerQueue.QueueName, token), cancellationToken);
+            var response = await pipeline.ExecuteAsync(async token => await managementClient.Value.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"
@@ -198,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 managementClient.GetQueues(page, 500, token), cancellationToken);
+        var (StatusCode, Reason, Value, MorePages) = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetQueues(page, 500, token), cancellationToken);
 
         ValidateResponse((StatusCode, Reason, Value));
 
@@ -227,7 +227,7 @@ static List<RabbitMQBrokerQueueDetails> MaterializeQueueDetails(List<Queue> item
     {
         try
         {
-            var (statusCode, reason, value) = await managementClient.GetOverview(cancellationToken);
+            var (statusCode, reason, value) = await managementClient.Value.GetOverview(cancellationToken);
 
             if (value is not null)
             {

Original file line number	Diff line number	Diff line change
`@@ -1,8 +1,9 @@`
`1`	`1`	`namespace ServiceControl.Transports.RabbitMQ;`
`2`	`2`
	`3`	`+using System;`
`3`	`4`	`using NServiceBus.Transport.RabbitMQ.ManagementApi;`
`4`	`5`
`5`	`6`	`interface IManagementClientProvider`
`6`	`7`	`{`
`7`		`- ManagementClient ManagementClient { get; }`
	`8`	`+ Lazy<ManagementClient> GetManagementClient();`
`8`	`9`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ public QueueLengthProvider(TransportSettings settings, Action<QueueLengthEntry[]`
`13`	`13`	`{`
`14`	`14`	`if (transportCustomization is IManagementClientProvider provider)`
`15`	`15`	`{`
`16`		`- managementClient = provider.ManagementClient;`
	`16`	`+ managementClient = provider.GetManagementClient();`
`17`	`17`	`}`
`18`	`18`	`else`
`19`	`19`	`{`
`@@ -84,7 +84,7 @@ async Task FetchQueueLengths(CancellationToken cancellationToken)`
`84`	`84`
`85`	`85`	`try`
`86`	`86`	`{`
`87`		`- var (statusCode, reason, queue) = await managementClient.GetQueue(queueName, cancellationToken);`
	`87`	`+ var (statusCode, reason, queue) = await managementClient.Value.GetQueue(queueName, cancellationToken);`
`88`	`88`
`89`	`89`	`if (queue is not null)`
`90`	`90`	`{`
`@@ -111,6 +111,6 @@ async Task FetchQueueLengths(CancellationToken cancellationToken)`
`111`	`111`
`112`	`112`	`static readonly ILog Logger = LogManager.GetLogger<QueueLengthProvider>();`
`113`	`113`
`114`		`- readonly ManagementClient managementClient;`
	`114`	`+ readonly Lazy<ManagementClient> managementClient;`
`115`	`115`	`}`
`116`	`116`	`}`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ public class RabbitMQQuery : BrokerThroughputQuery`
`20`	`20`	`{`
`21`	`21`	`readonly ILogger<RabbitMQQuery> logger;`
`22`	`22`	`readonly TimeProvider timeProvider;`
`23`		`- readonly ManagementClient managementClient;`
	`23`	`+ readonly Lazy<ManagementClient> managementClient;`
`24`	`24`
`25`	`25`	`readonly ResiliencePipeline pipeline = new ResiliencePipelineBuilder()`
`26`	`26`	`.AddRetry(new RetryStrategyOptions()) // Add retry using the default options`
`@@ -34,7 +34,7 @@ public RabbitMQQuery(ILogger<RabbitMQQuery> logger, TimeProvider timeProvider, I`
`34`	`34`
`35`	`35`	`if (transportCustomization is IManagementClientProvider provider)`
`36`	`36`	`{`
`37`		`- managementClient = provider.ManagementClient;`
	`37`	`+ managementClient = provider.GetManagementClient();`
`38`	`38`	`}`
`39`	`39`	`else`
`40`	`40`	`{`
`@@ -63,7 +63,7 @@ void CheckLegacySettings(ReadOnlyDictionary<string, string> settings, string key`
`63`	`63`	`public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBrokerQueue brokerQueue, DateOnly startDate, [EnumeratorCancellation] CancellationToken cancellationToken = default)`
`64`	`64`	`{`
`65`	`65`	`var queue = (RabbitMQBrokerQueueDetails)brokerQueue;`
`66`		`- var response = await pipeline.ExecuteAsync(async token => await managementClient.GetQueue(queue.QueueName, token), cancellationToken);`
	`66`	`+ var response = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetQueue(queue.QueueName, token), cancellationToken);`
`67`	`67`
`68`	`68`	`if (response.Value is null)`
`69`	`69`	`{`
`@@ -79,7 +79,7 @@ public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBro`
`79`	`79`	`{`
`80`	`80`	`await Task.Delay(TimeSpan.FromMinutes(15), timeProvider, cancellationToken);`
`81`	`81`
`82`		`- response = await pipeline.ExecuteAsync(async token => await managementClient.GetQueue(queue.QueueName, token), cancellationToken);`
	`82`	`+ response = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetQueue(queue.QueueName, token), cancellationToken);`
`83`	`83`
`84`	`84`	`if (response.Value is null)`
`85`	`85`	`{`
`@@ -99,7 +99,7 @@ public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBro`
`99`	`99`
`100`	`100`	`async Task GetRabbitDetails(CancellationToken cancellationToken)`
`101`	`101`	`{`
`102`		`- var response = await pipeline.ExecuteAsync(async async => await managementClient.GetOverview(cancellationToken), cancellationToken);`
	`102`	`+ var response = await pipeline.ExecuteAsync(async async => await managementClient.Value.GetOverview(cancellationToken), cancellationToken);`
`103`	`103`
`104`	`104`	`ValidateResponse(response);`
`105`	`105`
`@@ -160,7 +160,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can`
`160`	`160`	`{`
`161`	`161`	`try`
`162`	`162`	`{`
`163`		`- var response = await pipeline.ExecuteAsync(async token => await managementClient.GetBindingsForQueue(brokerQueue.QueueName, token), cancellationToken);`
	`163`	`+ var response = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetBindingsForQueue(brokerQueue.QueueName, token), cancellationToken);`
`164`	`164`
`165`	`165`	`// Check if conventional binding is found`
`166`	`166`	`if (response.Value.Any(binding => binding?.Source == brokerQueue.QueueName`
`@@ -179,7 +179,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can`
`179`	`179`
`180`	`180`	`try`
`181`	`181`	`{`
`182`		`- var response = await pipeline.ExecuteAsync(async token => await managementClient.GetBindingsForExchange(brokerQueue.QueueName, token), cancellationToken);`
	`182`	`+ var response = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetBindingsForExchange(brokerQueue.QueueName, token), cancellationToken);`
`183`	`183`
`184`	`184`	`// Check if delayed binding is found`
`185`	`185`	`if (response.Value.Any(binding => binding?.Source is "nsb.v2.delay-delivery" or "nsb.delay-delivery"`
`@@ -198,7 +198,7 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can`
`198`	`198`
`199`	`199`	`internal async Task<(List<RabbitMQBrokerQueueDetails>?, bool morePages)> GetPage(int page, CancellationToken cancellationToken)`
`200`	`200`	`{`
`201`		`- var (StatusCode, Reason, Value, MorePages) = await pipeline.ExecuteAsync(async token => await managementClient.GetQueues(page, 500, token), cancellationToken);`
	`201`	`+ var (StatusCode, Reason, Value, MorePages) = await pipeline.ExecuteAsync(async token => await managementClient.Value.GetQueues(page, 500, token), cancellationToken);`
`202`	`202`
`203`	`203`	`ValidateResponse((StatusCode, Reason, Value));`
`204`	`204`
`@@ -227,7 +227,7 @@ static List<RabbitMQBrokerQueueDetails> MaterializeQueueDetails(List<Queue> item`
`227`	`227`	`{`
`228`	`228`	`try`
`229`	`229`	`{`
`230`		`- var (statusCode, reason, value) = await managementClient.GetOverview(cancellationToken);`
	`230`	`+ var (statusCode, reason, value) = await managementClient.Value.GetOverview(cancellationToken);`
`231`	`231`
`232`	`232`	`if (value is not null)`
`233`	`233`	`{`