AzureServiceBus ensure client reuse

src/HealthChecks.AzureServiceBus/AzureServiceBusHealthCheck.cs

@@ -1,4 +1,3 @@
-using System.Collections.Concurrent;
 using Azure.Messaging.ServiceBus;
 using Azure.Messaging.ServiceBus.Administration;
 using HealthChecks.AzureServiceBus.Configuration;
@@ -7,21 +6,19 @@ namespace HealthChecks.AzureServiceBus;
 
 public abstract class AzureServiceBusHealthCheck<TOptions> where TOptions : AzureServiceBusHealthCheckOptions
 {
-    protected static readonly ConcurrentDictionary<string, ServiceBusClient> ClientConnections = new();
-
-    protected static readonly ConcurrentDictionary<string, ServiceBusAdministrationClient> ManagementClientConnections = new();
-
-    protected static readonly ConcurrentDictionary<string, ServiceBusReceiver> ServiceBusReceivers = new();
-
     protected TOptions Options { get; }
 
+    private readonly ServiceBusClientProvider _clientProvider;
+
+    [Obsolete("Please, override ConnectionKey property instead.")]
     protected string Prefix => Options.ConnectionString ?? Options.FullyQualifiedNamespace!;
 
-    protected abstract string ConnectionKey { get; }
+    protected virtual string ConnectionKey => Options.ConnectionString ?? Options.FullyQualifiedNamespace!;
 
-    protected AzureServiceBusHealthCheck(TOptions options)
+    protected AzureServiceBusHealthCheck(TOptions options, ServiceBusClientProvider clientProvider)
     {
         Options = options;
+        _clientProvider = clientProvider;
 
         if (!string.IsNullOrWhiteSpace(options.ConnectionString))
             return;
@@ -37,11 +34,11 @@ protected AzureServiceBusHealthCheck(TOptions options)
 
     protected ServiceBusClient CreateClient() =>
         Options.Credential is null
-            ? new ServiceBusClient(Options.ConnectionString)
-            : new ServiceBusClient(Options.FullyQualifiedNamespace, Options.Credential);
+            ? _clientProvider.CreateClient(Options.ConnectionString)
+            : _clientProvider.CreateClient(Options.FullyQualifiedNamespace, Options.Credential);
 
     protected ServiceBusAdministrationClient CreateManagementClient() =>
         Options.Credential is null
-            ? new ServiceBusAdministrationClient(Options.ConnectionString)
-            : new ServiceBusAdministrationClient(Options.FullyQualifiedNamespace, Options.Credential);
+            ? _clientProvider.CreateManagementClient(Options.ConnectionString)
+            : _clientProvider.CreateManagementClient(Options.FullyQualifiedNamespace, Options.Credential);
 }

src/HealthChecks.AzureServiceBus/AzureServiceBusQueueHealthCheck.cs

@@ -5,16 +5,20 @@ namespace HealthChecks.AzureServiceBus;
 
 public class AzureServiceBusQueueHealthCheck : AzureServiceBusHealthCheck<AzureServiceBusQueueHealthCheckOptions>, IHealthCheck
 {
-    private string? _connectionKey;
+    private readonly string _queueKey;
 
-    protected override string ConnectionKey => _connectionKey ??= $"{Prefix}_{Options.QueueName}";
-
-    public AzureServiceBusQueueHealthCheck(AzureServiceBusQueueHealthCheckOptions options)
-        : base(options)
+    public AzureServiceBusQueueHealthCheck(AzureServiceBusQueueHealthCheckOptions options, ServiceBusClientProvider clientProvider)
+        : base(options, clientProvider)
     {
         Guard.ThrowIfNull(options.QueueName, true);
+
+        _queueKey = $"{nameof(AzureServiceBusQueueHealthCheck)}_{ConnectionKey}_{Options.QueueName}";
     }
 
+    public AzureServiceBusQueueHealthCheck(AzureServiceBusQueueHealthCheckOptions options)
+        : this(options, new ServiceBusClientProvider())
+    { }
+
     /// <inheritdoc />
     public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
     {
@@ -32,19 +36,20 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
             return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
         }
 
-        Task CheckWithReceiver()
+        async Task CheckWithReceiver()
         {
-            var client = ClientConnections.GetOrAdd(ConnectionKey, _ => CreateClient());
-            var receiver = ServiceBusReceivers.GetOrAdd(
-                $"{nameof(AzureServiceBusQueueHealthCheck)}_{ConnectionKey}",
-                client.CreateReceiver(Options.QueueName));
+            var client = await ClientCache.GetOrAddAsyncDisposableAsync(ConnectionKey, _ => CreateClient()).ConfigureAwait(false);
+            var receiver = await ClientCache.GetOrAddAsyncDisposableAsync(
+                _queueKey,
+                _ => client.CreateReceiver(Options.QueueName))
+                .ConfigureAwait(false);
 
-            return receiver.PeekMessageAsync(cancellationToken: cancellationToken);
+            await receiver.PeekMessageAsync(cancellationToken: cancellationToken).ConfigureAwait(false);
         }
 
         Task CheckWithManagement()
         {
-            var managementClient = ManagementClientConnections.GetOrAdd(ConnectionKey, _ => CreateManagementClient());
+            var managementClient = ClientCache.GetOrAdd(ConnectionKey, _ => CreateManagementClient());
 
             return managementClient.GetQueueRuntimePropertiesAsync(Options.QueueName, cancellationToken);
         }

src/HealthChecks.AzureServiceBus/AzureServiceBusQueueMessageCountThresholdHealthCheck.cs

@@ -9,22 +9,24 @@ public class AzureServiceBusQueueMessageCountThresholdHealthCheck : AzureService
     private readonly AzureServiceBusQueueMessagesCountThreshold? _activeMessagesThreshold;
     private readonly AzureServiceBusQueueMessagesCountThreshold? _deadLetterMessagesThreshold;
 
-    public AzureServiceBusQueueMessageCountThresholdHealthCheck(AzureServiceBusQueueMessagesCountThresholdHealthCheckOptions options)
-        : base(options)
+    public AzureServiceBusQueueMessageCountThresholdHealthCheck(AzureServiceBusQueueMessagesCountThresholdHealthCheckOptions options, ServiceBusClientProvider clientProvider)
+        : base(options, clientProvider)
     {
         _queueName = Guard.ThrowIfNull(options.QueueName);
         _activeMessagesThreshold = options.ActiveMessages;
         _deadLetterMessagesThreshold = options.DeadLetterMessages;
     }
 
-    protected override string ConnectionKey => $"{Prefix}_{_queueName}";
+    public AzureServiceBusQueueMessageCountThresholdHealthCheck(AzureServiceBusQueueMessagesCountThresholdHealthCheckOptions options)
+        : this(options, new ServiceBusClientProvider())
+    { }
 
     /// <inheritdoc />
     public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
     {
         try
         {
-            var managementClient = ManagementClientConnections.GetOrAdd(ConnectionKey, CreateManagementClient());
+            var managementClient = ClientCache.GetOrAdd(ConnectionKey, _ => CreateManagementClient());
 
             var properties = await managementClient.GetQueueRuntimePropertiesAsync(_queueName, cancellationToken).ConfigureAwait(false);
 

src/HealthChecks.AzureServiceBus/AzureServiceBusSubscriptionHealthCheck.cs

@@ -5,18 +5,21 @@ namespace HealthChecks.AzureServiceBus;
 
 public class AzureServiceBusSubscriptionHealthCheck : AzureServiceBusHealthCheck<AzureServiceBusSubscriptionHealthCheckHealthCheckOptions>, IHealthCheck
 {
-    private string? _connectionKey;
+    private readonly string _subscriptionKey;
 
-    protected override string ConnectionKey =>
-        _connectionKey ??= $"{Prefix}_{Options.TopicName}_{Options.SubscriptionName}";
-
-    public AzureServiceBusSubscriptionHealthCheck(AzureServiceBusSubscriptionHealthCheckHealthCheckOptions options)
-        : base(options)
+    public AzureServiceBusSubscriptionHealthCheck(AzureServiceBusSubscriptionHealthCheckHealthCheckOptions options, ServiceBusClientProvider clientProvider)
+        : base(options, clientProvider)
     {
         Guard.ThrowIfNull(options.TopicName, true);
         Guard.ThrowIfNull(options.SubscriptionName, true);
+
+        _subscriptionKey = $"{nameof(AzureServiceBusSubscriptionHealthCheck)}_{ConnectionKey}_{Options.TopicName}_{Options.SubscriptionName}";
     }
 
+    public AzureServiceBusSubscriptionHealthCheck(AzureServiceBusSubscriptionHealthCheckHealthCheckOptions options)
+        : this(options, new ServiceBusClientProvider())
+    { }
+
     /// <inheritdoc />
     public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
     {
@@ -34,19 +37,20 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
             return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
         }
 
-        Task CheckWithReceiver()
+        async Task CheckWithReceiver()
         {
-            var client = ClientConnections.GetOrAdd(ConnectionKey, _ => CreateClient());
-            var receiver = ServiceBusReceivers.GetOrAdd(
-                $"{nameof(AzureServiceBusSubscriptionHealthCheck)}_{ConnectionKey}",
-                client.CreateReceiver(Options.TopicName, Options.SubscriptionName));
+            var client = await ClientCache.GetOrAddAsyncDisposableAsync(ConnectionKey, _ => CreateClient()).ConfigureAwait(false);
+            var receiver = await ClientCache.GetOrAddAsyncDisposableAsync(
+                _subscriptionKey,
+                _ => client.CreateReceiver(Options.TopicName, Options.SubscriptionName))
+                .ConfigureAwait(false);
 
-            return receiver.PeekMessageAsync(cancellationToken: cancellationToken);
+            await receiver.PeekMessageAsync(cancellationToken: cancellationToken).ConfigureAwait(false);
         }
 
         Task CheckWithManagement()
         {
-            var managementClient = ManagementClientConnections.GetOrAdd(ConnectionKey, _ => CreateManagementClient());
+            var managementClient = ClientCache.GetOrAdd(ConnectionKey, _ => CreateManagementClient());
 
             return managementClient.GetSubscriptionRuntimePropertiesAsync(
                 Options.TopicName, Options.SubscriptionName, cancellationToken);

src/HealthChecks.AzureServiceBus/AzureServiceBusTopicHealthCheck.cs

@@ -5,23 +5,23 @@ namespace HealthChecks.AzureServiceBus;
 
 public class AzureServiceBusTopicHealthCheck : AzureServiceBusHealthCheck<AzureServiceBusTopicHealthCheckOptions>, IHealthCheck
 {
-    private string? _connectionKey;
-
-    protected override string ConnectionKey => _connectionKey ??= $"{Prefix}_{Options.TopicName}";
-
-    public AzureServiceBusTopicHealthCheck(AzureServiceBusTopicHealthCheckOptions options)
-        : base(options)
+    public AzureServiceBusTopicHealthCheck(AzureServiceBusTopicHealthCheckOptions options, ServiceBusClientProvider clientProvider)
+        : base(options, clientProvider)
     {
         Guard.ThrowIfNull(options.TopicName, true);
     }
 
+    public AzureServiceBusTopicHealthCheck(AzureServiceBusTopicHealthCheckOptions options)
+        : this(options, new ServiceBusClientProvider())
+    { }
+
     /// <inheritdoc />
     public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context,
         CancellationToken cancellationToken = default)
     {
         try
         {
-            var managementClient = ManagementClientConnections.GetOrAdd(ConnectionKey, _ => CreateManagementClient());
+            var managementClient = ClientCache.GetOrAdd(ConnectionKey, _ => CreateManagementClient());
 
             _ = await managementClient.GetTopicRuntimePropertiesAsync(Options.TopicName, cancellationToken).ConfigureAwait(false);
 

src/HealthChecks.AzureServiceBus/ServiceBusClientProvider.cs

@@ -0,0 +1,20 @@
+using Azure.Core;
+using Azure.Messaging.ServiceBus;
+using Azure.Messaging.ServiceBus.Administration;
+
+namespace HealthChecks.AzureServiceBus;
+
+public class ServiceBusClientProvider
+{
+    public virtual ServiceBusClient CreateClient(string? connectionString)
+        => new ServiceBusClient(connectionString);
+
+    public virtual ServiceBusClient CreateClient(string? fullyQualifiedName, TokenCredential credential)
+        => new ServiceBusClient(fullyQualifiedName, credential);
+
+    public virtual ServiceBusAdministrationClient CreateManagementClient(string? connectionString)
+        => new ServiceBusAdministrationClient(connectionString);
+
+    public virtual ServiceBusAdministrationClient CreateManagementClient(string? fullyQualifiedName, TokenCredential credential)
+        => new ServiceBusAdministrationClient(fullyQualifiedName, credential);
+}