dev: Prepare ISession for switching implementation (#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

@@ -1,6 +1,7 @@
 using System.Data;
 using System.Data.Common;
 using System.Diagnostics.CodeAnalysis;
+using Ydb.Sdk.Ado.Session;
 using Ydb.Sdk.Services.Query;
 using static System.Data.IsolationLevel;
 
@@ -24,7 +25,7 @@ private YdbConnectionStringBuilder ConnectionStringBuilder
         [param: AllowNull] init => _connectionStringBuilder = value;
     }
 
-    internal Services.Query.Session Session
+    internal ISession Session
     {
         get
         {
@@ -35,7 +36,7 @@ internal Services.Query.Session Session
         private set => _session = value;
     }
 
-    private Services.Query.Session _session = null!;
+    private ISession _session = null!;
 
     public YdbConnection()
     {
@@ -91,8 +92,18 @@ public override void ChangeDatabase(string databaseName)
     public override async Task OpenAsync(CancellationToken cancellationToken)
     {
         ThrowIfConnectionOpen();
-
-        Session = await PoolManager.GetSession(ConnectionStringBuilder, cancellationToken);
+        try
+        {
+            Session = await PoolManager.GetSession(ConnectionStringBuilder, cancellationToken);
+        }
+        catch (OperationCanceledException e)
+        {
+            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
+            );
+        }
 
         OnStateChange(ClosedToOpenEventArgs);
 
@@ -124,7 +135,7 @@ public override async Task CloseAsync()
         }
         finally
         {
-            await _session.Release();
+            _session.Close();
         }
     }
 
@@ -152,7 +163,7 @@ internal void OnNotSuccessStatusCode(StatusCode code)
     {
         _session.OnNotSuccessStatusCode(code);
 
-        if (!_session.IsActive)
+        if (_session.IsBroken)
         {
             ConnectionState = ConnectionState.Broken;
         }

src/Ydb.Sdk/src/Ado/YdbConnection.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;