fix tests

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);
+        using var ctsGetSession = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _disposeCts.Token);
         if (_createSessionTimeout > 0)
             ctsGetSession.CancelAfter(TimeSpan.FromSeconds(_createSessionTimeout));
 
@@ -133,10 +133,6 @@ private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
                 () => waiterTcs.TrySetCanceled(),
                 useSynchronizationContext: false
             );
-            await using var disposedCancellationTokenRegistration = _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))
@@ -300,7 +296,7 @@ internal abstract class PoolingSessionBase<T> : ISession where T : PoolingSessio
 {
     private readonly PoolingSessionSource<T> _source;
 
-    private int _state = (int)PoolingSessionState.In;
+    private int _state = (int)PoolingSessionState.Out;
 
     protected PoolingSessionBase(PoolingSessionSource<T> source)
     {

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

@@ -104,7 +104,52 @@ public async Task HighContention_OpenClose_NotCanceledException()
     }
 
     [Fact]
-    public async Task DisposeAsync_Close_Driver()
+    public async Task DisposeAsync_Cancel_WaitersSession()
+    {
+        const int maxSessionSize = 10;
+        var mockFactory = new MockPoolingSessionFactory(maxSessionSize);
+        var sessionSource = new PoolingSessionSource<MockPoolingSession>(
+            mockFactory, new YdbConnectionStringBuilder { MaxSessionPool = maxSessionSize }
+        );
+
+        var openSessions = new List<ISession>();
+        var waitingSessionTasks = new List<Task>();
+        for (var i = 0; i < maxSessionSize; i++)
+        {
+            openSessions.Add(await sessionSource.OpenSession());
+        }
+
+        for (var i = 0; i < maxSessionSize; i++)
+        {
+            waitingSessionTasks.Add(Task.Run(async () =>
+            {
+                var session = await sessionSource.OpenSession();
+                await session.Close();
+            }));
+        }
+
+        await sessionSource.DisposeAsync();
+        Assert.Equal(maxSessionSize, mockFactory.NumSession);
+        Assert.Equal("Session Source is disposed.",
+            (await Assert.ThrowsAsync<YdbException>(async () => await sessionSource.OpenSession())).Message);
+
+        for (var i = 0; i < maxSessionSize; i++)
+        {
+            await Assert.ThrowsAsync<TaskCanceledException>(() => waitingSessionTasks[i]);
+        }
+
+        for (var i = 0; i < maxSessionSize; i++)
+        {
+            await openSessions[i].Close();
+        }
+
+        Assert.Equal(0, mockFactory.NumSession);
+        Assert.Equal("Session Source is disposed.",
+            (await Assert.ThrowsAsync<YdbException>(async () => await sessionSource.OpenSession())).Message);
+    }
+
+    [Fact]
+    public async Task StressTest_DisposeAsync_Close_Driver()
     {
         const int contentionTasks = 200;
         const int maxSessionSize = 100;
@@ -113,7 +158,6 @@ public async Task DisposeAsync_Close_Driver()
             var disposeCalled = false;
             var mockFactory = new MockPoolingSessionFactory(maxSessionSize)
             {
-                Open = async _ => await Task.Yield(),
                 Dispose = () =>
                 {
                     Volatile.Write(ref disposeCalled, true);
@@ -146,6 +190,9 @@ public async Task DisposeAsync_Close_Driver()
                     {
                         Assert.Equal("Session Source is disposed.", e.Message);
                     }
+                    catch (OperationCanceledException)
+                    {
+                    }
                 }));
             }