YdbConnection.OpenAsync throws OperationCancelException when the CancellationToken is cancelled. (#501)

Author: KirillKurdyukov

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,4 @@
+- ADO.NET: `YdbConnection.OpenAsync` throws OperationCancelException when the `CancellationToken` is cancelled.
 - Feat ADO.NET: decimal type with arbitrary precision/scale ([#498](https://github.com/ydb-platform/ydb-dotnet-sdk/issues/498)).
 - Fixed bug: interval value parsing in microseconds and double instead of ticks ([#497](https://github.com/ydb-platform/ydb-dotnet-sdk/issues/497)).
 - ADO.NET: Changed `IBulkUpsertImporter.AddRowAsync` signature: `object?[] row` → `params object[]`.

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

@@ -126,9 +126,11 @@ private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
                 continue;
             }
 
-            await using var _ = finalToken.Register(
-                () => waiterTcs.TrySetCanceled(),
-                useSynchronizationContext: false
+            await using var _ = finalToken.Register(() => waiterTcs.TrySetException(
+                    new YdbException($"The connection pool has been exhausted, either raise 'MaxSessionPool' " +
+                                     $"(currently {_maxSessionSize}) or 'CreateSessionTimeout' " +
+                                     $"(currently {_createSessionTimeout} seconds) in your connection string.")
+                ), useSynchronizationContext: false
             );
             await using var disposeRegistration = _disposeCts.Token.Register(
                 () => waiterTcs.TrySetException(new YdbException("The session source has been shut down.")),

src/Ydb.Sdk/src/Ado/YdbConnection.cs

@@ -110,18 +110,8 @@ public override void ChangeDatabase(string databaseName)
     public override async Task OpenAsync(CancellationToken cancellationToken)
     {
         ThrowIfConnectionOpen();
-        try
-        {
-            Session = await PoolManager.GetSession(ConnectionStringBuilder, cancellationToken);
-        }
-        catch (OperationCanceledException e)
-        {
-            throw new YdbException(StatusCode.ClientTransportTimeout,
-                $"The connection pool has been exhausted, either raise 'MaxSessionPool' " +
-                $"(currently {ConnectionStringBuilder.MaxSessionPool}) or 'CreateSessionTimeout' " +
-                $"(currently {ConnectionStringBuilder.CreateSessionTimeout} seconds) in your connection string.", e
-            );
-        }
+
+        Session = await PoolManager.GetSession(ConnectionStringBuilder, cancellationToken);
 
         OnStateChange(ClosedToOpenEventArgs);
 

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Specification.Tests/YdbConnectionTests.cs

@@ -28,13 +28,4 @@ public override void ServerVersion_returns_value()
     {
         base.ServerVersion_returns_value();
     }
-
-    public override async Task OpenAsync_is_canceled()
-    {
-        await using var connection = CreateConnection();
-        connection.ConnectionString = ConnectionString;
-        var task = connection.OpenAsync(CanceledToken);
-        await Assert.ThrowsAnyAsync<YdbException>(() => task);
-        Assert.True(task.IsFaulted);
-    }
 }

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

@@ -295,6 +295,9 @@ public async Task StressTest_HighContention_OpenClose()
                 catch (OperationCanceledException)
                 {
                 }
+                catch (YdbException)
+                {
+                }
             }, cts.Token));
         }
 
@@ -316,7 +319,9 @@ public async Task Get_Session_From_Exhausted_Pool()
         var cts = new CancellationTokenSource();
         cts.CancelAfter(500);
 
-        await Assert.ThrowsAsync<TaskCanceledException>(async () => await sessionSource.OpenSession(cts.Token));
+        Assert.Equal("The connection pool has been exhausted, either raise 'MaxSessionPool' (currently 1) " +
+                     "or 'CreateSessionTimeout' (currently 5 seconds) in your connection string.",
+            (await Assert.ThrowsAsync<YdbException>(async () => await sessionSource.OpenSession(cts.Token))).Message);
         session.Close();
 
         Assert.Equal(1, mockFactory.NumSession);

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/YdbConnectionTests.cs

@@ -201,9 +201,7 @@ public async Task OpenAsync_WhenCancelTokenIsCanceled_ThrowYdbException()
         connection.ConnectionString = ConnectionString + ";MinSessionPool=1";
         using var cts = new CancellationTokenSource();
         cts.Cancel();
-        Assert.Equal("The connection pool has been exhausted, either raise 'MaxSessionPool' (currently 10) " +
-                     "or 'CreateSessionTimeout' (currently 5 seconds) in your connection string.",
-            (await Assert.ThrowsAsync<YdbException>(async () => await connection.OpenAsync(cts.Token))).Message);
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await connection.OpenAsync(cts.Token));
         Assert.Equal(ConnectionState.Closed, connection.State);
     }