Use recommended NATS.Net client for Nats package

Directory.Packages.props

@@ -71,7 +71,7 @@
     <PackageVersion Include="MongoDB.Driver" Version="3.0.0" />
     <PackageVersion Include="MySqlConnector" Version="2.3.1" />
     <PackageVersion Include="MySqlConnector.DependencyInjection" Version="2.3.1" />
-    <PackageVersion Include="NATS.Client" Version="1.1.1" />
+    <PackageVersion Include="NATS.Net" Version="2.5.4" />
     <PackageVersion Include="Npgsql" Version="8.0.3" />
     <PackageVersion Include="Npgsql.DependencyInjection" Version="8.0.1" />
     <PackageVersion Include="Npgsql.EntityFrameworkCore.PostgreSQL" Version="8.0.11" />

src/HealthChecks.Nats/DependencyInjection/NatsHealthCheckBuilderExtensions.cs

@@ -1,5 +1,6 @@
 using HealthChecks.Nats;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
+using NATS.Client.Core;
 
 namespace Microsoft.Extensions.DependencyInjection;
 
@@ -11,81 +12,41 @@ public static class NatsHealthCheckBuilderExtensions
     internal const string NAME = "nats";
 
     /// <summary>
-    /// Add a health check for Nats.
+    /// Add a health check for Nats services.
     /// </summary>
     /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
-    /// <param name="setup">The action to configure the Nats setup.</param>
-    /// <param name="name">
-    /// The health check name.
-    /// Optional.
-    /// If <see langword="null"/> the type name 'nats' will be used for the name.
-    /// </param>
+    /// <param name="clientFactory">
+    /// An optional factory to obtain <see cref="INatsConnection" /> instance.
+    /// When not provided, <see cref="INatsConnection" /> is simply resolved from <see cref="IServiceProvider"/>.</param>
+    /// <param name="name">The health check name. Optional. If <c>null</c> the type name 'nats' will be used for the name.</param>
     /// <param name="failureStatus">
-    /// The <see cref="HealthStatus"/> that should be reported when the health check fails.
-    /// Optional.
-    /// If <see langword="null"/> then the default status of <see cref="HealthStatus.Unhealthy"/> will be reported.
-    /// </param>
-    /// <param name="tags">A list of tags that can be used to filter sets of health checks. Optional.</param>
-    /// <param name="timeout">An optional System.TimeSpan representing the timeout of the check.</param>
-    /// <returns>The <see cref="IHealthChecksBuilder"/>.</returns>
-    public static IHealthChecksBuilder AddNats(
-        this IHealthChecksBuilder builder,
-        Action<NatsOptions>? setup,
-        string? name = default,
-        HealthStatus? failureStatus = default,
-        IEnumerable<string>? tags = default,
-        TimeSpan? timeout = default)
-    {
-        var options = new NatsOptions();
-        setup?.Invoke(options);
-
-        builder.Services.AddSingleton(_ => new NatsHealthCheck(options));
-
-        return builder.Add(new HealthCheckRegistration(
-            name ?? NAME,
-            sp => sp.GetRequiredService<NatsHealthCheck>(),
-            failureStatus,
-            tags,
-            timeout));
-    }
-
-    /// <summary>
-    /// Add a health check for Nats.
-    /// </summary>
-    /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
-    /// <param name="setup">The factory to configure the Nats setup.</param>
-    /// <param name="name">
-    /// The health check name.
-    /// Optional.
-    /// If <see langword="null"/> the type name 'nats' will be used for the name.
-    /// </param>
-    /// <param name="failureStatus">
-    /// The <see cref="HealthStatus"/> that should be reported when the health check fails.
-    /// Optional.
-    /// If <see langword="null"/> then the default status of <see cref="HealthStatus.Unhealthy"/> will be reported.
+    /// The <see cref="HealthStatus"/> that should be reported when the health check fails. Optional. If <c>null</c> then
+    /// the default status of <see cref="HealthStatus.Unhealthy"/> will be reported.
     /// </param>
     /// <param name="tags">A list of tags that can be used to filter sets of health checks. Optional.</param>
     /// <param name="timeout">An optional <see cref="TimeSpan"/> representing the timeout of the check.</param>
-    /// <returns>The <see cref="IHealthChecksBuilder"/>.</returns>
+    /// <returns>The specified <paramref name="builder"/>.</returns>
     public static IHealthChecksBuilder AddNats(
         this IHealthChecksBuilder builder,
-        Action<IServiceProvider, NatsOptions>? setup,
-        string? name = default,
+        Func<IServiceProvider, INatsConnection>? clientFactory = default,
+        string? name = NAME,
         HealthStatus? failureStatus = default,
         IEnumerable<string>? tags = default,
         TimeSpan? timeout = default)
     {
         return builder.Add(new HealthCheckRegistration(
             name ?? NAME,
-            sp =>
-            {
-                var options = new NatsOptions();
-                setup?.Invoke(sp, options);
-
-                return new NatsHealthCheck(options);
-            },
+            sp => Factory(clientFactory, sp),
             failureStatus,
             tags,
             timeout));
+
+        static NatsHealthCheck Factory(Func<IServiceProvider, INatsConnection>? clientFactory, IServiceProvider sp)
+        {
+            // The user might have registered a factory for NatsConnection type, but not for the abstraction (INatsConnection).
+            // That is why we try to resolve NatsConnection first.
+            INatsConnection client = clientFactory?.Invoke(sp) ?? sp.GetService<NatsConnection>() ?? sp.GetRequiredService<INatsConnection>();
+            return new(client);
+        }
     }
 }

src/HealthChecks.Nats/HealthChecks.Nats.csproj

