Update RabbitMQQuery to use the transport management client

Travis Nickels · bording · commit a226649f2424 · 2025-02-26T11:26:47.000-05:00
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQBrokerQueueDetails.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQBrokerQueueDetails.cs
@@ -2,18 +2,18 @@
 namespace ServiceControl.Transports.RabbitMQ;
 
 using System.Collections.Generic;
-using System.Text.Json;
 using ServiceControl.Transports.BrokerThroughput;
+using NServiceBus.Transport.RabbitMQ.ManagementApi;
 
-public class RabbitMQBrokerQueueDetails(JsonElement token) : IBrokerQueue
+class RabbitMQBrokerQueueDetails(Queue queue) : IBrokerQueue
 {
-    public string QueueName { get; } = token.GetProperty("name").GetString()!;
+    public string QueueName { get; } = queue.Name;
     public string SanitizedName => QueueName;
     public string Scope => VHost;
-    public string VHost { get; } = token.GetProperty("vhost").GetString()!;
+    public string VHost { get; } = queue.Vhost;
     public List<string> EndpointIndicators { get; } = [];
-    long? AckedMessages { get; set; } = FromToken(token);
-    long Baseline { get; set; } = FromToken(token) ?? 0;
+    long? AckedMessages { get; set; } = queue.MessageStats?.Ack;
+    long Baseline { get; set; } = queue.MessageStats?.Ack ?? 0;
 
     public long CalculateThroughputFrom(RabbitMQBrokerQueueDetails newReading)
     {
@@ -32,9 +32,4 @@ public long CalculateThroughputFrom(RabbitMQBrokerQueueDetails newReading)
 
         return newlyAckedMessages;
     }
-
-    static long? FromToken(JsonElement jsonElement) =>
-        jsonElement.TryGetProperty("message_stats", out var stats) && stats.TryGetProperty("ack", out var val)
-            ? val.GetInt64()
-            : null;
 }
diff --git a/src/ServiceControl.Transports.RabbitMQ/RabbitMQQuery.cs b/src/ServiceControl.Transports.RabbitMQ/RabbitMQQuery.cs
@@ -7,10 +7,7 @@ namespace ServiceControl.Transports.RabbitMQ;
 using System.Linq;
 using System.Net;
 using System.Net.Http;
-using System.Net.Http.Json;
 using System.Runtime.CompilerServices;
-using System.Text.Json;
-using System.Text.Json.Nodes;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Web;
@@ -19,6 +16,7 @@ namespace ServiceControl.Transports.RabbitMQ;
 using Polly;
 using Polly.Retry;
 using ServiceControl.Transports.BrokerThroughput;
