FIFO LOCK FREE POOL

Author: KirillKurdyukov

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

@@ -1,5 +1,4 @@
 using System.Collections.Concurrent;
-using Microsoft.Extensions.Logging.Abstractions;
 using Ydb.Sdk.Ado.Session;
 
 namespace Ydb.Sdk.Ado;
@@ -49,8 +48,8 @@ internal static async Task ClearPool(string connectionString)
             try
             {
                 await SemaphoreSlim.WaitAsync();
-                //
-                // await sessionPool.DisposeAsync();
+                
+                await sessionPool.DisposeAsync();
             }
             finally
             {

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

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

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

@@ -9,7 +9,7 @@ internal class PoolingSessionFactory : IPoolingSessionFactory<PoolingSession>
     private readonly bool _disableServerBalancer;
     private readonly ILogger<PoolingSession> _logger;
 
-    private PoolingSessionFactory(IDriver driver, YdbConnectionStringBuilder settings, ILoggerFactory loggerFactory)
+    internal PoolingSessionFactory(IDriver driver, YdbConnectionStringBuilder settings, ILoggerFactory loggerFactory)
     {
         _driver = driver;
         _disableServerBalancer = settings.DisableServerBalancer;
@@ -21,4 +21,6 @@ public static async Task<PoolingSessionFactory> Create(YdbConnectionStringBuilde
 
     public PoolingSession NewSession(PoolingSessionSource<PoolingSession> source) =>
         new(_driver, source, _disableServerBalancer, _logger);
+
+    public ValueTask DisposeAsync() => _driver.DisposeAsync();
 }

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

@@ -10,6 +10,7 @@ internal sealed class PoolingSessionSource<T> : ISessionSource where T : Pooling
 {
     private readonly ConcurrentStack<T> _idleSessions = new();
     private readonly ConcurrentQueue<TaskCompletionSource<T?>> _waiters = new();
+    private readonly CancellationTokenSource _disposeCts = new();
 
     private readonly IPoolingSessionFactory<T> _sessionFactory;
     private readonly int _minSessionSize;
@@ -20,7 +21,9 @@ internal sealed class PoolingSessionSource<T> : ISessionSource where T : Pooling
     private readonly Timer _cleanerTimer;
 
     private volatile int _numSessions;
-    private volatile int _waiterSize;
+    private volatile int _disposed;
+
+    private bool IsDisposed => _disposed == 1;
 
     public PoolingSessionSource(
         IPoolingSessionFactory<T> sessionFactory,
@@ -43,10 +46,15 @@ YdbConnectionStringBuilder settings
         _cleanerTimer = new Timer(CleanIdleSessions, this, _sessionIdleTimeout, _sessionIdleTimeout);
     }
 
-    public ValueTask<ISession> OpenSession(CancellationToken cancellationToken = default) =>
-        TryGetIdleSession(out var session)
+    public ValueTask<ISession> OpenSession(CancellationToken cancellationToken = default)
+    {
+        if (IsDisposed)
+            throw new YdbException("Session Source is disposed.");
+
+        return TryGetIdleSession(out var session)
             ? new ValueTask<ISession>(session)
             : RentAsync(cancellationToken);
+    }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private bool TryGetIdleSession([NotNullWhen(true)] out T? session)
@@ -119,7 +127,14 @@ private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
                 continue;
             }
 
-            await using var _ = finalToken.Register(() => waiterTcs.TrySetCanceled(), useSynchronizationContext: false);
+            await using var _ = finalToken.Register(
+                () => waiterTcs.TrySetCanceled(),
+                useSynchronizationContext: false
+            );
+            await using var disposedCancellationTokenRegistration = _disposeCts.Token.Register(
+                () => waiterTcs.TrySetException(new YdbException("Session Source is disposed.")),
+                useSynchronizationContext: false
+            );
             session = await waiterTcs.Task.ConfigureAwait(false);
 
             if (CheckIdleSession(session) || TryGetIdleSession(out session))
@@ -174,10 +189,15 @@ private void WakeUpWaiter()
 
     public void Return(T session)
     {
-        if (session.IsBroken)
+        if (session.IsBroken || IsDisposed)
         {
             CloseSession(session);
 
+            if (IsDisposed)
+            {
+                _ = TryDisposeCore();
+            }
+
             return;
         }
 
@@ -237,9 +257,27 @@ private static void CleanIdleSessions(object? state)
             }
         }
     }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (Interlocked.CompareExchange(ref _disposed, 1, 0) != 0)
+        {
+            return;
+        }
+
+        await _cleanerTimer.DisposeAsync();
+        _disposeCts.Cancel();
+
+        CleanIdleSessions(_idleSessions);
+
+        await TryDisposeCore();
+    }
+
+    private ValueTask TryDisposeCore() =>
+        _numSessions == 0 ? _sessionFactory.DisposeAsync() : ValueTask.CompletedTask;
 }
 
