Feat ADO.NET: Added dispose timeout (10 seconds) to PoolingSessionSource

Author: KirillKurdyukov

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,4 @@
+- Feat ADO.NET: Added dispose timeout (10 seconds) to `PoolingSessionSource`.
 - Feat ADO.NET: Added `EnableImplicitSession` to support implicit sessions.
 
 ## v0.23.1

src/Ydb.Sdk/src/Ado/Session/ImplicitSessionSource.cs

@@ -25,7 +25,7 @@ public ValueTask<ISession> OpenSession(CancellationToken cancellationToken)
 
         return TryAcquireLease()
             ? new ValueTask<ISession>(new ImplicitSession(_driver, this))
-            : throw new ObjectDisposedException(nameof(ImplicitSessionSource));
+            : throw new YdbException("The implicit session source has been shut down.");
     }
 
     private bool TryAcquireLease()

src/Ydb.Sdk/src/Ado/Session/PoolingSessionSource.cs

@@ -1,12 +1,16 @@
 using System.Collections.Concurrent;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
+using Microsoft.Extensions.Logging;
 using Ydb.Query;
 
 namespace Ydb.Sdk.Ado.Session;
 
 internal sealed class PoolingSessionSource<T> : ISessionSource where T : PoolingSessionBase<T>
 {
+    private const int DisposeTimeoutSeconds = 10;
+
     private readonly ConcurrentStack<T> _idleSessions = new();
     private readonly ConcurrentQueue<TaskCompletionSource<T?>> _waiters = new();
     private readonly CancellationTokenSource _disposeCts = new();
@@ -18,16 +22,14 @@ internal sealed class PoolingSessionSource<T> : ISessionSource where T : Pooling
     private readonly int _createSessionTimeout;
     private readonly TimeSpan _sessionIdleTimeout;
     private readonly Timer _cleanerTimer;
+    private readonly ILogger _logger;
 
     private volatile int _numSessions;
     private volatile int _disposed;
 
     private bool IsDisposed => _disposed == 1;
 
-    public PoolingSessionSource(
-        IPoolingSessionFactory<T> sessionFactory,
-        YdbConnectionStringBuilder settings
-    )
+    public PoolingSessionSource(IPoolingSessionFactory<T> sessionFactory, YdbConnectionStringBuilder settings)
     {
         _sessionFactory = sessionFactory;
         _minSessionSize = settings.MinSessionPool;
@@ -43,6 +45,7 @@ YdbConnectionStringBuilder settings
         _createSessionTimeout = settings.CreateSessionTimeout;
         _sessionIdleTimeout = TimeSpan.FromSeconds(settings.SessionIdleTimeout);
         _cleanerTimer = new Timer(CleanIdleSessions, this, _sessionIdleTimeout, _sessionIdleTimeout);
+        _logger = settings.LoggerFactory.CreateLogger<PoolingSessionSource<T>>();
     }
 
     public ValueTask<ISession> OpenSession(CancellationToken cancellationToken = default)
@@ -156,7 +159,7 @@ private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
             try
             {
                 if (IsDisposed)
-                    throw new YdbException("The session source has been shut down.");
+                    throw new YdbException("The session source has been shutdown.");
 
                 var session = _sessionFactory.NewSession(this);
                 await session.Open(cancellationToken);
@@ -167,8 +170,7 @@ private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
                         return session;
                 }
 
-                throw new YdbException(
-                    $"Could not find free slot in {_sessions} when opening. Please report a bug.");
+                throw new YdbException($"Could not find free slot in {_sessions} when opening. Please report a bug.");
             }
             catch
             {
@@ -266,6 +268,7 @@ public async ValueTask DisposeAsync()
         await _cleanerTimer.DisposeAsync();
         _disposeCts.Cancel();
 
+        var sw = Stopwatch.StartNew();
         var spinWait = new SpinWait();
         do
         {
@@ -280,9 +283,28 @@ public async ValueTask DisposeAsync()
             }
 
             spinWait.SpinOnce();
-        } while (_numSessions > 0);
+        } while (_numSessions > 0 && sw.Elapsed < TimeSpan.FromSeconds(DisposeTimeoutSeconds));
+
+        try
+        {
+            await _sessionFactory.DisposeAsync();
+        }
+        catch (Exception e)
+        {
+            _logger.LogError(e, "Failed to dispose the transport driver");
+        }
 
