Add: hint "session-balancer" (#462)

* Add: hint "session-balancer"
* Added DisableServerBalancer: support for server-side load balancing of sessions.
* Added server balancer disabling check in Query/SessionPool

Author: LiamHamsters

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,5 @@
+- Added DisableServerBalancer option to ADO.NET session creation; default false.
+
 ## v0.20.1
 
 - Fixed bug ADO.NET: `YdbSchema.SchemaObjects` and `Ydb.DescribeTable`methods are public for `EntityFrameworkCore.Ydb`.

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

@@ -38,6 +38,7 @@ private void InitDefaultValues()
         _maxReceiveMessageSize = GrpcDefaultSettings.MaxReceiveMessageSize;
         _disableDiscovery = GrpcDefaultSettings.DisableDiscovery;
         _createSessionTimeout = SessionPoolDefaultSettings.CreateSessionTimeoutSeconds;
+        _disableServerBalancer = false;
     }
 
     public string Host
@@ -236,6 +237,18 @@ public int MaxReceiveMessageSize
 
     private int _maxReceiveMessageSize;
 
+    public bool DisableServerBalancer
+    {
+        get => _disableServerBalancer;
+        set
+        {
+            _disableServerBalancer = value;
+            SaveValue(nameof(DisableServerBalancer), value);
+        }
+    }
+
+    private bool _disableServerBalancer;
+
     public bool DisableDiscovery
     {
         get => _disableDiscovery;
@@ -433,6 +446,9 @@ static YdbConnectionOption()
             AddOption(new YdbConnectionOption<int>(IntExtractor,
                     (builder, createSessionTimeout) => builder.CreateSessionTimeout = createSessionTimeout),
                 "CreateSessionTimeout", "Create Session Timeout");
+            AddOption(new YdbConnectionOption<bool>(BoolExtractor, (builder, disableServerBalancer) =>
+                    builder.DisableServerBalancer = disableServerBalancer),
+                "DisableServerBalancer", "Disable Server Balancer");
         }
 
         private static void AddOption(YdbConnectionOption option, params string[] keys)

src/Ydb.Sdk/src/GrpcRequestSettings.cs

@@ -1,5 +1,4 @@
-using System.Collections.Immutable;
-using Google.Protobuf.WellKnownTypes;
+using Google.Protobuf.WellKnownTypes;
 using Ydb.Operations;
 
 namespace Ydb.Sdk;
@@ -8,9 +7,9 @@ public class GrpcRequestSettings
 {
     public string TraceId { get; set; } = string.Empty;
     public TimeSpan TransportTimeout { get; set; } = TimeSpan.Zero;
-    public ImmutableArray<string> CustomClientHeaders { get; } = new();
     public CancellationToken CancellationToken = default;
 
+    internal List<string> ClientCapabilities { get; } = new();
     internal long NodeId { get; set; }
     internal Action<Grpc.Core.Metadata> TrailersHandler { get; set; } = _ => { };
 }

src/Ydb.Sdk/src/IDriver.cs

@@ -182,6 +182,11 @@ protected async ValueTask<CallOptions> GetCallOptions(GrpcRequestSettings settin
             meta.Add(Metadata.RpcTraceIdHeader, settings.TraceId);
         }
 
+        foreach (var clientCapabilitiesHeader in settings.ClientCapabilities)
+        {
+            meta.Add(Metadata.RpcClientCapabilitiesHeader, clientCapabilitiesHeader);
+        }
+
         var options = new CallOptions(headers: meta, cancellationToken: settings.CancellationToken);
 
         if (settings.TransportTimeout != TimeSpan.Zero)

src/Ydb.Sdk/src/Metadata.cs

@@ -8,6 +8,7 @@ internal static class Metadata
     public const string RpcTraceIdHeader = "x-ydb-trace-id";
     public const string RpcSdkInfoHeader = "x-ydb-sdk-build-info";
     public const string RpcServerHintsHeader = "x-ydb-server-hints";
+    public const string RpcClientCapabilitiesHeader = "x-ydb-client-capabilities";
 
     public const string GracefulShutdownHint = "session-close";
 }

src/Ydb.Sdk/src/Pool/SessionPool.cs

@@ -12,16 +12,18 @@ internal abstract class SessionPool<TSession> where TSession : SessionBase<TSess
     private readonly int _createSessionTimeoutMs;
     private readonly int _size;
 
+    protected readonly SessionPoolConfig Config;
     protected readonly ILogger<SessionPool<TSession>> Logger;
 
     private volatile int _waitingCount;
     private volatile bool _disposed;
 
-    protected SessionPool(ILogger<SessionPool<TSession>> logger, SessionPoolConfig sessionPoolConfig)
+    protected SessionPool(ILogger<SessionPool<TSession>> logger, SessionPoolConfig config)
     {
         Logger = logger;
-        _size = sessionPoolConfig.MaxSessionPool;
-        _createSessionTimeoutMs = sessionPoolConfig.CreateSessionTimeout * 1000;
+        Config = config;
+        _size = config.MaxSessionPool;
+        _createSessionTimeoutMs = config.CreateSessionTimeout * 1000;
         _semaphore = new SemaphoreSlim(_size);
     }
 
