Harden MSSQL pool

Author: vkuttyp

Directory.Build.props

@@ -7,6 +7,6 @@
     <Authors>vkuttyp</Authors>
     <PackageLicenseExpression>MIT</PackageLicenseExpression>
     <RepositoryUrl>https://github.com/vkuttyp/CosmoSQLClient-Dotnet</RepositoryUrl>
-    <Version>1.9.42</Version>
+    <Version>1.9.43</Version>
   </PropertyGroup>
 </Project>

src/CosmoSQLClient.MsSql/AssemblyInfo.cs

@@ -0,0 +1,3 @@
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("CosmoSQLClient.MsSql.Tests")]

src/CosmoSQLClient.MsSql/MsSqlConnectionPool.cs

@@ -26,6 +26,8 @@ public sealed class MsSqlConnectionPool : ISqlDatabase
     private readonly int _maxConnections;
     private readonly int _minIdle;
     private readonly Channel<MsSqlConnection> _idle;
+    private readonly Func<MsSqlConnection, CancellationToken, ValueTask<bool>> _validateConnection;
+    private readonly Func<CancellationToken, Task<MsSqlConnection>> _connectionFactory;
     private int _count;
     private bool _disposed;
     private CancellationTokenSource? _keepAliveCts;
@@ -50,7 +52,9 @@ public sealed class MsSqlConnectionPool : ISqlDatabase
     public MsSqlConnectionPool(
         MsSqlConfiguration config,
         int maxConnections = 10,
-        int minIdle = 0)
+        int minIdle = 0,
+        Func<MsSqlConnection, CancellationToken, ValueTask<bool>>? validateConnection = null,
+        Func<CancellationToken, Task<MsSqlConnection>>? connectionFactory = null)
     {
         _config = config;
         _maxConnections = Math.Max(1, maxConnections);
@@ -88,6 +92,9 @@ public MsSqlConnectionPool(
             _sslOptions = new object(); // Placeholder for netstandard2.0, ConnectAsync handles it
 #endif
         }
+
+        _validateConnection = validateConnection ?? DefaultValidateConnectionAsync;
+        _connectionFactory = connectionFactory ?? OpenConnectionAsync;
     }
 
     public bool IsOpen => !_disposed;
