feat: FIFO PoolingSessionSource (#487)

Author: KirillKurdyukov

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

@@ -11,11 +11,11 @@ internal static class PoolManager
     private static readonly ConcurrentDictionary<string, SessionPool> Pools = new();
 
     internal static async Task<ISession> GetSession(
-        YdbConnectionStringBuilder connectionString,
+        YdbConnectionStringBuilder settings,
         CancellationToken cancellationToken
     )
     {
-        if (Pools.TryGetValue(connectionString.ConnectionString, out var sessionPool))
+        if (Pools.TryGetValue(settings.ConnectionString, out var sessionPool))
         {
             return await sessionPool.GetSession(cancellationToken);
         }
@@ -24,21 +24,21 @@ CancellationToken cancellationToken
         {
             await SemaphoreSlim.WaitAsync(cancellationToken);
 
-            if (Pools.TryGetValue(connectionString.ConnectionString, out var pool))
+            if (Pools.TryGetValue(settings.ConnectionString, out var pool))
             {
                 return await pool.GetSession(cancellationToken);
             }
 
             var newSessionPool = new SessionPool(
-                await connectionString.BuildDriver(),
+                await settings.BuildDriver(),
                 new SessionPoolConfig(
-                    MaxSessionPool: connectionString.MaxSessionPool,
-                    CreateSessionTimeout: connectionString.CreateSessionTimeout,
+                    MaxSessionPool: settings.MaxSessionPool,
+                    CreateSessionTimeout: settings.CreateSessionTimeout,
                     DisposeDriver: true
                 )
             );
 
-            Pools[connectionString.ConnectionString] = newSessionPool;
+            Pools[settings.ConnectionString] = newSessionPool;
 
             return await newSessionPool.GetSession(cancellationToken);
         }

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

@@ -1,8 +1,6 @@
 namespace Ydb.Sdk.Ado.Session;
 
-internal interface ISessionSource<TSession> where TSession : ISession
+internal interface ISessionSource : IAsyncDisposable
 {
-    ValueTask<TSession> OpenSession(CancellationToken cancellationToken);
-
-    void Return(TSession session);
+    ValueTask<ISession> OpenSession(CancellationToken cancellationToken);
 }

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

@@ -8,42 +8,40 @@
 
 namespace Ydb.Sdk.Ado.Session;
 
-internal class PoolingSession : IPoolingSession
+internal class PoolingSession : PoolingSessionBase<PoolingSession>
 {
     private const string SessionBalancer = "session-balancer";
 
     private static readonly TimeSpan DeleteSessionTimeout = TimeSpan.FromSeconds(5);
     private static readonly CreateSessionRequest CreateSessionRequest = new();
 
-    private readonly PoolingSessionSource _poolingSessionSource;
     private readonly ILogger<PoolingSession> _logger;
+    private readonly bool _disableServerBalancer;
+
     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;
+    public override IDriver Driver { get; }
+    public override bool IsBroken => _isBroken;
 
     internal PoolingSession(
         IDriver driver,
-        PoolingSessionSource poolingSessionSource,
+        PoolingSessionSource<PoolingSession> poolingSessionSource,
         bool disableServerBalancer,
         ILogger<PoolingSession> logger
-    )
+    ) : base(poolingSessionSource)
     {
-        _poolingSessionSource = poolingSessionSource;
         _disableServerBalancer = disableServerBalancer;
         _logger = logger;
         Driver = driver;
     }
 
-    public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
+    public override ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
         string query,
         Dictionary<string, YdbValue> parameters,
         GrpcRequestSettings settings,
@@ -65,10 +63,7 @@ public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
         return Driver.ServerStreamCall(QueryService.ExecuteQueryMethod, request, settings);
     }
 
-    public async Task CommitTransaction(
-        string txId,
-        CancellationToken cancellationToken = default
-    )
+    public override async Task CommitTransaction(string txId, CancellationToken cancellationToken = default)
     {
         var response = await Driver.UnaryCall(
             QueryService.CommitTransactionMethod,
@@ -82,10 +77,7 @@ public async Task CommitTransaction(
         }
     }
 
-    public async Task RollbackTransaction(
-        string txId,
-        CancellationToken cancellationToken = default
-    )
+    public override async Task RollbackTransaction(string txId, CancellationToken cancellationToken = default)
     {
         var response = await Driver.UnaryCall(
             QueryService.RollbackTransactionMethod,
@@ -99,7 +91,7 @@ public async Task RollbackTransaction(
         }
     }
 
-    public void OnNotSuccessStatusCode(StatusCode statusCode)
+    public override void OnNotSuccessStatusCode(StatusCode statusCode)
     {
         _isBadSession = _isBadSession || statusCode is StatusCode.BadSession;
 
@@ -116,7 +108,7 @@ StatusCode.ClientTransportTimeout or
         }
     }
 
-    public async Task Open(CancellationToken cancellationToken)
+    internal override async Task Open(CancellationToken cancellationToken)
     {
         var requestSettings = new GrpcRequestSettings { CancellationToken = cancellationToken };
 
@@ -217,7 +209,7 @@ public async Task Open(CancellationToken cancellationToken)
         await completeTask.Task;
     }
 
-    public async Task DeleteSession()
+    internal override async Task DeleteSession()
     {
         try
         {
@@ -248,6 +240,4 @@ public async Task DeleteSession()
                 SessionId, NodeId);
         }
     }
-
-    public void Close() => _poolingSessionSource.Return(this);
 }

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

@@ -1,20 +1,26 @@
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Abstractions;
 
 namespace Ydb.Sdk.Ado.Session;
 
-internal class PoolingSessionFactory : IPoolingSessionFactory
+internal class PoolingSessionFactory : IPoolingSessionFactory<PoolingSession>
 {
     private readonly IDriver _driver;
     private readonly bool _disableServerBalancer;
     private readonly ILogger<PoolingSession> _logger;
 
-    public PoolingSessionFactory(IDriver driver, YdbConnectionStringBuilder settings, ILoggerFactory loggerFactory)
+    internal PoolingSessionFactory(IDriver driver, YdbConnectionStringBuilder settings, ILoggerFactory loggerFactory)
     {
         _driver = driver;
         _disableServerBalancer = settings.DisableServerBalancer;
         _logger = loggerFactory.CreateLogger<PoolingSession>();
     }
 
-    public IPoolingSession NewSession(PoolingSessionSource source) =>
-        new PoolingSession(_driver, source, _disableServerBalancer, _logger);
+    public static async Task<PoolingSessionFactory> Create(YdbConnectionStringBuilder settings) =>
+        new(await settings.BuildDriver(), settings, settings.LoggerFactory ?? NullLoggerFactory.Instance);
+
+    public PoolingSession NewSession(PoolingSessionSource<PoolingSession> source) =>
+        new(_driver, source, _disableServerBalancer, _logger);
+
+    public ValueTask DisposeAsync() => _driver.DisposeAsync();
 }