dev: Prepare ISession for switching implementation (ydb-platform#481)

- Optimization: On BadSession, do not invoke the `DeleteSession()` method.
- Canceling AttachStream after calling the `DeleteSession` method.
- Fixed bug: fixed issue where session was not deleted (`ClientTransportTimeout`).

Author: KirillKurdyukov

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,6 @@
+- Optimization: On BadSession, do not invoke the `DeleteSession()` method.
+- Canceling AttachStream after calling the `DeleteSession` method.
+- Fixed bug: fixed issue where session was not deleted (`ClientTransportTimeout`).
 - Fixed bug: Grpc.Core.StatusCode.Cancelled was mapped to server's Canceled status.
 - ADO.NET: PoolingSessionSource 2.0 based on Npgsql pooling algorithm.
 - Added new ADO.NET options:

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

@@ -1,4 +1,5 @@
 using System.Collections.Concurrent;
+using Ydb.Sdk.Ado.Session;
 using Ydb.Sdk.Pool;
 using Ydb.Sdk.Services.Query;
 
@@ -9,7 +10,7 @@ internal static class PoolManager
     private static readonly SemaphoreSlim SemaphoreSlim = new(1); // async mutex
     private static readonly ConcurrentDictionary<string, SessionPool> Pools = new();
 
-    internal static async Task<Services.Query.Session> GetSession(
+    internal static async Task<ISession> GetSession(
         YdbConnectionStringBuilder connectionString,
         CancellationToken cancellationToken
     )

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

@@ -6,6 +6,8 @@ namespace Ydb.Sdk.Ado.Session;
 
 internal interface ISession
 {
+    IDriver Driver { get; }
+
     bool IsBroken { get; }
 
     ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(

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

@@ -6,13 +6,12 @@ namespace Ydb.Sdk.Ado.Session;
 
 internal class ImplicitSession : ISession
 {
-    private readonly IDriver _driver;
-
     public ImplicitSession(IDriver driver)
     {
-        _driver = driver;
+        Driver = driver;
     }
 
+    public IDriver Driver { get; }
     public bool IsBroken => false;
 
     public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
@@ -36,7 +35,7 @@ public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
         };
         request.Parameters.Add(parameters.ToDictionary(p => p.Key, p => p.Value.GetProto()));
 
-        return _driver.ServerStreamCall(QueryService.ExecuteQueryMethod, request, settings);
+        return Driver.ServerStreamCall(QueryService.ExecuteQueryMethod, request, settings);
     }
 
     public Task CommitTransaction(string txId, CancellationToken cancellationToken = default) =>

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

@@ -15,17 +15,19 @@ internal class PoolingSession : IPoolingSession
     private static readonly TimeSpan DeleteSessionTimeout = TimeSpan.FromSeconds(5);
     private static readonly CreateSessionRequest CreateSessionRequest = new();
 
-    private readonly IDriver _driver;
     private readonly PoolingSessionSource _poolingSessionSource;
     private readonly ILogger<PoolingSession> _logger;
+    private readonly CancellationTokenSource _attachStreamLifecycleCts = new();
 
     private volatile bool _isBroken = true;
+    private volatile bool _isBadSession;
 
     private readonly bool _disableServerBalancer;
 
     private string SessionId { get; set; } = string.Empty;
     private long NodeId { get; set; }
 
+    public IDriver Driver { get; }
     public bool IsBroken => _isBroken;
 
     internal PoolingSession(
@@ -35,10 +37,10 @@ internal PoolingSession(
         ILogger<PoolingSession> logger
     )
     {
-        _driver = driver;
         _poolingSessionSource = poolingSessionSource;
         _disableServerBalancer = disableServerBalancer;
         _logger = logger;
+        Driver = driver;
     }
 
     public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
@@ -60,15 +62,15 @@ public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
         };
         request.Parameters.Add(parameters.ToDictionary(p => p.Key, p => p.Value.GetProto()));
 
-        return _driver.ServerStreamCall(QueryService.ExecuteQueryMethod, request, settings);
+        return Driver.ServerStreamCall(QueryService.ExecuteQueryMethod, request, settings);
     }
 
     public async Task CommitTransaction(
         string txId,
         CancellationToken cancellationToken = default
     )
     {
-        var response = await _driver.UnaryCall(
+        var response = await Driver.UnaryCall(
             QueryService.CommitTransactionMethod,
             new CommitTransactionRequest { SessionId = SessionId, TxId = txId },
             new GrpcRequestSettings { CancellationToken = cancellationToken, NodeId = NodeId }
@@ -85,7 +87,7 @@ public async Task RollbackTransaction(
         CancellationToken cancellationToken = default
     )
     {
-        var response = await _driver.UnaryCall(
+        var response = await Driver.UnaryCall(
             QueryService.RollbackTransactionMethod,
             new RollbackTransactionRequest { SessionId = SessionId, TxId = txId },
             new GrpcRequestSettings { CancellationToken = cancellationToken, NodeId = NodeId }
@@ -99,6 +101,8 @@ public async Task RollbackTransaction(
 
     public void OnNotSuccessStatusCode(StatusCode statusCode)
     {
+        _isBadSession = _isBadSession || statusCode is StatusCode.BadSession;
+
         if (statusCode is
             StatusCode.BadSession or
             StatusCode.SessionBusy or
@@ -121,7 +125,7 @@ public async Task Open(CancellationToken cancellationToken)
             requestSettings.ClientCapabilities.Add(SessionBalancer);
         }
 
-        var response = await _driver.UnaryCall(QueryService.CreateSessionMethod, CreateSessionRequest, requestSettings);
+        var response = await Driver.UnaryCall(QueryService.CreateSessionMethod, CreateSessionRequest, requestSettings);
 
         if (response.Status.IsNotSuccess())
         {
@@ -138,7 +142,7 @@ public async Task Open(CancellationToken cancellationToken)
         {
             try
             {
-                using var stream = await _driver.ServerStreamCall(
+                using var stream = await Driver.ServerStreamCall(
                     QueryService.AttachSessionMethod,
                     new AttachSessionRequest { SessionId = SessionId },
                     new GrpcRequestSettings { NodeId = NodeId }
@@ -161,10 +165,12 @@ public async Task Open(CancellationToken cancellationToken)
 
                 completeTask.SetResult();
 
+                var lifecycleAttachToken = _attachStreamLifecycleCts.Token;
+
                 try
                 {
                     // ReSharper disable once MethodSupportsCancellation
-                    while (await stream.MoveNextAsync())
+                    while (await stream.MoveNextAsync(lifecycleAttachToken))
                     {
                         var sessionState = stream.Current;
 
@@ -215,14 +221,16 @@ public async Task DeleteSession()
     {
         try
         {
-            if (_isBroken)
+            _isBroken = true;
+            _attachStreamLifecycleCts.CancelAfter(DeleteSessionTimeout);
+
+            if (_isBadSession)
             {
                 return;
             }
 
-            _isBroken = true;
-
-            var deleteSessionResponse = await _driver.UnaryCall(
+            _isBadSession = true;
+            var deleteSessionResponse = await Driver.UnaryCall(
                 QueryService.DeleteSessionMethod,
                 new DeleteSessionRequest { SessionId = SessionId },
                 new GrpcRequestSettings { TransportTimeout = DeleteSessionTimeout, NodeId = NodeId }

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

@@ -1,8 +1,3 @@
-// This file contains session pooling algorithms adapted from Npgsql
-// Original source: https://github.com/npgsql/npgsql
-// Copyright (c) 2002-2025, Npgsql
-// Licence https://github.com/npgsql/npgsql?tab=PostgreSQL-1-ov-file
-
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 using System.Threading.Channels;
@@ -92,45 +87,34 @@ private async ValueTask<IPoolingSession> RentAsync(CancellationToken cancellatio
 
         var finalToken = ctsGetSession.Token;
 
-        try
+        var session = await OpenNewSession(finalToken).ConfigureAwait(false);
+        if (session != null)
+            return session;
+
+        while (true)
         {
-            var session = await OpenNewSession(finalToken).ConfigureAwait(false);
-            if (session != null)
-                return session;
+            session = await _idleSessionReader.ReadAsync(finalToken).ConfigureAwait(false);
 
-            while (true)
+            if (CheckIdleSession(session))
             {
-                session = await _idleSessionReader.ReadAsync(finalToken).ConfigureAwait(false);
-
-                if (CheckIdleSession(session))
-                {
-                    return session;
-                }
+                return session;
+            }
 
-                // If we're here, our waiting attempt on the idle session channel was released with a null
-                // (or bad session), or we're in sync mode. Check again if a new idle session has appeared since we last checked.
-                if (TryGetIdleSession(out session))
-                {
-                    return session;
-                }
+            // If we're here, our waiting attempt on the idle session channel was released with a null
+            // (or bad session), or we're in sync mode. Check again if a new idle session has appeared since we last checked.
+            if (TryGetIdleSession(out session))
+            {
+                return session;
+            }
 
-                // We might have closed a session in the meantime and no longer be at max capacity
-                // so try to open a new session and if that fails, loop again.
-                session = await OpenNewSession(finalToken).ConfigureAwait(false);
-                if (session != null)
-                {
-                    return session;
-                }
+            // We might have closed a session in the meantime and no longer be at max capacity
+            // so try to open a new session and if that fails, loop again.
+            session = await OpenNewSession(finalToken).ConfigureAwait(false);
+            if (session != null)
+            {
+                return session;
             }
         }
-        catch (OperationCanceledException e)
-        {
-            throw new YdbException(StatusCode.Cancelled,
-                $"The connection pool has been exhausted, either raise 'MaxSessionPool' " +
-                $"(currently {_maxSessionSize}) or 'CreateSessionTimeout' " +
-                $"(currently {_createSessionTimeout} seconds) in your connection string.", e
-            );
-        }
     }
 
     private async ValueTask<IPoolingSession?> OpenNewSession(CancellationToken cancellationToken)

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

@@ -3,7 +3,6 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Text;
 using Ydb.Sdk.Ado.Internal;
-using Ydb.Sdk.Services.Query;
 using Ydb.Sdk.Value;
 
 namespace Ydb.Sdk.Ado;
@@ -206,8 +205,8 @@ protected override async Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBeha
         preparedSql.Append(sql);
 
         var execSettings = CommandTimeout > 0
-            ? new ExecuteQuerySettings { TransportTimeout = TimeSpan.FromSeconds(CommandTimeout) }
-            : new ExecuteQuerySettings();
+            ? new GrpcRequestSettings { TransportTimeout = TimeSpan.FromSeconds(CommandTimeout) }
+            : new GrpcRequestSettings();
 
         var transaction = YdbConnection.CurrentTransaction;
 

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

@@ -25,15 +25,15 @@ public YdbConnectionStringBuilder(string connectionString)
     // Init default connection string
     private void InitDefaultValues()
     {
-        _host = YdbAdoDefaultSettings.Host;
-        _port = YdbAdoDefaultSettings.Port;
-        _database = YdbAdoDefaultSettings.Database;
+        _host = "localhost";
+        _port = 2136;
+        _database = "/local";
         _minSessionPool = 0;
         _maxSessionPool = SessionPoolDefaultSettings.MaxSessionPool;
         _createSessionTimeout = SessionPoolDefaultSettings.CreateSessionTimeoutSeconds;
         _sessionIdleTimeout = 300;
         _sessionPruningInterval = 10;
-        _useTls = YdbAdoDefaultSettings.UseTls;
+        _useTls = false;
         _connectTimeout = GrpcDefaultSettings.ConnectTimeoutSeconds;
         _keepAlivePingDelay = GrpcDefaultSettings.KeepAlivePingSeconds;
         _keepAlivePingTimeout = GrpcDefaultSettings.KeepAlivePingTimeoutSeconds;

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

@@ -3,8 +3,11 @@
 using System.Globalization;
 using Ydb.Scheme;
 using Ydb.Scheme.V1;
+using Ydb.Sdk.Ado.Internal;
 using Ydb.Sdk.Ado.Schema;
 using Ydb.Sdk.Services.Table;
+using Ydb.Table;
+using Ydb.Table.V1;
 
 namespace Ydb.Sdk.Ado;
 
@@ -59,10 +62,28 @@ public static async Task<YdbTable> DescribeTable(
     {
         try
         {
-            var describeResponse = await ydbConnection.Session
-                .DescribeTable(WithSuffix(ydbConnection.Database) + tableName, describeTableSettings);
+            describeTableSettings ??= new DescribeTableSettings();
 
-            return new YdbTable(tableName, describeResponse);
+            var describeResponse = await ydbConnection.Session.Driver.UnaryCall(
+                TableService.DescribeTableMethod,
+                new DescribeTableRequest
+                {
+                    Path = WithSuffix(ydbConnection.Database) + tableName,
+                    IncludeTableStats = describeTableSettings.IncludeTableStats,
+                    IncludePartitionStats = describeTableSettings.IncludePartitionStats,
+                    IncludeShardKeyBounds = describeTableSettings.IncludeShardKeyBounds
+                },
+                describeTableSettings
+            );
+
+            if (describeResponse.Operation.Status.IsNotSuccess())
+            {
+                throw YdbException.FromServer(describeResponse.Operation.Status, describeResponse.Operation.Issues);
+            }
+
+            var describeResult = describeResponse.Operation.Result.Unpack<DescribeTableResult>();
+
+            return new YdbTable(tableName, describeResult);
         }
         catch (YdbException e)
         {
@@ -241,31 +262,20 @@ private static async Task AppendDescribeTable(
         string? tableType,
         Action<YdbTable, string> appendInTable)
     {
-        try
-        {
-            var describeResponse = await ydbConnection.Session
-                .DescribeTable(WithSuffix(ydbConnection.Database) + tableName, describeTableSettings);
-            var ydbTable = new YdbTable(tableName, describeResponse);
+        var ydbTable = await DescribeTable(ydbConnection, tableName, describeTableSettings);
 
-            var type = ydbTable.IsSystem
-                ? "SYSTEM_TABLE"
-                : ydbTable.Type switch
-                {
-                    YdbTable.TableType.Table => "TABLE",
-                    YdbTable.TableType.ColumnTable => "COLUMN_TABLE",
-                    YdbTable.TableType.ExternalTable => "EXTERNAL_TABLE",
-                    _ => throw new ArgumentOutOfRangeException(nameof(tableType))
-                };
-            if (type.IsPattern(tableType))
+        var type = ydbTable.IsSystem
+            ? "SYSTEM_TABLE"
+            : ydbTable.Type switch
             {
-                appendInTable(ydbTable, type);
-            }
-        }
-        catch (YdbException e)
+                YdbTable.TableType.Table => "TABLE",
+                YdbTable.TableType.ColumnTable => "COLUMN_TABLE",
+                YdbTable.TableType.ExternalTable => "EXTERNAL_TABLE",
+                _ => throw new ArgumentOutOfRangeException(nameof(tableType))
+            };
+        if (type.IsPattern(tableType))
         {
-            ydbConnection.OnNotSuccessStatusCode(e.Code);
-
-            throw;
+            appendInTable(ydbTable, type);
         }
     }