fix(MsSql): don't leak idle slot when validation throws OperationCanceledException

The previous AcquireAsync removed an idle connection from the channel,
called validateConnection on it, and only disposed + decremented _count
when validation returned false. When validation threw — most commonly
OperationCanceledException from the request CT, but anything propagating
out — the connection was already pulled from _idle and never put back,
disposed, or counted down. Each cancelled validate leaks one slot.

After enough leaks _count exceeds _maxConnections even though the
channel is empty, and AcquireAsync deadlocks at _idle.Reader.ReadAsync
waiting for a write that will never come. This wedged CosmoS3 PUTs
through cosmoproxy: every aws s3 cp triggered parallel cred + user
lookups whose CT got cancelled when AWS CLI hit its read timeout.

Wrap both validation sites in TryValidateOrDiscardAsync, which always
disposes and decrements before propagating an exception. Add a
regression test that fires N cancelled acquires and asserts _count
returns to zero.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: vkuttyp

src/CosmoSQLClient.MsSql/MsSqlConnectionPool.cs

@@ -205,11 +205,8 @@ public async Task<MsSqlConnection> AcquireAsync(CancellationToken ct = default)
             {
                 if (_idle.Reader.TryRead(out var idleConn))
                 {
-                    if (await _validateConnection(idleConn, ct).ConfigureAwait(false))
+                    if (await TryValidateOrDiscardAsync(idleConn, ct).ConfigureAwait(false))
                         return idleConn;
-
-                    await idleConn.DisposeAsync().ConfigureAwait(false);
-                    Interlocked.Decrement(ref _count);
                     continue;
                 }
 
@@ -222,12 +219,42 @@ public async Task<MsSqlConnection> AcquireAsync(CancellationToken ct = default)
 
                 Interlocked.Decrement(ref _count);
                 var pooledConn = await _idle.Reader.ReadAsync(ct).ConfigureAwait(false);
-                if (await _validateConnection(pooledConn, ct).ConfigureAwait(false))
+                if (await TryValidateOrDiscardAsync(pooledConn, ct).ConfigureAwait(false))
                     return pooledConn;
+            }
+        }
 
-                await pooledConn.DisposeAsync().ConfigureAwait(false);
+        // Validate a connection that's been pulled out of the idle channel.
+        // Returns true if the connection is healthy and the caller may keep it;
+        // returns false when validation rejected it (already disposed + decremented).
+        // CRITICAL: cancellation during _validateConnection propagates as
+        // OperationCanceledException, so the connection is already removed
+        // from the channel. We MUST dispose it and decrement _count before
+        // re-throwing — otherwise the slot is leaked and, after enough
+        // cancellations, the pool deadlocks waiting on an empty channel
+        // because _count says we're at capacity.
+        private async ValueTask<bool> TryValidateOrDiscardAsync(MsSqlConnection conn, CancellationToken ct)
+        {
+            try
+            {
+                if (await _validateConnection(conn, ct).ConfigureAwait(false))
+                    return true;
+            }
+            catch
+            {
+                await SafeDisposeAsync(conn).ConfigureAwait(false);
                 Interlocked.Decrement(ref _count);
+                throw;
             }
+
+            await SafeDisposeAsync(conn).ConfigureAwait(false);
+            Interlocked.Decrement(ref _count);
+            return false;
+        }
+
+        private static async ValueTask SafeDisposeAsync(MsSqlConnection conn)
+        {
+            try { await conn.DisposeAsync().ConfigureAwait(false); } catch { /* best effort */ }
         }
 
         internal async Task InjectIdleConnectionAsync(MsSqlConnection conn)

tests/CosmoSQLClient.MsSql.Tests/MsSqlConnectionPoolTests.cs

@@ -1,3 +1,5 @@
+using System;
+using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
 using CosmoSQLClient.MsSql;
@@ -7,6 +9,14 @@ namespace CosmoSQLClient.MsSql.Tests;
 
 public sealed class MsSqlConnectionPoolTests
 {
+    private static int ReadCount(MsSqlConnectionPool pool)
+    {
+        // _count is a private int field accessed via Interlocked. Reflection
+        // is the only way to read it from a test without changing visibility.
+        var f = typeof(MsSqlConnectionPool).GetField("_count", BindingFlags.NonPublic | BindingFlags.Instance);
+        return (int)f!.GetValue(pool)!;
+    }
+
     [Fact]
     public async Task AcquireAsync_RecreatesConnectionWhenValidationFails()
     {
@@ -42,4 +52,58 @@ public async Task AcquireAsync_RecreatesConnectionWhenValidationFails()
         Assert.Equal(1, factoryCalls);
         Assert.Equal(1, validatorCalls);
     }
+
+    [Fact]
+    public async Task AcquireAsync_CancelledValidation_DoesNotLeakIdleSlot()
+    {
+        // Regression for the bug that wedged CosmoS3 PUTs through cosmoproxy:
+        // when validateConnection threw OperationCanceledException, the
+        // already-claimed idle connection was neither disposed nor counted
+        // back. After enough cancellations, _count exceeded _maxConnections
+        // even though _idle was empty, and AcquireAsync deadlocked on
+        // _idle.Reader.ReadAsync waiting for a write that would never come.
+
+        var config = new MsSqlConfiguration
+        {
+            Host = "localhost",
+            Database = "master",
+            TrustServerCertificate = true
+        };
+
+        var pool = new MsSqlConnectionPool(
+            config,
+            maxConnections: 2,
+            minIdle: 0,
+            validateConnection: (conn, ct) =>
+            {
+                ct.ThrowIfCancellationRequested();
+                throw new OperationCanceledException(ct);
+            },
+            connectionFactory: ct =>
+            {
+                ct.ThrowIfCancellationRequested();
+                return Task.FromResult(new MsSqlConnection("Server=localhost;Database=master;User Id=sa;Password=pass;Pooling=false;"));
+            });
+
+        // Pre-warm two idle connections so the validation path is exercised.
+        await pool.InjectIdleConnectionAsync(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;"));
+        Assert.Equal(2, ReadCount(pool));
+
+        var alreadyCancelled = new CancellationToken(canceled: true);
+
+        for (int i = 0; i < 5; i++)
+        {
+            await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+            {
+                await pool.AcquireAsync(alreadyCancelled);
+            });
+        }
+
+        // After all the cancelled acquires, the pool's accounting should still
+        // be correct — every connection that was pulled from _idle must have
+        // been disposed and decremented. With the old code _count would still
+        // be 2, leaving the pool unable to open new connections.
+        Assert.Equal(0, ReadCount(pool));
+    }
 }