tried to make a “two-phase”

Author: LiamHamsters

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

@@ -2,6 +2,8 @@ namespace Ydb.Sdk.Ado.Session;
 
 internal sealed class ImplicitSessionSource : ISessionSource
 {
+    private enum State { Open = 0, Closing = 1, Closed = 2 }
+
     private readonly IDriver _driver;
     private readonly ManualResetEventSlim _allReleased = new(false);
 
@@ -25,12 +27,12 @@ public ValueTask<ISession> OpenSession(CancellationToken cancellationToken)
 
     private bool TryAcquireLease()
     {
-        if (Volatile.Read(ref _state) == 2)
+        if (Volatile.Read(ref _state) == (int)State.Closed)
             return false;
 
         Interlocked.Increment(ref _activeLeaseCount);
 
-        if (Volatile.Read(ref _state) == 2)
+        if (Volatile.Read(ref _state) == (int)State.Closed)
         {
             Interlocked.Decrement(ref _activeLeaseCount);
             return false;
@@ -41,21 +43,30 @@ private bool TryAcquireLease()
 
     internal void ReleaseLease()
     {
-        if (Interlocked.Decrement(ref _activeLeaseCount) == 0 && Volatile.Read(ref _state) == 1)
+        if (Interlocked.Decrement(ref _activeLeaseCount) == 0 &&
+            Volatile.Read(ref _state) != (int)State.Open)
+        {
             _allReleased.Set();
+        }
     }
 
     public async ValueTask DisposeAsync()
     {
-        if (Interlocked.CompareExchange(ref _state, 1, 0) != 0)
-            return;
+        var prev = Interlocked.CompareExchange(ref _state, (int)State.Closing, (int)State.Open);
+        switch (prev)
+        {
+            case (int)State.Closed:
+                return;
+            case (int)State.Closing:
+                break;
+        }
 
         if (Volatile.Read(ref _activeLeaseCount) != 0)
             _allReleased.Wait();
 
         try
         {
-            Volatile.Write(ref _state, 2);
+            Volatile.Write(ref _state, (int)State.Closed);
             await _driver.DisposeAsync();
         }
         finally