-        await _sessionFactory.DisposeAsync();
+        for (var i = 0; i < _maxSessionSize; i++)
+        {
+            var session = Volatile.Read(ref _sessions[i]);
+
+            if (session == null || session.CompareAndSet(PoolingSessionState.Clean, PoolingSessionState.Clean))
+                continue;
+            _logger.LogCritical("Disposal timed out: Some sessions are still active");
+
+            throw new YdbException("Timeout while disposing of the pool: some sessions are still active. " +
+                                   "This may indicate a connection leak or suspended operations.");
+        }
     }
 }
 

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/Session/PoolingSessionSourceMockTests.cs

@@ -185,6 +185,7 @@ public async Task StressTest_DisposeAsync_Close_Driver()
                     {
                         using var session = await sessionSource.OpenSession();
                         await Task.Yield();
+                        Assert.False(disposeCalled);
                     }
                     catch (YdbException e)
                     {
@@ -201,6 +202,34 @@ public async Task StressTest_DisposeAsync_Close_Driver()
             Assert.True(disposeCalled);
         }
     }
+    
+    [Fact]
+    public async Task DisposeAsync_WhenSessionIsLeaked_ThrowsYdbExceptionWithTimeoutMessage()
+    {
+        var disposeCalled = false;
+        const int maxSessionSize = 10;
+        var mockFactory = new MockPoolingSessionFactory(maxSessionSize)
+        {
+            Dispose = () =>
+            {
+                Volatile.Write(ref disposeCalled, true);
+                return ValueTask.CompletedTask;
+            }
+        };
+        var settings = new YdbConnectionStringBuilder {  MaxSessionPool = maxSessionSize };
+        var sessionSource = new PoolingSessionSource<MockPoolingSession>(mockFactory, settings);
+        
+#pragma warning disable CA2012
+        _ = sessionSource.OpenSession(CancellationToken.None);
+#pragma warning restore CA2012
+
+        Assert.Equal("Timeout while disposing of the pool: some sessions are still active. " +
+                     "This may indicate a connection leak or suspended operations.",
+            (await Assert.ThrowsAsync<YdbException>(async () => await sessionSource.DisposeAsync())).Message);
+        Assert.True(disposeCalled);
+        Assert.Equal("The session source has been shut down.", (await Assert.ThrowsAsync<YdbException>(
+            () => sessionSource.OpenSession(CancellationToken.None).AsTask())).Message);
+    }
 
     [Fact]
     public async Task IdleTimeout_MinSessionSize_CloseNumSessionsMinusMinSessionCount()

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/Session/YdbImplicitStressTests.cs

@@ -32,6 +32,7 @@ public async Task StressTest_OpenSession_RaceWithDispose_SuccessfulOpensAreNotDi
                 try
                 {
                     using var s = await source.OpenSession(CancellationToken.None);
+                    await Task.Yield();
                     Assert.False(_isDisposed);
                 }
                 catch (ObjectDisposedException)
@@ -65,6 +66,7 @@ public async Task DisposeAsync_WhenSessionIsLeaked_ThrowsYdbExceptionWithTimeout
                      "This may indicate a connection leak or suspended operations.",
             (await Assert.ThrowsAsync<YdbException>(async () => await source.DisposeAsync())).Message);
         Assert.True(_isDisposed);
-        await Assert.ThrowsAsync<ObjectDisposedException>(() => source.OpenSession(CancellationToken.None).AsTask());
+        Assert.Equal("The implicit session source has been shut down.", (await Assert.ThrowsAsync<YdbException>(
+            () => source.OpenSession(CancellationToken.None).AsTask())).Message);
     }
 }