Add sessions per pool to DbPool config

DbPool: Add session per pool argument to config

7 files changed, 23 insertions, 16 deletions

diff --git a/ydb/core/yq/libs/control_plane_config/control_plane_config.cpp b/ydb/core/yq/libs/control_plane_config/control_plane_config.cpp
index 170c3a56aa6..8330a201265 100644
--- a/ydb/core/yq/libs/control_plane_config/control_plane_config.cpp
+++ b/ydb/core/yq/libs/control_plane_config/control_plane_config.cpp
@@ -69,7 +69,7 @@ public:
         Become(&TControlPlaneConfigActor::StateFunc);
         if (Config.GetUseDbMapping()) {
             YdbConnection = NewYdbConnection(Config.GetStorage(), CredProviderFactory, YqSharedResources->CoreYdbDriver);
-            DbPool = YqSharedResources->DbPoolHolder->GetOrCreate(static_cast<ui32>(EDbPoolId::MAIN), 10);
+            DbPool = YqSharedResources->DbPoolHolder->GetOrCreate(static_cast<ui32>(EDbPoolId::MAIN));
             TablePathPrefix = YdbConnection->TablePathPrefix;
             Schedule(TDuration::Zero(), new NActors::TEvents::TEvWakeup());
         } else {
diff --git a/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage.cpp b/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage.cpp
index 6cc1856deb4..0612ed48ed4 100644
--- a/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage.cpp
+++ b/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage.cpp
@@ -37,7 +37,7 @@ void TYdbControlPlaneStorageActor::Bootstrap() {
     NLwTraceMonPage::ProbeRegistry().AddProbesList(LWTRACE_GET_PROBES(YQ_CONTROL_PLANE_STORAGE_PROVIDER));
 
     YdbConnection = NewYdbConnection(Config->Proto.GetStorage(), CredProviderFactory, YqSharedResources->CoreYdbDriver);
-    DbPool = YqSharedResources->DbPoolHolder->GetOrCreate(static_cast<ui32>(EDbPoolId::MAIN), 10);
+    DbPool = YqSharedResources->DbPoolHolder->GetOrCreate(static_cast<ui32>(EDbPoolId::MAIN));
     TablePathPrefix = YdbConnection->TablePathPrefix;
     CreateDirectory();
 
diff --git a/ydb/core/yq/libs/quota_manager/quota_manager.cpp b/ydb/core/yq/libs/quota_manager/quota_manager.cpp
index 5c3b3d665c1..aedfd6b5cb1 100644
--- a/ydb/core/yq/libs/quota_manager/quota_manager.cpp
+++ b/ydb/core/yq/libs/quota_manager/quota_manager.cpp
@@ -189,7 +189,7 @@ public:
 
         YdbConnection = NewYdbConnection(StorageConfig, CredProviderFactory, YqSharedResources->CoreYdbDriver);
         TablePathPrefix = YdbConnection->TablePathPrefix;
-        DbPool = YqSharedResources->DbPoolHolder->GetOrCreate(static_cast<ui32>(EDbPoolId::MAIN), 10);
+        DbPool = YqSharedResources->DbPoolHolder->GetOrCreate(static_cast<ui32>(EDbPoolId::MAIN));
         Send(NActors::GetNameserviceActorId(), new NActors::TEvInterconnect::TEvListNodes());
         Become(&TQuotaManagementService::StateFunc);
         LOG_I("STARTED");
diff --git a/ydb/core/yq/libs/shared_resources/shared_resources.cpp b/ydb/core/yq/libs/shared_resources/shared_resources.cpp
index 4ab3cbd379d..1066763ff1e 100644
--- a/ydb/core/yq/libs/shared_resources/shared_resources.cpp
+++ b/ydb/core/yq/libs/shared_resources/shared_resources.cpp
@@ -38,7 +38,8 @@ struct TYqSharedResourcesImpl : public TActorSystemPtrMixin, public TYqSharedRes
     NDbPool::TConfig PrepareDbPoolConfig(const NYq::NConfig::TConfig& config) {
         NDbPool::TConfig dbPoolConfig;
         const auto& storageConfig = config.GetDbPool().GetStorage();
-        dbPoolConfig.SetMaxSessionCount(config.GetDbPool().GetMaxSessionCount());
+        auto maxSessionCount = config.GetDbPool().GetMaxSessionCount();
+        (*dbPoolConfig.MutablePools())[0] = maxSessionCount ? maxSessionCount : 10;
         dbPoolConfig.SetEndpoint(storageConfig.GetEndpoint());
         dbPoolConfig.SetDatabase(storageConfig.GetDatabase());
         dbPoolConfig.SetOAuthFile(storageConfig.GetOAuthFile());
diff --git a/ydb/library/db_pool/db_pool.cpp b/ydb/library/db_pool/db_pool.cpp
index 8e1c7e22cd4..50cd7030592 100644
--- a/ydb/library/db_pool/db_pool.cpp
+++ b/ydb/library/db_pool/db_pool.cpp
@@ -218,10 +218,6 @@ static void PrepareConfig(NDbPool::TConfig& config) {
     if (!config.GetToken() && config.GetOAuthFile()) {
         config.SetToken(StripString(TFileInput(config.GetOAuthFile()).ReadAll()));
     }
-
-    if (!config.GetMaxSessionCount()) {
-        config.SetMaxSessionCount(10);
-    }
 }
 
 TDbPoolMap::TDbPoolMap(
@@ -267,25 +263,35 @@ void TDbPoolMap::Reset(const NDbPool::TConfig& config) {
     TableClient = nullptr;
 }
 
-TDbPool::TPtr TDbPoolHolder::GetOrCreate(ui32 dbPoolId, ui32 sessionsCount) {
-    return Pools->GetOrCreate(dbPoolId, sessionsCount);
+TDbPool::TPtr TDbPoolHolder::GetOrCreate(ui32 dbPoolId) {
+    return Pools->GetOrCreate(dbPoolId);
 }
 
-TDbPool::TPtr TDbPoolMap::GetOrCreate(ui32 dbPoolId, ui32 sessionsCount) {
+TDbPool::TPtr TDbPoolMap::GetOrCreate(ui32 dbPoolId) {
     TGuard<TMutex> lock(Mutex);
     auto it = Pools.find(dbPoolId);
     if (it != Pools.end()) {
         return it->second;
     }
 
+    auto it_pool = Config.GetPools().find(dbPoolId);
+    if (it_pool == Config.GetPools().end()) {
+        return nullptr;
+    }
+
     if (!Config.GetEndpoint()) {
         return nullptr;
     }
 
     if (!TableClient) {
+        auto maxSessionCount = 0;
+        for (const auto& pool : Config.GetPools())
+        {
+            maxSessionCount += pool.second;
+        }
         auto clientSettings = NYdb::NTable::TClientSettings()
             .UseQueryCache(false)
-            .SessionPoolSettings(NYdb::NTable::TSessionPoolSettings().MaxActiveSessions(1 + Config.GetMaxSessionCount()))
+            .SessionPoolSettings(NYdb::NTable::TSessionPoolSettings().MaxActiveSessions(1 + maxSessionCount))
             .Database(Config.GetDatabase())
             .DiscoveryEndpoint(Config.GetEndpoint())
             .DiscoveryMode(NYdb::EDiscoveryMode::Async);
@@ -300,7 +306,7 @@ TDbPool::TPtr TDbPoolMap::GetOrCreate(ui32 dbPoolId, ui32 sessionsCount) {
         TableClient = MakeHolder<NYdb::NTable::TTableClient>(Driver, clientSettings);
     }
 
-    TDbPool::TPtr dbPool = new TDbPool(sessionsCount, *TableClient, Counters);
+    TDbPool::TPtr dbPool = new TDbPool(it_pool->second, *TableClient, Counters);
     Pools.emplace(dbPoolId, dbPool);
     return dbPool;
 }
diff --git a/ydb/library/db_pool/db_pool.h b/ydb/library/db_pool/db_pool.h
index c321f5a400c..2aeebac4a5b 100644
--- a/ydb/library/db_pool/db_pool.h
+++ b/ydb/library/db_pool/db_pool.h
@@ -39,7 +39,7 @@ class TDbPoolMap: public TThrRefBase {
 public:
     using TPtr = TIntrusivePtr<TDbPoolMap>;
 
-    TDbPool::TPtr GetOrCreate(ui32 poolId, ui32 sessionsCount);
+    TDbPool::TPtr GetOrCreate(ui32 poolId);
 
 private:
     friend class TDbPoolHolder;
@@ -75,7 +75,7 @@ public:
     ~TDbPoolHolder();
 
     void Reset(const NDbPool::TConfig& config);
-    TDbPool::TPtr GetOrCreate(ui32 poolId, ui32 sessionsCount);
+    TDbPool::TPtr GetOrCreate(ui32 poolId);
     NYdb::TDriver& GetDriver();
     TDbPoolMap::TPtr Get();
 
diff --git a/ydb/library/db_pool/protos/config.proto b/ydb/library/db_pool/protos/config.proto
index f8c0825d971..d2fac3a8a65 100644
--- a/ydb/library/db_pool/protos/config.proto
+++ b/ydb/library/db_pool/protos/config.proto
@@ -4,7 +4,7 @@ option cc_enable_arenas = true;
 package NDbPool;
 
 message TConfig {
-    uint32 MaxSessionCount = 1;
+    map<uint32, uint32> Pools = 1; // sessions per pool
     string Endpoint = 2;
     string Database = 3;
     string OAuthFile = 4;