Added server balancer disabling check in Query/SessionPool

Author: LiamHamsters

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

src/Ydb.Sdk/src/DriverConfig.cs

@@ -1,6 +1,7 @@
 using System.Reflection;
 using System.Security.Cryptography.X509Certificates;
 using Ydb.Sdk.Auth;
+using Ydb.Sdk.Pool;
 
 namespace Ydb.Sdk;
 
@@ -26,7 +27,7 @@ public class DriverConfig
 
     public int MaxSendMessageSize { get; init; } = GrpcDefaultSettings.MaxSendMessageSize;
     public int MaxReceiveMessageSize { get; init; } = GrpcDefaultSettings.MaxReceiveMessageSize;
-    public bool DisableServerBalancer { get; init; } = GrpcDefaultSettings.DisableServerBalancer;
+    public bool DisableServerBalancer { get; init; } = SessionPoolDefaultSettings.DisableServerBalancer;
 
     internal X509Certificate2Collection CustomServerCertificates { get; } = new();
     internal TimeSpan EndpointDiscoveryInterval = TimeSpan.FromMinutes(1);

src/Ydb.Sdk/src/IDriver.cs

@@ -182,8 +182,7 @@ protected async ValueTask<CallOptions> GetCallOptions(GrpcRequestSettings settin
             meta.Add(Metadata.RpcTraceIdHeader, settings.TraceId);
         }
 
-        if (!settings.List.Contains("session-balancer"))
-        {
+        foreach (var setting in settings.List) {
             meta.Add(Metadata.RpcClientCapabilitiesHeader, "session-balancer");
         }
 

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

@@ -11,6 +11,7 @@ internal abstract class SessionPool<TSession> where TSession : SessionBase<TSess
     private readonly ConcurrentQueue<TSession> _idleSessions = new();
     private readonly int _createSessionTimeoutMs;
     private readonly int _size;
+    private protected readonly SessionPoolConfig Config;
 
     protected readonly ILogger<SessionPool<TSession>> Logger;
 
@@ -253,12 +254,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.List.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;
 
@@ -17,7 +18,7 @@ public SessionPoolConfig(uint? sizeLimit = null)
     public TimeSpan KeepAliveTimeout { get; set; } = TimeSpan.FromSeconds(1);
     public TimeSpan CreateSessionTimeout { get; set; } = TimeSpan.FromSeconds(1);
 
-    public bool DisableServerBalancer { get; set; } = GrpcDefaultSettings.DisableServerBalancer;
+    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

@@ -26,7 +26,7 @@ private protected override async Task<GetSessionResponse> CreateSession()
         var createSessionSettings = new CreateSessionSettings
         {
             TransportTimeout = Config.CreateSessionTimeout,
-            OperationTimeout = Config.CreateSessionTimeout,
+            OperationTimeout = Config.CreateSessionTimeout
         };
 
         if (!Config.DisableServerBalancer)