Reuse Event Hubs client for health checks (#7628)

Co-authored-by: Sebastien Ros <[email protected]>

Author: github-actions[bot]

src/Components/Aspire.Azure.Messaging.EventHubs/EventHubProducerClientComponent.cs

@@ -4,10 +4,8 @@
 using Aspire.Azure.Messaging.EventHubs;
 using Azure.Core.Extensions;
 using Azure.Messaging.EventHubs.Producer;
-using HealthChecks.Azure.Messaging.EventHubs;
 using Microsoft.Extensions.Azure;
 using Microsoft.Extensions.Configuration;
-using Microsoft.Extensions.Diagnostics.HealthChecks;
 
 namespace Microsoft.Extensions.Hosting;
 
@@ -51,9 +49,4 @@ protected override IAzureClientBuilder<EventHubProducerClient, EventHubProducerC
 
         }, requiresCredential: false);
     }
-
-    protected override IHealthCheck CreateHealthCheck(EventHubProducerClient client, AzureMessagingEventHubsProducerSettings settings)
-    {
-        return new AzureEventHubHealthCheck(client);
-    }
 }

src/Components/Aspire.Azure.Messaging.EventHubs/EventHubsComponent.cs

@@ -14,10 +14,12 @@ namespace Microsoft.Extensions.Hosting;
 
 internal abstract class EventHubsComponent<TSettings, TClient, TClientOptions> :
     AzureComponent<TSettings, TClient, TClientOptions>
-    where TClientOptions: class
+    where TClientOptions : class
     where TClient : class
     where TSettings : AzureMessagingEventHubsSettings, new()
 {
+    private EventHubProducerClient? _healthCheckClient;
+
     // each EventHub client class is in a different namespace, so the base AzureComponent.ActivitySourceNames logic doesn't work
     protected override string[] ActivitySourceNames => ["Azure.Messaging.EventHubs.*"];
 
@@ -26,15 +28,26 @@ protected override IHealthCheck CreateHealthCheck(TClient client, TSettings sett
         // HealthChecks.Azure.Messaging.EventHubs currently only supports EventHubProducerClient.
         // https://github.com/Xabaril/AspNetCore.Diagnostics.HealthChecks/issues/2258 tracks supporting other client types.
 
-        var producerClientOptions = new EventHubProducerClientOptions
+        // Reuse the client if it's an EventHubProducerClient
+        if (client is EventHubProducerClient producerClient)
+        {
+            _healthCheckClient = producerClient;
+        }
+
+        // Create a custom EventHubProducerClient otherwise
+        if (_healthCheckClient == null)
         {
-            Identifier = $"AspireEventHubHealthCheck-{settings.EventHubName}",
-        };
-        var producerClient = !string.IsNullOrEmpty(settings.ConnectionString) ?
-            new EventHubProducerClient(settings.ConnectionString, producerClientOptions) :
-            new EventHubProducerClient(settings.FullyQualifiedNamespace, settings.EventHubName, settings.Credential ?? new DefaultAzureCredential(), producerClientOptions);
+            var producerClientOptions = new EventHubProducerClientOptions
+            {
+                Identifier = $"AspireEventHubHealthCheck-{settings.EventHubName}",
+            };
+
+            _healthCheckClient = !string.IsNullOrEmpty(settings.ConnectionString) ?
+                new EventHubProducerClient(settings.ConnectionString, producerClientOptions) :
+                new EventHubProducerClient(settings.FullyQualifiedNamespace, settings.EventHubName, settings.Credential ?? new DefaultAzureCredential(), producerClientOptions);
+        }
 
-        return new AzureEventHubHealthCheck(producerClient);
+        return new AzureEventHubHealthCheck(_healthCheckClient);
     }
 
     protected override bool GetHealthCheckEnabled(TSettings settings)

src/Components/Aspire.Azure.Messaging.EventHubs/PartitionReceiverClientComponent.cs

@@ -69,7 +69,7 @@ protected override IAzureClientBuilder<PartitionReceiver, PartitionReceiverOptio
                 settings.EventPosition,
                 settings.ConnectionString,
                 settings.EventHubName,
-                options);            
+                options);
 
         }, requiresCredential: false);
     }

tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTests.ConsumerClient.cs renamed to tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTests.EventHubConsumerClient.cs

@@ -1,10 +1,16 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Reflection;
 using Aspire.Components.ConformanceTests;
+using HealthChecks.Azure.Messaging.EventHubs;
 using Microsoft.DotNet.RemoteExecutor;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Options;
+using Xunit;
 
 namespace Aspire.Azure.Messaging.EventHubs.Tests;
 
@@ -60,4 +66,32 @@ protected static RemoteInvokeOptions EnableTracingForAzureSdk()
      {
          RuntimeConfigurationOptions = { { "Azure.Experimental.EnableActivitySource", true } }
      };
+
+    [ConditionalFact]
+    public void HealthChecksClientsAreReused()
+    {
+        SkipIfHealthChecksAreNotSupported();
+
+        // DisableRetries so the test doesn't take so long retrying when the server isn't available.
+        using IHost host = CreateHostWithComponent(configureComponent: DisableRetries);
+
+        HealthCheckServiceOptions healthCheckServiceOptions = host.Services.GetRequiredService<IOptions<HealthCheckServiceOptions>>().Value;
+
+        var registration = healthCheckServiceOptions.Registrations.First();
+
+        var healthCheck1 = registration.Factory(host.Services) as AzureEventHubHealthCheck;
+        var healthCheck2 = registration.Factory(host.Services) as AzureEventHubHealthCheck;
+
+        Assert.NotNull(healthCheck1);
+        Assert.NotNull(healthCheck2);
+
+        var clientAccessor = typeof(AzureEventHubHealthCheck).GetField("_client", BindingFlags.NonPublic | BindingFlags.Instance);
+
+        Assert.NotNull(clientAccessor);
+
+        var client1 = clientAccessor?.GetValue(healthCheck1);
+        var client2 = clientAccessor?.GetValue(healthCheck2);
+
+        Assert.Same(client1, client2);
+    }
 }