@@ -101,7 +108,7 @@ public MsSqlConnectionPool(
     public async Task WarmUpAsync(int? count = null, CancellationToken ct = default)
     {
         int n = Math.Min(count ?? _minIdle, _maxConnections);
-        var tasks = Enumerable.Range(0, n).Select(_ => OpenConnectionAsync(ct));
+        var tasks = Enumerable.Range(0, n).Select(_ => _connectionFactory(ct));
         var conns = await Task.WhenAll(tasks).ConfigureAwait(false);
         foreach (var c in conns)
             await _idle.Writer.WriteAsync(c, ct).ConfigureAwait(false);
@@ -192,35 +199,72 @@ private async Task PingIdleAsync(CancellationToken ct)
 
     // ── Pool acquire/release ───────────────────────────────────────────────────
 
-    public async Task<MsSqlConnection> AcquireAsync(CancellationToken ct = default)
-    {
-        // Try an idle connection first (non-blocking).
-        if (_idle.Reader.TryRead(out var conn))
-            return conn;
+        public async Task<MsSqlConnection> AcquireAsync(CancellationToken ct = default)
+        {
+            while (true)
+            {
+                if (_idle.Reader.TryRead(out var idleConn))
+                {
+                    if (await _validateConnection(idleConn, ct).ConfigureAwait(false))
+                        return idleConn;
+
+                    await idleConn.DisposeAsync().ConfigureAwait(false);
+                    Interlocked.Decrement(ref _count);
+                    continue;
+                }
+
+                int current = Interlocked.Increment(ref _count);
+                if (current <= _maxConnections)
+                {
+                    try { return await _connectionFactory(ct).ConfigureAwait(false); }
+                    catch { Interlocked.Decrement(ref _count); throw; }
+                }
 
-        // If below capacity, create a new connection.
-        int current = Interlocked.Increment(ref _count);
-        if (current <= _maxConnections)
+                Interlocked.Decrement(ref _count);
+                var pooledConn = await _idle.Reader.ReadAsync(ct).ConfigureAwait(false);
+                if (await _validateConnection(pooledConn, ct).ConfigureAwait(false))
+                    return pooledConn;
+
+                await pooledConn.DisposeAsync().ConfigureAwait(false);
+                Interlocked.Decrement(ref _count);
+            }
+        }
+
+        internal async Task InjectIdleConnectionAsync(MsSqlConnection conn)
         {
-            try { return await OpenConnectionAsync(ct).ConfigureAwait(false); }
-            catch { Interlocked.Decrement(ref _count); throw; }
+            await _idle.Writer.WriteAsync(conn).ConfigureAwait(false);
+            Interlocked.Increment(ref _count);
         }
 
-        // Already at max — wait for an idle one.
-        Interlocked.Decrement(ref _count);
-        return await _idle.Reader.ReadAsync(ct).ConfigureAwait(false);
-    }
+        public async Task ReleaseAsync(MsSqlConnection conn)
+        {
+            if (conn.IsOpen && !_disposed)
+                await _idle.Writer.WriteAsync(conn).ConfigureAwait(false);
+            else
+            {
+                await conn.DisposeAsync().ConfigureAwait(false);
+                Interlocked.Decrement(ref _count);
+            }
+        }
 
-    public async Task ReleaseAsync(MsSqlConnection conn)
-    {
-        if (conn.IsOpen && !_disposed)
-            await _idle.Writer.WriteAsync(conn).ConfigureAwait(false);
-        else
+        private async ValueTask<bool> DefaultValidateConnectionAsync(MsSqlConnection conn, CancellationToken ct)
         {
-            await conn.DisposeAsync().ConfigureAwait(false);
-            Interlocked.Decrement(ref _count);
+            if (!conn.IsOpen)
+                return false;
+            try
+            {
+                await conn.QueryAsync("SELECT 1", ct: ct).ConfigureAwait(false);
+                return true;
+            }
+            catch (OperationCanceledException)
+            {
+                throw;
+            }
+            catch
+            {
+                return false;
+            }
         }
-    }
 
     // ── ISqlDatabase ───────────────────────────────────────────────────────────
 

tests/CosmoSQLClient.MsSql.Tests/MsSqlConnectionPoolTests.cs

@@ -0,0 +1,45 @@
+using System.Threading;
+using System.Threading.Tasks;
+using CosmoSQLClient.MsSql;
+using Xunit;
+
+namespace CosmoSQLClient.MsSql.Tests;
+
+public sealed class MsSqlConnectionPoolTests
+{
+    [Fact]
+    public async Task AcquireAsync_RecreatesConnectionWhenValidationFails()
+    {
+        var config = new MsSqlConfiguration
+        {
+            Host = "localhost",
+            Database = "master",
+            TrustServerCertificate = true
+        };
+
+        int validatorCalls = 0;
+        int factoryCalls = 0;
+
+        var pool = new MsSqlConnectionPool(
+            config,
+            maxConnections: 1,
+            minIdle: 0,
+            validateConnection: (conn, ct) =>
+            {
+                Interlocked.Increment(ref validatorCalls);
+                return new ValueTask<bool>(false);
+            },
+            connectionFactory: ct =>
+            {
+                Interlocked.Increment(ref factoryCalls);
+                return Task.FromResult(new MsSqlConnection("Server=localhost;Database=master;User Id=sa;Password=pass;Pooling=false;"));
+            });
+
+        await pool.InjectIdleConnectionAsync(new MsSqlConnection("Server=localhost;Database=master;User Id=sa;Password=pass;Pooling=false;"));
+        var connection = await pool.AcquireAsync();
+        await pool.ReleaseAsync(connection);
+
+        Assert.Equal(1, factoryCalls);
+        Assert.Equal(1, validatorCalls);
+    }
+}