feat ADO.NET: create session timeout & cancellation token propagation

Author: KirillKurdyukov

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,8 @@
+- `YdbCommand`: Cancellation token propagation in `Execute*` methods
+- `YdbConnection`: Cancellation token propagation in `OpenAsync` method.
+- `YdbDataReader`: Cancellation token propagation in `ReadAsync` and `NextResultAsync` methods.
+- Added `CreateSessionTimeout` to ADO.NET when creating a session. (Default 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/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,7 +63,7 @@ private YdbDataReader(
     }
 
     internal static async Task<YdbDataReader> CreateYdbDataReader(
-        IAsyncEnumerator<ExecuteQueryResponsePart> resultSetStream,
+        IServerStream<ExecuteQueryResponsePart> resultSetStream,
         Action<Status> onStatus,
         YdbTransaction? ydbTransaction = null)
     {
@@ -75,7 +75,7 @@ internal static async Task<YdbDataReader> CreateYdbDataReader(
 
     private async Task Init()
     {
-        if (State.IsConsumed == await NextExecPart())
+        if (State.IsConsumed == await NextExecPart(CancellationToken.None))
         {
             throw new YdbException("YDB server closed the stream");
         }
@@ -448,7 +448,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 +475,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 +521,7 @@ public override async Task CloseAsync()
         }
 
         _onNotSuccessStatus(new Status(StatusCode.SessionBusy));
-        await _stream.DisposeAsync();
+        _stream.Dispose();
 
         if (_ydbTransaction != null)
         {
@@ -542,13 +542,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 +561,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/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;

src/Ydb.Sdk/src/GrpcDefaultSettings.cs

@@ -5,16 +5,20 @@ internal static class GrpcDefaultSettings
     /// <summary>
     /// Default interval (in seconds) for sending keep-alive ping messages.
     /// </summary>
-    internal const int DefaultKeepAlivePingSeconds = 10;
+    internal const int KeepAlivePingSeconds = 10;
 
     /// <summary>
     /// Default timeout (in seconds) for receiving a response to a keep-alive ping.
     /// </summary>
-    internal const int DefaultKeepAlivePingTimeoutSeconds = 10;
+    internal const int KeepAlivePingTimeoutSeconds = 10;
 
     internal const int MaxSendMessageSize = 64 * 1024 * 1024; // 64 Mb
 
     internal const int MaxReceiveMessageSize = 64 * 1024 * 1024; // 64 Mb
 
     internal const int ConnectTimeoutSeconds = 5;
+
+    internal const bool EnableMultipleHttp2Connections = false;
+
+    internal const bool DisableDiscovery = false;
 }

src/Ydb.Sdk/src/IDriver.cs

@@ -45,6 +45,13 @@ public interface IBidirectionalStream<in TRequest, out TResponse> : IDisposable
     public Task RequestStreamComplete();
 }
 
+public interface IServerStream<out TResponse> : IDisposable
+{
+    public ValueTask<bool> MoveNextAsync(CancellationToken cancellationToken = default);
+
+    public TResponse Current { get; }
+}
+
 public abstract class BaseDriver : IDriver
 {
     private readonly ICredentialsProvider? _credentialsProvider;
@@ -200,7 +207,7 @@ public async ValueTask DisposeAsync()
     }
 }
 
-public sealed class ServerStream<TResponse> : IAsyncEnumerator<TResponse>, IAsyncEnumerable<TResponse>
+public sealed class ServerStream<TResponse> : IServerStream<TResponse>
 {
     private readonly AsyncServerStreamingCall<TResponse> _stream;
     private readonly Action<RpcException> _rpcErrorAction;
@@ -211,18 +218,11 @@ internal ServerStream(AsyncServerStreamingCall<TResponse> stream, Action<RpcExce
         _rpcErrorAction = rpcErrorAction;
     }
 
-    public ValueTask DisposeAsync()
-    {
-        _stream.Dispose();
-
-        return default;
-    }
-
-    public async ValueTask<bool> MoveNextAsync()
+    public async ValueTask<bool> MoveNextAsync(CancellationToken cancellationToken = default)
     {
         try
         {
-            return await _stream.ResponseStream.MoveNext(CancellationToken.None);
+            return await _stream.ResponseStream.MoveNext(cancellationToken);
         }
         catch (RpcException e)
         {
@@ -234,7 +234,7 @@ public async ValueTask<bool> MoveNextAsync()
 
     public TResponse Current => _stream.ResponseStream.Current;
 
-    public IAsyncEnumerator<TResponse> GetAsyncEnumerator(CancellationToken cancellationToken = new()) => this;
+    public void Dispose() => _stream.Dispose();
 }
 
 internal class BidirectionalStream<TRequest, TResponse> : IBidirectionalStream<TRequest, TResponse>