From d93a17b29196ed7a7b1f03495ff6001b954e269c Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Fri, 7 Feb 2025 13:43:11 +0500 Subject: [PATCH] YQ-4092 KqpRun improve tenants flags (#14279) --- ydb/tests/tools/kqprun/kqprun.cpp | 53 ++++- ydb/tests/tools/kqprun/src/common.h | 6 +- .../tools/kqprun/src/proto/storage_meta.proto | 14 ++ ydb/tests/tools/kqprun/src/ydb_setup.cpp | 193 ++++++++++++------ 4 files changed, 192 insertions(+), 74 deletions(-) diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index ec8ce224a64c..1ddbd4adc78d 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -95,7 +95,7 @@ struct TExecutionOptions { .TraceId = DefaultTraceId, .PoolId = "", .UserSID = BUILTIN_ACL_ROOT, - .Database = "", + .Database = GetValue(0, Databases, TString()), .Timeout = TDuration::Zero() }; } @@ -135,15 +135,15 @@ struct TExecutionOptions { private: void ValidateOptionsSizes(const TRunnerOptions& runnerOptions) const { - const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName) { - if (checkSize > numberQueries) { + const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName, bool useInSchemeQuery = false) { + if (checkSize > std::max(numberQueries, static_cast(useInSchemeQuery ? 1 : 0))) { ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of script queries is " << numberQueries; } }; checker(ExecutionCases.size(), "execution cases"); checker(ScriptQueryActions.size(), "script query actions"); - checker(Databases.size(), "databases"); + checker(Databases.size(), "databases", true); checker(TraceIds.size(), "trace ids"); checker(PoolIds.size(), "pool ids"); checker(UserSIDs.size(), "user SIDs"); @@ -257,7 +257,7 @@ struct TExecutionOptions { static void ValidateStorageSettings(const TYdbSetupSettings& ydbSettings) { if (ydbSettings.DisableDiskMock) { - if (ydbSettings.NodeCount + ydbSettings.SharedTenants.size() + ydbSettings.DedicatedTenants.size() > 1) { + if (ydbSettings.NodeCount + ydbSettings.Tenants.size() > 1) { ythrow yexception() << "Disable disk mock cannot be used for multi node clusters (already disabled)"; } else if (ydbSettings.PDisksPath) { ythrow yexception() << "Disable disk mock cannot be used with real PDisks (already disabled)"; @@ -876,17 +876,50 @@ class TMain : public TMainClassArgs { .DefaultValue(RunnerOptions.YdbSettings.DomainName) .StoreResult(&RunnerOptions.YdbSettings.DomainName); - options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain") + const auto addTenant = [this](const TString& type, TStorageMeta::TTenant::EType protoType, const NLastGetopt::TOptsParser* option) { + TStringBuf tenant; + TStringBuf nodesCountStr; + TStringBuf(option->CurVal()).Split(':', tenant, nodesCountStr); + if (tenant.empty()) { + ythrow yexception() << type << " tenant name should not be empty"; + } + + TStorageMeta::TTenant tenantInfo; + tenantInfo.SetType(protoType); + tenantInfo.SetNodesCount(nodesCountStr ? FromString(nodesCountStr) : 1); + if (tenantInfo.GetNodesCount() == 0) { + ythrow yexception() << type << " tenant should have at least one node"; + } + + if (!RunnerOptions.YdbSettings.Tenants.emplace(tenant, tenantInfo).second) { + ythrow yexception() << "Got duplicated tenant name: " << tenant; + } + }; + options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain (for node count use dedicated-name:node-count)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.DedicatedTenants); + .Handler1(std::bind(addTenant, "Dedicated", TStorageMeta::TTenant::DEDICATED, std::placeholders::_1)); - options.AddLongOption("shared", "Shared tenant path, relative inside domain") + options.AddLongOption("shared", "Shared tenant path, relative inside domain (for node count use dedicated-name:node-count)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.SharedTenants); + .Handler1(std::bind(addTenant, "Shared", TStorageMeta::TTenant::SHARED, std::placeholders::_1)); options.AddLongOption("serverless", "Serverless tenant path, relative inside domain (use string serverless-name@shared-name to specify shared database)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.ServerlessTenants); + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TStringBuf serverless; + TStringBuf shared; + TStringBuf(option->CurVal()).Split('@', serverless, shared); + if (serverless.empty()) { + ythrow yexception() << "Serverless tenant name should not be empty"; + } + + TStorageMeta::TTenant tenantInfo; + tenantInfo.SetType(TStorageMeta::TTenant::SERVERLESS); + tenantInfo.SetSharedTenant(TString(shared)); + if (!RunnerOptions.YdbSettings.Tenants.emplace(serverless, tenantInfo).second) { + ythrow yexception() << "Got duplicated tenant name: " << serverless; + } + }); options.AddLongOption("storage-size", TStringBuilder() << "Domain storage size in gigabytes (" << NKikimr::NBlobDepot::FormatByteSize(DEFAULT_STORAGE_SIZE) << " by default)") .RequiredArgument("uint") diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index e85378c8667d..4f4f4ac36913 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -9,6 +9,8 @@ #include #include +#include + #include #include @@ -48,9 +50,7 @@ struct TYdbSetupSettings { ui32 NodeCount = 1; TString DomainName = "Root"; - std::unordered_set DedicatedTenants; - std::unordered_set SharedTenants; - std::unordered_set ServerlessTenants; + std::map Tenants; TDuration HealthCheckTimeout = TDuration::Seconds(10); EHealthCheck HealthCheckLevel = EHealthCheck::NodesCount; bool SameSession = false; diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto index 98401b8200d0..ee4a7c4162a2 100644 --- a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto +++ b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto @@ -3,6 +3,20 @@ syntax = "proto3"; package NKqpRun; message TStorageMeta { + message TTenant { + enum EType { + DEDICATED = 0; + SHARED = 1; + SERVERLESS = 2; + } + + EType Type = 1; + uint32 NodesCount = 2; + string SharedTenant = 3; // Only for serverless tenants + } + uint64 StorageGeneration = 1; uint64 StorageSize = 2; + string DomainName = 3; + map Tenants = 4; } diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 2f646d06c675..d754d22e254c 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -177,7 +177,7 @@ class TYdbSetup::TImpl { serverSettings.SetFrFactory(functionRegistryFactory); } - void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const { + void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) { TFsPath diskPath; if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") { diskPath = *Settings_.PDisksPath; @@ -195,32 +195,33 @@ class TYdbSetup::TImpl { } bool formatDisk = true; - NKqpRun::TStorageMeta storageMeta; if (diskPath) { - const auto storageMetaPath = TFsPath(diskPath).Child("kqprun_storage_meta.conf"); - if (storageMetaPath.Exists() && !Settings_.FormatStorage) { - if (!google::protobuf::TextFormat::ParseFromString(TFileInput(storageMetaPath.GetPath()).ReadAll(), &storageMeta)) { + StorageMetaPath_ = TFsPath(diskPath).Child("kqprun_storage_meta.conf"); + if (StorageMetaPath_.Exists() && !Settings_.FormatStorage) { + if (!google::protobuf::TextFormat::ParseFromString(TFileInput(StorageMetaPath_.GetPath()).ReadAll(), &StorageMeta_)) { ythrow yexception() << "Storage meta is corrupted, please use --format-storage"; } - storageMeta.SetStorageGeneration(storageMeta.GetStorageGeneration() + 1); + StorageMeta_.SetStorageGeneration(StorageMeta_.GetStorageGeneration() + 1); formatDisk = false; } - if (Settings_.DiskSize && storageMeta.GetStorageSize() != *Settings_.DiskSize) { + if (Settings_.DiskSize && StorageMeta_.GetStorageSize() != *Settings_.DiskSize) { if (!formatDisk) { - ythrow yexception() << "Cannot change disk size without formatting storage, current disk size " << NKikimr::NBlobDepot::FormatByteSize(storageMeta.GetStorageSize()) << ", please use --format-storage"; + ythrow yexception() << "Cannot change disk size without formatting storage, current disk size " << NKikimr::NBlobDepot::FormatByteSize(StorageMeta_.GetStorageSize()) << ", please use --format-storage"; } - storageMeta.SetStorageSize(*Settings_.DiskSize); - } else if (!storageMeta.GetStorageSize()) { - storageMeta.SetStorageSize(DEFAULT_STORAGE_SIZE); + StorageMeta_.SetStorageSize(*Settings_.DiskSize); + } else if (!StorageMeta_.GetStorageSize()) { + StorageMeta_.SetStorageSize(DEFAULT_STORAGE_SIZE); } - TString storageMetaStr; - google::protobuf::TextFormat::PrintToString(storageMeta, &storageMetaStr); + const TString& domainName = NKikimr::CanonizePath(Settings_.DomainName); + if (!StorageMeta_.GetDomainName()) { + StorageMeta_.SetDomainName(domainName); + } else if (StorageMeta_.GetDomainName() != domainName) { + ythrow yexception() << "Cannot change domain name without formatting storage, current name " << StorageMeta_.GetDomainName() << ", please use --format-storage"; + } - TFileOutput storageMetaOutput(storageMetaPath.GetPath()); - storageMetaOutput.Write(storageMetaStr); - storageMetaOutput.Finish(); + UpdateStorageMeta(); } TString storagePath = diskPath.GetPath(); @@ -237,7 +238,7 @@ class TYdbSetup::TImpl { serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.PDisksPath); serverSettings.SetCustomDiskParams(storage); - serverSettings.SetStorageGeneration(storageMeta.GetStorageGeneration()); + serverSettings.SetStorageGeneration(StorageMeta_.GetStorageGeneration()); } NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) { @@ -278,28 +279,54 @@ class TYdbSetup::TImpl { serverSettings.SetGrpcPort(grpcPort); } - if (!Settings_.SharedTenants.empty() || !Settings_.DedicatedTenants.empty()) { - serverSettings.SetDynamicNodeCount(Settings_.SharedTenants.size() + Settings_.DedicatedTenants.size()); - for (const TString& dedicatedTenant : Settings_.DedicatedTenants) { - serverSettings.AddStoragePoolType(dedicatedTenant); - } - for (const auto& sharedTenant : Settings_.SharedTenants) { - serverSettings.AddStoragePoolType(sharedTenant); + for (const auto& [tenantPath, tenantInfo] : StorageMeta_.GetTenants()) { + Settings_.Tenants.emplace(tenantPath, tenantInfo); + } + + ui32 dynNodesCount = 0; + for (const auto& [tenantPath, tenantInfo] : Settings_.Tenants) { + if (tenantInfo.GetType() != TStorageMeta::TTenant::SERVERLESS) { + serverSettings.AddStoragePoolType(tenantPath); + dynNodesCount += tenantInfo.GetNodesCount(); } } + serverSettings.SetDynamicNodeCount(dynNodesCount); return serverSettings; } - void CreateTenant(Ydb::Cms::CreateDatabaseRequest&& request, const TString& type) const { - const auto path = request.path(); - Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating " << type << " tenant " << path << "..." << CoutColors_.Default() << Endl; - Tenants_->CreateTenant(std::move(request)); + void CreateTenant(Ydb::Cms::CreateDatabaseRequest&& request, const TString& relativePath, const TString& type, TStorageMeta::TTenant tenantInfo) { + const auto absolutePath = request.path(); + const auto [it, inserted] = StorageMeta_.MutableTenants()->emplace(relativePath, tenantInfo); + if (inserted) { + if (Settings_.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating " << type << " tenant " << absolutePath << "..." << CoutColors_.Default() << Endl; + } + Tenants_->CreateTenant(std::move(request), tenantInfo.GetNodesCount()); + UpdateStorageMeta(); + } else { + if (it->second.GetType() != tenantInfo.GetType()) { + ythrow yexception() << "Can not change tenant " << absolutePath << " type without formatting storage, current type " << TStorageMeta::TTenant::EType_Name(it->second.GetType()) << ", please use --format-storage"; + } + if (it->second.GetSharedTenant() != tenantInfo.GetSharedTenant()) { + ythrow yexception() << "Can not change tenant " << absolutePath << " shared resources without formatting storage from '" << it->second.GetSharedTenant() << "', please use --format-storage"; + } + if (it->second.GetNodesCount() != tenantInfo.GetNodesCount()) { + it->second.SetNodesCount(tenantInfo.GetNodesCount()); + UpdateStorageMeta(); + } + if (Settings_.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Starting " << type << " tenant " << absolutePath << "..." << CoutColors_.Default() << Endl; + } + if (!request.has_serverless_resources()) { + Tenants_->Run(absolutePath, tenantInfo.GetNodesCount()); + } + } if (Settings_.MonitoringEnabled) { - ui32 nodeIndex = GetNodeIndexForDatabase(path); + ui32 nodeIndex = GetNodeIndexForDatabase(absolutePath); NActors::TActorId edgeActor = GetRuntime()->AllocateEdgeActor(nodeIndex); - GetRuntime()->Register(NKikimr::CreateBoardPublishActor(NKikimr::MakeEndpointsBoardPath(path), "", edgeActor, 0, true), nodeIndex, GetRuntime()->GetAppData(nodeIndex).UserPoolId); + GetRuntime()->Register(NKikimr::CreateBoardPublishActor(NKikimr::MakeEndpointsBoardPath(absolutePath), "", edgeActor, 0, true), nodeIndex, GetRuntime()->GetAppData(nodeIndex).UserPoolId); } } @@ -309,39 +336,50 @@ class TYdbSetup::TImpl { } void CreateTenants() { - for (const TString& dedicatedTenant : Settings_.DedicatedTenants) { + std::set sharedTenants; + std::map serverlessTenants; + for (const auto& [tenantPath, tenantInfo] : Settings_.Tenants) { Ydb::Cms::CreateDatabaseRequest request; - request.set_path(GetTenantPath(dedicatedTenant)); - AddTenantStoragePool(request.mutable_resources()->add_storage_units(), dedicatedTenant); - CreateTenant(std::move(request), "dedicated"); + request.set_path(GetTenantPath(tenantPath)); + + switch (tenantInfo.GetType()) { + case TStorageMeta::TTenant::DEDICATED: + AddTenantStoragePool(request.mutable_resources()->add_storage_units(), tenantPath); + CreateTenant(std::move(request), tenantPath, "dedicated", tenantInfo); + break; + + case TStorageMeta::TTenant::SHARED: + sharedTenants.emplace(tenantPath); + AddTenantStoragePool(request.mutable_shared_resources()->add_storage_units(), tenantPath); + CreateTenant(std::move(request), tenantPath, "shared", tenantInfo); + break; + + case TStorageMeta::TTenant::SERVERLESS: + serverlessTenants.emplace(tenantPath, tenantInfo); + break; + + default: + ythrow yexception() << "Unexpected tenant type: " << TStorageMeta::TTenant::EType_Name(tenantInfo.GetType()); + break; + } } - for (const TString& sharedTenant : Settings_.SharedTenants) { - Ydb::Cms::CreateDatabaseRequest request; - request.set_path(GetTenantPath(sharedTenant)); - AddTenantStoragePool(request.mutable_shared_resources()->add_storage_units(), sharedTenant); - CreateTenant(std::move(request), "shared"); - } + for (auto [tenantPath, tenantInfo] : serverlessTenants) { + if (!tenantInfo.GetSharedTenant()) { + if (sharedTenants.empty()) { + ythrow yexception() << "Can not create serverless tenant, there is no shared tenants, please use `--shared `"; + } + if (sharedTenants.size() > 1) { + ythrow yexception() << "Can not create serverless tenant, there is more than one shared tenant, please use `--serverless " << tenantPath << "@`"; + } + tenantInfo.SetSharedTenant(*sharedTenants.begin()); + } - ServerlessToShared_.reserve(Settings_.ServerlessTenants.size()); - for (const TString& serverlessTenant : Settings_.ServerlessTenants) { Ydb::Cms::CreateDatabaseRequest request; - if (serverlessTenant.Contains('@')) { - TStringBuf serverless; - TStringBuf shared; - TStringBuf(serverlessTenant).Split('@', serverless, shared); - - request.set_path(GetTenantPath(TString(serverless))); - request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(TString(shared))); - } else if (!Settings_.SharedTenants.empty()) { - request.set_path(GetTenantPath(serverlessTenant)); - request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(*Settings_.SharedTenants.begin())); - } else { - ythrow yexception() << "Can not create serverless tenant " << serverlessTenant << ", there is no shared tenants"; - } + request.set_path(GetTenantPath(tenantPath)); + request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(tenantInfo.GetSharedTenant())); ServerlessToShared_[request.path()] = request.serverless_resources().shared_database_path(); - - CreateTenant(std::move(request), "serverless"); + CreateTenant(std::move(request), tenantPath, "serverless", tenantInfo); } } @@ -402,8 +440,18 @@ class TYdbSetup::TImpl { if (Settings_.MonitoringEnabled && Settings_.VerboseLevel >= EVerbose::Info) { for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount; ++nodeIndex) { - Cout << CoutColors_.Cyan() << "Monitoring port" << (Settings_.NodeCount > 1 ? TStringBuilder() << " for node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; + Cout << CoutColors_.Cyan() << "Monitoring port" << (Server_->StaticNodes() + Server_->DynamicNodes() > 1 ? TStringBuilder() << " for static node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; } + const auto printTenantNodes = [this](const std::pair& tenantInfo) { + if (tenantInfo.second.GetType() == TStorageMeta::TTenant::SERVERLESS) { + return; + } + const auto& nodes = Tenants_->List(GetTenantPath(tenantInfo.first)); + for (auto it = nodes.rbegin(); it != nodes.rend(); ++it) { + Cout << CoutColors_.Cyan() << "Monitoring port for dynamic node " << *it + 1 << " [" << tenantInfo.first << "]: " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(*it) << Endl; + } + }; + std::for_each(Settings_.Tenants.rbegin(), Settings_.Tenants.rend(), std::bind(printTenantNodes, std::placeholders::_1)); } if (Settings_.GrpcEnabled && Settings_.VerboseLevel >= EVerbose::Info) { @@ -590,12 +638,12 @@ class TYdbSetup::TImpl { } TString GetDatabasePath(const TString& database) const { - return NKikimr::CanonizePath(database ? database : Settings_.DomainName); + return NKikimr::CanonizePath(database ? database : GetDefaultDatabase()); } ui32 GetNodeIndexForDatabase(const TString& path) const { - auto canonizedPath = NKikimr::CanonizePath(path); - if (canonizedPath.empty() || canonizedPath == NKikimr::CanonizePath(Settings_.DomainName)) { + auto canonizedPath = NKikimr::CanonizePath(path ? path : GetDefaultDatabase()); + if (canonizedPath == NKikimr::CanonizePath(Settings_.DomainName)) { return RandomNumber(Settings_.NodeCount); } @@ -610,6 +658,27 @@ class TYdbSetup::TImpl { ythrow yexception() << "Unknown tenant '" << canonizedPath << "'"; } + TString GetDefaultDatabase() const { + if (StorageMeta_.TenantsSize() > 1) { + ythrow yexception() << "Can not choose default database, there is more than one tenants, please use `-D `"; + } + if (StorageMeta_.TenantsSize() == 1) { + return StorageMeta_.GetTenants().begin()->first; + } + return Settings_.DomainName; + } + + void UpdateStorageMeta() const { + if (StorageMetaPath_) { + TString storageMetaStr; + google::protobuf::TextFormat::PrintToString(StorageMeta_, &storageMetaStr); + + TFileOutput storageMetaOutput(StorageMetaPath_.GetPath()); + storageMetaOutput.Write(storageMetaStr); + storageMetaOutput.Finish(); + } + } + private: TYdbSetupSettings Settings_; NColorizer::TColors CoutColors_; @@ -622,6 +691,8 @@ class TYdbSetup::TImpl { std::unordered_map ServerlessToShared_; std::optional AsyncQueryRunnerActorId_; std::optional SessionState_; + TFsPath StorageMetaPath_; + NKqpRun::TStorageMeta StorageMeta_; };