Improve Npgsql health check design

src/HealthChecks.NpgSql/DependencyInjection/NpgSqlHealthCheckBuilderExtensions.cs

@@ -37,7 +37,13 @@ public static IHealthChecksBuilder AddNpgSql(
         IEnumerable<string>? tags = default,
         TimeSpan? timeout = default)
     {
-        return builder.AddNpgSql(_ => connectionString, healthQuery, configure, name, failureStatus, tags, timeout);
+        Guard.ThrowIfNull(connectionString, throwOnEmptyString: true);
+
+        return builder.AddNpgSql(new NpgSqlHealthCheckOptions(connectionString)
+        {
+            CommandText = healthQuery,
+            Configure = configure
+        }, name, failureStatus, tags, timeout);
     }
 
     /// <summary>
@@ -65,14 +71,36 @@ public static IHealthChecksBuilder AddNpgSql(
         IEnumerable<string>? tags = default,
         TimeSpan? timeout = default)
     {
-        return builder.AddNpgSql(sp => new NpgsqlDataSourceBuilder(connectionStringFactory(sp)).Build(), healthQuery, configure, name, failureStatus, tags, timeout);
+        // This instance is captured in lambda closure, so it can be reused (perf)
+        NpgSqlHealthCheckOptions options = new()
+        {
+            CommandText = healthQuery,
+            Configure = configure,
+        };
+
+        return builder.Add(new HealthCheckRegistration(
+            name ?? NAME,
+            sp =>
+            {
+                options.ConnectionString ??= Guard.ThrowIfNull(connectionStringFactory.Invoke(sp), throwOnEmptyString: true, paramName: nameof(connectionStringFactory));
+
+                ResolveDataSourceIfPossible(options, sp);
+
+                return new NpgSqlHealthCheck(options);
+            },
+            failureStatus,
+            tags,
+            timeout));
     }
 
     /// <summary>
     /// Add a health check for Postgres databases.
     /// </summary>
     /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
-    /// <param name="dbDataSourceFactory">A factory to build the NpgsqlDataSource to use.</param>
+    /// <param name="dbDataSourceFactory">
+    /// An optional factory to obtain <see cref="NpgsqlDataSource" /> instance.
+    /// When not provided, <see cref="NpgsqlDataSource" /> is simply resolved from <see cref="IServiceProvider"/>.
+    /// </param>
     /// <param name="healthQuery">The query to be used in check.</param>
     /// <param name="configure">An optional action to allow additional Npgsql specific configuration.</param>
     /// <param name="name">The health check name. Optional. If <c>null</c> the type name 'npgsql' will be used for the name.</param>
