type safe generic

Author: KirillKurdyukov

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

@@ -7,7 +7,7 @@ namespace Ydb.Sdk.Ado;
 internal static class PoolManager
 {
     private static readonly SemaphoreSlim SemaphoreSlim = new(1); // async mutex
-    private static readonly ConcurrentDictionary<string, PoolingSessionSource> Pools = new();
+    private static readonly ConcurrentDictionary<string, ISessionSource> Pools = new();
 
     internal static async Task<ISession> GetSession(
         YdbConnectionStringBuilder settings,
@@ -28,7 +28,7 @@ CancellationToken cancellationToken
                 return await pool.OpenSession(cancellationToken);
             }
 
-            var newSessionPool = new PoolingSessionSource(
+            var newSessionPool = new PoolingSessionSource<PoolingSession>(
                 new PoolingSessionFactory(
                     await settings.BuildDriver(),
                     settings,

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
 {
-    ValueTask<TSession> OpenSession(CancellationToken cancellationToken);
-
-    void Return(TSession session);
+    ValueTask<ISession> OpenSession(CancellationToken cancellationToken);
 }

src/Ydb.Sdk/src/Ado/Session/PoolingSession.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 Microsoft.Extensions.Logging;
 using Ydb.Query;
 using Ydb.Query.V1;
@@ -13,7 +8,7 @@
 
 namespace Ydb.Sdk.Ado.Session;
 
-internal class PoolingSession : PoolingSessionBase
+internal class PoolingSession : PoolingSessionBase<PoolingSession>
 {
     private const string SessionBalancer = "session-balancer";
 
@@ -36,7 +31,7 @@ internal class PoolingSession : PoolingSessionBase
 
     internal PoolingSession(
         IDriver driver,
-        PoolingSessionSource poolingSessionSource,
+        PoolingSessionSource<PoolingSession> poolingSessionSource,
         bool disableServerBalancer,
         ILogger<PoolingSession> logger
     ) : base(poolingSessionSource)

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

@@ -2,7 +2,7 @@
 
 namespace Ydb.Sdk.Ado.Session;
 
-internal class PoolingSessionFactory : IPoolingSessionFactory
+internal class PoolingSessionFactory : IPoolingSessionFactory<PoolingSession>
 {
     private readonly IDriver _driver;
     private readonly bool _disableServerBalancer;
@@ -15,6 +15,6 @@ public PoolingSessionFactory(IDriver driver, YdbConnectionStringBuilder settings
         _logger = loggerFactory.CreateLogger<PoolingSession>();
     }
 
-    public PoolingSessionBase NewSession(PoolingSessionSource source) =>
-        new PoolingSession(_driver, source, _disableServerBalancer, _logger);
+    public PoolingSession NewSession(PoolingSessionSource<PoolingSession> source) =>
+        new(_driver, source, _disableServerBalancer, _logger);
 }

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

@@ -6,26 +6,23 @@
 
 namespace Ydb.Sdk.Ado.Session;
 
-internal sealed class PoolingSessionSource : ISessionSource<PoolingSessionBase>
+internal sealed class PoolingSessionSource<T> : ISessionSource where T : PoolingSessionBase<T>
 {
-    private readonly ConcurrentStack<PoolingSessionBase> _idleSessions = new();
-    private readonly ConcurrentQueue<TaskCompletionSource<PoolingSessionBase?>> _waiters = new();
-
-    private readonly IPoolingSessionFactory _sessionFactory;
+    private readonly ConcurrentStack<T> _idleSessions = new();
+    private readonly ConcurrentQueue<TaskCompletionSource<T?>> _waiters = new();
 
+    private readonly IPoolingSessionFactory<T> _sessionFactory;
     private readonly int _minSessionSize;
     private readonly int _maxSessionSize;
-
-    private readonly PoolingSessionBase?[] _sessions;
-
+    private readonly T?[] _sessions;
     private readonly int _createSessionTimeout;
     private readonly TimeSpan _sessionIdleTimeout;
     private readonly Timer _cleanerTimer;
 
     private volatile int _numSessions;
 
     public PoolingSessionSource(
-        IPoolingSessionFactory sessionFactory,
+        IPoolingSessionFactory<T> sessionFactory,
         YdbConnectionStringBuilder settings
     )
     {
@@ -39,19 +36,19 @@ YdbConnectionStringBuilder settings
                 $"Connection can't have 'Max Session Pool' {_maxSessionSize} under 'Min Session Pool' {_minSessionSize}");
         }
 
-        _sessions = new PoolingSessionBase?[_maxSessionSize];
+        _sessions = new T?[_maxSessionSize];
         _createSessionTimeout = settings.CreateSessionTimeout;
         _sessionIdleTimeout = TimeSpan.FromSeconds(settings.SessionIdleTimeout);
         _cleanerTimer = new Timer(CleanIdleSessions, this, _sessionIdleTimeout, _sessionIdleTimeout);
     }
 
-    public ValueTask<PoolingSessionBase> OpenSession(CancellationToken cancellationToken = default) =>
+    public ValueTask<ISession> OpenSession(CancellationToken cancellationToken = default) =>
         TryGetIdleSession(out var session)
-            ? new ValueTask<PoolingSessionBase>(session)
+            ? new ValueTask<ISession>(session)
             : RentAsync(cancellationToken);
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private bool TryGetIdleSession([NotNullWhen(true)] out PoolingSessionBase? session)
+    private bool TryGetIdleSession([NotNullWhen(true)] out T? session)
     {
         while (_idleSessions.TryPop(out session))
         {
@@ -65,7 +62,7 @@ private bool TryGetIdleSession([NotNullWhen(true)] out PoolingSessionBase? sessi
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private bool CheckIdleSession([NotNullWhen(true)] PoolingSessionBase? session)
+    private bool CheckIdleSession([NotNullWhen(true)] T? session)
     {
         if (session == null || session.State == PoolingSessionState.Clean)
         {
@@ -82,7 +79,7 @@ private bool CheckIdleSession([NotNullWhen(true)] PoolingSessionBase? session)
         return session.CompareAndSet(PoolingSessionState.In, PoolingSessionState.Out);
     }
 
-    private async ValueTask<PoolingSessionBase> RentAsync(CancellationToken cancellationToken)
+    private async ValueTask<ISession> RentAsync(CancellationToken cancellationToken)
     {
         using var ctsGetSession = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
         if (_createSessionTimeout > 0)
@@ -96,8 +93,7 @@ private async ValueTask<PoolingSessionBase> RentAsync(CancellationToken cancella
 
         while (true)
         {
-            var waiterTcs =
-                new TaskCompletionSource<PoolingSessionBase?>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var waiterTcs = new TaskCompletionSource<T?>(TaskCreationOptions.RunContinuationsAsynchronously);
             _waiters.Enqueue(waiterTcs);
             await using var _ = finalToken.Register(() => waiterTcs.TrySetCanceled(), useSynchronizationContext: false);
             session = await waiterTcs.Task.ConfigureAwait(false);
@@ -111,9 +107,8 @@ private async ValueTask<PoolingSessionBase> RentAsync(CancellationToken cancella
         }
     }
 
-    private async ValueTask<PoolingSessionBase?> OpenNewSession(CancellationToken cancellationToken)
+    private async ValueTask<T?> OpenNewSession(CancellationToken cancellationToken)
     {
-        // As long as we're under max capacity, attempt to increase the session count and open a new session.
         for (var numSessions = _numSessions; numSessions < _maxSessionSize; numSessions = _numSessions)
         {
             if (Interlocked.CompareExchange(ref _numSessions, numSessions + 1, numSessions) != numSessions)
@@ -135,12 +130,8 @@ private async ValueTask<PoolingSessionBase> RentAsync(CancellationToken cancella
             }
             catch
             {
-                // RPC open failed, decrement the open and busy counter back down.
                 Interlocked.Decrement(ref _numSessions);
 
-                // In case there's a waiting attempt on the waiters queue, we write a null to the idle connector channel
-                // to wake it up, so it will try opening (and probably throw immediately)
-                // Statement order is important since we have synchronous completions on the channel.
                 WakeUpWaiter();
 
                 throw;
@@ -156,7 +147,7 @@ private void WakeUpWaiter()
             waiter.TrySetResult(null); // wake up waiter!
     }
 
-    public void Return(PoolingSessionBase session)
+    public void Return(T session)
     {
         if (session.IsBroken)
         {
@@ -181,7 +172,7 @@ public void Return(PoolingSessionBase session)
         WakeUpWaiter();
     }
 
-    private void CloseSession(PoolingSessionBase session)
+    private void CloseSession(T session)
     {
         var i = 0;
         for (; i < _maxSessionSize; i++)
@@ -202,7 +193,7 @@ private void CloseSession(PoolingSessionBase session)
 
     private static void CleanIdleSessions(object? state)
     {
-        var pool = (PoolingSessionSource)state!;
+        var pool = (PoolingSessionSource<T>)state!;
         var now = DateTime.Now;
 
         for (var i = 0; i < pool._maxSessionSize; i++)
@@ -222,9 +213,9 @@ private static void CleanIdleSessions(object? state)
     }
 }
 
-internal interface IPoolingSessionFactory
+internal interface IPoolingSessionFactory<T> where T : PoolingSessionBase<T>
 {
-    PoolingSessionBase NewSession(PoolingSessionSource source);
+    T NewSession(PoolingSessionSource<T> source);
 }
 
 internal enum PoolingSessionState
@@ -234,13 +225,13 @@ internal enum PoolingSessionState
     Clean
 }
 
-internal abstract class PoolingSessionBase : ISession
+internal abstract class PoolingSessionBase<T> : ISession where T : PoolingSessionBase<T>
 {
-    private readonly PoolingSessionSource _source;
+    private readonly PoolingSessionSource<T> _source;
 
     private int _state = (int)PoolingSessionState.In;
 
-    protected PoolingSessionBase(PoolingSessionSource source)
+    protected PoolingSessionBase(PoolingSessionSource<T> source)
     {
         _source = source;
     }
@@ -272,5 +263,5 @@ public abstract ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
 
     public abstract void OnNotSuccessStatusCode(StatusCode code);
 
-    public void Close() => _source.Return(this);
+    public void Close() => _source.Return((T)this);
 }

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

@@ -9,7 +9,7 @@ namespace Ydb.Sdk.Ado.Benchmarks;
 [ThreadingDiagnoser]
 public class SessionSourceBenchmark
 {
-    private PoolingSessionSource _poolingSessionSource = null!;
+    private PoolingSessionSource<MockPoolingSession> _poolingSessionSource = null!;
     private const int SessionPoolSize = 50;
     private const int ConcurrentTasks = 20;
 
@@ -18,7 +18,7 @@ public void Setup()
     {
         var settings = new YdbConnectionStringBuilder { MaxSessionPool = SessionPoolSize };
 
-        _poolingSessionSource = new PoolingSessionSource(new MockSessionFactory(), settings);
+        _poolingSessionSource = new PoolingSessionSource<MockPoolingSession>(new MockSessionFactory(), settings);
     }
 
     [Benchmark]
@@ -85,12 +85,13 @@ public async Task SessionReuse_Pattern()
     }
 }
 
-internal class MockSessionFactory : IPoolingSessionFactory
+internal class MockSessionFactory : IPoolingSessionFactory<MockPoolingSession>
 {
-    public PoolingSessionBase NewSession(PoolingSessionSource source) => new MockPoolingSession(source);
+    public MockPoolingSession NewSession(PoolingSessionSource<MockPoolingSession> source) => new(source);
 }
 
-internal class MockPoolingSession(PoolingSessionSource source) : PoolingSessionBase(source)
+internal class MockPoolingSession(PoolingSessionSource<MockPoolingSession> source)
+    : PoolingSessionBase<MockPoolingSession>(source)
 {
     public override IDriver Driver => null!;
     public override bool IsBroken => false;

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

@@ -9,31 +9,44 @@ public class PoolingSessionSourceMockTests
 {
     [Fact]
     public void MinSessionPool_bigger_than_MaxSessionPool_throws() => Assert.Throws<ArgumentException>(() =>
-        new PoolingSessionSource(new MockPoolingSessionFactory(),
+        new PoolingSessionSource<MockPoolingSession>(new MockPoolingSessionFactory(),
             new YdbConnectionStringBuilder { MaxSessionPool = 1, MinSessionPool = 2 })
     );
 
     [Fact]
     public async Task Reuse_Session_Before_Creating_new()
     {
-        var sessionSource = new PoolingSessionSource(new MockPoolingSessionFactory(), new YdbConnectionStringBuilder());
-        var session = (MockPoolingSession)await sessionSource.OpenSession();
-        var sessionId = session.SessionId;
+        var sessionSource =
+            new PoolingSessionSource<MockPoolingSession>(new MockPoolingSessionFactory(),
+                new YdbConnectionStringBuilder());
+        var session = await sessionSource.OpenSession();
+        var sessionId = session.SessionId();
         session.Close();
-        session = (MockPoolingSession)await sessionSource.OpenSession();
-        Assert.Equal(sessionId, session.SessionId);
+        session = await sessionSource.OpenSession();
+        Assert.Equal(sessionId, session.SessionId());
     }
+
+    [Fact]
+    public async Task Creating_Session_Throw_Exception()
+    {
+    }
+}
+
+internal static class ISessionExtension
+{
+    internal static string SessionId(this ISession session) => ((MockPoolingSession)session).SessionId;
 }
 
-internal class MockPoolingSessionFactory : IPoolingSessionFactory
+internal class MockPoolingSessionFactory : IPoolingSessionFactory<MockPoolingSession>
 {
     private int _sessionNum;
 
-    public PoolingSessionBase NewSession(PoolingSessionSource source) =>
-        new MockPoolingSession(source, Interlocked.Increment(ref _sessionNum));
+    public MockPoolingSession NewSession(PoolingSessionSource<MockPoolingSession> source) =>
+        new(source, Interlocked.Increment(ref _sessionNum));
 }
 
-internal class MockPoolingSession(PoolingSessionSource source, int sessionNum) : PoolingSessionBase(source)
+internal class MockPoolingSession(PoolingSessionSource<MockPoolingSession> source, int sessionNum)
+    : PoolingSessionBase<MockPoolingSession>(source)
 {
     public string SessionId => $"session_{sessionNum}";
     public override IDriver Driver => null!;
@@ -44,7 +57,8 @@ internal class MockPoolingSession(PoolingSessionSource source, int sessionNum) :
 
     public override ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
         string query,
-        Dictionary<string, YdbValue> parameters, GrpcRequestSettings settings,
+        Dictionary<string, YdbValue> parameters,
+        GrpcRequestSettings settings,
         TransactionControl? txControl
     ) => throw new NotImplementedException();
 

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

@@ -17,7 +17,7 @@ public class PoolingSessionTests
     private readonly Mock<IDriver> _mockIDriver;
     private readonly Mock<IServerStream<SessionState>> _mockAttachStream = new(MockBehavior.Strict);
     private readonly PoolingSessionFactory _poolingSessionFactory;
-    private readonly PoolingSessionSource _poolingSessionSource;
+    private readonly PoolingSessionSource<PoolingSession> _poolingSessionSource;
 
     public PoolingSessionTests()
     {
@@ -32,7 +32,7 @@ public PoolingSessionTests()
         ).ReturnsAsync(_mockAttachStream.Object);
         _mockAttachStream.Setup(stream => stream.Dispose());
         _poolingSessionFactory = new PoolingSessionFactory(_mockIDriver.Object, settings, TestUtils.LoggerFactory);
-        _poolingSessionSource = new PoolingSessionSource(_poolingSessionFactory, settings);
+        _poolingSessionSource = new PoolingSessionSource<PoolingSession>(_poolingSessionFactory, settings);
     }
 
     [Theory]