@@ -8,7 +8,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="NATS.Client" />
+    <PackageReference Include="NATS.Net" />
   </ItemGroup>
 
 </Project>

src/HealthChecks.Nats/NatsHealthCheck.cs

@@ -1,111 +1,36 @@
-using System.Text;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
-using NATS.Client;
+using NATS.Client.Core;
 
 namespace HealthChecks.Nats;
 
 /// <summary>
 /// Health check for Nats Server.
 /// </summary>
-/// <remarks>
-/// Relies on a static <see cref="ConnectionFactory"/> which provides factory methods to create
-/// connections to NATS Servers, and a <see cref="IConnection"/> object connected to the NATS server.
-/// </remarks>
-public sealed class NatsHealthCheck : IHealthCheck, IDisposable
+public sealed class NatsHealthCheck(INatsConnection connection) : IHealthCheck
 {
-    private static readonly ConnectionFactory _connectionFactory = new();
-
-    private readonly NatsOptions _options;
-
-    private IConnection? _connection;
-
-    public NatsHealthCheck(NatsOptions natsOptions)
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
     {
-        _options = Guard.ThrowIfNull(natsOptions);
+        var result = connection.ConnectionState switch
+        {
+            NatsConnectionState.Open => HealthCheckResult.Healthy(),
+            NatsConnectionState.Connecting or NatsConnectionState.Reconnecting => HealthCheckResult.Degraded(),
+            NatsConnectionState.Closed => await TryConnectAsync(connection).ConfigureAwait(false),
+            _ => new HealthCheckResult(context.Registration.FailureStatus)
+        };
+
+        return result;
     }
 
-    /// <inheritdoc />
-    public Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    private static async Task<HealthCheckResult> TryConnectAsync(INatsConnection natsConnection)
     {
         try
         {
-            // Create new connection if there is no existing one
-            IConnection? connection = _connection;
-            if (connection == null)
-            {
-#pragma warning disable IDISP001 // Dispose created [false positive, https://github.com/DotNetAnalyzers/IDisposableAnalyzers/issues/515]
-                connection = CreateConnection(_options);
-#pragma warning restore IDISP001 // Dispose created
-                var exchanged = Interlocked.CompareExchange(ref _connection, connection, null);
-                if (exchanged != null) // was set by other thread
-                {
-                    connection.Dispose();
-                    connection = exchanged;
-                }
-            }
-
-            // reset connection in case of stuck so the next HC call will establish it again
-            // https://github.com/Xabaril/AspNetCore.Diagnostics.HealthChecks/issues/1544
-            if (connection.State == ConnState.DISCONNECTED || connection.State == ConnState.CLOSED)
-                _connection = null;
-
-            var healthCheckResult = GetHealthCheckResultFromState(connection);
-            return Task.FromResult(healthCheckResult);
+            await natsConnection.ConnectAsync().ConfigureAwait(false);
+            return HealthCheckResult.Healthy();
         }
-        catch (Exception ex)
-        {
-            var unhealthy = new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
-            return Task.FromResult(unhealthy);
-        }
-
-        IConnection CreateConnection(NatsOptions options)
+        catch (Exception)
         {
-            if (!string.IsNullOrWhiteSpace(options.CredentialsPath))
-                return _connectionFactory.CreateConnection(options.Url, options.CredentialsPath);
-            if (!string.IsNullOrWhiteSpace(options.Jwt) && !string.IsNullOrWhiteSpace(options.PrivateNKey))
-                return _connectionFactory.CreateConnection(options.Url, options.Jwt, options.PrivateNKey);
-            return _connectionFactory.CreateConnection(options.Url);
+            return HealthCheckResult.Unhealthy();
         }
-
-        HealthCheckResult GetHealthCheckResultFromState(IConnection connection)
-        {
-            string description = GetDescription(connection);
-
-            return connection.State switch
-            {
-                ConnState.CONNECTED => HealthCheckResult.Healthy(description, GetStatsData(connection)),
-                ConnState.CONNECTING
-                or ConnState.RECONNECTING
-                or ConnState.DRAINING_SUBS
-                or ConnState.DRAINING_PUBS => HealthCheckResult.Degraded(description),
-                ConnState.CLOSED
-                or ConnState.DISCONNECTED => HealthCheckResult.Unhealthy(description),
-                _ => new HealthCheckResult(context.Registration.FailureStatus, description),
-            };
-        }
-
-        static string GetDescription(IConnection connection)
-        {
-            var sb = new StringBuilder();
-            sb.AppendFormat("{0}: {1}; ", nameof(connection.ClientIP), connection.ClientIP);
-            if (!string.IsNullOrWhiteSpace(connection.ConnectedUrl))
-                sb.AppendFormat("{0}: {1}; ", nameof(connection.ConnectedUrl), connection.ConnectedUrl);
-            sb.AppendFormat("{0}: {1}; ", nameof(connection.State), connection.State);
-            if (connection.SubscriptionCount != default)
-                sb.AppendFormat("{0}: {1}", nameof(connection.SubscriptionCount), connection.SubscriptionCount);
-            return sb.ToString();
-        }
-
-        static IReadOnlyDictionary<string, object> GetStatsData(IConnection connection) =>
-            new Dictionary<string, object>
-            {
-                [nameof(connection.Stats.InMsgs)] = connection.Stats.InMsgs,
-                [nameof(connection.Stats.OutMsgs)] = connection.Stats.OutMsgs,
-                [nameof(connection.Stats.InBytes)] = connection.Stats.InBytes,
-                [nameof(connection.Stats.OutBytes)] = connection.Stats.OutBytes,
-                [nameof(connection.Stats.Reconnects)] = connection.Stats.Reconnects
-            };
     }
-
-    public void Dispose() => _connection?.Dispose();
 }