feat: PoolingSessionBase

Author: KirillKurdyukov

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

@@ -13,53 +13,40 @@
 
 namespace Ydb.Sdk.Ado.Session;
 
-internal class PoolingSession : ISession
+internal class PoolingSession : PoolingSessionBase
 {
     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; }
 
-    private int _state = (int)PoolingSessionState.In;
-
-    public IDriver Driver { get; }
-    public bool IsBroken => _isBroken;
+    public override IDriver Driver { get; }
+    public override bool IsBroken => _isBroken;
 
     internal PoolingSession(
         IDriver driver,
         PoolingSessionSource poolingSessionSource,
         bool disableServerBalancer,
         ILogger<PoolingSession> logger
-    )
+    ) : base(poolingSessionSource)
     {
-        _poolingSessionSource = poolingSessionSource;
         _disableServerBalancer = disableServerBalancer;
         _logger = logger;
         Driver = driver;
     }
 
-    internal bool CompareAndSet(PoolingSessionState expected, PoolingSessionState actual) =>
-        Interlocked.CompareExchange(ref _state, (int)expected, (int)actual) == (int)expected;
-
-    internal void Set(PoolingSessionState state) => Interlocked.Exchange(ref _state, (int)state);
-
-    internal PoolingSessionState State => (PoolingSessionState)Volatile.Read(ref _state);
-
-    internal DateTime IdleStartTime { get; set; }
-
-    public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
+    public override ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
         string query,
         Dictionary<string, YdbValue> parameters,
         GrpcRequestSettings settings,
@@ -81,10 +68,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,
@@ -98,10 +82,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,
@@ -115,7 +96,7 @@ public async Task RollbackTransaction(
         }
     }
 
-    public void OnNotSuccessStatusCode(StatusCode statusCode)
+    public override void OnNotSuccessStatusCode(StatusCode statusCode)
     {
         _isBadSession = _isBadSession || statusCode is StatusCode.BadSession;
 
@@ -132,7 +113,7 @@ StatusCode.ClientTransportTimeout or
         }
     }
 
-    public async Task Open(CancellationToken cancellationToken)
+    internal override async Task Open(CancellationToken cancellationToken)
     {
         var requestSettings = new GrpcRequestSettings { CancellationToken = cancellationToken };
 
@@ -233,7 +214,7 @@ public async Task Open(CancellationToken cancellationToken)
         await completeTask.Task;
     }
 
-    public async Task DeleteSession()
+    internal override async Task DeleteSession()
     {
         try
         {
@@ -264,6 +245,4 @@ public async Task DeleteSession()
                 SessionId, NodeId);
         }
     }
-
-    public void Close() => _poolingSessionSource.Return(this);
 }

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

@@ -15,6 +15,6 @@ public PoolingSessionFactory(IDriver driver, YdbConnectionStringBuilder settings
         _logger = loggerFactory.CreateLogger<PoolingSession>();
     }
 
-    public PoolingSession NewSession(PoolingSessionSource source) =>
-        new(_driver, source, _disableServerBalancer, _logger);
+    public PoolingSessionBase NewSession(PoolingSessionSource source) =>
+        new PoolingSession(_driver, source, _disableServerBalancer, _logger);
 }

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

@@ -1,22 +1,22 @@
 using System.Collections.Concurrent;
-using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
-using System.Threading.Channels;
+using Ydb.Query;
+using Ydb.Sdk.Value;
 
 namespace Ydb.Sdk.Ado.Session;
 
-internal sealed class PoolingSessionSource : ISessionSource<PoolingSession>
+internal sealed class PoolingSessionSource : ISessionSource<PoolingSessionBase>
 {
-    private readonly ConcurrentStack<PoolingSession> _idleSessions = new();
-    private readonly ConcurrentQueue<TaskCompletionSource<PoolingSession?>> _waiters = new();
+    private readonly ConcurrentStack<PoolingSessionBase> _idleSessions = new();
+    private readonly ConcurrentQueue<TaskCompletionSource<PoolingSessionBase?>> _waiters = new();
 
     private readonly IPoolingSessionFactory _sessionFactory;
 
     private readonly int _minSessionSize;
     private readonly int _maxSessionSize;
 
-    private readonly PoolingSession?[] _sessions;
+    private readonly PoolingSessionBase?[] _sessions;
 
     private readonly int _createSessionTimeout;
     private readonly TimeSpan _sessionIdleTimeout;
@@ -39,19 +39,19 @@ YdbConnectionStringBuilder settings
                 $"Connection can't have 'Max Session Pool' {_maxSessionSize} under 'Min Session Pool' {_minSessionSize}");
         }
 
