feat ADO.NET: create session timeout & cancellation token propagation (#459)

- ADO.NET: session is now deactivated when cancelled.
- Fixed bug ADO.NET: throws an `InvalidOperationException` if the connection is broken during the next invocation.
- Fixed bug `YdbCommand`: `Execute*` methods now propagate the cancellation token only for initializing YdbDataReader; the token is not passed to the server stream.
- `YdbCommand`: Improved cancellation token propagation in `Execute*` methods.
- `YdbConnection`: Added cancellation token propagation support in `OpenAsync`.
- `YdbDataReader`: Added cancellation token propagation support in `ReadAsync` and `NextResultAsync`.
- Added `CreateSessionTimeout` option to ADO.NET session creation; default is 5 seconds.

Author: KirillKurdyukov

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,11 @@
+- ADO.NET: session is now deactivated when cancelled.
+- Fixed bug ADO.NET: throws an `InvalidOperationException` if the connection is broken during the next invocation.
+- Fixed bug `YdbCommand`: `Execute*` methods now propagate the cancellation token only for initializing YdbDataReader; the token is not passed to the server stream.
+- `YdbCommand`: Improved cancellation token propagation in `Execute*` methods.
+- `YdbConnection`: Added cancellation token propagation support in `OpenAsync`.
+- `YdbDataReader`: Added cancellation token propagation support in `ReadAsync` and `NextResultAsync`.
+- Added `CreateSessionTimeout` option to ADO.NET session creation; default is 5 seconds.
+
 ## v0.18.3
 
 - Added `ConnectTimeout`: time to wait (in seconds) while trying to establish a connection.

src/Ydb.Sdk/src/Ado/PoolManager.cs

@@ -1,4 +1,5 @@
 using System.Collections.Concurrent;
+using Ydb.Sdk.Pool;
 using Ydb.Sdk.Services.Query;
 
 namespace Ydb.Sdk.Ado;
@@ -15,7 +16,7 @@ CancellationToken cancellationToken
     {
         if (Pools.TryGetValue(connectionString.ConnectionString, out var sessionPool))
         {
-            return await sessionPool.GetSession();
+            return await sessionPool.GetSession(cancellationToken);
         }
 
         try
@@ -24,15 +25,21 @@ CancellationToken cancellationToken
 
             if (Pools.TryGetValue(connectionString.ConnectionString, out var pool))
             {
-                return await pool.GetSession();
+                return await pool.GetSession(cancellationToken);
             }
 
-            var newSessionPool = new SessionPool(await connectionString.BuildDriver(), connectionString.MaxSessionPool,
-                disposingDriver: true);
+            var newSessionPool = new SessionPool(
+                await connectionString.BuildDriver(),
+                new SessionPoolConfig(
+                    MaxSessionPool: connectionString.MaxSessionPool,
+                    CreateSessionTimeout: connectionString.CreateSessionTimeout,
+                    DisposeDriver: true
+                )
+            );
 
             Pools[connectionString.ConnectionString] = newSessionPool;
 
-            return await newSessionPool.GetSession();
+            return await newSessionPool.GetSession(cancellationToken);
         }
         finally
         {

src/Ydb.Sdk/src/Ado/YdbAdoDefaultSettings.cs

@@ -0,0 +1,12 @@
+namespace Ydb.Sdk.Ado;
+
+internal static class YdbAdoDefaultSettings
+{
+    internal const string Host = "localhost";
+
+    internal const int Port = 2136;
+
+    internal const string Database = "/local";
+
+    internal const bool UseTls = false;
+}

src/Ydb.Sdk/src/Ado/YdbCommand.cs

@@ -174,7 +174,7 @@ protected override async Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBeha
             throw new YdbOperationInProgressException(YdbConnection);
         }
 
-        YdbConnection.EnsureConnectionOpen();
+        YdbConnection.ThrowIfConnectionClosed();
 
         var ydbParameters = DbParameterCollection.YdbParameters;
         var (sql, paramNames) = SqlParser.Parse(
@@ -201,7 +201,6 @@ protected override async Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBeha
         var execSettings = CommandTimeout > 0
             ? new ExecuteQuerySettings { TransportTimeout = TimeSpan.FromSeconds(CommandTimeout) }
             : new ExecuteQuerySettings();
-        execSettings.CancellationToken = cancellationToken;
 
         var transaction = YdbConnection.CurrentTransaction;
 
@@ -211,9 +210,9 @@ protected override async Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBeha
         }
 
         var ydbDataReader = await YdbDataReader.CreateYdbDataReader(
-            await YdbConnection.Session.ExecuteQuery(
-                preparedSql.ToString(), ydbParameters, execSettings, transaction?.TransactionControl
-            ), YdbConnection.OnStatus, transaction
+            await YdbConnection.Session
+                .ExecuteQuery(preparedSql.ToString(), ydbParameters, execSettings, transaction?.TransactionControl),
+            YdbConnection.OnStatus, transaction, cancellationToken
         );
 
         YdbConnection.LastReader = ydbDataReader;

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

@@ -28,7 +28,7 @@ internal Session Session
     {
         get
         {
-            EnsureConnectionOpen();
+            ThrowIfConnectionClosed();
 
             return _session;
         }
@@ -53,7 +53,7 @@ public YdbConnection(YdbConnectionStringBuilder connectionStringBuilder)
 
     protected override YdbTransaction BeginDbTransaction(IsolationLevel isolationLevel)
     {
-        EnsureConnectionOpen();
+        ThrowIfConnectionClosed();
 
         return BeginTransaction(isolationLevel switch
         {
@@ -66,7 +66,7 @@ protected override YdbTransaction BeginDbTransaction(IsolationLevel isolationLev
 
     public YdbTransaction BeginTransaction(TxMode txMode = TxMode.SerializableRw)
     {
-        EnsureConnectionOpen();
+        ThrowIfConnectionClosed();
 
         if (CurrentTransaction is { Completed: false })
         {
@@ -90,7 +90,7 @@ public override void ChangeDatabase(string databaseName)
 
     public override async Task OpenAsync(CancellationToken cancellationToken)
     {
-        EnsureConnectionClosed();
+        ThrowIfConnectionOpen();
 
         try
         {
@@ -100,6 +100,11 @@ public override async Task OpenAsync(CancellationToken cancellationToken)
         {
             throw e switch
             {
+                OperationCanceledException => throw new YdbException(StatusCode.Cancelled,
+                    $"The connection pool has been exhausted, either raise 'MaxSessionPool' " +
+                    $"(currently {ConnectionStringBuilder.MaxSessionPool}) or 'CreateSessionTimeout' " +
+                    $"(currently {ConnectionStringBuilder.CreateSessionTimeout} seconds) in your connection string.", e
+                ),
                 Driver.TransportException transportException => new YdbException(transportException),
                 StatusUnsuccessfulException unsuccessfulException => new YdbException(unsuccessfulException.Status),
                 _ => e
@@ -147,7 +152,7 @@ public override string ConnectionString
         set
 #pragma warning restore CS8765 // Nullability of type of parameter doesn't match overridden member (possibly because of nullability attributes).
         {
-            EnsureConnectionClosed();
+            ThrowIfConnectionOpen();
 
             // ReSharper disable once ConditionIsAlwaysTrueOrFalseAccordingToNullableAPIContract
             _connectionStringBuilder = value != null ? new YdbConnectionStringBuilder(value) : null;
@@ -181,7 +186,7 @@ public override string ServerVersion
     {
         get
         {
-            EnsureConnectionOpen();
+            ThrowIfConnectionClosed();
 
             return string.Empty; // TODO ServerVersion
         }
@@ -212,17 +217,17 @@ public override Task<DataTable> GetSchemaAsync(
         CancellationToken cancellationToken = default
     ) => YdbSchema.GetSchemaAsync(this, collectionName, restrictionValues, cancellationToken);
 
-    internal void EnsureConnectionOpen()
+    internal void ThrowIfConnectionClosed()
     {
-        if (ConnectionState == ConnectionState.Closed)
+        if (ConnectionState is ConnectionState.Closed or ConnectionState.Broken)
         {
             throw new InvalidOperationException("Connection is closed");
         }
     }
 
-    private void EnsureConnectionClosed()
+    private void ThrowIfConnectionOpen()
     {
-        if (ConnectionState != ConnectionState.Closed)
+        if (ConnectionState == ConnectionState.Open)
         {
             throw new InvalidOperationException("Connection already open");
         }

src/Ydb.Sdk/src/Ado/YdbConnectionStringBuilder.cs

@@ -4,6 +4,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using Ydb.Sdk.Auth;
+using Ydb.Sdk.Pool;
 using Ydb.Sdk.Transport;
 
 namespace Ydb.Sdk.Ado;
@@ -24,18 +25,19 @@ public YdbConnectionStringBuilder(string connectionString)
     // Init default connection string
     private void InitDefaultValues()
     {
-        _host = "localhost";
-        _port = 2136;
-        _database = "/local";
-        _maxSessionPool = 100;
-        _useTls = false;
+        _host = YdbAdoDefaultSettings.Host;
+        _port = YdbAdoDefaultSettings.Port;
+        _database = YdbAdoDefaultSettings.Database;
+        _maxSessionPool = SessionPoolDefaultSettings.MaxSessionPool;
+        _useTls = YdbAdoDefaultSettings.UseTls;
         _connectTimeout = GrpcDefaultSettings.ConnectTimeoutSeconds;
-        _keepAlivePingDelay = GrpcDefaultSettings.DefaultKeepAlivePingSeconds;
-        _keepAlivePingTimeout = GrpcDefaultSettings.DefaultKeepAlivePingTimeoutSeconds;
-        _enableMultipleHttp2Connections = false;
+        _keepAlivePingDelay = GrpcDefaultSettings.KeepAlivePingSeconds;
+        _keepAlivePingTimeout = GrpcDefaultSettings.KeepAlivePingTimeoutSeconds;
+        _enableMultipleHttp2Connections = GrpcDefaultSettings.EnableMultipleHttp2Connections;
         _maxSendMessageSize = GrpcDefaultSettings.MaxSendMessageSize;
         _maxReceiveMessageSize = GrpcDefaultSettings.MaxReceiveMessageSize;
-        _disableDiscovery = false;
+        _disableDiscovery = GrpcDefaultSettings.DisableDiscovery;
+        _createSessionTimeout = SessionPoolDefaultSettings.CreateSessionTimeoutSeconds;
     }
 
     public string Host
@@ -246,6 +248,24 @@ public bool DisableDiscovery
 
     private bool _disableDiscovery;
 
+    public int CreateSessionTimeout
+    {
+        get => _createSessionTimeout;
+        set
+        {
+            if (value < 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(value), value,
+                    "Invalid create session timeout: " + value);
+            }
+
+            _createSessionTimeout = value;
+            SaveValue(nameof(CreateSessionTimeout), value);
+        }
+    }
+
+    private int _createSessionTimeout;
+
     public ILoggerFactory? LoggerFactory { get; init; }
 
     public ICredentialsProvider? CredentialsProvider { get; init; }
@@ -410,6 +430,9 @@ static YdbConnectionOption()
                 "MaxReceiveMessageSize", "Max Receive Message Size");
             AddOption(new YdbConnectionOption<bool>(BoolExtractor, (builder, disableDiscovery) =>
                 builder.DisableDiscovery = disableDiscovery), "DisableDiscovery", "Disable Discovery");
+            AddOption(new YdbConnectionOption<int>(IntExtractor,
+                    (builder, createSessionTimeout) => builder.CreateSessionTimeout = createSessionTimeout),
+                "CreateSessionTimeout", "Create Session Timeout");
         }
 
         private static void AddOption(YdbConnectionOption option, params string[] keys)

src/Ydb.Sdk/src/Ado/YdbDataReader.cs

@@ -8,7 +8,7 @@ namespace Ydb.Sdk.Ado;
 
 public sealed class YdbDataReader : DbDataReader, IAsyncEnumerable<YdbDataRecord>
 {
-    private readonly IAsyncEnumerator<ExecuteQueryResponsePart> _stream;
+    private readonly IServerStream<ExecuteQueryResponsePart> _stream;
     private readonly YdbTransaction? _ydbTransaction;
     private readonly RepeatedField<IssueMessage> _issueMessagesInStream = new();
     private readonly Action<Status> _onNotSuccessStatus;
@@ -53,7 +53,7 @@ private enum State
     internal bool IsOpen => ReaderState is State.NewResultSet or State.ReadResultSet;
 
     private YdbDataReader(
-        IAsyncEnumerator<ExecuteQueryResponsePart> resultSetStream,
+        IServerStream<ExecuteQueryResponsePart> resultSetStream,
         Action<Status> onNotSuccessStatus,
         YdbTransaction? ydbTransaction)
     {
@@ -63,19 +63,21 @@ private YdbDataReader(
     }
 
     internal static async Task<YdbDataReader> CreateYdbDataReader(
-        IAsyncEnumerator<ExecuteQueryResponsePart> resultSetStream,
+        IServerStream<ExecuteQueryResponsePart> resultSetStream,
         Action<Status> onStatus,
-        YdbTransaction? ydbTransaction = null)
+        YdbTransaction? ydbTransaction = null,
+        CancellationToken cancellationToken = default
+    )
     {
         var ydbDataReader = new YdbDataReader(resultSetStream, onStatus, ydbTransaction);
-        await ydbDataReader.Init();
+        await ydbDataReader.Init(cancellationToken);
 
         return ydbDataReader;
     }
 
-    private async Task Init()
+    private async Task Init(CancellationToken cancellationToken)
     {
-        if (State.IsConsumed == await NextExecPart())
+        if (State.IsConsumed == await NextExecPart(cancellationToken))
         {
             throw new YdbException("YDB server closed the stream");
         }
@@ -448,7 +450,7 @@ public override async Task<bool> NextResultAsync(CancellationToken cancellationT
             State.ReadResultSet => await new Func<Task<State>>(async () =>
             {
                 State state;
-                while ((state = await NextExecPart()) == State.ReadResultSet)
+                while ((state = await NextExecPart(cancellationToken)) == State.ReadResultSet)
                 {
                 }
 
@@ -475,7 +477,7 @@ public override async Task<bool> ReadAsync(CancellationToken cancellationToken)
             return true;
         }
 
-        while ((ReaderState = await NextExecPart()) == State.ReadResultSet) // reset _currentRowIndex
+        while ((ReaderState = await NextExecPart(cancellationToken)) == State.ReadResultSet) // reset _currentRowIndex
         {
             if (++_currentRowIndex < RowsCount)
             {
@@ -521,7 +523,7 @@ public override async Task CloseAsync()
         }
 
         _onNotSuccessStatus(new Status(StatusCode.SessionBusy));
-        await _stream.DisposeAsync();
+        _stream.Dispose();
 
         if (_ydbTransaction != null)
         {
@@ -542,13 +544,13 @@ private YdbValue GetFieldYdbValue(int ordinal)
             : ydbValue;
     }
 
-    private async ValueTask<State> NextExecPart()
+    private async ValueTask<State> NextExecPart(CancellationToken cancellationToken)
     {
         try
         {
             _currentRowIndex = -1;
 
-            if (!await _stream.MoveNextAsync())
+            if (!await _stream.MoveNextAsync(cancellationToken))
             {
                 return State.IsConsumed;
             }
@@ -561,7 +563,7 @@ private async ValueTask<State> NextExecPart()
             {
                 OnFailReadStream();
 
-                while (await _stream.MoveNextAsync())
+                while (await _stream.MoveNextAsync(cancellationToken))
                 {
                     _issueMessagesInStream.AddRange(_stream.Current.Issues);
                 }

src/Ydb.Sdk/src/Ado/YdbException.cs

@@ -14,10 +14,15 @@ internal YdbException(Driver.TransportException transportException)
     }
 
     internal YdbException(Status status, Exception? innerException = null)
-        : base(status.ToString(), innerException)
+        : this(status.StatusCode, status.ToString(), innerException)
     {
-        Code = status.StatusCode;
-        var policy = RetrySettings.DefaultInstance.GetRetryRule(status.StatusCode).Policy;
+    }
+
+    internal YdbException(StatusCode statusCode, string message, Exception? innerException = null)
+        : base(message, innerException)
+    {
+        Code = statusCode;
+        var policy = RetrySettings.DefaultInstance.GetRetryRule(statusCode).Policy;
 
         IsTransient = policy == RetryPolicy.Unconditional;
         IsTransientWhenIdempotent = policy != RetryPolicy.None;

src/Ydb.Sdk/src/DriverConfig.cs

@@ -14,15 +14,15 @@ public class DriverConfig
         TimeSpan.FromSeconds(GrpcDefaultSettings.ConnectTimeoutSeconds);
 
     public TimeSpan KeepAlivePingDelay { get; init; } =
-        TimeSpan.FromSeconds(GrpcDefaultSettings.DefaultKeepAlivePingSeconds);
+        TimeSpan.FromSeconds(GrpcDefaultSettings.KeepAlivePingSeconds);
 
     public TimeSpan KeepAlivePingTimeout { get; init; } =
-        TimeSpan.FromSeconds(GrpcDefaultSettings.DefaultKeepAlivePingTimeoutSeconds);
+        TimeSpan.FromSeconds(GrpcDefaultSettings.KeepAlivePingTimeoutSeconds);
 
     public string? User { get; init; }
     public string? Password { get; init; }
 
-    public bool EnableMultipleHttp2Connections { get; init; }
+    public bool EnableMultipleHttp2Connections { get; init; } = GrpcDefaultSettings.EnableMultipleHttp2Connections;
 
     public int MaxSendMessageSize { get; init; } = GrpcDefaultSettings.MaxSendMessageSize;
     public int MaxReceiveMessageSize { get; init; } = GrpcDefaultSettings.MaxReceiveMessageSize;