feat: add EnableImplicitSession support

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,5 @@
+- Added provider support for implicit sessions.
+
 ## v0.23.1
 
 - Fixed bug Topic Reader: NullReferenceException when handling StopPartitionSessionRequest ([#528](https://github.com/ydb-platform/ydb-dotnet-sdk/issues/528)).

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

@@ -33,10 +33,20 @@ 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.MaxSessionPool > 0)
+            {
+                var factory = new PoolingSessionFactory(driver, settings);
+                newSessionPool = new PoolingSessionSource<PoolingSession>(factory, settings);
+            }
+            else
+            {
+                newSessionPool = new ImplicitSessionSource(driver);
+            }
 
             Pools[settings.ConnectionString] = newSessionPool;
 

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

@@ -5,9 +5,12 @@ namespace Ydb.Sdk.Ado.Session;
 
 internal class ImplicitSession : ISession
 {
-    public ImplicitSession(IDriver driver)
+    private readonly ImplicitSessionSource _source;
+
+    public ImplicitSession(IDriver driver, ImplicitSessionSource source)
     {
         Driver = driver;
+        _source = source;
     }
 
     public IDriver Driver { get; }
@@ -47,9 +50,7 @@ public void OnNotSuccessStatusCode(StatusCode code)
     {
     }
 
-    public void Dispose()
-    {
-    }
+    public void Dispose() => _source.ReleaseLease();
 
     private static YdbException NotSupportedTransaction => new("Transactions are not supported in implicit sessions");
 }

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

@@ -0,0 +1,71 @@
+namespace Ydb.Sdk.Ado.Session;
+
+internal sealed class ImplicitSessionSource : ISessionSource
+{
+    private readonly IDriver _driver;
+    private readonly ManualResetEventSlim _allReleased = new(false);
+
+    private int _state;
+    private int _activeLeaseCount;
+
+    internal ImplicitSessionSource(IDriver driver)
+    {
+        _driver = driver;
+    }
+
+    public ValueTask<ISession> OpenSession(CancellationToken cancellationToken)
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+
+        if (!TryAcquireLease())
+            throw new ObjectDisposedException(nameof(ImplicitSessionSource));
+
+        return new ValueTask<ISession>(new ImplicitSession(_driver, this));
+    }
+
+    private bool TryAcquireLease()
+    {
+        if (Volatile.Read(ref _state) == 2)
+            return false;
+
+        var newCount = Interlocked.Increment(ref _activeLeaseCount);
+
+        var state = Volatile.Read(ref _state);
+
+        if (state == 2 || (state == 1 && newCount == 1))
+        {
+            Interlocked.Decrement(ref _activeLeaseCount);
+            return false;
+        }
+
+        return true;
+    }
+
+    internal void ReleaseLease()
+    {
+        if (Interlocked.Decrement(ref _activeLeaseCount) == 0 &&
+            Volatile.Read(ref _state) != 0)
+        {
+            _allReleased.Set();
+        }
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (Interlocked.CompareExchange(ref _state, 1, 0) != 0)
+            return;
+
+        if (Volatile.Read(ref _activeLeaseCount) != 0)
+            _allReleased.Wait();
+
+        try
+        {
+            Volatile.Write(ref _state, 2);
+            await _driver.DisposeAsync();
+        }
+        finally
+        {
+            _allReleased.Dispose();
+        }
+    }
+}

src/Ydb.Sdk/src/Ado/YdbConnectionStringBuilder.cs

@@ -40,6 +40,7 @@ private void InitDefaultValues()
         _maxReceiveMessageSize = GrpcDefaultSettings.MaxReceiveMessageSize;
         _disableDiscovery = GrpcDefaultSettings.DisableDiscovery;
         _disableServerBalancer = false;
+        _enableImplicitSession = false;
     }
 
     public string Host
@@ -314,6 +315,18 @@ public int CreateSessionTimeout
 
     private int _createSessionTimeout;
 
+    public bool EnableImplicitSession
+    {
+        get => _enableImplicitSession;
+        set
+        {
+            _enableImplicitSession = value;
+            SaveValue(nameof(EnableImplicitSession), value);
+        }
+    }
+
+    private bool _enableImplicitSession;
+
     public ILoggerFactory LoggerFactory { get; init; } = NullLoggerFactory.Instance;
 
     public ICredentialsProvider? CredentialsProvider { get; init; }
@@ -495,6 +508,9 @@ static YdbConnectionOption()
             AddOption(new YdbConnectionOption<bool>(BoolExtractor,
                     (builder, disableServerBalancer) => builder.DisableServerBalancer = disableServerBalancer),
                 "DisableServerBalancer", "Disable Server Balancer");
