fix test

Author: KirillKurdyukov

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

@@ -92,7 +92,7 @@ private bool CheckIdleSession([NotNullWhen(true)] T? session)
 
     private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
     {
-        using var ctsGetSession = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _disposeCts.Token);
+        using var ctsGetSession = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
         if (_createSessionTimeout > 0)
             ctsGetSession.CancelAfter(TimeSpan.FromSeconds(_createSessionTimeout));
 
@@ -133,6 +133,10 @@ private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
                 () => waiterTcs.TrySetCanceled(),
                 useSynchronizationContext: false
             );
+            await using var disposeRegistration = _disposeCts.Token.Register(
+                () => waiterTcs.TrySetException(new YdbException("Session Source is disposed.")),
+                useSynchronizationContext: false
+            );
             session = await waiterTcs.Task.ConfigureAwait(false);
 
             if (CheckIdleSession(session) || TryGetIdleSession(out session))
@@ -151,6 +155,11 @@ private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
             if (Interlocked.CompareExchange(ref _numSessions, numSessions + 1, numSessions) != numSessions)
                 continue;
 
+            if (IsDisposed)
+            {
+                break;
+            }
+
             try
             {
                 var session = _sessionFactory.NewSession(this);
@@ -261,16 +270,6 @@ public async ValueTask DisposeAsync()
         await _cleanerTimer.DisposeAsync();
         _disposeCts.Cancel();
 
-        for (var i = 0; i < _maxSessionSize; i++)
-        {
-            var session = Volatile.Read(ref _sessions[i]);
-
-            if (session != null && session.CompareAndSet(PoolingSessionState.In, PoolingSessionState.Clean))
-            {
-                CloseSession(session);
-            }
-        }
-
         var spinWait = new SpinWait();
         do
         {

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

@@ -130,7 +130,7 @@ public async Task DisposeAsync_Cancel_WaitersSession()
 
         var disposeTask = Task.Run(async () => await sessionSource.DisposeAsync());
         Assert.Equal(maxSessionSize, mockFactory.NumSession);
-
+        await Task.Delay(5_000);
         for (var i = 0; i < maxSessionSize; i++)
         {
             openSessions[i].Close();
@@ -140,7 +140,8 @@ public async Task DisposeAsync_Cancel_WaitersSession()
         Assert.Equal(0, mockFactory.NumSession);
         for (var i = 0; i < maxSessionSize; i++)
         {
-            await Assert.ThrowsAnyAsync<Exception>(() => waitingSessionTasks[i]);
+            Assert.Equal("Session Source is disposed.",
+                (await Assert.ThrowsAsync<YdbException>(() => waitingSessionTasks[i])).Message);
         }
 
         Assert.Equal("Session Source is disposed.",