@@ -83,19 +111,21 @@ public static IHealthChecksBuilder AddNpgSql(
     /// <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 specified <paramref name="builder"/>.</returns>
+    /// <remarks>
+    /// Depending on how the <see cref="NpgsqlDataSource" /> was configured, the connections it hands out may be pooled.
+    /// That is why it should be the exact same <see cref="NpgsqlDataSource" /> that is used by other parts of your app.
+    /// </remarks>
     public static IHealthChecksBuilder AddNpgSql(
         this IHealthChecksBuilder builder,
-        Func<IServiceProvider, NpgsqlDataSource> dbDataSourceFactory,
+        Func<IServiceProvider, NpgsqlDataSource>? dbDataSourceFactory = null,
         string healthQuery = HEALTH_QUERY,
         Action<NpgsqlConnection>? configure = null,
         string? name = default,
         HealthStatus? failureStatus = default,
         IEnumerable<string>? tags = default,
         TimeSpan? timeout = default)
     {
-        Guard.ThrowIfNull(dbDataSourceFactory);
-
-        NpgsqlDataSource? dataSource = null;
+        // This instance is captured in lambda closure, so it can be reused (perf)
         NpgSqlHealthCheckOptions options = new()
         {
             CommandText = healthQuery,
@@ -106,49 +136,7 @@ public static IHealthChecksBuilder AddNpgSql(
             name ?? NAME,
             sp =>
             {
-                // The Data Source needs to be created only once,
-                // as each instance has it's own connection pool.
-                // See https://github.com/Xabaril/AspNetCore.Diagnostics.HealthChecks/issues/1993 for more details.
-
-                // Perform an atomic read of the current value.
-                NpgsqlDataSource? existingDataSource = Volatile.Read(ref dataSource);
-                if (existingDataSource is null)
-                {
-                    // Create a new Data Source
-                    NpgsqlDataSource fromFactory = dbDataSourceFactory(sp);
-                    // Try to resolve the Data Source from DI.
-                    NpgsqlDataSource? fromDI = sp.GetService<NpgsqlDataSource>();
-
-                    if (fromDI is not null && fromDI.ConnectionString.Equals(fromFactory.ConnectionString))
-                    {
-                        // If they are using the same ConnectionString, we can reuse the instance from DI.
-                        // So there is only ONE NpgsqlDataSource per the whole app and ONE connection pool.
-
-                        if (!ReferenceEquals(fromDI, fromFactory))
-                        {
-                            // Dispose it, as long as it's not the same instance.
-                            fromFactory.Dispose();
-                        }
-                        Interlocked.Exchange(ref dataSource, fromDI);
-                        options.DataSource = fromDI;
-                    }
-                    else
-                    {
-                        // Perform an atomic exchange, but only if the value is still null.
-                        existingDataSource = Interlocked.CompareExchange(ref dataSource, fromFactory, null);
-                        if (existingDataSource is not null)
-                        {
-                            // Some other thread has created the data source in the meantime,
-                            // we dispose our own copy, and use the existing instance.
-                            fromFactory.Dispose();
-                            options.DataSource = existingDataSource;
-                        }
-                        else
-                        {
-                            options.DataSource = fromFactory;
-                        }
-                    }
-                }
+                options.DataSource ??= dbDataSourceFactory?.Invoke(sp) ?? sp.GetRequiredService<NpgsqlDataSource>();
 
                 return new NpgSqlHealthCheck(options);
             },
@@ -182,9 +170,29 @@ public static IHealthChecksBuilder AddNpgSql(
 
         return builder.Add(new HealthCheckRegistration(
             name ?? NAME,
-            _ => new NpgSqlHealthCheck(options),
+            sp =>
+            {
+                ResolveDataSourceIfPossible(options, sp);
+
+                return new NpgSqlHealthCheck(options);
+            },
             failureStatus,
             tags,
             timeout));
     }
+
+    private static void ResolveDataSourceIfPossible(NpgSqlHealthCheckOptions options, IServiceProvider sp)
+    {
+        if (options.DataSource is null && !options.TriedToResolveFromDI)
+        {
+            NpgsqlDataSource? fromDi = sp.GetService<NpgsqlDataSource>();
+            if (fromDi?.ConnectionString == options.ConnectionString)
+            {
+                // When it's possible, we reuse the DataSource registered in the DI.
+                // We do that to achieve best performance and avoid issues like https://github.com/Xabaril/AspNetCore.Diagnostics.HealthChecks/issues/1993
+                options.DataSource = fromDi;
+            }
+            options.TriedToResolveFromDI = true; // save the answer, so we don't do it more than once
+        }
+    }
 }

src/HealthChecks.NpgSql/NpgSqlHealthCheck.cs

@@ -1,4 +1,6 @@
+using System.Diagnostics;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Npgsql;
 
 namespace HealthChecks.NpgSql;
 
@@ -11,7 +13,7 @@ public class NpgSqlHealthCheck : IHealthCheck
 
     public NpgSqlHealthCheck(NpgSqlHealthCheckOptions options)
     {
-        Guard.ThrowIfNull(options.DataSource);
+        Debug.Assert(options.ConnectionString is not null || options.DataSource is not null);
         Guard.ThrowIfNull(options.CommandText, true);
         _options = options;
     }
@@ -21,7 +23,9 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
     {
         try
         {
-            await using var connection = _options.DataSource!.CreateConnection();
+            await using var connection = _options.DataSource is not null
+                ? _options.DataSource.CreateConnection()
+                : new NpgsqlConnection(_options.ConnectionString);
 
             _options.Configure?.Invoke(connection);
             await connection.OpenAsync(cancellationToken).ConfigureAwait(false);
@@ -33,7 +37,6 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
             return _options.HealthCheckResultBuilder == null
                 ? HealthCheckResult.Healthy()
                 : _options.HealthCheckResultBuilder(result);
-
         }
         catch (Exception ex)
         {

src/HealthChecks.NpgSql/NpgSqlHealthCheckOptions.cs

@@ -9,20 +9,63 @@ namespace HealthChecks.NpgSql;
 /// </summary>
 public class NpgSqlHealthCheckOptions
 {
+    internal NpgSqlHealthCheckOptions()
+    {
+        // This ctor is internal on purpose: those who want to use NpgSqlHealthCheckOptions
+        // need to specify either ConnectionString or DataSource when creating it.
+        // Making the ConnectionString and DataSource setters internal
+        // allows us to ensure that the customers don't try to mix both concepts.
+        // By encapsulating all of that, we ensure that all instances of this type are valid.
+    }
+
     /// <summary>
-    /// The Postgres connection string to be used.
-    /// Use <see cref="DataSource"/> property for advanced configuration.
+    /// Creates an instance of <see cref="NpgSqlHealthCheckOptions"/>.
+    /// </summary>
+    /// <param name="connectionString">The PostgreSQL connection string to be used.</param>
+    /// <remarks>
+    /// <see cref="NpgsqlDataSource"/> supports additional configuration beyond the connection string,
+    /// such as logging, advanced authentication options, type mapping management, and more.
+    /// To further configure a data source, use <see cref=" NpgsqlDataSourceBuilder"/> and
+    /// the <see cref="NpgSqlHealthCheckOptions(NpgsqlDataSource)"/> constructor.
+    /// </remarks>
+    public NpgSqlHealthCheckOptions(string connectionString)
+    {
+        ConnectionString = Guard.ThrowIfNull(connectionString, throwOnEmptyString: true);
+    }
+
+    /// <summary>
+    /// Creates an instance of <see cref="NpgSqlHealthCheckOptions"/>.
     /// </summary>
-    public string? ConnectionString
+    /// <param name="dataSource">The Postgres <see cref="NpgsqlDataSource" /> to be used.</param>
+    /// <remarks>
+    /// Depending on how the <see cref="NpgsqlDataSource" /> was configured, the connections it hands out may be pooled.
+    /// That is why it should be the exact same <see cref="NpgsqlDataSource" /> that is used by other parts of your app.
+    /// </remarks>
+    public NpgSqlHealthCheckOptions(NpgsqlDataSource dataSource)
     {
-        get => DataSource?.ConnectionString;
-        set => DataSource = value is not null ? NpgsqlDataSource.Create(value) : null;
+        DataSource = Guard.ThrowIfNull(dataSource);
     }
 
     /// <summary>
-    /// The Postgres data source to be used.
+    /// The Postgres connection string to be used.
     /// </summary>
-    public NpgsqlDataSource? DataSource { get; set; }
+    /// <remarks>
+    /// <see cref="NpgsqlDataSource"/> supports additional configuration beyond the connection string,
+    /// such as logging, advanced authentication options, type mapping management, and more.
+    /// To further configure a data source, use <see cref=" NpgsqlDataSourceBuilder"/> and the <see cref="NpgSqlHealthCheckOptions(NpgsqlDataSource)"/> constructor.
+    /// </remarks>
+    public string? ConnectionString { get; internal set; }
+
+    /// <summary>
+    /// The Postgres <see cref="NpgsqlDataSource" /> to be used.
+    /// </summary>
+    /// <remarks>
+    /// Depending on how the <see cref="NpgsqlDataSource" /> was configured, the connections it hands out may be pooled.
+    /// That is why it should be the exact same <see cref="NpgsqlDataSource" /> that is used by other parts of your app.
+    /// </remarks>
+    public NpgsqlDataSource? DataSource { get; internal set; }
+
+    internal bool TriedToResolveFromDI;
 
     /// <summary>
     /// The query to be executed.

src/HealthChecks.NpgSql/README.md

@@ -0,0 +1,47 @@
+## PostgreSQL Health Check
+
+This health check verifies the ability to communicate with [PostgreSQL](https://www.postgresql.org/). It uses the [Npgsql](https://www.npgsql.org/) library.
+
+## NpgsqlDataSource
+
+Starting with Npgsql 7.0 (and .NET 7), the starting point for any database operation is [NpgsqlDataSource](https://www.npgsql.org/doc/api/Npgsql.NpgsqlDataSource.html). The data source represents your PostgreSQL database, and can hand out connections to it, or support direct execution of SQL against it. The data source encapsulates the various Npgsql configuration needed to connect to PostgreSQL, as well as the **connection pooling which makes Npgsql efficient**.
+
+Npgsql's **data source supports additional configuration beyond the connection string**, such as logging, advanced authentication options, type mapping management, and more.
+
+## Recommended approach
+
+To take advantage of the performance `NpgsqlDataSource` has to offer, it should be used as a **singleton**. Otherwise, the app might end up with having multiple data source instances, all of which would have their own connection pools. This can lead to resources exhaustion and major performance issues (Example: [#1993](https://github.com/Xabaril/AspNetCore.Diagnostics.HealthChecks/issues/1993)).
+
+We encourage you to use [Npgsql.DependencyInjection](https://www.nuget.org/packages/Npgsql.DependencyInjection) package for registering a singleton factory for `NpgsqlDataSource`.  It allows easy configuration of your Npgsql connections and registers the appropriate services in your DI container.
+
+To make the shift to `NpgsqlDataSource` as easy as possible,  the `Npgsql.DependencyInjection` package registers not just a factory for the data source, but also factory for `NpgsqlConnection`  (and even `DbConnection`). So, your app does not need to suddenly start using `NpgsqlDataSource` everywhere.
+
+```csharp
+void Configure(IServiceCollection services)
+{
+    services.AddNpgsqlDataSource("Host=pg_server;Username=test;Password=test;Database=test");
+    services.AddHealthChecks().AddNpgSql();
+}
+```
+
+By default, the `NpgsqlDataSource` instance is resolved from service provider. If you need to access more than one PostgreSQL database, you can use keyed DI services to achieve that:
+
+```csharp
+void Configure(IServiceCollection services)
+{
+    services.AddNpgsqlDataSource("Host=pg_server;Username=test;Password=test;Database=first", serviceKey: "first");
+    services.AddHealthChecks().AddNpgSql(sp => sp.GetRequiredKeyedService<NpgsqlDataSource>("first"));
+
+    services.AddNpgsqlDataSource("Host=pg_server;Username=test;Password=test;Database=second", serviceKey: "second");
+    services.AddHealthChecks().AddNpgSql(sp => sp.GetRequiredKeyedService<NpgsqlDataSource>("second"));
+}
+```
+
+
+## Connection String
+
+Raw connection string is of course still supported:
+
+```csharp
+services.AddHealthChecks().AddNpgSql("Host=pg_server;Username=test;Password=test;Database=test");
+```

test/HealthChecks.Npgsql.Tests/DependencyInjection/RegistrationTests.cs

@@ -96,7 +96,7 @@ public void factory_is_called_only_once()
     [InlineData(false)]
     public void factory_reuses_pre_registered_datasource_when_possible(bool sameConnectionString)
     {
-        const string connectionString = "Server=localhost";
+        const string connectionString = "Host=localhost";
         ServiceCollection services = new();
 
         services.AddSingleton<NpgsqlDataSource>(serviceProvider =>
@@ -130,4 +130,19 @@ public void factory_reuses_pre_registered_datasource_when_possible(bool sameConn
 
         factoryCalls.ShouldBe(1);
     }
+
+    [Fact]
+    public void recommended_scenario_compiles_and_works_as_expected()
+    {
+        ServiceCollection services = new();
+        services.AddNpgsqlDataSource("Host=pg_server;Username=test;Password=test;Database=test");
+        services.AddHealthChecks().AddNpgSql();
+
+        using var serviceProvider = services.BuildServiceProvider();
+        var options = serviceProvider.GetRequiredService<IOptions<HealthCheckServiceOptions>>();
+        var registration = options.Value.Registrations.Single();
+        var healthCheck = registration.Factory(serviceProvider);
+
+        healthCheck.ShouldBeOfType<NpgSqlHealthCheck>();
+    }
 }