+            AddOption(new YdbConnectionOption<bool>(BoolExtractor,
+                    (builder, enableImplicitSession) => builder.EnableImplicitSession = enableImplicitSession),
+                "EnableImplicitSession", "ImplicitSession");
         }
 
         private static void AddOption(YdbConnectionOption option, params string[] keys)

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

@@ -0,0 +1,154 @@
+using Moq;
+using Xunit;
+using Ydb.Sdk.Ado.Session;
+
+namespace Ydb.Sdk.Ado.Tests.Session;
+
+public class YdbImplicitStressTests : TestBase
+{
+    private static IDriver DummyDriver()
+    {
+        var m = new Mock<IDriver>(MockBehavior.Loose);
+        m.Setup(d => d.DisposeAsync()).Returns(ValueTask.CompletedTask);
+        return m.Object;
+    }
+
+    private sealed class Counter
+    {
+        public int Value;
+        public void Inc() => Interlocked.Increment(ref Value);
+    }
+
+    [Fact(Timeout = 30_000)]
+    public async Task Dispose_WaitsForAllLeases_AndSignalsOnEmptyExactlyOnce()
+    {
+        var driver = DummyDriver();
+
+        var opened = new Counter();
+        var closed = new Counter();
+
+        var source = new ImplicitSessionSource(driver);
+
+        var workers = Enumerable.Range(0, 200).Select(async _ =>
+        {
+            var rnd = Random.Shared;
+            for (var j = 0; j < 10; j++)
+            {
+                ISession s;
+                try
+                {
+                    s = await source.OpenSession(CancellationToken.None);
+                    opened.Inc();
+
+                    await Task.Delay(rnd.Next(0, 5));
+                }
+                catch (ObjectDisposedException)
+                {
+                    return;
+                }
+
+                var s2 = await source.OpenSession(CancellationToken.None);
+                s2.Dispose();
+
+                s.Dispose();
+                closed.Inc();
+            }
+        }).ToArray();
+
+        var disposer = Task.Run(async () =>
+        {
+            await Task.Delay(10);
+            await source.DisposeAsync();
+        });
+
+        await Task.WhenAll(workers.Append(disposer));
+
+        Assert.True(opened.Value > 0);
+        Assert.Equal(opened.Value, closed.Value);
+
+        await Assert.ThrowsAsync<ObjectDisposedException>(() => source.OpenSession(CancellationToken.None).AsTask());
+    }
+
+    [Fact(Timeout = 30_000)]
+    public async Task Stress_Counts_AreBalanced()
+    {
+        var driver = DummyDriver();
+
+        var opened = new Counter();
+        var closed = new Counter();
+
+        var source = new ImplicitSessionSource(driver);
+
+        var workers = Enumerable.Range(0, 200).Select(async _ =>
+        {
+            var rnd = Random.Shared;
+            for (var j = 0; j < 10; j++)
+            {
+                ISession s;
+                try
+                {
+                    s = await source.OpenSession(CancellationToken.None);
+                    opened.Inc();
+
+                    await Task.Delay(rnd.Next(0, 3));
+                }
+                catch (ObjectDisposedException)
+                {
+                    return;
+                }
+
+                var s2 = await source.OpenSession(CancellationToken.None);
+                s2.Dispose();
+
+                s.Dispose();
+                closed.Inc();
+            }
+        }).ToArray();
+
+        var disposer = Task.Run(async () => await source.DisposeAsync());
+
+        await Task.WhenAll(workers.Append(disposer));
+
+        Assert.Equal(opened.Value, closed.Value);
+        Assert.True(opened.Value > 0);
+
+        await Assert.ThrowsAsync<ObjectDisposedException>(() => source.OpenSession(CancellationToken.None).AsTask());
+    }
+
+    [Fact(Timeout = 30_000)]
+    public async Task Open_RacingWithDispose_StateRemainsConsistent()
+    {
+        var driver = DummyDriver();
+
+        var source = new ImplicitSessionSource(driver);
+
+        var opens = Enumerable.Range(0, 1000).Select(async _ =>
+        {
+            ISession s;
+            try
+            {
+                s = await source.OpenSession(CancellationToken.None);
+            }
+            catch (ObjectDisposedException)
+            {
+                return 0;
+            }
+
+            var s2 = await source.OpenSession(CancellationToken.None);
+            s2.Dispose();
+
+            s.Dispose();
+            return 1;
+        }).ToArray();
+
+        var disposeTask = Task.Run(async () =>
+        {
+            await Task.Yield();
+            await source.DisposeAsync();
+        });
+
+        await Task.WhenAll(opens.Append(disposeTask));
+
+        await Assert.ThrowsAsync<ObjectDisposedException>(() => source.OpenSession(CancellationToken.None).AsTask());
+    }
+}