-        _sessions = new PoolingSession?[_maxSessionSize];
+        _sessions = new PoolingSessionBase?[_maxSessionSize];
         _createSessionTimeout = settings.CreateSessionTimeout;
         _sessionIdleTimeout = TimeSpan.FromSeconds(settings.SessionIdleTimeout);
         _cleanerTimer = new Timer(CleanIdleSessions, this, _sessionIdleTimeout, _sessionIdleTimeout);
     }
 
-    public ValueTask<PoolingSession> OpenSession(CancellationToken cancellationToken = default) =>
+    public ValueTask<PoolingSessionBase> OpenSession(CancellationToken cancellationToken = default) =>
         TryGetIdleSession(out var session)
-            ? new ValueTask<PoolingSession>(session)
+            ? new ValueTask<PoolingSessionBase>(session)
             : RentAsync(cancellationToken);
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private bool TryGetIdleSession([NotNullWhen(true)] out PoolingSession? session)
+    private bool TryGetIdleSession([NotNullWhen(true)] out PoolingSessionBase? session)
     {
         while (_idleSessions.TryPop(out session))
         {
@@ -65,7 +65,7 @@ private bool TryGetIdleSession([NotNullWhen(true)] out PoolingSession? session)
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private bool CheckIdleSession([NotNullWhen(true)] PoolingSession? session)
+    private bool CheckIdleSession([NotNullWhen(true)] PoolingSessionBase? session)
     {
         if (session == null || session.State == PoolingSessionState.Clean)
         {
@@ -82,7 +82,7 @@ private bool CheckIdleSession([NotNullWhen(true)] PoolingSession? session)
         return session.CompareAndSet(PoolingSessionState.In, PoolingSessionState.Out);
     }
 
-    private async ValueTask<PoolingSession> RentAsync(CancellationToken cancellationToken)
+    private async ValueTask<PoolingSessionBase> RentAsync(CancellationToken cancellationToken)
     {
         using var ctsGetSession = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
         if (_createSessionTimeout > 0)
@@ -97,7 +97,7 @@ private async ValueTask<PoolingSession> RentAsync(CancellationToken cancellation
         while (true)
         {
             var waiterTcs =
-                new TaskCompletionSource<PoolingSession?>(TaskCreationOptions.RunContinuationsAsynchronously);
+                new TaskCompletionSource<PoolingSessionBase?>(TaskCreationOptions.RunContinuationsAsynchronously);
             _waiters.Enqueue(waiterTcs);
             await using var _ = finalToken.Register(() => waiterTcs.TrySetCanceled(), useSynchronizationContext: false);
             session = await waiterTcs.Task.ConfigureAwait(false);
@@ -111,7 +111,7 @@ private async ValueTask<PoolingSession> RentAsync(CancellationToken cancellation
         }
     }
 
-    private async ValueTask<PoolingSession?> OpenNewSession(CancellationToken cancellationToken)
+    private async ValueTask<PoolingSessionBase?> 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)
@@ -156,7 +156,7 @@ private void WakeUpWaiter()
             waiter.TrySetResult(null); // wake up waiter!
     }
 
-    public void Return(PoolingSession session)
+    public void Return(PoolingSessionBase session)
     {
         if (session.IsBroken)
         {
@@ -181,7 +181,7 @@ public void Return(PoolingSession session)
         WakeUpWaiter();
     }
 
-    private void CloseSession(PoolingSession session)
+    private void CloseSession(PoolingSessionBase session)
     {
         var i = 0;
         for (; i < _maxSessionSize; i++)
@@ -224,7 +224,7 @@ private static void CleanIdleSessions(object? state)
 
 internal interface IPoolingSessionFactory
 {
-    PoolingSession NewSession(PoolingSessionSource source);
+    PoolingSessionBase NewSession(PoolingSessionSource source);
 }
 
 internal enum PoolingSessionState
@@ -233,3 +233,44 @@ internal enum PoolingSessionState
     Out,
     Clean
 }
+
+internal abstract class PoolingSessionBase : ISession
+{
+    private readonly PoolingSessionSource _source;
+
+    private int _state = (int)PoolingSessionState.In;
+
+    protected PoolingSessionBase(PoolingSessionSource source)
+    {
+        _source = source;
+    }
+
+    internal bool CompareAndSet(PoolingSessionState expected, PoolingSessionState actual) =>
+        Interlocked.CompareExchange(ref _state, (int)expected, (int)actual) == (int)expected;
+
+    internal void Set(PoolingSessionState state) => Interlocked.Exchange(ref _state, (int)state);
+
+    internal PoolingSessionState State => (PoolingSessionState)Volatile.Read(ref _state);
+
+    internal DateTime IdleStartTime { get; set; }
+
+    public abstract IDriver Driver { get; }
+
+    public abstract bool IsBroken { get; }
+
+    internal abstract Task Open(CancellationToken cancellationToken);
+
+    internal abstract Task DeleteSession();
+
+    public abstract ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(string query,
+        Dictionary<string, YdbValue> parameters, GrpcRequestSettings settings,
+        TransactionControl? txControl);
+
+    public abstract Task CommitTransaction(string txId, CancellationToken cancellationToken = default);
+
+    public abstract Task RollbackTransaction(string txId, CancellationToken cancellationToken = default);
+
+    public abstract void OnNotSuccessStatusCode(StatusCode code);
+
+    public void Close() => _source.Return(this);
+}

src/Ydb.Sdk/src/Ado/YdbSchema.cs

@@ -423,9 +423,7 @@ CancellationToken cancellationToken
             );
 
             var operation = response.Operation;
-            var status = Status.FromProto(operation.Status, operation.Issues);
-
-            if (status.IsNotSuccess)
+            if (operation.Status.IsNotSuccess())
             {
                 throw YdbException.FromServer(operation.Status, operation.Issues);
             }

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

@@ -87,5 +87,30 @@ public async Task SessionReuse_Pattern()
 
 internal class MockSessionFactory : IPoolingSessionFactory
 {
-    public PoolingSession NewSession(PoolingSessionSource source) => new PoolingSession(null, null, false, null);
+    public PoolingSessionBase NewSession(PoolingSessionSource source) => new MockPoolingSession(source);
+}
+
+internal class MockPoolingSession(PoolingSessionSource source) : PoolingSessionBase(source)
+{
+    public override IDriver Driver => null!;
+    public override bool IsBroken => false;
+
+    internal override Task Open(CancellationToken cancellationToken) => Task.CompletedTask;
+    internal override Task DeleteSession() => Task.CompletedTask;
+
+    public override ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
+        string query,
+        Dictionary<string, YdbValue> parameters, GrpcRequestSettings settings,
+        TransactionControl? txControl
+    ) => throw new NotImplementedException();
+
+    public override Task CommitTransaction(string txId, CancellationToken cancellationToken = default) =>
+        throw new NotImplementedException();
+
+    public override Task RollbackTransaction(string txId, CancellationToken cancellationToken = default) =>
+        throw new NotImplementedException();
+
+    public override void OnNotSuccessStatusCode(StatusCode code)
+    {
+    }
 }

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

@@ -29,33 +29,32 @@ internal class MockPoolingSessionFactory : IPoolingSessionFactory
 {
     private int _sessionNum;
 
-    public PoolingSession NewSession(PoolingSessionSource source) =>
+    public PoolingSessionBase NewSession(PoolingSessionSource source) =>
         new MockPoolingSession(source, Interlocked.Increment(ref _sessionNum));
 }
 
-internal class MockPoolingSession(PoolingSessionSource source, int sessionNum) : PoolingSession(null, null, true, null)
+internal class MockPoolingSession(PoolingSessionSource source, int sessionNum) : PoolingSessionBase(source)
 {
-    internal string SessionId { get; } = $"session_{sessionNum}";
+    public string SessionId => $"session_{sessionNum}";
+    public override IDriver Driver => null!;
+    public override bool IsBroken => false;
 
-    public IDriver Driver => throw new NotImplementedException();
-    public bool IsBroken { get; set; }
+    internal override Task Open(CancellationToken cancellationToken) => Task.CompletedTask;
+    internal override Task DeleteSession() => Task.CompletedTask;
 
-    public void Close() => source.Return(this);
-
-    public Task Open(CancellationToken cancellationToken) => Task.CompletedTask;
-
-    public Task DeleteSession() => Task.CompletedTask;
-
-    public Task CommitTransaction(string txId, CancellationToken cancellationToken = default) =>
-        throw new NotImplementedException();
+    public override ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(
+        string query,
+        Dictionary<string, YdbValue> parameters, GrpcRequestSettings settings,
+        TransactionControl? txControl
+    ) => throw new NotImplementedException();
 
-    public Task RollbackTransaction(string txId, CancellationToken cancellationToken = default) =>
+    public override Task CommitTransaction(string txId, CancellationToken cancellationToken = default) =>
         throw new NotImplementedException();
 
-    public ValueTask<IServerStream<ExecuteQueryResponsePart>> ExecuteQuery(string query,
-        Dictionary<string, YdbValue> parameters, GrpcRequestSettings settings,
-        TransactionControl? txControl) =>
+    public override Task RollbackTransaction(string txId, CancellationToken cancellationToken = default) =>
         throw new NotImplementedException();
 
-    public void OnNotSuccessStatusCode(StatusCode code) => throw new NotImplementedException();
+    public override void OnNotSuccessStatusCode(StatusCode code)
+    {
+    }
 }