Reuse Event Hubs client for health checks (#7625)

sebastienros · web-flow · commit b4304093c44a · 2025-02-14T21:41:08.000-08:00
diff --git a/src/Components/Aspire.Azure.Messaging.EventHubs/EventHubProducerClientComponent.cs b/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);
-    }
 }
diff --git a/src/Components/Aspire.Azure.Messaging.EventHubs/EventHubsComponent.cs b/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)
diff --git a/src/Components/Aspire.Azure.Messaging.EventHubs/PartitionReceiverClientComponent.cs b/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);
     }
diff --git a/tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTests.EventHubConsumerClient.cs b/tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTests.EventHubConsumerClient.cs
diff --git a/tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTests.EventHubProducerClient.cs b/tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTests.EventHubProducerClient.cs
diff --git a/tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTestsBase.cs b/tests/Aspire.Azure.Messaging.EventHubs.Tests/ConformanceTestsBase.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);
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ protected override IAzureClientBuilder<PartitionReceiver, PartitionReceiverOptio`
`69`	`69`	`settings.EventPosition,`
`70`	`70`	`settings.ConnectionString,`
`71`	`71`	`settings.EventHubName,`
`72`		`- options);`
	`72`	`+ options);`
`73`	`73`
`74`	`74`	`}, requiresCredential: false);`
`75`	`75`	`}`