Move implicit session creation to PoolManager by flag

Author: LiamHamsters

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

@@ -33,10 +33,22 @@ CancellationToken cancellationToken
                          !cacheDriver.IsDisposed
                 ? cacheDriver
                 : Drivers[settings.GrpcConnectionString] = await settings.BuildDriver();
-            driver.RegisterOwner();
 
-            var factory = new PoolingSessionFactory(driver, settings);
-            var newSessionPool = new PoolingSessionSource<PoolingSession>(factory, settings);
+            ISessionSource newSessionPool;
+            if (settings.EnableImplicitSession)
+            {
+                var key = settings.ConnectionString;
+                newSessionPool = new ImplicitSessionSource(
+                    driver,
+                    onEmpty: () => Pools.TryRemove(key, out _)
+                );
+            }
+            else
+            {
+                driver.RegisterOwner();
+                var factory = new PoolingSessionFactory(driver, settings);
+                newSessionPool = new PoolingSessionSource<PoolingSession>(factory, settings);
+            }
 
             Pools[settings.ConnectionString] = newSessionPool;
 

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

@@ -3,49 +3,57 @@ namespace Ydb.Sdk.Ado.Session;
 internal sealed class ImplicitSessionSource : ISessionSource
 {
     private readonly IDriver _driver;
-    private readonly Action? _onEmpty;
-    private int _leased;
-    private int _closed;
+    private readonly Action? _onBecameEmpty;
+    private int _isDisposed;
+    private int _activeLeaseCount;
 
     internal ImplicitSessionSource(IDriver driver, Action? onEmpty = null)
     {
-        _driver = driver;
-        _onEmpty = onEmpty;
+        _driver = driver ?? throw new ArgumentNullException(nameof(driver));
+        _onBecameEmpty = onEmpty;
     }
 
     public ValueTask<ISession> OpenSession(CancellationToken cancellationToken)
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        if (Volatile.Read(ref _closed) == 1)
+        if (!TryAcquireLease())
             throw new ObjectDisposedException(nameof(ImplicitSessionSource));
 
-        Interlocked.Increment(ref _leased);
+        return new ValueTask<ISession>(new ImplicitSession(_driver, ReleaseLease));
+    }
+
+    private bool TryAcquireLease()
+    {
+        if (Volatile.Read(ref _isDisposed) != 0)
+            return false;
 
-        if (Volatile.Read(ref _closed) == 1)
+        Interlocked.Increment(ref _activeLeaseCount);
+
+        if (Volatile.Read(ref _isDisposed) != 0)
         {
-            Interlocked.Decrement(ref _leased);
-            throw new ObjectDisposedException(nameof(ImplicitSessionSource));
+            Interlocked.Decrement(ref _activeLeaseCount);
+            return false;
         }
 
-        return new ValueTask<ISession>(new ImplicitSession(_driver, Release));
+        return true;
     }
 
-    private void Release()
+    private void ReleaseLease()
     {
-        if (Interlocked.Decrement(ref _leased) == 0)
+        if (Interlocked.Decrement(ref _activeLeaseCount) == 0)
         {
-            _onEmpty?.Invoke();
+            _onBecameEmpty?.Invoke();
         }
     }
 
     public ValueTask DisposeAsync()
     {
-        Interlocked.Exchange(ref _closed, 1);
+        Interlocked.Exchange(ref _isDisposed, 1);
 
-        if (Volatile.Read(ref _leased) == 0)
+        if (Volatile.Read(ref _activeLeaseCount) == 0)
         {
-            _onEmpty?.Invoke();
+            _onBecameEmpty?.Invoke();
         }
 
         return ValueTask.CompletedTask;

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

@@ -39,22 +39,6 @@ internal ISession Session
 
     private ISession _session = null!;
 
-    private ImplicitSessionSource? _implicitSessionSource;
-
-    internal bool EnableImplicitSession => ConnectionStringBuilder.EnableImplicitSession;
-
-    internal ISession GetExecutionSession(bool useImplicit)
-    {
-        ThrowIfConnectionClosed();
-
-        if (!useImplicit)
-            return Session;
-
-        _implicitSessionSource ??=
-            new ImplicitSessionSource(Session.Driver, onEmpty: () => _implicitSessionSource = null);
-        return _implicitSessionSource.OpenSession(CancellationToken.None).GetAwaiter().GetResult();
-    }
-
     public YdbConnection()
     {
     }
@@ -143,7 +127,6 @@ public override async Task CloseAsync()
         finally
         {
             _session.Close();
-            _implicitSessionSource = null;
         }
     }
 
@@ -165,7 +148,7 @@ public override string ConnectionString
 
     public override ConnectionState State => ConnectionState;
 
-    private ConnectionState ConnectionState { get; set; } = ConnectionState.Closed; // Invoke AsyncOpen()
+    private ConnectionState ConnectionState { get; set; } = ConnectionState.Closed;
 
     internal void OnNotSuccessStatusCode(StatusCode code)
     {