@@ -253,12 +255,15 @@ internal TS MakeGrpcRequestSettings<TS>(TS settings) where TS : GrpcRequestSetti
 internal record SessionPoolConfig(
     int MaxSessionPool = SessionPoolDefaultSettings.MaxSessionPool,
     int CreateSessionTimeout = SessionPoolDefaultSettings.CreateSessionTimeoutSeconds,
-    bool DisposeDriver = false
+    bool DisposeDriver = false,
+    bool DisableServerBalancer = SessionPoolDefaultSettings.DisableServerBalancer
 );
 
 internal static class SessionPoolDefaultSettings
 {
     internal const int MaxSessionPool = 100;
 
     internal const int CreateSessionTimeoutSeconds = 5;
+
+    internal const bool DisableServerBalancer = false;
 }

src/Ydb.Sdk/src/Services/Query/SessionPool.cs

@@ -35,10 +35,20 @@ protected override async Task<Session> CreateSession(
         CancellationToken cancellationToken = default
     )
     {
+        var requestSettings = new GrpcRequestSettings
+        {
+            CancellationToken = cancellationToken
+        };
+
+        if (!Config.DisableServerBalancer)
+        {
+            requestSettings.ClientCapabilities.Add("session-balancer");
+        }
+
         var response = await _driver.UnaryCall(
             QueryService.CreateSessionMethod,
             CreateSessionRequest,
-            new GrpcRequestSettings { CancellationToken = cancellationToken }
+            requestSettings
         );
 
         Status.FromProto(response.Status, response.Issues).EnsureSuccess();

src/Ydb.Sdk/src/Services/Sessions/SessionPoolBase.cs

@@ -1,5 +1,6 @@
 using Microsoft.Extensions.Logging;
 using Ydb.Sdk.Client;
+using Ydb.Sdk.Pool;
 
 namespace Ydb.Sdk.Services.Sessions;
 
@@ -16,6 +17,8 @@ public SessionPoolConfig(uint? sizeLimit = null)
     public TimeSpan PeriodicCheckInterval { get; set; } = TimeSpan.FromSeconds(10);
     public TimeSpan KeepAliveTimeout { get; set; } = TimeSpan.FromSeconds(1);
     public TimeSpan CreateSessionTimeout { get; set; } = TimeSpan.FromSeconds(1);
+
+    public bool DisableServerBalancer { get; set; } = SessionPoolDefaultSettings.DisableServerBalancer;
 }
 
 public class GetSessionResponse<TSession> : ResponseWithResultBase<TSession>, IDisposable where TSession : SessionBase

src/Ydb.Sdk/src/Services/Table/SessionPool.cs

@@ -23,11 +23,18 @@ public SessionPool(Driver driver, SessionPoolConfig config) :
 
     private protected override async Task<GetSessionResponse> CreateSession()
     {
-        var createSessionResponse = await _tableClient.CreateSession(new CreateSessionSettings
+        var createSessionSettings = new CreateSessionSettings
         {
             TransportTimeout = Config.CreateSessionTimeout,
             OperationTimeout = Config.CreateSessionTimeout
-        });
+        };
+
+        if (!Config.DisableServerBalancer)
+        {
+            createSessionSettings.ClientCapabilities.Add("session-balancer");
+        }
+
+        var createSessionResponse = await _tableClient.CreateSession(createSessionSettings);
 
         lock (Lock)
         {

src/Ydb.Sdk/tests/Ado/YdbConnectionStringBuilderTests.cs

@@ -24,6 +24,7 @@ public void InitDefaultValues_WhenEmptyConstructorInvoke_ReturnDefaultConnection
         Assert.Equal(64 * 1024 * 1024, ydbConnectionStringBuilder.MaxSendMessageSize);
         Assert.Equal(64 * 1024 * 1024, ydbConnectionStringBuilder.MaxReceiveMessageSize);
         Assert.False(ydbConnectionStringBuilder.DisableDiscovery);
+        Assert.False(ydbConnectionStringBuilder.DisableServerBalancer);
         Assert.Equal(5, ydbConnectionStringBuilder.CreateSessionTimeout);
     }
 
@@ -45,7 +46,7 @@ public void InitConnectionStringBuilder_WhenExpectedKeys_ReturnUpdatedConnection
             "ConnectTimeout=30;KeepAlivePingDelay=30;KeepAlivePingTimeout=60;" +
             "EnableMultipleHttp2Connections=true;CreateSessionTimeout=30;" +
             "MaxSendMessageSize=1000000;MaxReceiveMessageSize=1000000;" +
-            "DisableDiscovery=true"
+            "DisableDiscovery=true;DisableServerBalancer=true"
         );
 
         Assert.Equal(2135, connectionString.Port);
@@ -64,8 +65,9 @@ public void InitConnectionStringBuilder_WhenExpectedKeys_ReturnUpdatedConnection
                      "ConnectTimeout=30;KeepAlivePingDelay=30;KeepAlivePingTimeout=60;" +
                      "EnableMultipleHttp2Connections=True;CreateSessionTimeout=30;" +
                      "MaxSendMessageSize=1000000;MaxReceiveMessageSize=1000000;" +
-                     "DisableDiscovery=True", connectionString.ConnectionString);
+                     "DisableDiscovery=True;DisableServerBalancer=True", connectionString.ConnectionString);
         Assert.True(connectionString.DisableDiscovery);
+        Assert.True(connectionString.DisableServerBalancer);
         Assert.Equal(30, connectionString.CreateSessionTimeout);
     }