Npgsql: don't create more than one DataSource (#2045)

* add a failing test

* provide simple fix, non-thread safe fix

* provide thread-safe fix

* bump Npgsql version

* add .NET 7 to list of supported tfms to avoid "System.TypeLoadException : Could not load type 'System.Data.Common.DbDataSource' from assembly 'Npgsql, Version=7.0.6.0, Culture=neutral, PublicKeyToken=5d8b90d52f46fda7'." thrown for .NET 7 test runner

* bump the version

* add test for scenario when there is pre-registered DataSource in the DI

* use NpgsqlDataSource registered in the DI when the connection string match

* reorder the usings to satisfy dotnet format

* one more fix

build/versions.props

@@ -39,7 +39,7 @@
     <HealthCheckMySql>7.0.0</HealthCheckMySql>
     <HealthCheckNetwork>7.0.0</HealthCheckNetwork>
     <HealthCheckNats>7.0.0</HealthCheckNats>
-    <HealthCheckNpgSql>7.0.0</HealthCheckNpgSql>
+    <HealthCheckNpgSql>7.1.0</HealthCheckNpgSql>
     <HealthCheckOpenIdConnectServer>7.1.0</HealthCheckOpenIdConnectServer>
     <HealthCheckOracle>7.0.0</HealthCheckOracle>
     <HealthCheckPrometheusMetrics>7.0.0</HealthCheckPrometheusMetrics>

src/HealthChecks.NpgSql/DependencyInjection/NpgSqlHealthCheckBuilderExtensions.cs

@@ -95,16 +95,61 @@ public static IHealthChecksBuilder AddNpgSql(
     {
         Guard.ThrowIfNull(dbDataSourceFactory);
 
+        NpgsqlDataSource? dataSource = null;
+        NpgSqlHealthCheckOptions options = new()
+        {
+            CommandText = healthQuery,
+            Configure = configure,
+        };
+
         return builder.Add(new HealthCheckRegistration(
             name ?? NAME,
             sp =>
             {
-                var options = new NpgSqlHealthCheckOptions
+                // 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)
                 {
-                    DataSource = dbDataSourceFactory(sp),
-                    CommandText = healthQuery,
-                    Configure = configure,
-                };
+                    // 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;
+                        }
+                    }
+                }
+
                 return new NpgSqlHealthCheck(options);
             },
             failureStatus,

src/HealthChecks.NpgSql/HealthChecks.NpgSql.csproj

@@ -1,14 +1,14 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
-    <TargetFramework>netstandard2.0</TargetFramework>
+    <TargetFrameworks>netstandard2.0;net7.0</TargetFrameworks>
     <PackageTags>$(PackageTags);Beat;Postgress</PackageTags>
     <Description>HealthChecks.NpgSql is a health check for Postgress Sql.</Description>
     <VersionPrefix>$(HealthCheckNpgSql)</VersionPrefix>
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Npgsql" Version="7.0.4" />
+    <PackageReference Include="Npgsql" Version="7.0.6" />
     <PackageReference Include="Microsoft.Extensions.Diagnostics.HealthChecks" Version="7.0.9" />
   </ItemGroup>
 

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

@@ -1,4 +1,6 @@
+using System.Reflection;
 using HealthChecks.NpgSql;
+using Npgsql;
 
 namespace HealthChecks.Npgsql.Tests.DependencyInjection;
 
@@ -62,4 +64,70 @@ public void add_health_check_with_connection_string_factory_when_properly_config
         check.ShouldBeOfType<NpgSqlHealthCheck>();
         factoryCalled.ShouldBeTrue();
     }
+
+    [Fact]
+    public void factory_is_called_only_once()
+    {
+        ServiceCollection services = new();
+        int factoryCalls = 0;
+        services.AddHealthChecks()
+            .AddNpgSql(_ =>
+            {
+                Interlocked.Increment(ref factoryCalls);
+                return "Server=localhost";
+            }, name: "my-npg-1");
+
+        using var serviceProvider = services.BuildServiceProvider();
+
+        var options = serviceProvider.GetRequiredService<IOptions<HealthCheckServiceOptions>>();
+
+        var registration = options.Value.Registrations.Single();
+
+        for (int i = 0; i < 10; i++)
+        {
+            _ = registration.Factory(serviceProvider);
+        }
+
+        factoryCalls.ShouldBe(1);
+    }
+
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public void factory_reuses_pre_registered_datasource_when_possible(bool sameConnectionString)
+    {
+        const string connectionString = "Server=localhost";
+        ServiceCollection services = new();
+
+        services.AddSingleton<NpgsqlDataSource>(serviceProvider =>
+        {
+            var dataSourceBuilder = new NpgsqlDataSourceBuilder(connectionString);
+            return dataSourceBuilder.Build();
+        });
+
+        int factoryCalls = 0;
+        services.AddHealthChecks()
+            .AddNpgSql(_ =>
+            {
+                Interlocked.Increment(ref factoryCalls);
+                return sameConnectionString ? connectionString : $"{connectionString}2";
+            }, name: "my-npg-1");
+
+        using var serviceProvider = services.BuildServiceProvider();
+
+        var options = serviceProvider.GetRequiredService<IOptions<HealthCheckServiceOptions>>();
+
+        var registration = options.Value.Registrations.Single();
+
+        for (int i = 0; i < 10; i++)
+        {
+            var healthCheck = (NpgSqlHealthCheck)registration.Factory(serviceProvider);
+            var fieldInfo = typeof(NpgSqlHealthCheck).GetField("_options", BindingFlags.Instance | BindingFlags.NonPublic);
+            var npgSqlHealthCheckOptions = (NpgSqlHealthCheckOptions)fieldInfo!.GetValue(healthCheck)!;
+
+            Assert.Equal(sameConnectionString, ReferenceEquals(serviceProvider.GetRequiredService<NpgsqlDataSource>(), npgSqlHealthCheckOptions.DataSource));
+        }
+
+        factoryCalls.ShouldBe(1);
+    }
 }