-internal interface IPoolingSessionFactory<T> where T : PoolingSessionBase<T>
+internal interface IPoolingSessionFactory<T> : IAsyncDisposable where T : PoolingSessionBase<T>
 {
     T NewSession(PoolingSessionSource<T> source);
 }

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

@@ -98,7 +98,15 @@ public override async Task OpenAsync(CancellationToken cancellationToken)
         }
         catch (OperationCanceledException e)
         {
-            throw new YdbException(StatusCode.Cancelled,
+            throw new YdbException(StatusCode.ClientTransportTimeout,
+                $"The connection pool has been exhausted, either raise 'MaxSessionPool' " +
+                $"(currently {ConnectionStringBuilder.MaxSessionPool}) or 'CreateSessionTimeout' " +
+                $"(currently {ConnectionStringBuilder.CreateSessionTimeout} seconds) in your connection string.", e
+            );
+        }
+        catch (YdbException e) when (e.Code == StatusCode.ClientTransportTimeout)
+        {
+            throw new YdbException(e.Code,
                 $"The connection pool has been exhausted, either raise 'MaxSessionPool' " +
                 $"(currently {ConnectionStringBuilder.MaxSessionPool}) or 'CreateSessionTimeout' " +
                 $"(currently {ConnectionStringBuilder.CreateSessionTimeout} seconds) in your connection string.", e

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Benchmarks/README.md

@@ -38,3 +38,15 @@ Allocated            : Allocated memory per single operation (managed only, incl
 | SessionReuse_Pattern                |    159,954.05 ns |   2,820.324 ns |   4,044.825 ns |             220.0000 |           6.0381 |   0.7324 |       - |    7307 B |
 | SessionReuse_HighContention_Pattern |  8,914,529.81 ns |  46,900.448 ns |  41,576.026 ns |           19756.6563 |         149.0469 | 625.0000 | 93.7500 | 5289794 B |
 | SessionReuse_HighIterations_Pattern | 81,211,792.96 ns | 749,115.160 ns | 664,071.077 ns |          200020.0000 |         614.8571 |        - |       - |    7458 B |
+
+# YDB .NET SDK Session Pool Benchmarks (FIFO lock-free)
+
+| Method                              |             Mean |            Error |           StdDev |           Median | Completed Work Items | Lock Contentions |     Gen0 |     Gen1 | Allocated |
+|-------------------------------------|-----------------:|-----------------:|-----------------:|-----------------:|---------------------:|-----------------:|---------:|---------:|----------:|
+| SingleThreaded_OpenClose            |         64.75 ns |         1.004 ns |         0.986 ns |         64.31 ns |                    - |                - |   0.0038 |        - |      32 B |                                        
+| MultiThreaded_OpenClose             |     22,547.55 ns |       443.931 ns |       511.232 ns |     22,613.77 ns |              40.0020 |                - |   0.9460 |        - |    7888 B |
+| HighContention_OpenClose            |     83,230.11 ns |     1,642.445 ns |     1,536.344 ns |     83,232.22 ns |             205.5402 |           0.0016 |   5.0049 |        - |   42337 B |
+| SessionReuse_Pattern                |    102,998.59 ns |     1,990.337 ns |     2,854.481 ns |    103,245.75 ns |             220.0000 |           0.0002 |   1.5869 |        - |   13678 B |
+| SessionReuse_HighContention_Pattern |  3,735,379.14 ns |    73,499.927 ns |   116,578.682 ns |  3,694,595.05 ns |           19845.2422 |           5.7344 | 812.5000 | 132.8125 | 6750890 B |
+| SessionReuse_HighIterations_Pattern | 69,402,122.32 ns | 1,382,232.565 ns | 4,075,544.202 ns | 70,535,363.07 ns |          200020.0000 |                - | 714.2857 |        - | 6407458 B |
+

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Benchmarks/SessionSourceBenchmark.cs

@@ -136,6 +136,8 @@ public async Task SessionReuse_HighIterations_Pattern()
 internal class MockSessionFactory : IPoolingSessionFactory<MockPoolingSession>
 {
     public MockPoolingSession NewSession(PoolingSessionSource<MockPoolingSession> source) => new(source);
+
+    public ValueTask DisposeAsync() => ValueTask.CompletedTask;
 }
 
 internal class MockPoolingSession(PoolingSessionSource<MockPoolingSession> source)

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/Session/PoolingSessionSourceMockTests.cs

@@ -127,6 +127,8 @@ internal class MockPoolingSessionFactory : IPoolingSessionFactory<MockPoolingSes
 
     public MockPoolingSession NewSession(PoolingSessionSource<MockPoolingSession> source) =>
         new(source, Open, DeleteSession, IsBroken, Interlocked.Increment(ref _sessionNum));
+
+    public ValueTask DisposeAsync() => ValueTask.CompletedTask;
 }
 
 internal class MockPoolingSession(