+using NServiceBus.Transport.RabbitMQ.ManagementApi;
 
 public class RabbitMQQuery : BrokerThroughputQuery
 {
@@ -29,7 +27,6 @@ public class RabbitMQQuery : BrokerThroughputQuery
         .Build();
     readonly ILogger<RabbitMQQuery> logger;
     readonly TimeProvider timeProvider;
-    readonly ConnectionConfiguration connectionConfiguration;
     readonly RabbitMQTransport rabbitMQTransport;
 
     public RabbitMQQuery(ILogger<RabbitMQQuery> logger,
@@ -39,81 +36,38 @@ public RabbitMQQuery(ILogger<RabbitMQQuery> logger,
     {
         this.logger = logger;
         this.timeProvider = timeProvider;
-        if (transportCustomization is IRabbitMQTransportExtensions rabbitMQTransportCustomization)
-        {
-            rabbitMQTransport = rabbitMQTransportCustomization.GetTransport();
-            _ = rabbitMQTransport;
-        }
-        else
-        {
-            throw new InvalidOperationException($"Expected a RabbitMQTransport but received {transportCustomization.GetType().Name}.");
-        }
-
-        connectionConfiguration = ConnectionConfiguration.Create(transportSettings.ConnectionString, string.Empty);
+        rabbitMQTransport = GetRabbitMQTransport(transportCustomization);
     }
 
     protected override void InitializeCore(ReadOnlyDictionary<string, string> settings)
     {
-        // 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, rabbitMQTransport.ManagementApiConfiguration.UserName);
-        var password = GetSettingsValue(settings, RabbitMQSettings.Password, rabbitMQTransport.ManagementApiConfiguration.Password);
-        var apiUrl = GetSettingsValue(settings, RabbitMQSettings.API, rabbitMQTransport.ManagementApiConfiguration.Url);
-
-        if (userName != rabbitMQTransport.ManagementApiConfiguration.UserName)
-        {
-            _ = Diagnostics.AppendLine($"UserName in settings is different from Management API URL: {userName} != {rabbitMQTransport.ManagementApiConfiguration.UserName}");
-        }
-
-        if (password != rabbitMQTransport.ManagementApiConfiguration.Password)
-        {
-            _ = Diagnostics.AppendLine($"Password in settings is different from Management API URL.");
-        }
-
-        if (apiUrl != rabbitMQTransport.ManagementApiConfiguration.Url)
-        {
-            _ = Diagnostics.AppendLine($"API URL in settings is different from Management API URL: {apiUrl} != {rabbitMQTransport.ManagementApiConfiguration.Url}");
-        }
+        ////  TODO: Update documentation
+        //// https://docs.particular.net/servicecontrol/servicecontrol-instances/configuration#usage-reporting-when-using-the-rabbitmq-transport
+        CheckLegacySettings(settings, RabbitMQSettings.UserName);
+        CheckLegacySettings(settings, RabbitMQSettings.Password);
+        CheckLegacySettings(settings, RabbitMQSettings.API);
+    }
 
-        if (!Uri.TryCreate(apiUrl, UriKind.Absolute, out _))
+    static RabbitMQTransport GetRabbitMQTransport(ITransportCustomization transportCustomization)
+    {
+        if (transportCustomization is IRabbitMQTransportExtensions rabbitMQTransportCustomization)
         {
-            InitialiseErrors.Add("API url configured is invalid");
+            return rabbitMQTransportCustomization.GetTransport();
         }
 
-        var defaultCredential = new NetworkCredential(userName, password);
-
-        if (InitialiseErrors.Count == 0)
-        {
-            // ideally we would use the HttpClientFactory, but it would be a bit more involved to set that up
-            // so for now we are using a virtual method that can be overriden in tests
-            // https://github.com/Particular/ServiceControl/issues/4493
-            httpClient = CreateHttpClient(defaultCredential, apiUrl);
-        }
+        throw new InvalidOperationException($"Expected a RabbitMQTransport but received {transportCustomization.GetType().Name}.");
     }
 
-    string GetSettingsValue(ReadOnlyDictionary<string, string> settings, string key, string defaultValue)
+    void CheckLegacySettings(ReadOnlyDictionary<string, string> settings, string key)
     {
-        if (!settings.TryGetValue(key, out string? value) ||
-            string.IsNullOrEmpty(value))
-        {
-            logger.LogInformation($"Using {key} from connection string");
-            value = defaultValue;
-            _ = Diagnostics.AppendLine(
-                $"{key} not set, defaulted to using {key} from the ConnectionString used by instance");
-        }
-        else
+        if (settings.TryGetValue(key, out _))
         {
-            if (key == RabbitMQSettings.Password)
-            {
-                _ = Diagnostics.AppendLine($"{key} set.");
-            }
-
-            _ = Diagnostics.AppendLine($"{key} set to {value}.");
+            logger.LogInformation($"The legacy LicensingComponent/{key} is still defined in the app.config or environment variables");
+            _ = Diagnostics.AppendLine($"LicensingComponent/{key} is still defined in the app.config or environment variables");
         }
-
-        return value;
     }
 
+    // TODO: Determine if this needs to be updated in the RabbitMQ Transport
     protected virtual HttpClient CreateHttpClient(NetworkCredential defaultCredential, string apiUrl) =>
         new(new SocketsHttpHandler
         {
@@ -130,16 +84,31 @@ public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBro
         var url = $"/api/queues/{HttpUtility.UrlEncode(queue.VHost)}/{HttpUtility.UrlEncode(queue.QueueName)}";
 
         logger.LogDebug($"Querying {url}");
-        var newReading = await pipeline.ExecuteAsync(async token => new RabbitMQBrokerQueueDetails(await httpClient!.GetFromJsonAsync<JsonElement>(url, token)), cancellationToken);
+
+        var response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetQueue(queue.QueueName, cancellationToken), cancellationToken);
+
+        if (!response.HasValue)
+        {
+            throw new InvalidOperationException($"Could not access RabbitMQ Management API. ({response.StatusCode}: {response.Reason})");
+        }
+
+        var newReading = new RabbitMQBrokerQueueDetails(response.Value);
+
         _ = queue.CalculateThroughputFrom(newReading);
 
         // looping for 24hrs, in 4 increments of 15 minutes
         for (var i = 0; i < 24 * 4; i++)
         {
             await Task.Delay(TimeSpan.FromMinutes(15), timeProvider, cancellationToken);
             logger.LogDebug($"Querying {url}");
-            newReading = await pipeline.ExecuteAsync(async token => new RabbitMQBrokerQueueDetails(await httpClient!.GetFromJsonAsync<JsonElement>(url, token)), cancellationToken);
+            response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetQueue(queue.QueueName, cancellationToken), cancellationToken);
 
+            if (!response.HasValue)
+            {
+                throw new InvalidOperationException($"Could not access RabbitMQ Management API. ({response.StatusCode}: {response.Reason})");
+            }
+
+            newReading = new RabbitMQBrokerQueueDetails(response.Value);
             var newTotalThroughput = queue.CalculateThroughputFrom(newReading);
             yield return new QueueThroughput
             {
@@ -151,31 +120,32 @@ public override async IAsyncEnumerable<QueueThroughput> GetThroughputPerDay(IBro
 
     async Task<(string rabbitVersion, string managementVersion)> GetRabbitDetails(bool skipResiliencePipeline, CancellationToken cancellationToken)
     {
-        var overviewUrl = "/api/overview";
+        Response<Overview?> response = skipResiliencePipeline
+            ? await rabbitMQTransport.ManagementClient.GetOverview(cancellationToken)
+            : await pipeline.ExecuteAsync(async async => await rabbitMQTransport.ManagementClient.GetOverview(cancellationToken), cancellationToken);
 
-        JsonObject obj;
+        var overview = GetResponseValue(response);
 
-        if (skipResiliencePipeline)
+        if (overview.DisableStats)
         {
-            obj = (await httpClient!.GetFromJsonAsync<JsonObject>(overviewUrl, cancellationToken))!;
-        }
-        else
-        {
-            obj = (await pipeline.ExecuteAsync(async token =>
-                await httpClient!.GetFromJsonAsync<JsonObject>(overviewUrl, token), cancellationToken))!;
+            throw new Exception("The RabbitMQ broker is configured with 'management.disable_stats = true' or 'management_agent.disable_metrics_collector = true' " +
+                "and as a result queue statistics cannot be collected using this tool. Consider changing the configuration of the RabbitMQ broker.");
         }
 
-        var statsDisabled = obj["disable_stats"]?.GetValue<bool>() ?? false;
+        var rabbitVersion = response.Value?.BrokerVersion ?? response.Value?.ProductVersion;
+        var mgmtVersion = response.Value?.ManagementVersion;
 
-        if (statsDisabled)
+        return (rabbitVersion?.ToString() ?? "Unknown", mgmtVersion?.ToString() ?? "Unknown");
+    }
+
+    static T GetResponseValue<T>(Response<T?> response) where T : class
+    {
+        if (!response.HasValue || response.Value is null)
         {
-            throw new Exception("The RabbitMQ broker is configured with 'management.disable_stats = true' or 'management_agent.disable_metrics_collector = true' and as a result queue statistics cannot be collected using this tool. Consider changing the configuration of the RabbitMQ broker.");
+            throw new InvalidOperationException($"Could not access RabbitMQ Management API. ({response.StatusCode}: {response.Reason})");
         }
 
-        var rabbitVersion = obj["rabbitmq_version"] ?? obj["product_version"];
-        var mgmtVersion = obj["management_version"];
-
-        return (rabbitVersion?.GetValue<string>() ?? "Unknown", mgmtVersion?.GetValue<string>() ?? "Unknown");
+        return response.Value;
     }
 
     public override async IAsyncEnumerable<IBrokerQueue> GetQueueNames(
@@ -219,16 +189,15 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
     {
         try
         {
-            var bindingsUrl = $"/api/queues/{HttpUtility.UrlEncode(brokerQueue.VHost)}/{HttpUtility.UrlEncode(brokerQueue.QueueName)}/bindings";
-            var bindings = await pipeline.ExecuteAsync(async token => await httpClient!.GetFromJsonAsync<JsonArray>(bindingsUrl, token), cancellationToken);
-            var conventionalBindingFound = bindings?.Any(binding => binding!["source"]?.GetValue<string>() == brokerQueue.QueueName
-                                                                    && binding["vhost"]?.GetValue<string>() == brokerQueue.VHost
-                                                                    && binding["destination"]?.GetValue<string>() == brokerQueue.QueueName
-                                                                    && binding["destination_type"]?.GetValue<string>() == "queue"
-                                                                    && binding["routing_key"]?.GetValue<string>() == string.Empty
-                                                                    && binding["properties_key"]?.GetValue<string>() == "~") ?? false;
-
-            if (conventionalBindingFound)
+            var response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetQueueBindings(brokerQueue.QueueName, cancellationToken), cancellationToken);
+
+            // Check if conventional binding is found
+            if (response.Value.Any(binding => binding?.Source == brokerQueue.QueueName
+                && binding?.Vhost == brokerQueue.VHost
+                && binding?.Destination == brokerQueue.QueueName
+                && binding?.DestinationType == "queue"
+                && binding?.RoutingKey == string.Empty
+                && binding?.PropertiesKey == "~"))
             {
                 brokerQueue.EndpointIndicators.Add("ConventionalTopologyBinding");
             }
@@ -240,20 +209,14 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
 
         try
         {
-            var exchangeUrl = $"/api/exchanges/{HttpUtility.UrlEncode(brokerQueue.VHost)}/{HttpUtility.UrlEncode(brokerQueue.QueueName)}/bindings/destination";
-            var bindings = await pipeline.ExecuteAsync(async token => await httpClient!.GetFromJsonAsync<JsonArray>(exchangeUrl, token), cancellationToken);
-            var delayBindingFound = bindings?.Any(binding =>
-            {
-                var source = binding!["source"]?.GetValue<string>();
-
-                return source is "nsb.v2.delay-delivery" or "nsb.delay-delivery"
-                    && binding["vhost"]?.GetValue<string>() == brokerQueue.VHost
-                    && binding["destination"]?.GetValue<string>() == brokerQueue.QueueName
-                    && binding["destination_type"]?.GetValue<string>() == "exchange"
-                    && binding["routing_key"]?.GetValue<string>() == $"#.{brokerQueue.QueueName}";
-            }) ?? false;
-
-            if (delayBindingFound)
+            var response = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetExchangeBindingsDestination(brokerQueue.QueueName, cancellationToken), cancellationToken);
+
+            // Check if delayed binding is found
+            if (response.Value.Any(binding => binding?.Source is "nsb.v2.delay-delivery" or "nsb.delay-delivery"
+                    && binding?.Vhost == brokerQueue.VHost
+                    && binding?.Destination == brokerQueue.QueueName
+                    && binding?.DestinationType == "exchange"
+                    && binding?.RoutingKey == $"#.{brokerQueue.QueueName}"))
             {
                 brokerQueue.EndpointIndicators.Add("DelayBinding");
             }
@@ -264,41 +227,41 @@ async Task AddAdditionalQueueDetails(RabbitMQBrokerQueueDetails brokerQueue, Can
         }
     }
 
-    public async Task<(RabbitMQBrokerQueueDetails[]?, bool morePages)> GetPage(int page, CancellationToken cancellationToken)
+    internal async Task<(List<RabbitMQBrokerQueueDetails>?, bool morePages)> GetPage(int page, CancellationToken cancellationToken)
     {
-        var url = $"/api/queues/{HttpUtility.UrlEncode(connectionConfiguration.VirtualHost)}?page={page}&page_size=500&name=&use_regex=false&pagination=true";
-
-        var container = await pipeline.ExecuteAsync(async token => await httpClient!.GetFromJsonAsync<JsonNode>(url, token), cancellationToken);
-        switch (container)
+        var pagination = await pipeline.ExecuteAsync(async token => await rabbitMQTransport.ManagementClient.GetPage(page, cancellationToken), cancellationToken);
+        switch (pagination.Value)
         {
-            case JsonObject obj:
+            case Pagination obj:
                 {
-                    var pageCount = obj["page_count"]!.GetValue<int>();
-                    var pageReturned = obj["page"]!.GetValue<int>();
+                    var pageCount = obj.PageCount;
+                    var pageReturned = obj.Page;
 
-                    if (obj["items"] is not JsonArray items)
+                    if (obj.Items is null) //is not JsonArray items
                     {
                         return (null, false);
                     }
 
-                    return (MaterializeQueueDetails(items), pageCount > pageReturned);
+                    return (MaterializeQueueDetails(obj.Items), pageCount > pageReturned);
                 }
             // Older versions of RabbitMQ API did not have paging and returned the array of items directly
-            case JsonArray arr:
-                {
-                    return (MaterializeQueueDetails(arr), false);
-                }
+            //case JsonArray arr:
+            //    {
+            //        return (MaterializeQueueDetails(arr), false);
+            //    }
             default:
                 throw new Exception("Was not able to get list of queues from RabbitMQ broker.");
         }
     }
 
-    static RabbitMQBrokerQueueDetails[] MaterializeQueueDetails(JsonArray items)
+    static List<RabbitMQBrokerQueueDetails> MaterializeQueueDetails(List<Queue> items)
     {
-        // It is not possible to directly operated on the JsonNode. When the JsonNode is a JObject
-        // and the indexer is access the internal dictionary is initialized which can cause key not found exceptions
-        // when the payload contains the same key multiple times (which happened in the past).
-        var queues = items.Select(item => new RabbitMQBrokerQueueDetails(item!.Deserialize<JsonElement>())).ToArray();
+        var queues = new List<RabbitMQBrokerQueueDetails>();
+        foreach (var item in items)
+        {
+            queues.Add(new RabbitMQBrokerQueueDetails(item));
+        }
+
         return queues;
     }
 
