From b230ebaf3c83909c65de7f12b9eabd5e5965f0b6 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Wed, 25 Sep 2024 17:59:00 +0300 Subject: [PATCH 1/5] YQ-3689 added kqp proxy database cache (#9644) --- ydb/core/kqp/common/events/events.h | 46 ++++ ydb/core/kqp/common/events/query.h | 9 + .../kqp/common/events/script_executions.h | 52 ++-- ydb/core/kqp/common/events/workload_service.h | 13 +- ydb/core/kqp/common/events/ya.make | 1 + ydb/core/kqp/common/kqp_event_impl.cpp | 5 + ydb/core/kqp/common/simple/kqp_event_ids.h | 6 +- .../resource_pool_classifier/checker.cpp | 30 ++- .../kqp_proxy_databases_cache.cpp | 249 ++++++++++++++++++ .../kqp/proxy_service/kqp_proxy_service.cpp | 100 +++++-- .../proxy_service/kqp_proxy_service_impl.h | 59 +++++ ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp | 110 ++++++++ ydb/core/kqp/proxy_service/ut/ya.make | 1 + ydb/core/kqp/proxy_service/ya.make | 1 + .../workload_service/actors/scheme_actors.cpp | 4 +- ydb/core/kqp/workload_service/common/events.h | 32 --- .../workload_service/kqp_workload_service.cpp | 4 +- .../kqp_workload_service_impl.h | 6 +- .../ut/kqp_workload_service_ut.cpp | 4 +- ydb/core/protos/kqp.proto | 1 + 20 files changed, 644 insertions(+), 89 deletions(-) create mode 100644 ydb/core/kqp/proxy_service/kqp_proxy_databases_cache.cpp diff --git a/ydb/core/kqp/common/events/events.h b/ydb/core/kqp/common/events/events.h index 45c833ef88d3..426bdfd22841 100644 --- a/ydb/core/kqp/common/events/events.h +++ b/ydb/core/kqp/common/events/events.h @@ -111,6 +111,18 @@ struct TEvKqp { struct TEvScriptRequest : public TEventLocal { TEvScriptRequest() = default; + const TString& GetDatabase() const { + return Record.GetRequest().GetDatabase(); + } + + const TString& GetDatabaseId() const { + return Record.GetRequest().GetDatabaseId(); + } + + void SetDatabaseId(const TString& databaseId) { + Record.MutableRequest()->SetDatabaseId(databaseId); + } + mutable NKikimrKqp::TEvQueryRequest Record; TDuration ForgetAfter; TDuration ResultsTtl; @@ -164,6 +176,40 @@ struct TEvKqp { return issues; } }; + + struct TEvUpdateDatabaseInfo : public TEventLocal { + TEvUpdateDatabaseInfo(const TString& database, Ydb::StatusIds::StatusCode status, NYql::TIssues issues) + : Status(status) + , Database(database) + , Issues(std::move(issues)) + {} + + TEvUpdateDatabaseInfo(const TString& database, const TString& databaseId, bool serverless) + : Status(Ydb::StatusIds::SUCCESS) + , Database(database) + , DatabaseId(databaseId) + , Serverless(serverless) + , Issues({}) + {} + + Ydb::StatusIds::StatusCode Status; + TString Database; + TString DatabaseId; + bool Serverless = false; + NYql::TIssues Issues; + }; + + struct TEvDelayedRequestError : public TEventLocal { + TEvDelayedRequestError(THolder requestEvent, Ydb::StatusIds::StatusCode status, NYql::TIssues issues) + : RequestEvent(std::move(requestEvent)) + , Status(status) + , Issues(std::move(issues)) + {} + + THolder RequestEvent; + Ydb::StatusIds::StatusCode Status; + NYql::TIssues Issues; + }; }; } // namespace NKikimr::NKqp diff --git a/ydb/core/kqp/common/events/query.h b/ydb/core/kqp/common/events/query.h index b1d01c8b5159..b180f665492d 100644 --- a/ydb/core/kqp/common/events/query.h +++ b/ydb/core/kqp/common/events/query.h @@ -353,6 +353,14 @@ struct TEvQueryRequest: public NActors::TEventLocal Token_; TActorId RequestActorId; TString Database; + TString DatabaseId; TString SessionId; TString YqlText; TString QueryId; diff --git a/ydb/core/kqp/common/events/script_executions.h b/ydb/core/kqp/common/events/script_executions.h index f6bf6d424101..62e2d767c744 100644 --- a/ydb/core/kqp/common/events/script_executions.h +++ b/ydb/core/kqp/common/events/script_executions.h @@ -22,13 +22,34 @@ enum EFinalizationStatus : i32 { FS_ROLLBACK, }; -struct TEvForgetScriptExecutionOperation : public NActors::TEventLocal { - TEvForgetScriptExecutionOperation(const TString& database, const NOperationId::TOperationId& id) +template +struct TEventWithDatabaseId : public NActors::TEventLocal { + TEventWithDatabaseId(const TString& database) : Database(database) - , OperationId(id) {} + const TString& GetDatabase() const { + return Database; + } + + const TString& GetDatabaseId() const { + return DatabaseId; + } + + void SetDatabaseId(const TString& databaseId) { + DatabaseId = databaseId; + } + const TString Database; + TString DatabaseId; +}; + +struct TEvForgetScriptExecutionOperation : public TEventWithDatabaseId { + TEvForgetScriptExecutionOperation(const TString& database, const NOperationId::TOperationId& id) + : TEventWithDatabaseId(database) + , OperationId(id) + {} + const NOperationId::TOperationId OperationId; }; @@ -43,14 +64,12 @@ struct TEvForgetScriptExecutionOperationResponse : public NActors::TEventLocal { - explicit TEvGetScriptExecutionOperation(const TString& database, const NOperationId::TOperationId& id) - : Database(database) +struct TEvGetScriptExecutionOperation : public TEventWithDatabaseId { + TEvGetScriptExecutionOperation(const TString& database, const NOperationId::TOperationId& id) + : TEventWithDatabaseId(database) , OperationId(id) - { - } + {} - TString Database; NOperationId::TOperationId OperationId; }; @@ -97,14 +116,13 @@ struct TEvGetScriptExecutionOperationResponse : public NActors::TEventLocal Metadata; }; -struct TEvListScriptExecutionOperations : public NActors::TEventLocal { +struct TEvListScriptExecutionOperations : public TEventWithDatabaseId { TEvListScriptExecutionOperations(const TString& database, const ui64 pageSize, const TString& pageToken) - : Database(database) + : TEventWithDatabaseId(database) , PageSize(pageSize) , PageToken(pageToken) {} - TString Database; ui64 PageSize; TString PageToken; }; @@ -151,14 +169,12 @@ struct TEvCheckAliveRequest : public NActors::TEventPB { }; -struct TEvCancelScriptExecutionOperation : public NActors::TEventLocal { - explicit TEvCancelScriptExecutionOperation(const TString& database, const NOperationId::TOperationId& id) - : Database(database) +struct TEvCancelScriptExecutionOperation : public TEventWithDatabaseId { + TEvCancelScriptExecutionOperation(const TString& database, const NOperationId::TOperationId& id) + : TEventWithDatabaseId(database) , OperationId(id) - { - } + {} - TString Database; NOperationId::TOperationId OperationId; }; diff --git a/ydb/core/kqp/common/events/workload_service.h b/ydb/core/kqp/common/events/workload_service.h index c1d36a957a76..9781b951e2d5 100644 --- a/ydb/core/kqp/common/events/workload_service.h +++ b/ydb/core/kqp/common/events/workload_service.h @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -90,14 +91,20 @@ struct TEvUpdatePoolInfo : public NActors::TEventLocal SecurityObject; }; -struct TEvUpdateDatabaseInfo : public NActors::TEventLocal { - TEvUpdateDatabaseInfo(const TString& database, bool serverless) - : Database(database) +struct TEvFetchDatabaseResponse : public NActors::TEventLocal { + TEvFetchDatabaseResponse(Ydb::StatusIds::StatusCode status, const TString& database, bool serverless, TPathId pathId, NYql::TIssues issues) + : Status(status) + , Database(database) , Serverless(serverless) + , PathId(pathId) + , Issues(std::move(issues)) {} + const Ydb::StatusIds::StatusCode Status; const TString Database; const bool Serverless; + const TPathId PathId; + const NYql::TIssues Issues; }; } // NKikimr::NKqp::NWorkload diff --git a/ydb/core/kqp/common/events/ya.make b/ydb/core/kqp/common/events/ya.make index 76d8e0fe6274..4ac5ea520cd1 100644 --- a/ydb/core/kqp/common/events/ya.make +++ b/ydb/core/kqp/common/events/ya.make @@ -16,6 +16,7 @@ PEERDIR( ydb/core/kqp/common/shutdown ydb/core/kqp/common/compilation ydb/core/resource_pools + ydb/core/scheme ydb/library/yql/dq/actors ydb/public/api/protos diff --git a/ydb/core/kqp/common/kqp_event_impl.cpp b/ydb/core/kqp/common/kqp_event_impl.cpp index 5396e7c0d6f3..d63b6f821fe2 100644 --- a/ydb/core/kqp/common/kqp_event_impl.cpp +++ b/ydb/core/kqp/common/kqp_event_impl.cpp @@ -91,6 +91,11 @@ void TEvKqp::TEvQueryRequest::PrepareRemote() const { } Record.MutableRequest()->SetUsePublicResponseDataFormat(true); + + if (!DatabaseId.empty()) { + Record.MutableRequest()->SetDatabaseId(DatabaseId); + } + Record.MutableRequest()->SetSessionId(SessionId); Record.MutableRequest()->SetAction(QueryAction); Record.MutableRequest()->SetType(QueryType); diff --git a/ydb/core/kqp/common/simple/kqp_event_ids.h b/ydb/core/kqp/common/simple/kqp_event_ids.h index 8c4b3fcab29a..571944731486 100644 --- a/ydb/core/kqp/common/simple/kqp_event_ids.h +++ b/ydb/core/kqp/common/simple/kqp_event_ids.h @@ -44,7 +44,9 @@ struct TKqpEvents { EvListSessionsRequest, EvListSessionsResponse, EvListProxyNodesRequest, - EvListProxyNodesResponse + EvListProxyNodesResponse, + EvUpdateDatabaseInfo, + EvDelayedRequestError }; static_assert (EvCompileInvalidateRequest + 1 == EvAbortExecution); @@ -175,8 +177,8 @@ struct TKqpWorkloadServiceEvents { EvCleanupRequest, EvCleanupResponse, EvUpdatePoolInfo, - EvUpdateDatabaseInfo, EvSubscribeOnPoolChanges, + EvFetchDatabaseResponse, }; }; diff --git a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp index 5b6bae22b411..7200999c7ca1 100644 --- a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp +++ b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp @@ -4,7 +4,6 @@ #include #include #include -#include #include #include @@ -20,6 +19,31 @@ using namespace NResourcePool; using namespace NWorkload; +struct TEvPrivate { + // Event ids + enum EEv : ui32 { + EvRanksCheckerResponse = EventSpaceBegin(TEvents::ES_PRIVATE), + + EvEnd + }; + + static_assert(EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE)"); + + struct TEvRanksCheckerResponse : public TEventLocal { + TEvRanksCheckerResponse(Ydb::StatusIds::StatusCode status, i64 maxRank, ui64 numberClassifiers, NYql::TIssues issues) + : Status(status) + , MaxRank(maxRank) + , NumberClassifiers(numberClassifiers) + , Issues(std::move(issues)) + {} + + const Ydb::StatusIds::StatusCode Status; + const i64 MaxRank; + const ui64 NumberClassifiers; + const NYql::TIssues Issues; + }; +}; + class TRanksCheckerActor : public NKikimr::TQueryBase { using TBase = NKikimr::TQueryBase; @@ -177,7 +201,7 @@ class TResourcePoolClassifierPreparationActor : public TActorBootstrappedGet()->Status != Ydb::StatusIds::SUCCESS) { FailAndPassAway("Database check failed", ev->Get()->Status, ev->Get()->Issues); return; @@ -223,7 +247,7 @@ class TResourcePoolClassifierPreparationActor : public TActorBootstrapped + +#include + + +namespace NKikimr::NKqp { + +namespace { + + +struct TEvPrivate { + // Event ids + enum EEv : ui32 { + EvSubscribeOnDatabase = EventSpaceBegin(TEvents::ES_PRIVATE), + EvPingDatabaseSubscription, + + EvEnd + }; + + static_assert(EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE)"); + + struct TEvSubscribeOnDatabase : public TEventLocal { + explicit TEvSubscribeOnDatabase(const TString& database) + : Database(database) + {} + + const TString Database; + }; + + struct TEvPingDatabaseSubscription : public TEventLocal { + explicit TEvPingDatabaseSubscription(const TString& database) + : Database(database) + {} + + const TString Database; + }; +}; + +class TDatabaseSubscriberActor : public TActor { + struct TDatabaseState { + TString Database; + TString DatabaseId = ""; + bool Serverless = false; + + bool FetchRequestIsRunning = true; + TInstant LastUpdateTime = TInstant::Now(); + ui32 WatchKey = 0; + }; + + using TBase = TActor; + +public: + TDatabaseSubscriberActor(TDuration idleTimeout) + : TBase(&TDatabaseSubscriberActor::StateFunc) + , IdleTimeout(idleTimeout) + , DatabaseStates(std::numeric_limits::max()) + {} + + void Registered(TActorSystem* sys, const TActorId& owner) { + TBase::Registered(sys, owner); + Owner = owner; + } + + void Handle(TEvPrivate::TEvSubscribeOnDatabase::TPtr& ev) { + const TString& database = ev->Get()->Database; + auto databaseStateIt = DatabaseStates.Find(database); + + if (databaseStateIt == DatabaseStates.End()) { + DatabaseStates.Insert({database, TDatabaseState{.Database = database}}); + Register(NWorkload::CreateDatabaseFetcherActor(SelfId(), database)); + StartIdleCheck(); + return; + } + + databaseStateIt->LastUpdateTime = TInstant::Now(); + if (databaseStateIt->DatabaseId) { + SendSubscriberInfo(*databaseStateIt, Ydb::StatusIds::SUCCESS); + } + } + + void Handle(TEvPrivate::TEvPingDatabaseSubscription::TPtr& ev) { + auto databaseStateIt = DatabaseStates.Find(ev->Get()->Database); + if (databaseStateIt != DatabaseStates.End()) { + databaseStateIt->LastUpdateTime = TInstant::Now(); + } + } + + void Handle(NWorkload::TEvFetchDatabaseResponse::TPtr& ev) { + auto databaseStateIt = DatabaseStates.Find(ev->Get()->Database); + if (databaseStateIt == DatabaseStates.End()) { + return; + } + + databaseStateIt->FetchRequestIsRunning = false; + UpdateDatabaseState(*databaseStateIt, ev->Get()->PathId, ev->Get()->Serverless); + SendSubscriberInfo(*databaseStateIt, ev->Get()->Status, ev->Get()->Issues); + + if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) { + FreeWatchKey++; + databaseStateIt->WatchKey = FreeWatchKey; + Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvWatchPathId(ev->Get()->PathId, FreeWatchKey)); + } + } + + void Handle(TEvTxProxySchemeCache::TEvWatchNotifyDeleted::TPtr& ev) { + auto databaseStateIt = DatabaseStates.Find(ev->Get()->Path); + if (databaseStateIt == DatabaseStates.End()) { + return; + } + + UnsubscribeFromSchemeCache(*databaseStateIt); + SendSubscriberInfo(*databaseStateIt, Ydb::StatusIds::NOT_FOUND, {NYql::TIssue{"Database was dropped"}}); + DatabaseStates.Erase(databaseStateIt); + } + + void HandlePoison() { + Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvWatchRemove(0)); + TBase::PassAway(); + } + + void HandleWakeup() { + IdleCheckStarted = false; + const auto minimalTime = TInstant::Now() - IdleTimeout; + while (!DatabaseStates.Empty()) { + auto oldestIt = DatabaseStates.FindOldest(); + if (oldestIt->LastUpdateTime > minimalTime) { + break; + } + + UnsubscribeFromSchemeCache(*oldestIt); + SendSubscriberInfo(*oldestIt, Ydb::StatusIds::ABORTED, {NYql::TIssue{"Database subscription was dropped by idle timeout"}}); + DatabaseStates.Erase(oldestIt); + } + + if (!DatabaseStates.Empty()) { + StartIdleCheck(); + } + } + + STRICT_STFUNC(StateFunc, + hFunc(TEvPrivate::TEvSubscribeOnDatabase, Handle); + hFunc(TEvPrivate::TEvPingDatabaseSubscription, Handle); + hFunc(NWorkload::TEvFetchDatabaseResponse, Handle); + sFunc(TEvents::TEvPoison, HandlePoison); + sFunc(TEvents::TEvWakeup, HandleWakeup); + + hFunc(TEvTxProxySchemeCache::TEvWatchNotifyDeleted, Handle); + IgnoreFunc(TEvTxProxySchemeCache::TEvWatchNotifyUpdated); + ) + +private: + static void UpdateDatabaseState(TDatabaseState& databaseState, TPathId pathId, bool serverless) { + databaseState.LastUpdateTime = TInstant::Now(); + databaseState.DatabaseId = (serverless ? TStringBuilder() << pathId.OwnerId << ":" << pathId.LocalPathId << ":" : TStringBuilder()) << databaseState.Database; + databaseState.Serverless = serverless; + } + + void UnsubscribeFromSchemeCache(TDatabaseState& databaseState) const { + if (databaseState.WatchKey) { + Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvWatchRemove(databaseState.WatchKey)); + databaseState.WatchKey = 0; + } + } + + void SendSubscriberInfo(const TDatabaseState& databaseState, Ydb::StatusIds::StatusCode status, NYql::TIssues issues = {}) { + if (status == Ydb::StatusIds::SUCCESS || status == Ydb::StatusIds::UNSUPPORTED) { + Send(Owner, new TEvKqp::TEvUpdateDatabaseInfo(databaseState.Database, databaseState.DatabaseId, databaseState.Serverless)); + } else { + NYql::TIssue rootIssue(TStringBuilder() << "Failed to describe database " << databaseState.Database); + for (const auto& issue : issues) { + rootIssue.AddSubIssue(MakeIntrusive(issue)); + } + Send(Owner, new TEvKqp::TEvUpdateDatabaseInfo(databaseState.Database, status, {rootIssue})); + } + } + + void StartIdleCheck() { + if (!IdleCheckStarted) { + IdleCheckStarted = true; + Schedule(IdleTimeout, new TEvents::TEvWakeup()); + } + } + +private: + const TDuration IdleTimeout; + TActorId Owner; + bool IdleCheckStarted = false; + + TLRUCache DatabaseStates; + ui32 FreeWatchKey = 0; +}; + +} // anonymous namespace + +TDatabasesCache::TDatabasesCache(TDuration idleTimeout) + : IdleTimeout(idleTimeout) +{} + +const TString& TDatabasesCache::GetTenantName() { + if (!TenantName) { + TenantName = CanonizePath(AppData()->TenantName); + } + return TenantName; +} + +void TDatabasesCache::UpdateDatabaseInfo(TEvKqp::TEvUpdateDatabaseInfo::TPtr& event, TActorContext actorContext) { + auto it = DatabasesCache.find(event->Get()->Database); + if (it == DatabasesCache.end()) { + return; + } + it->second.DatabaseId = event->Get()->DatabaseId; + + const bool success = event->Get()->Status == Ydb::StatusIds::SUCCESS; + for (auto& delayedEvent : it->second.DelayedEvents) { + if (success) { + actorContext.Send(std::move(delayedEvent.Event)); + } else { + actorContext.Send(actorContext.SelfID, new TEvKqp::TEvDelayedRequestError(std::move(delayedEvent.Event), event->Get()->Status, event->Get()->Issues), 0, delayedEvent.RequestType); + } + } + it->second.DelayedEvents.clear(); + + if (!success) { + DatabasesCache.erase(it); + } +} + +void TDatabasesCache::SubscribeOnDatabase(const TString& database, TActorContext actorContext) { + if (!SubscriberActor) { + SubscriberActor = actorContext.Register(new TDatabaseSubscriberActor(IdleTimeout)); + } + actorContext.Send(SubscriberActor, new TEvPrivate::TEvSubscribeOnDatabase(database)); +} + +void TDatabasesCache::PingDatabaseSubscription(const TString& database, TActorContext actorContext) const { + if (SubscriberActor) { + actorContext.Send(SubscriberActor, new TEvPrivate::TEvPingDatabaseSubscription(database)); + } +} + +void TDatabasesCache::StopSubscriberActor(TActorContext actorContext) const { + if (SubscriberActor) { + actorContext.Send(SubscriberActor, new TEvents::TEvPoison()); + } +} + +} // namespace NKikimr::NKqp diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp index 2e226519807c..126165cf434a 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp +++ b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp @@ -176,6 +176,15 @@ class TKqpProxyService : public TActorBootstrapped { }; }; + enum class EDelayedRequestType { + QueryRequest, + ScriptRequest, + ForgetScriptExecutionOperation, + GetScriptExecutionOperation, + ListScriptExecutionOperations, + CancelScriptExecutionOperation, + }; + public: static constexpr NKikimrServices::TActivity::EType ActorActivityType() { return NKikimrServices::TActivity::KQP_PROXY_ACTOR; @@ -486,6 +495,7 @@ class TKqpProxyService : public TActorBootstrapped { }); ResourcePoolsCache.UnsubscribeFromResourcePoolClassifiers(ActorContext()); + DatabasesCache.StopSubscriberActor(ActorContext()); return TActor::PassAway(); } @@ -635,6 +645,10 @@ class TKqpProxyService : public TActorBootstrapped { } void Handle(TEvKqp::TEvQueryRequest::TPtr& ev) { + if (!DatabasesCache.SetDatabaseIdOrDefer(ev, static_cast(EDelayedRequestType::QueryRequest), ActorContext())) { + return; + } + const TString& database = ev->Get()->GetDatabase(); const TString& traceId = ev->Get()->GetTraceId(); const auto queryType = ev->Get()->GetType(); @@ -735,7 +749,7 @@ class TKqpProxyService : public TActorBootstrapped { } void Handle(TEvKqp::TEvScriptRequest::TPtr& ev) { - if (CheckScriptExecutionsTablesReady(ev)) { + if (CheckScriptExecutionsTablesReady(ev, EDelayedRequestType::ScriptRequest)) { auto req = ev->Get()->Record.MutableRequest(); auto maxRunTime = GetQueryTimeout(req->GetType(), req->GetTimeoutMs(), TableServiceConfig, QueryServiceConfig); req->SetTimeoutMs(maxRunTime.MilliSeconds()); @@ -1360,7 +1374,8 @@ class TKqpProxyService : public TActorBootstrapped { hFunc(TEvKqp::TEvListSessionsRequest, Handle); hFunc(TEvKqp::TEvListProxyNodesRequest, Handle); hFunc(NWorkload::TEvUpdatePoolInfo, Handle); - hFunc(NWorkload::TEvUpdateDatabaseInfo, Handle); + hFunc(TEvKqp::TEvUpdateDatabaseInfo, Handle); + hFunc(TEvKqp::TEvDelayedRequestError, Handle); hFunc(NMetadata::NProvider::TEvRefreshSubscriberData, Handle); default: Y_ABORT("TKqpProxyService: unexpected event type: %" PRIx32 " event: %s", @@ -1643,12 +1658,53 @@ class TKqpProxyService : public TActorBootstrapped { NYql::NDq::SetYqlLogLevels(yqlPriority); } - template - bool CheckScriptExecutionsTablesReady(TEvent& ev) { + void HanleDelayedRequestError(EDelayedRequestType requestType, THolder requestEvent, Ydb::StatusIds::StatusCode status, NYql::TIssues issues) { + switch (requestType) { + case EDelayedRequestType::QueryRequest: { + auto response = std::make_unique(); + response->Record.GetRef().SetYdbStatus(status); + NYql::IssuesToMessage(issues, response->Record.GetRef().MutableResponse()->MutableQueryIssues()); + Send(requestEvent->Sender, std::move(response), 0, requestEvent->Cookie); + break; + } + + case EDelayedRequestType::ScriptRequest: + HanleDelayedScriptRequestError(std::move(requestEvent), status, std::move(issues)); + break; + + case EDelayedRequestType::ForgetScriptExecutionOperation: + HanleDelayedScriptRequestError(std::move(requestEvent), status, std::move(issues)); + break; + + case EDelayedRequestType::GetScriptExecutionOperation: + HanleDelayedScriptRequestError(std::move(requestEvent), status, std::move(issues)); + break; + + case EDelayedRequestType::ListScriptExecutionOperations: + HanleDelayedScriptRequestError(std::move(requestEvent), status, std::move(issues)); + break; + + case EDelayedRequestType::CancelScriptExecutionOperation: + HanleDelayedScriptRequestError(std::move(requestEvent), status, std::move(issues)); + break; + } + } + + template + void HanleDelayedScriptRequestError(THolder requestEvent, Ydb::StatusIds::StatusCode status, NYql::TIssues issues) const { + Send(requestEvent->Sender, new TResponse(status, std::move(issues)), 0, requestEvent->Cookie); + } + + template + bool CheckScriptExecutionsTablesReady(TEvent& ev, EDelayedRequestType requestType) { if (!AppData()->FeatureFlags.GetEnableScriptExecutionOperations()) { NYql::TIssues issues; issues.AddIssue("ExecuteScript feature is not enabled"); - Send(ev->Sender, new TResponse(Ydb::StatusIds::UNSUPPORTED, std::move(issues))); + HanleDelayedRequestError(requestType, std::move(ev), Ydb::StatusIds::UNSUPPORTED, std::move(issues)); + return false; + } + + if (!DatabasesCache.SetDatabaseIdOrDefer(ev, static_cast(requestType), ActorContext())) { return false; } @@ -1661,14 +1717,12 @@ class TKqpProxyService : public TActorBootstrapped { if (DelayedEventsQueue.size() < 10000) { DelayedEventsQueue.push_back({ .Event = std::move(ev), - .ResponseBuilder = [](Ydb::StatusIds::StatusCode status, NYql::TIssues issues) { - return new TResponse(status, std::move(issues)); - } + .RequestType = static_cast(requestType) }); } else { NYql::TIssues issues; issues.AddIssue("Too many queued requests"); - Send(ev->Sender, new TResponse(Ydb::StatusIds::OVERLOADED, std::move(issues))); + HanleDelayedRequestError(requestType, std::move(ev), Ydb::StatusIds::OVERLOADED, std::move(issues)); } return false; case EScriptExecutionsCreationStatus::Finished: @@ -1693,32 +1747,32 @@ class TKqpProxyService : public TActorBootstrapped { if (ev->Get()->Success) { Send(std::move(delayedEvent.Event)); } else { - Send(delayedEvent.Event->Sender, delayedEvent.ResponseBuilder(Ydb::StatusIds::INTERNAL_ERROR, {rootIssue})); + HanleDelayedRequestError(static_cast(delayedEvent.RequestType), std::move(delayedEvent.Event), Ydb::StatusIds::INTERNAL_ERROR, {rootIssue}); } DelayedEventsQueue.pop_front(); } } void Handle(NKqp::TEvForgetScriptExecutionOperation::TPtr& ev) { - if (CheckScriptExecutionsTablesReady(ev)) { + if (CheckScriptExecutionsTablesReady(ev, EDelayedRequestType::ForgetScriptExecutionOperation)) { Register(CreateForgetScriptExecutionOperationActor(std::move(ev)), TMailboxType::HTSwap, AppData()->SystemPoolId); } } void Handle(NKqp::TEvGetScriptExecutionOperation::TPtr& ev) { - if (CheckScriptExecutionsTablesReady(ev)) { + if (CheckScriptExecutionsTablesReady(ev, EDelayedRequestType::GetScriptExecutionOperation)) { Register(CreateGetScriptExecutionOperationActor(std::move(ev)), TMailboxType::HTSwap, AppData()->SystemPoolId); } } void Handle(NKqp::TEvListScriptExecutionOperations::TPtr& ev) { - if (CheckScriptExecutionsTablesReady(ev)) { + if (CheckScriptExecutionsTablesReady(ev, EDelayedRequestType::ListScriptExecutionOperations)) { Register(CreateListScriptExecutionOperationsActor(std::move(ev)), TMailboxType::HTSwap, AppData()->SystemPoolId); } } void Handle(NKqp::TEvCancelScriptExecutionOperation::TPtr& ev) { - if (CheckScriptExecutionsTablesReady(ev)) { + if (CheckScriptExecutionsTablesReady(ev, EDelayedRequestType::CancelScriptExecutionOperation)) { Register(CreateCancelScriptExecutionOperationActor(std::move(ev)), TMailboxType::HTSwap, AppData()->SystemPoolId); } } @@ -1822,8 +1876,15 @@ class TKqpProxyService : public TActorBootstrapped { ResourcePoolsCache.UpdatePoolInfo(ev->Get()->Database, ev->Get()->PoolId, ev->Get()->Config, ev->Get()->SecurityObject, ActorContext()); } - void Handle(NWorkload::TEvUpdateDatabaseInfo::TPtr& ev) { - ResourcePoolsCache.UpdateDatabaseInfo(ev->Get()->Database, ev->Get()->Serverless); + void Handle(TEvKqp::TEvUpdateDatabaseInfo::TPtr& ev) { + if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) { + ResourcePoolsCache.UpdateDatabaseInfo(ev->Get()->Database, ev->Get()->Serverless); + } + DatabasesCache.UpdateDatabaseInfo(ev, ActorContext()); + } + + void Handle(TEvKqp::TEvDelayedRequestError::TPtr& ev) { + HanleDelayedRequestError(static_cast(ev->Cookie), std::move(ev->Get()->RequestEvent), ev->Get()->Status, std::move(ev->Get()->Issues)); } void Handle(NMetadata::NProvider::TEvRefreshSubscriberData::TPtr& ev) { @@ -1883,16 +1944,13 @@ class TKqpProxyService : public TActorBootstrapped { Pending, Finished, }; - struct TDelayedEvent { - THolder Event; - std::function ResponseBuilder; - }; EScriptExecutionsCreationStatus ScriptExecutionsCreationStatus = EScriptExecutionsCreationStatus::NotStarted; - std::deque DelayedEventsQueue; + std::deque DelayedEventsQueue; bool IsLookupByRmScheduled = false; TActorId KqpTempTablesAgentActor; TResourcePoolsCache ResourcePoolsCache; + TDatabasesCache DatabasesCache; }; } // namespace diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h b/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h index 4947091cf7d3..323458630d15 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h +++ b/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h @@ -642,4 +642,63 @@ class TResourcePoolsCache { bool SubscribedOnResourcePoolClassifiers = false; }; +class TDatabasesCache { +public: + struct TDelayedEvent { + THolder Event; + i32 RequestType; + }; + +private: + struct TDatabaseInfo { + TString DatabaseId; // string "::" + std::vector DelayedEvents; + }; + +public: + TDatabasesCache(TDuration idleTimeout = TDuration::Seconds(60)); + + template + bool SetDatabaseIdOrDefer(TEvent& event, i32 requestType, TActorContext actorContext) { + if (!event->Get()->GetDatabaseId().empty()) { + return true; + } + + const auto& database = CanonizePath(event->Get()->GetDatabase()); + if (database.empty() || database == GetTenantName()) { + event->Get()->SetDatabaseId(GetTenantName()); + return true; + } + + auto& databaseInfo = DatabasesCache[database]; + if (databaseInfo.DatabaseId) { + PingDatabaseSubscription(database, actorContext); + event->Get()->SetDatabaseId(databaseInfo.DatabaseId); + return true; + } + + SubscribeOnDatabase(database, actorContext); + databaseInfo.DelayedEvents.push_back(TDelayedEvent{ + .Event = std::move(event), + .RequestType = requestType + }); + + return false; + } + + void UpdateDatabaseInfo(TEvKqp::TEvUpdateDatabaseInfo::TPtr& event, TActorContext actorContext); + void StopSubscriberActor(TActorContext actorContext) const; + +private: + const TString& GetTenantName(); + void SubscribeOnDatabase(const TString& database, TActorContext actorContext); + void PingDatabaseSubscription(const TString& database, TActorContext actorContext) const; + +private: + const TDuration IdleTimeout; + std::unordered_map DatabasesCache; + TActorId SubscriberActor; + TString TenantName; +}; + } // namespace NKikimr::NKqp diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp index 883ec7d9198e..1d008a72be80 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp +++ b/ydb/core/kqp/proxy_service/kqp_proxy_ut.cpp @@ -2,7 +2,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -65,6 +67,89 @@ TString CreateSession(TTestActorRuntime* runtime, const TActorId& kqpProxy, cons return sessionId; } +class TDatabaseCacheTestActor : public TActorBootstrapped { +public: + TDatabaseCacheTestActor(const TString& database, const TString& expectedDatabaseId, TDuration idleTimeout, NThreading::TPromise promise) + : IdleTimeout(idleTimeout) + , Database(database) + , ExpectedDatabaseId(expectedDatabaseId) + , Cache(idleTimeout) + , Promise(promise) + {} + + void Bootstrap() { + Become(&TDatabaseCacheTestActor::StateFunc); + + auto event = MakeHolder(); + event->Record.MutableRequest()->SetDatabase(Database); + Send(SelfId(), event.Release()); + + Schedule(3 * IdleTimeout, new TEvents::TEvWakeup()); + } + + void Handle(TEvKqp::TEvUpdateDatabaseInfo::TPtr& ev) { + if (!CacheUpdated) { + UNIT_ASSERT_VALUES_EQUAL_C(ev->Get()->Status, Ydb::StatusIds::SUCCESS, TStringBuilder() << GetErrorString() << ev->Get()->Issues.ToString()); + Cache.UpdateDatabaseInfo(ev, ActorContext()); + CacheUpdated = true; + } else { + UNIT_ASSERT_VALUES_EQUAL_C(ev->Get()->Status, Ydb::StatusIds::ABORTED, TStringBuilder() << GetErrorString() << ev->Get()->Issues.ToString()); + UNIT_ASSERT_STRING_CONTAINS_C(ev->Get()->Issues.ToString(), "Database subscription was dropped by idle timeout", GetErrorString()); + Finish(); + } + } + + void Handle(TEvKqp::TEvDelayedRequestError::TPtr& ev) { + UNIT_ASSERT_C(false, TStringBuilder() << "Unexpected fail, status: " << ev->Get()->Status << ", " << GetErrorString() << ev->Get()->Issues.ToString()); + } + + void Handle(TEvKqp::TEvQueryRequest::TPtr& ev) { + auto success = Cache.SetDatabaseIdOrDefer(ev, 0, ActorContext()); + + bool dedicated = Database == ExpectedDatabaseId; + if (CacheUpdated || dedicated) { + UNIT_ASSERT_C(success, TStringBuilder() << "Expected database id from cache, " << GetErrorString()); + UNIT_ASSERT_STRING_CONTAINS_C(ev->Get()->GetDatabaseId(), ExpectedDatabaseId, GetErrorString()); + if (dedicated) { + Finish(); + } + } else { + UNIT_ASSERT_C(!success, TStringBuilder() << "Unexpected database id from cache, " << GetErrorString()); + } + } + + void HandleWakeup() { + UNIT_ASSERT_C(false, TStringBuilder() << "Test cache timeout, " << GetErrorString()); + Finish(); + } + + STRICT_STFUNC(StateFunc, + hFunc(TEvKqp::TEvUpdateDatabaseInfo, Handle); + hFunc(TEvKqp::TEvDelayedRequestError, Handle); + hFunc(TEvKqp::TEvQueryRequest, Handle); + sFunc(TEvents::TEvWakeup, HandleWakeup); + ) + +private: + TString GetErrorString() const { + return TStringBuilder() << "cache updated: " << CacheUpdated << ", database: " << Database << "\n"; + } + + void Finish() { + Promise.SetValue(); + PassAway(); + } + +private: + const TDuration IdleTimeout; + const TString Database; + const TString ExpectedDatabaseId; + TDatabasesCache Cache; + NThreading::TPromise Promise; + + bool CacheUpdated = false; +}; + } Y_UNIT_TEST_SUITE(KqpProxy) { @@ -542,5 +627,30 @@ Y_UNIT_TEST_SUITE(KqpProxy) { UNIT_ASSERT(allDoneOk); } + + Y_UNIT_TEST(DatabasesCacheForServerless) { + auto ydb = NWorkload::TYdbSetupSettings() + .CreateSampleTenants(true) + .Create(); + + auto& runtime = *ydb->GetRuntime(); + TDuration idleTimeout = TDuration::Seconds(5); + + auto checkCache = [&](const TString& database, const TString& expectedDatabaseId, ui32 nodeIndex) { + auto promise = NThreading::NewPromise(); + runtime.Register(new TDatabaseCacheTestActor(database, expectedDatabaseId, idleTimeout, promise), nodeIndex); + promise.GetFuture().GetValueSync(); + }; + + const auto& dedicatedTennant = ydb->GetSettings().GetDedicatedTenantName(); + checkCache(dedicatedTennant, dedicatedTennant, 2); + + const auto& sharedTennant = ydb->GetSettings().GetSharedTenantName(); + checkCache(sharedTennant, sharedTennant, 1); + + const auto& serverlessTennant = ydb->GetSettings().GetServerlessTenantName(); + checkCache(serverlessTennant, TStringBuilder() << ":4:" << serverlessTennant, 1); + } + } // namspace NKqp } // namespace NKikimr diff --git a/ydb/core/kqp/proxy_service/ut/ya.make b/ydb/core/kqp/proxy_service/ut/ya.make index 730f59a3fcae..fc6d9e7c89cb 100644 --- a/ydb/core/kqp/proxy_service/ut/ya.make +++ b/ydb/core/kqp/proxy_service/ut/ya.make @@ -13,6 +13,7 @@ PEERDIR( ydb/core/kqp/run_script_actor ydb/core/kqp/proxy_service ydb/core/kqp/ut/common + ydb/core/kqp/workload_service/ut/common ydb/library/yql/sql/pg_dummy ydb/public/sdk/cpp/client/ydb_query ydb/public/sdk/cpp/client/ydb_driver diff --git a/ydb/core/kqp/proxy_service/ya.make b/ydb/core/kqp/proxy_service/ya.make index e1c2b9e1b76a..8a143789e701 100644 --- a/ydb/core/kqp/proxy_service/ya.make +++ b/ydb/core/kqp/proxy_service/ya.make @@ -2,6 +2,7 @@ LIBRARY() SRCS( kqp_proxy_service.cpp + kqp_proxy_databases_cache.cpp kqp_proxy_peer_stats_calculator.cpp kqp_script_executions.cpp kqp_session_info.cpp diff --git a/ydb/core/kqp/workload_service/actors/scheme_actors.cpp b/ydb/core/kqp/workload_service/actors/scheme_actors.cpp index 781d9f4a6eca..e7cb25220418 100644 --- a/ydb/core/kqp/workload_service/actors/scheme_actors.cpp +++ b/ydb/core/kqp/workload_service/actors/scheme_actors.cpp @@ -489,6 +489,7 @@ class TDatabaseFetcherActor : public TSchemeActorBase { } if (result.DomainInfo) { Serverless = result.DomainInfo->IsServerless(); + PathId = result.DomainInfo->DomainKey; } Reply(Ydb::StatusIds::SUCCESS); return; @@ -537,7 +538,7 @@ class TDatabaseFetcherActor : public TSchemeActorBase { } Issues.AddIssues(std::move(issues)); - Send(ReplyActorId, new TEvPrivate::TEvFetchDatabaseResponse(status, Database, Serverless, std::move(Issues))); + Send(ReplyActorId, new TEvFetchDatabaseResponse(status, Database, Serverless, PathId, std::move(Issues))); PassAway(); } @@ -560,6 +561,7 @@ class TDatabaseFetcherActor : public TSchemeActorBase { const NACLib::EAccessRights CheckAccess; bool Serverless = false; + TPathId PathId; }; } // anonymous namespace diff --git a/ydb/core/kqp/workload_service/common/events.h b/ydb/core/kqp/workload_service/common/events.h index 48643582d7f0..0c6f7d1d8bd6 100644 --- a/ydb/core/kqp/workload_service/common/events.h +++ b/ydb/core/kqp/workload_service/common/events.h @@ -22,7 +22,6 @@ struct TEvPrivate { EvRefreshPoolState = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), EvResolvePoolResponse, EvFetchPoolResponse, - EvFetchDatabaseResponse, EvCreatePoolResponse, EvPrepareTablesRequest, EvPlaceRequestIntoPoolResponse, @@ -46,8 +45,6 @@ struct TEvPrivate { EvStartRequestResponse, EvCleanupRequestsResponse, - EvRanksCheckerResponse, - EvEnd }; @@ -93,20 +90,6 @@ struct TEvPrivate { const NYql::TIssues Issues; }; - struct TEvFetchDatabaseResponse : public NActors::TEventLocal { - TEvFetchDatabaseResponse(Ydb::StatusIds::StatusCode status, const TString& database, bool serverless, NYql::TIssues issues) - : Status(status) - , Database(database) - , Serverless(serverless) - , Issues(std::move(issues)) - {} - - const Ydb::StatusIds::StatusCode Status; - const TString Database; - const bool Serverless; - const NYql::TIssues Issues; - }; - struct TEvCreatePoolResponse : public NActors::TEventLocal { TEvCreatePoolResponse(Ydb::StatusIds::StatusCode status, NYql::TIssues issues) : Status(status) @@ -321,21 +304,6 @@ struct TEvPrivate { const std::vector SesssionIds; const NYql::TIssues Issues; }; - - // Resource pool classifier events - struct TEvRanksCheckerResponse : public TEventLocal { - TEvRanksCheckerResponse(Ydb::StatusIds::StatusCode status, i64 maxRank, ui64 numberClassifiers, NYql::TIssues issues) - : Status(status) - , MaxRank(maxRank) - , NumberClassifiers(numberClassifiers) - , Issues(std::move(issues)) - {} - - const Ydb::StatusIds::StatusCode Status; - const i64 MaxRank; - const ui64 NumberClassifiers; - const NYql::TIssues Issues; - }; }; } // NKikimr::NKqp::NWorkload diff --git a/ydb/core/kqp/workload_service/kqp_workload_service.cpp b/ydb/core/kqp/workload_service/kqp_workload_service.cpp index 24fdd21e3d63..94e2571d6099 100644 --- a/ydb/core/kqp/workload_service/kqp_workload_service.cpp +++ b/ydb/core/kqp/workload_service/kqp_workload_service.cpp @@ -207,7 +207,7 @@ class TKqpWorkloadService : public TActorBootstrapped { hFunc(TEvCleanupRequest, Handle); hFunc(TEvents::TEvWakeup, Handle); - hFunc(TEvPrivate::TEvFetchDatabaseResponse, Handle); + hFunc(TEvFetchDatabaseResponse, Handle); hFunc(TEvPrivate::TEvFetchPoolResponse, Handle); hFunc(TEvPrivate::TEvResolvePoolResponse, Handle); hFunc(TEvPrivate::TEvPlaceRequestIntoPoolResponse, Handle); @@ -223,7 +223,7 @@ class TKqpWorkloadService : public TActorBootstrapped { ) private: - void Handle(TEvPrivate::TEvFetchDatabaseResponse::TPtr& ev) { + void Handle(TEvFetchDatabaseResponse::TPtr& ev) { GetOrCreateDatabaseState(ev->Get()->Database)->UpdateDatabaseInfo(ev); } diff --git a/ydb/core/kqp/workload_service/kqp_workload_service_impl.h b/ydb/core/kqp/workload_service/kqp_workload_service_impl.h index 9ae115235a25..b4a748dc2de8 100644 --- a/ydb/core/kqp/workload_service/kqp_workload_service_impl.h +++ b/ydb/core/kqp/workload_service/kqp_workload_service_impl.h @@ -67,17 +67,13 @@ struct TDatabaseState { subscribers.clear(); } - void UpdateDatabaseInfo(const TEvPrivate::TEvFetchDatabaseResponse::TPtr& ev) { + void UpdateDatabaseInfo(const TEvFetchDatabaseResponse::TPtr& ev) { DatabaseUnsupported = ev->Get()->Status == Ydb::StatusIds::UNSUPPORTED; if (ev->Get()->Status != Ydb::StatusIds::SUCCESS) { ReplyContinueError(ev->Get()->Status, GroupIssues(ev->Get()->Issues, "Failed to fetch database info")); return; } - if (Serverless != ev->Get()->Serverless) { - ActorContext.Send(MakeKqpProxyID(ActorContext.SelfID.NodeId()), new TEvUpdateDatabaseInfo(ev->Get()->Database, ev->Get()->Serverless)); - } - LastUpdateTime = TInstant::Now(); Serverless = ev->Get()->Serverless; StartPendingRequests(); diff --git a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp index 1813f8b41500..bb0ee347e6be 100644 --- a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp +++ b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp @@ -635,7 +635,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { } void WaitForFail(TIntrusivePtr ydb, const TQueryRunnerSettings& settings, const TString& poolId) { - ydb->WaitFor(TDuration::Seconds(5), "Resource pool classifier fail", [ydb, settings, poolId](TString& errorString) { + ydb->WaitFor(TDuration::Seconds(10), "Resource pool classifier fail", [ydb, settings, poolId](TString& errorString) { auto result = ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, settings); errorString = result.GetIssues().ToOneLineString(); @@ -644,7 +644,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { } void WaitForSuccess(TIntrusivePtr ydb, const TQueryRunnerSettings& settings) { - ydb->WaitFor(TDuration::Seconds(5), "Resource pool classifier success", [ydb, settings](TString& errorString) { + ydb->WaitFor(TDuration::Seconds(10), "Resource pool classifier success", [ydb, settings](TString& errorString) { auto result = ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, settings); errorString = result.GetIssues().ToOneLineString(); diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto index d5ccc68639f8..7399ec1cc164 100644 --- a/ydb/core/protos/kqp.proto +++ b/ydb/core/protos/kqp.proto @@ -124,6 +124,7 @@ message TQueryRequest { optional string UserSID = 33; optional uint64 OutputChunkMaxSize = 34; optional string PoolId = 35; + optional string DatabaseId = 36; } message TKqpPathIdProto { From cd24dc82342a50a73163045ad1199bdde43844f5 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy <79596613+GrigoriyPA@users.noreply.github.com> Date: Fri, 20 Sep 2024 21:11:55 +0300 Subject: [PATCH 2/5] YQ-3684 fixed error duplicate session id (#9583) --- .../kqp/session_actor/kqp_session_actor.cpp | 2 + .../actors/pool_handlers_acors.cpp | 7 +-- ydb/core/kqp/workload_service/common/events.h | 15 +++++- .../workload_service/kqp_workload_service.cpp | 47 ++++++++++++++++--- .../kqp_workload_service_impl.h | 26 ++++++++++ 5 files changed, 86 insertions(+), 11 deletions(-) diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp index e46dbcf51ac6..fe5f31be2fc1 100644 --- a/ydb/core/kqp/session_actor/kqp_session_actor.cpp +++ b/ydb/core/kqp/session_actor/kqp_session_actor.cpp @@ -255,6 +255,7 @@ class TKqpSessionActor : public TActorBootstrapped { QueryState->UserToken ), IEventHandle::FlagTrackDelivery); + QueryState->PoolHandlerActor = MakeKqpWorkloadServiceId(SelfId().NodeId()); Become(&TKqpSessionActor::ExecuteState); } @@ -2387,6 +2388,7 @@ class TKqpSessionActor : public TActorBootstrapped { hFunc(TEvTxProxySchemeCache::TEvNavigateKeySetResult, HandleNoop); hFunc(TEvents::TEvUndelivered, HandleNoop); hFunc(TEvTxUserProxy::TEvAllocateTxIdResult, HandleNoop); + hFunc(TEvKqpExecuter::TEvStreamData, HandleNoop); hFunc(NWorkload::TEvContinueRequest, HandleNoop); // always come from WorkerActor diff --git a/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp b/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp index 2ff6e6912abe..77512a4d529f 100644 --- a/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp +++ b/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp @@ -161,6 +161,7 @@ class TPoolHandlerActorBase : public TActor { } SendPoolInfoUpdate(std::nullopt, std::nullopt, Subscribers); + this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvStopPoolHandlerResponse(Database, PoolId)); Counters.OnCleanup(ResetCountersOnStrop); @@ -184,16 +185,16 @@ class TPoolHandlerActorBase : public TActor { } void Handle(TEvPrivate::TEvResolvePoolResponse::TPtr& ev) { - this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPlaceRequestIntoPoolResponse(Database, PoolId)); - auto event = std::move(ev->Get()->Event); + const TString& sessionId = event->Get()->SessionId; + this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPlaceRequestIntoPoolResponse(Database, PoolId, sessionId)); + const TActorId& workerActorId = event->Sender; if (!InFlightLimit) { this->Send(workerActorId, new TEvContinueRequest(Ydb::StatusIds::PRECONDITION_FAILED, PoolId, PoolConfig, {NYql::TIssue(TStringBuilder() << "Resource pool " << PoolId << " was disabled due to zero concurrent query limit")})); return; } - const TString& sessionId = event->Get()->SessionId; if (LocalSessions.contains(sessionId)) { this->Send(workerActorId, new TEvContinueRequest(Ydb::StatusIds::INTERNAL_ERROR, PoolId, PoolConfig, {NYql::TIssue(TStringBuilder() << "Got duplicate session id " << sessionId << " for pool " << PoolId)})); return; diff --git a/ydb/core/kqp/workload_service/common/events.h b/ydb/core/kqp/workload_service/common/events.h index 0c6f7d1d8bd6..57e332ac66b3 100644 --- a/ydb/core/kqp/workload_service/common/events.h +++ b/ydb/core/kqp/workload_service/common/events.h @@ -28,6 +28,7 @@ struct TEvPrivate { EvFinishRequestInPool, EvResignPoolHandler, EvStopPoolHandler, + EvStopPoolHandlerResponse, EvCancelRequest, EvUpdatePoolSubscription, @@ -111,13 +112,15 @@ struct TEvPrivate { }; struct TEvPlaceRequestIntoPoolResponse : public NActors::TEventLocal { - TEvPlaceRequestIntoPoolResponse(const TString& database, const TString& poolId) + TEvPlaceRequestIntoPoolResponse(const TString& database, const TString& poolId, const TString& sessionId) : Database(database) , PoolId(poolId) + , SessionId(sessionId) {} const TString Database; const TString PoolId; + const TString SessionId; }; struct TEvFinishRequestInPool : public NActors::TEventLocal { @@ -156,6 +159,16 @@ struct TEvPrivate { const bool ResetCounters; }; + struct TEvStopPoolHandlerResponse : public NActors::TEventLocal { + TEvStopPoolHandlerResponse(const TString& database, const TString& poolId) + : Database(database) + , PoolId(poolId) + {} + + const TString Database; + const TString PoolId; + }; + struct TEvCancelRequest : public NActors::TEventLocal { explicit TEvCancelRequest(const TString& sessionId) : SessionId(sessionId) diff --git a/ydb/core/kqp/workload_service/kqp_workload_service.cpp b/ydb/core/kqp/workload_service/kqp_workload_service.cpp index 94e2571d6099..9bcfe49f4a4e 100644 --- a/ydb/core/kqp/workload_service/kqp_workload_service.cpp +++ b/ydb/core/kqp/workload_service/kqp_workload_service.cpp @@ -169,6 +169,13 @@ class TKqpWorkloadService : public TActorBootstrapped { void Handle(TEvCleanupRequest::TPtr& ev) { const TString& database = ev->Get()->Database; const TString& poolId = ev->Get()->PoolId; + const TString& sessionId = ev->Get()->SessionId; + if (GetOrCreateDatabaseState(database)->PendingSessionIds.contains(sessionId)) { + LOG_D("Finished request with worker actor " << ev->Sender << ", wait for place request, Database: " << database << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId); + GetOrCreateDatabaseState(database)->PendingCancelRequests[sessionId].emplace_back(std::move(ev)); + return; + } + auto poolState = GetPoolState(database, poolId); if (!poolState) { ReplyCleanupError(ev->Sender, Ydb::StatusIds::NOT_FOUND, TStringBuilder() << "Pool " << poolId << " not found"); @@ -176,7 +183,7 @@ class TKqpWorkloadService : public TActorBootstrapped { } LOG_D("Finished request with worker actor " << ev->Sender << ", Database: " << database << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId); - Send(ev->Forward(poolState->PoolHandler)); + poolState->DoCleanupRequest(std::move(ev)); } void Handle(TEvents::TEvWakeup::TPtr& ev) { @@ -220,6 +227,7 @@ class TKqpWorkloadService : public TActorBootstrapped { hFunc(TEvPrivate::TEvTablesCreationFinished, Handle); hFunc(TEvPrivate::TEvCpuLoadResponse, Handle); hFunc(TEvPrivate::TEvResignPoolHandler, Handle); + hFunc(TEvPrivate::TEvStopPoolHandlerResponse, Handle); ) private: @@ -245,12 +253,16 @@ class TKqpWorkloadService : public TActorBootstrapped { void Handle(TEvPrivate::TEvResolvePoolResponse::TPtr& ev) { const auto& event = ev->Get()->Event; const TString& database = event->Get()->Database; + auto databaseState = GetOrCreateDatabaseState(database); if (ev->Get()->DefaultPoolCreated) { - GetOrCreateDatabaseState(database)->HasDefaultPool = true; + databaseState->HasDefaultPool = true; } const TString& poolId = event->Get()->PoolId; if (ev->Get()->Status != Ydb::StatusIds::SUCCESS) { + databaseState->RemovePendingSession(event->Get()->SessionId, [this](TEvCleanupRequest::TPtr event) { + ReplyCleanupError(event->Sender, Ydb::StatusIds::NOT_FOUND, TStringBuilder() << "Pool " << event->Get()->PoolId << " not found"); + }); ReplyContinueError(event->Sender, ev->Get()->Status, ev->Get()->Issues); return; } @@ -265,9 +277,19 @@ class TKqpWorkloadService : public TActorBootstrapped { void Handle(TEvPrivate::TEvPlaceRequestIntoPoolResponse::TPtr& ev) { const TString& database = ev->Get()->Database; const TString& poolId = ev->Get()->PoolId; - LOG_T("Request placed into pool, Database: " << database << ", PoolId: " << poolId); + const TString& sessionId = ev->Get()->SessionId; + LOG_T("Request placed into pool, Database: " << database << ", PoolId: " << poolId << ", SessionId: " << sessionId); - if (auto poolState = GetPoolState(database, poolId)) { + auto poolState = GetPoolState(database, poolId); + GetOrCreateDatabaseState(database)->RemovePendingSession(sessionId, [this, poolState](TEvCleanupRequest::TPtr event) { + if (poolState) { + poolState->DoCleanupRequest(std::move(event)); + } else { + ReplyCleanupError(event->Sender, Ydb::StatusIds::NOT_FOUND, TStringBuilder() << "Pool " << event->Get()->PoolId << " not found"); + } + }); + + if (poolState) { poolState->PlaceRequestRunning = false; poolState->UpdateHandler(); poolState->StartPlaceRequest(); @@ -388,6 +410,17 @@ class TKqpWorkloadService : public TActorBootstrapped { } } + void Handle(TEvPrivate::TEvStopPoolHandlerResponse::TPtr& ev) { + const TString& database = ev->Get()->Database; + const TString& poolId = ev->Get()->PoolId; + LOG_T("Got stop pool handler response, Database: " << database << ", PoolId: " << poolId); + + Counters.ActivePools->Dec(); + if (auto poolState = GetPoolState(database, poolId)) { + poolState->PreviousPoolHandlers.erase(ev->Sender); + } + } + private: void InitializeWorkloadService() { if (ServiceInitialized) { @@ -441,7 +474,7 @@ class TKqpWorkloadService : public TActorBootstrapped { std::vector poolsToDelete; poolsToDelete.reserve(PoolIdToState.size()); for (const auto& [poolKey, poolState] : PoolIdToState) { - if (!poolState.InFlightRequests && TInstant::Now() - poolState.LastUpdateTime > IDLE_DURATION) { + if (!poolState.InFlightRequests && TInstant::Now() - poolState.LastUpdateTime > IDLE_DURATION && poolState.PendingRequests.empty()) { CpuQuotaManager->CleanupHandler(poolState.PoolHandler); Send(poolState.PoolHandler, new TEvPrivate::TEvStopPoolHandler(true)); poolsToDelete.emplace_back(poolKey); @@ -449,7 +482,6 @@ class TKqpWorkloadService : public TActorBootstrapped { } for (const auto& poolKey : poolsToDelete) { PoolIdToState.erase(poolKey); - Counters.ActivePools->Dec(); } if (!PoolIdToState.empty()) { @@ -512,7 +544,8 @@ class TKqpWorkloadService : public TActorBootstrapped { Send(replyActorId, new TEvCleanupResponse(status, {NYql::TIssue(message)})); } - TDatabaseState* GetOrCreateDatabaseState(const TString& database) { + TDatabaseState* GetOrCreateDatabaseState(TString database) { + database = CanonizePath(database); auto databaseIt = DatabaseToState.find(database); if (databaseIt != DatabaseToState.end()) { return &databaseIt->second; diff --git a/ydb/core/kqp/workload_service/kqp_workload_service_impl.h b/ydb/core/kqp/workload_service/kqp_workload_service_impl.h index b4a748dc2de8..e9e292d81dfc 100644 --- a/ydb/core/kqp/workload_service/kqp_workload_service_impl.h +++ b/ydb/core/kqp/workload_service/kqp_workload_service_impl.h @@ -19,6 +19,8 @@ struct TDatabaseState { bool& EnabledResourcePoolsOnServerless; std::vector PendingRequersts = {}; + std::unordered_set PendingSessionIds = {}; + std::unordered_map> PendingCancelRequests = {}; std::unordered_map> PendingSubscriptions = {}; bool HasDefaultPool = false; bool Serverless = false; @@ -38,6 +40,7 @@ struct TDatabaseState { void DoPlaceRequest(TEvPlaceRequestIntoPool::TPtr ev) { TString database = ev->Get()->Database; + PendingSessionIds.emplace(ev->Get()->SessionId); PendingRequersts.emplace_back(std::move(ev)); if (!EnabledResourcePoolsOnServerless && (TInstant::Now() - LastUpdateTime) > IDLE_DURATION) { @@ -79,6 +82,14 @@ struct TDatabaseState { StartPendingRequests(); } + void RemovePendingSession(const TString& sessionId, std::function callback) { + for (auto& event : PendingCancelRequests[sessionId]) { + callback(std::move(event)); + } + PendingCancelRequests.erase(sessionId); + PendingSessionIds.erase(sessionId); + } + private: void StartPendingRequests() { if (!EnabledResourcePoolsOnServerless && Serverless) { @@ -94,6 +105,9 @@ struct TDatabaseState { void ReplyContinueError(Ydb::StatusIds::StatusCode status, NYql::TIssues issues) { for (const auto& ev : PendingRequersts) { + RemovePendingSession(ev->Get()->SessionId, [this](TEvCleanupRequest::TPtr event) { + ActorContext.Send(event->Sender, new TEvCleanupResponse(Ydb::StatusIds::NOT_FOUND, {NYql::TIssue(TStringBuilder() << "Pool " << event->Get()->PoolId << " not found")})); + }); ActorContext.Send(ev->Sender, new TEvContinueRequest(status, {}, {}, issues)); } PendingRequersts.clear(); @@ -108,6 +122,7 @@ struct TPoolState { bool WaitingInitialization = false; bool PlaceRequestRunning = false; std::optional NewPoolHandler = std::nullopt; + std::unordered_set PreviousPoolHandlers = {}; ui64 InFlightRequests = 0; TInstant LastUpdateTime = TInstant::Now(); @@ -118,6 +133,7 @@ struct TPoolState { } ActorContext.Send(PoolHandler, new TEvPrivate::TEvStopPoolHandler(false)); + PreviousPoolHandlers.insert(PoolHandler); PoolHandler = *NewPoolHandler; NewPoolHandler = std::nullopt; InFlightRequests = 0; @@ -139,6 +155,16 @@ struct TPoolState { InFlightRequests--; LastUpdateTime = TInstant::Now(); } + + void DoCleanupRequest(TEvCleanupRequest::TPtr event) { + for (const auto& poolHandler : PreviousPoolHandlers) { + ActorContext.Send(poolHandler, new TEvCleanupRequest( + event->Get()->Database, event->Get()->SessionId, event->Get()->PoolId, + event->Get()->Duration, event->Get()->CpuConsumed + )); + } + ActorContext.Send(event->Forward(PoolHandler)); + } }; struct TCpuQuotaManagerState { From a93af9bbabac7468e496e61c84b519ad4c4a3c68 Mon Sep 17 00:00:00 2001 From: Mikhail Surin Date: Tue, 3 Sep 2024 08:09:33 +0300 Subject: [PATCH 3/5] Use actual cores count (#8576) --- .../kqp/runtime/kqp_compute_scheduler.cpp | 32 +++++++++++++++---- ydb/core/kqp/runtime/kqp_compute_scheduler.h | 2 ++ 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp b/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp index d24bc60e870a..8d22b554a671 100644 --- a/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp +++ b/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp @@ -126,6 +126,8 @@ class TSchedulerEntity { std::atomic DelayedSumBatches = 0; std::atomic DelayedCount = 0; + double Share; + TMultithreadPublisher MutableStats; }; @@ -215,13 +217,16 @@ struct TComputeScheduler::TImpl { TDuration MaxDelay = TDuration::Seconds(10); - void AssignWeights() { } + void AssignWeights() { + for (auto& record : Records) { + record->MutableStats.Next()->Weight = SumCores * record->Share; + } + } - void CreateGroup(TString groupName, double maxShare, NMonotonic::TMonotonic now) { + void CreateGroup(TString groupName, double maxShare) { PoolId[groupName] = Records.size(); auto group = std::make_unique(); - group->MutableStats.Next()->LastNowRecalc = now; - group->MutableStats.Next()->Weight = maxShare; + group->Share = maxShare; Records.push_back(std::move(group)); } }; @@ -369,14 +374,18 @@ bool TComputeScheduler::Disable(TString group, TMonotonic now) { void TComputeScheduler::UpdateMaxShare(TString group, double share, TMonotonic now) { auto ptr = Impl->PoolId.FindPtr(group); if (!ptr) { - Impl->CreateGroup(group, share, now); + Impl->CreateGroup(group, share); } else { auto& record = Impl->Records[*ptr]; - record->MutableStats.Next()->Weight = share; + record->Share = share; } + Impl->AssignWeights(); AdvanceTime(now); } +void TComputeScheduler::SetCapacity(ui64 cores) { + Impl->SumCores = cores; +} ::NMonitoring::TDynamicCounters::TCounterPtr TComputeScheduler::GetGroupUsageCounter(TString group) const { return Impl->Counters @@ -418,6 +427,16 @@ class TSchedulerActor : public TActorBootstrapped { IEventHandle::FlagTrackDelivery); Become(&TSchedulerActor::State); + SetCapacity(SelfId().PoolID()); + } + + void SetCapacity(ui32 pool) { + NActors::TExecutorPoolStats poolStats; + TVector threadsStats; + TlsActivationContext->ActorSystem()->GetPoolStats(pool, poolStats, threadsStats); + ui64 threads = Max(poolStats.MaxThreadCount, 1); + Opts.Counters->SchedulerCapacity->Set(threads); + Opts.Scheduler->SetCapacity(threads); } STATEFN(State) { @@ -467,6 +486,7 @@ class TSchedulerActor : public TActorBootstrapped { } void Handle(TEvents::TEvWakeup::TPtr&) { + SetCapacity(SelfId().PoolID()); Opts.Scheduler->AdvanceTime(TlsActivationContext->Monotonic()); Schedule(Opts.AdvanceTimeInterval, new TEvents::TEvWakeup()); } diff --git a/ydb/core/kqp/runtime/kqp_compute_scheduler.h b/ydb/core/kqp/runtime/kqp_compute_scheduler.h index fc58b4e8dc98..fd66610fbd33 100644 --- a/ydb/core/kqp/runtime/kqp_compute_scheduler.h +++ b/ydb/core/kqp/runtime/kqp_compute_scheduler.h @@ -62,6 +62,8 @@ class TComputeScheduler { void ReportCounters(TIntrusivePtr); void UpdateMaxShare(TString, double, TMonotonic now); + + void SetCapacity(ui64 cores); void SetMaxDeviation(TDuration); void SetForgetInterval(TDuration); From a6127fbf08056754cba18580e942a11366832ecd Mon Sep 17 00:00:00 2001 From: Mikhail Surin Date: Mon, 16 Sep 2024 13:23:18 +0300 Subject: [PATCH 4/5] Per query cpu limits (#8520) --- ydb/core/kqp/counters/kqp_counters.cpp | 2 + ydb/core/kqp/counters/kqp_counters.h | 2 + ydb/core/kqp/executer_actor/kqp_planner.cpp | 5 +- .../kqp/node_service/kqp_node_service.cpp | 39 +- .../kqp/runtime/kqp_compute_scheduler.cpp | 607 ++++++++++++++---- ydb/core/kqp/runtime/kqp_compute_scheduler.h | 17 +- ydb/core/protos/kqp.proto | 4 +- .../resource_pools/resource_pool_settings.cpp | 3 +- .../resource_pools/resource_pool_settings.h | 1 + 9 files changed, 526 insertions(+), 154 deletions(-) diff --git a/ydb/core/kqp/counters/kqp_counters.cpp b/ydb/core/kqp/counters/kqp_counters.cpp index 4c39ee6d89d6..be0d9fa8d1a0 100644 --- a/ydb/core/kqp/counters/kqp_counters.cpp +++ b/ydb/core/kqp/counters/kqp_counters.cpp @@ -830,6 +830,8 @@ TKqpCounters::TKqpCounters(const ::NMonitoring::TDynamicCounterPtr& counters, co FullScansExecuted = KqpGroup->GetCounter("FullScans", true); SchedulerThrottled = KqpGroup->GetCounter("NodeScheduler/ThrottledUs", true); + SchedulerGroupsCount = KqpGroup->GetCounter("NodeScheduler/GroupsCount", false); + SchedulerValuesCount = KqpGroup->GetCounter("NodeScheduler/ValuesCount", false); SchedulerCapacity = KqpGroup->GetCounter("NodeScheduler/Capacity"); ComputeActorExecutions = KqpGroup->GetHistogram("NodeScheduler/BatchUs", NMonitoring::ExponentialHistogram(20, 2, 1)); ComputeActorDelays = KqpGroup->GetHistogram("NodeScheduler/Delays", NMonitoring::ExponentialHistogram(20, 2, 1)); diff --git a/ydb/core/kqp/counters/kqp_counters.h b/ydb/core/kqp/counters/kqp_counters.h index 7d6f38d05fe9..0c4042a14ef1 100644 --- a/ydb/core/kqp/counters/kqp_counters.h +++ b/ydb/core/kqp/counters/kqp_counters.h @@ -416,6 +416,8 @@ class TKqpCounters : public TKqpCountersBase, public NYql::NDq::TSpillingCounter NMonitoring::THistogramPtr ComputeActorDelays; ::NMonitoring::TDynamicCounters::TCounterPtr ThrottledActorsSpuriousActivations; NMonitoring::THistogramPtr SchedulerDelays; + NMonitoring::TDynamicCounters::TCounterPtr SchedulerGroupsCount; + NMonitoring::TDynamicCounters::TCounterPtr SchedulerValuesCount; // Sequences counters ::NMonitoring::TDynamicCounters::TCounterPtr SequencerActorsCount; diff --git a/ydb/core/kqp/executer_actor/kqp_planner.cpp b/ydb/core/kqp/executer_actor/kqp_planner.cpp index 7ad66678f3cd..d9eaf4619cf9 100644 --- a/ydb/core/kqp/executer_actor/kqp_planner.cpp +++ b/ydb/core/kqp/executer_actor/kqp_planner.cpp @@ -238,7 +238,10 @@ std::unique_ptr TKqpPlanner::SerializeReque request.SetDatabase(Database); if (UserRequestContext->PoolConfig.has_value()) { request.SetMemoryPoolPercent(UserRequestContext->PoolConfig->QueryMemoryLimitPercentPerNode); - request.SetMaxCpuShare(UserRequestContext->PoolConfig->TotalCpuLimitPercentPerNode / 100.0); + request.SetPoolMaxCpuShare(UserRequestContext->PoolConfig->TotalCpuLimitPercentPerNode / 100.0); + if (UserRequestContext->PoolConfig->QueryCpuLimitPercentPerNode >= 0) { + request.SetQueryCpuShare(UserRequestContext->PoolConfig->QueryCpuLimitPercentPerNode / 100.0); + } } return result; diff --git a/ydb/core/kqp/node_service/kqp_node_service.cpp b/ydb/core/kqp/node_service/kqp_node_service.cpp index 0b2952e0e286..ac482ab9a1d3 100644 --- a/ydb/core/kqp/node_service/kqp_node_service.cpp +++ b/ydb/core/kqp/node_service/kqp_node_service.cpp @@ -195,6 +195,26 @@ class TKqpNodeService : public TActorBootstrapped { auto schedulerNow = TlsActivationContext->Monotonic(); + TString schedulerGroup = msg.GetSchedulerGroup(); + + if (SchedulerOptions.Scheduler->Disabled(schedulerGroup)) { + auto share = msg.GetPoolMaxCpuShare(); + if (share <= 0 && msg.HasQueryCpuShare()) { + share = 1.0; + } + if (share > 0) { + Scheduler->UpdateGroupShare(schedulerGroup, share, schedulerNow); + Send(SchedulerActorId, new TEvSchedulerNewPool(msg.GetDatabase(), schedulerGroup)); + } else { + schedulerGroup = ""; + } + } + + std::optional querySchedulerGroup; + if (msg.HasQueryCpuShare() && schedulerGroup) { + querySchedulerGroup = Scheduler->MakePerQueryGroup(schedulerNow, msg.GetQueryCpuShare(), schedulerGroup); + } + // start compute actors TMaybe rlPath = Nothing(); if (msgRtSettings.HasRlPath()) { @@ -208,30 +228,21 @@ class TKqpNodeService : public TActorBootstrapped { const ui32 tasksCount = msg.GetTasks().size(); for (auto& dqTask: *msg.MutableTasks()) { - TString group = msg.GetSchedulerGroup(); - TComputeActorSchedulingOptions schedulingTaskOptions { .Now = schedulerNow, .SchedulerActorId = SchedulerActorId, .Scheduler = Scheduler.get(), - .Group = group, + .Group = schedulerGroup, .Weight = 1, - .NoThrottle = false, + .NoThrottle = schedulerGroup.empty(), .Counters = Counters }; - if (SchedulerOptions.Scheduler->Disabled(group)) { - auto share = msg.GetMaxCpuShare(); - if (share > 0) { - Scheduler->UpdateMaxShare(group, share, schedulerNow); - Send(SchedulerActorId, new TEvSchedulerNewPool(msg.GetDatabase(), group, share)); - } else { - schedulingTaskOptions.NoThrottle = true; - } - } - if (!schedulingTaskOptions.NoThrottle) { schedulingTaskOptions.Handle = SchedulerOptions.Scheduler->Enroll(schedulingTaskOptions.Group, schedulingTaskOptions.Weight, schedulingTaskOptions.Now); + if (querySchedulerGroup) { + Scheduler->AddToGroup(schedulerNow, *querySchedulerGroup, schedulingTaskOptions.Handle); + } } auto result = CaFactory_->CreateKqpComputeActor({ diff --git a/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp b/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp index 8d22b554a671..db3a9b844400 100644 --- a/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp +++ b/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp @@ -17,14 +17,259 @@ namespace { return TDuration::MicroSeconds(t); } - static constexpr double MinEntitiesWeight = 1e-8; - static constexpr TDuration AvgBatch = TDuration::MicroSeconds(100); } namespace NKikimr { namespace NKqp { +class IObservable : TNonCopyable { +public: + virtual bool Update() = 0; + + void AddDependency(IObservable* dep) { + Depth = Max(Depth, dep->Depth + 1); + Dependencies.insert(dep); + dep->Dependents.insert(this); + } + + bool HasDependents() { + return !Dependents.empty(); + } + + virtual ~IObservable() { + for (auto& dep : Dependencies) { + dep->Dependents.erase(this); + } + for (auto& dep : Dependents) { + dep->Dependencies.erase(this); + } + } + + size_t GetDepth() { + return Depth; + } + + template + void ForAllDependents(T&& f) { + for (auto* dep : Dependents) { + f(dep); + } + } + +protected: + TSet CutAllDependents() { + TSet res; + Dependents.swap(res); + for (auto* dep : res) { + dep->Dependencies.erase(this); + } + return res; + } + +private: + size_t Depth = 0; + + TSet Dependencies; + TSet Dependents; +}; + +template +class IObservableValue : public IObservable { +protected: + virtual double DoUpdateValue() = 0; + +public: + bool Update() override { + if (auto val = DoUpdateValue()) { + Value = val; + return true; + } else { + return false; + } + } + + T GetValue() { + return Value; + } + +private: + T Value; +}; + +class TShare : public IObservableValue { +protected: + double DoUpdateValue() override { + return Base->GetValue() * Share->GetValue(); + } + +public: + TShare(IObservableValue* base, IObservableValue* share) + : Base(base) + , Share(share) + { + AddDependency(base); + AddDependency(share); + Update(); + } + +private: + IObservableValue* Base; + IObservableValue* Share; +}; + +template +class TParameter; + +class TObservableUpdater : IObservable { +private: + bool Update() override { + return false; + } + +public: + void UpdateAll() { + TVector> queue; + auto deps = CutAllDependents(); + for (auto* dep : deps) { + queue.resize(Max(queue.size(), dep->GetDepth() + 1)); + queue[dep->GetDepth()].insert(dep); + } + + for (size_t i = 0; i < queue.size(); ++i) { + TSet cur; + queue[i].swap(cur); + for (auto* node : cur) { + if (node->Update()) { + node->ForAllDependents([&](auto* dep){ + queue.resize(Max(queue.size(), dep->GetDepth() + 1)); + queue[dep->GetDepth()].insert(dep); + }); + } + } + } + + } + + void ToUpdate(IObservable* dep) { + dep->AddDependency(this); + } + + using TParameterKey = std::pair; + + template + T* FindValue(TParameterKey key) { + if (auto ptr = Params.FindPtr(key)) { + return ptr->Get(); + } + return nullptr; + } + + template + void AddValue(TParameterKey key, THolder value) { + Params.emplace(key, std::move(value)); + } + + template + TParameter* FindOrAddParameter(TParameterKey key, double def); + + ui64 ValuesCount() { + return Params.size(); + } + + void CollectValues() { + std::vector toerase; + for (auto& [k, v] : Params) { + if (!v.Holder->HasDependents()) { + toerase.push_back(k); + } + } + for (auto& key : toerase) { + Params.erase(key); + } + } + +private: + struct TValueContainer { + TParameter* AsDoubleParameter = nullptr; + TParameter* AsBoolParameter = nullptr; + + THolder Holder; + + TValueContainer() = default; + TValueContainer(TValueContainer&&) = default; + + TValueContainer(THolder value) { + Holder = std::move(value); + } + + TValueContainer(THolder> value); + TValueContainer(THolder> value); + + template + T* Get(); + }; + + THashMap Params; +}; + +template +class TParameter : public IObservableValue { +public: + TParameter(TObservableUpdater* engine, double initialValue) + : Value_(initialValue) + , Updater_(engine) + { + Updater_->ToUpdate(this); + } + + void SetValue(T val) { + Value_ = val; + Updater_->ToUpdate(this); + } + +protected: + double DoUpdateValue() override { + return Value_; + } + +private: + T Value_; + TObservableUpdater* Updater_; +}; + +template +TParameter* TObservableUpdater::FindOrAddParameter(TParameterKey key, double def) { + if (auto* ptr = FindValue>(key)) { + return ptr; + } + auto value = MakeHolder>(this, def); + auto* result = value.Get(); + AddValue>(key, std::move(value)); + return result; +} + +TObservableUpdater::TValueContainer::TValueContainer(THolder> value) { + AsDoubleParameter = value.Get(); + Holder = THolder(value.Release()); +} + +TObservableUpdater::TValueContainer::TValueContainer(THolder> value) { + AsBoolParameter = value.Get(); + Holder = THolder(value.Release()); +} + +template +T* TObservableUpdater::TValueContainer::Get() { + if constexpr (std::is_same_v>) { + return AsDoubleParameter; + } else if constexpr (std::is_same_v>) { + return AsBoolParameter; + } else { + return Holder.Get(); + } +} + template class TMultiThreadView { public: @@ -34,7 +279,6 @@ class TMultiThreadView { { Usage->fetch_add(1); } - const T* get() { return Slot; } @@ -107,17 +351,17 @@ class TSchedulerEntity { ~TSchedulerEntity() {} struct TGroupMutableStats { - double Weight = 0; + double Capacity = 0; TMonotonic LastNowRecalc; bool Disabled = false; - double EntitiesWeight = 0; + i64 EntitiesWeight = 0; double MaxDeviation = 0; double MaxLimitDeviation = 0; ssize_t TrackedBefore = 0; double Limit(TMonotonic now) const { - return FromDuration(now - LastNowRecalc) * Weight + MaxLimitDeviation + TrackedBefore; + return FromDuration(now - LastNowRecalc) * Capacity + MaxLimitDeviation + TrackedBefore; } }; @@ -126,13 +370,43 @@ class TSchedulerEntity { std::atomic DelayedSumBatches = 0; std::atomic DelayedCount = 0; - double Share; + THolder> Share; + + ::NMonitoring::TDynamicCounters::TCounterPtr Vtime; + ::NMonitoring::TDynamicCounters::TCounterPtr EntitiesWeight; + ::NMonitoring::TDynamicCounters::TCounterPtr Limit; + ::NMonitoring::TDynamicCounters::TCounterPtr Weight; + + ::NMonitoring::TDynamicCounters::TCounterPtr SchedulerClock; + ::NMonitoring::TDynamicCounters::TCounterPtr SchedulerLimitUs; + ::NMonitoring::TDynamicCounters::TCounterPtr SchedulerTrackedUs; + + TString Name; + + void AssignWeight() { + MutableStats.Next()->Capacity = Share->GetValue(); + } + + void InitCounters(const TIntrusivePtr& counters) { + if (Vtime || !Name) { + return; + } + + auto group = counters->GetKqpCounters()->GetSubgroup("NodeScheduler/Group", Name); + Vtime = group->GetCounter("VTime", true); + EntitiesWeight = group->GetCounter("Entities", false); + Limit = group->GetCounter("Limit", true); + Weight = group->GetCounter("Weight", false); + SchedulerClock = group->GetCounter("Clock", false); + SchedulerTrackedUs = group->GetCounter("Tracked", true); + SchedulerLimitUs = group->GetCounter("AbsoluteLimit", true); + } TMultithreadPublisher MutableStats; }; - TGroupRecord* Group; - double Weight; + TStackVec> Groups; + i64 Weight; double Vruntime = 0; double Vstart; @@ -152,8 +426,10 @@ class TSchedulerEntity { bool isThrottled = false; void TrackTime(TDuration time, TMonotonic) { - auto group = Group->MutableStats.Current(); - Group->TrackedMicroSeconds.fetch_add(time.MicroSeconds()); + for (auto group : Groups) { + //auto current = group->MutableStats.Current(); + group->TrackedMicroSeconds.fetch_add(time.MicroSeconds()); + } } void UpdateBatchTime(TDuration time) { @@ -168,48 +444,62 @@ class TSchedulerEntity { } } - TMaybe GroupDelay(TMonotonic now) { - auto group = Group->MutableStats.Current(); - auto limit = group.get()->Limit(now); - auto tracked = Group->TrackedMicroSeconds.load(); + TMaybe GroupDelay(TMonotonic now, TGroupRecord* group) { + auto current = group->MutableStats.Current(); + auto limit = current.get()->Limit(now); + auto tracked = group->TrackedMicroSeconds.load(); //double Coeff = pow(WakeupDelay, Wakeups); if (limit > tracked) { return {}; } else { return Min(MaxDelay, ToDuration(/*Coeff * */(tracked - limit + - Max(0, Group->DelayedSumBatches.load()) + BatchTime.MicroSeconds() + - ActivationPenalty.MicroSeconds() * (Group->DelayedCount.load() + 1) + - group.get()->MaxLimitDeviation) / group.get()->Weight)); + Max(0, group->DelayedSumBatches.load()) + BatchTime.MicroSeconds() + + ActivationPenalty.MicroSeconds() * (group->DelayedCount.load() + 1) + + current.get()->MaxLimitDeviation) / current.get()->Capacity)); } } + TMaybe GroupDelay(TMonotonic now) { + TMaybe result; + for (auto group : Groups) { + auto groupResult = GroupDelay(now, group); + if (!result) { + result = groupResult; + } else if (groupResult && *result < *groupResult) { + result = groupResult; + } + } + return result; + } + void MarkThrottled() { isThrottled = true; - Group->DelayedSumBatches.fetch_add(BatchTime.MicroSeconds()); - Group->DelayedCount.fetch_add(1); + for (auto group : Groups) { + group->DelayedSumBatches.fetch_add(BatchTime.MicroSeconds()); + group->DelayedCount.fetch_add(1); + } } void MarkResumed() { isThrottled = false; - Group->DelayedSumBatches.fetch_sub(BatchTime.MicroSeconds()); - Group->DelayedCount.fetch_sub(1); + for (auto group : Groups) { + group->DelayedSumBatches.fetch_sub(BatchTime.MicroSeconds()); + group->DelayedCount.fetch_sub(1); + } } }; struct TComputeScheduler::TImpl { - TVector<::NMonitoring::TDynamicCounters::TCounterPtr> VtimeCounters; - TVector<::NMonitoring::TDynamicCounters::TCounterPtr> EntitiesWeightCounters; - TVector<::NMonitoring::TDynamicCounters::TCounterPtr> LimitCounters; - TVector<::NMonitoring::TDynamicCounters::TCounterPtr> WeightCounters; - - TVector<::NMonitoring::TDynamicCounters::TCounterPtr> SchedulerClock; - TVector<::NMonitoring::TDynamicCounters::TCounterPtr> SchedulerLimitUs; - TVector<::NMonitoring::TDynamicCounters::TCounterPtr> SchedulerTrackedUs; - - THashMap PoolId; + THashMap GroupId; std::vector> Records; - double SumCores; + TObservableUpdater WeightsUpdater; + TParameter SumCores{&WeightsUpdater, 1}; + + enum : ui32 { + Share = 1, + PerQueryShare = 2, + }; TIntrusivePtr Counters; TDuration SmoothPeriod = TDuration::MilliSeconds(100); @@ -217,18 +507,52 @@ struct TComputeScheduler::TImpl { TDuration MaxDelay = TDuration::Seconds(10); - void AssignWeights() { - for (auto& record : Records) { - record->MutableStats.Next()->Weight = SumCores * record->Share; + void CreateGroup(THolder> share, NMonotonic::TMonotonic now, std::optional groupName = std::nullopt) { + auto group = std::make_unique(); + group->Share = std::move(share); + if (groupName) { + group->Name = *groupName; + GroupId[*groupName] = Records.size(); } + AdvanceTime(now, group.get()); + Records.push_back(std::move(group)); } - void CreateGroup(TString groupName, double maxShare) { - PoolId[groupName] = Records.size(); - auto group = std::make_unique(); - group->Share = maxShare; - Records.push_back(std::move(group)); + void CollectGroups() { + std::vector remap; + std::vector> records; + + for (size_t i = 0; i < Records.size(); ++i) { + auto record = Records[i]->MutableStats.Current(); + if (record.get()->EntitiesWeight > 0 || Records[i]->Share->HasDependents()) { + remap.push_back(records.size()); + records.emplace_back(Records[i].release()); + } else { + // to delete + remap.push_back(-1); + } + } + + Records.swap(records); + + { + std::vector toerase; + for (auto& [k, v] : GroupId) { + if (remap[v] >= 0) { + v = remap[v]; + } else { + toerase.push_back(k); + } + } + for (auto& k: toerase) { + GroupId.erase(k); + } + } + + WeightsUpdater.CollectValues(); } + + void AdvanceTime(TMonotonic now, TSchedulerEntity::TGroupRecord* record); }; TComputeScheduler::TComputeScheduler() { @@ -237,86 +561,93 @@ TComputeScheduler::TComputeScheduler() { TComputeScheduler::~TComputeScheduler() = default; -TSchedulerEntityHandle TComputeScheduler::Enroll(TString groupName, double weight, TMonotonic now) { - Y_ENSURE(Impl->PoolId.contains(groupName), "unknown scheduler group"); - auto* groupEntry = Impl->Records[Impl->PoolId.at(groupName)].get(); - groupEntry->MutableStats.Next()->EntitiesWeight += weight; - Impl->AssignWeights(); - AdvanceTime(now); +void TComputeScheduler::AddToGroup(TMonotonic now, ui64 id, TSchedulerEntityHandle& handle) { + auto group = Impl->Records[id].get(); + (*handle).Groups.push_back(group); + group->MutableStats.Next()->EntitiesWeight += (*handle).Weight; + Impl->AdvanceTime(now, group); +} + +TSchedulerEntityHandle TComputeScheduler::Enroll(TString groupName, i64 weight, TMonotonic now) { + Y_ENSURE(Impl->GroupId.contains(groupName), "unknown scheduler group"); + auto id = Impl->GroupId.at(groupName); - auto result = std::make_unique(); - result->Group = groupEntry; - result->Weight = weight; - result->MaxDelay = Impl->MaxDelay; + TSchedulerEntityHandle result{new TSchedulerEntity()}; + (*result).Weight = weight; + (*result).MaxDelay = Impl->MaxDelay; - return TSchedulerEntityHandle(result.release()); + AddToGroup(now, id, result); + return result; } -void TComputeScheduler::AdvanceTime(TMonotonic now) { - if (Impl->Counters) { - if (Impl->VtimeCounters.size() < Impl->Records.size()) { - Impl->VtimeCounters.resize(Impl->Records.size()); - Impl->EntitiesWeightCounters.resize(Impl->Records.size()); - Impl->LimitCounters.resize(Impl->Records.size()); - Impl->WeightCounters.resize(Impl->Records.size()); - Impl->SchedulerClock.resize(Impl->Records.size()); - Impl->SchedulerLimitUs.resize(Impl->Records.size()); - Impl->SchedulerTrackedUs.resize(Impl->Records.size()); - - for (auto& [k, i] : Impl->PoolId) { - auto group = Impl->Counters->GetKqpCounters()->GetSubgroup("NodeScheduler/Group", k); - Impl->VtimeCounters[i] = group->GetCounter("VTime", true); - Impl->EntitiesWeightCounters[i] = group->GetCounter("Entities", false); - Impl->LimitCounters[i] = group->GetCounter("Limit", true); - Impl->WeightCounters[i] = group->GetCounter("Weight", false); - Impl->SchedulerClock[i] = group->GetCounter("Clock", false); - Impl->SchedulerTrackedUs[i] = group->GetCounter("Tracked", true); - Impl->SchedulerLimitUs[i] = group->GetCounter("AbsoluteLimit", true); - } +void TComputeScheduler::TImpl::AdvanceTime(TMonotonic now, TSchedulerEntity::TGroupRecord* record) { + if (Counters) { + record->InitCounters(Counters); + } + record->MutableStats.Next()->Capacity = record->Share->GetValue(); + auto& v = record->MutableStats; + { + auto group = v.Current(); + if (group.get()->LastNowRecalc > now) { + return; + } + double delta = 0; + + auto tracked = record->TrackedMicroSeconds.load(); + v.Next()->MaxLimitDeviation = SmoothPeriod.MicroSeconds() * v.Next()->Capacity; + v.Next()->LastNowRecalc = now; + v.Next()->TrackedBefore = + Max( + tracked - FromDuration(ForgetInteval) * group.get()->Capacity, + Min(group.get()->Limit(now) - group.get()->MaxLimitDeviation, tracked)); + + v.Next()->MaxDeviation = (FromDuration(SmoothPeriod) * v.Next()->Capacity) / v.Next()->Capacity; + + //if (group.get()->EntitiesWeight > 0) { + // delta = FromDuration(now - group.get()->LastNowRecalc) * group.get()->Capacity / group.get()->EntitiesWeight; + //} + + if (record->Vtime) { + record->SchedulerLimitUs->Set(group.get()->Limit(now)); + record->SchedulerTrackedUs->Set(record->TrackedMicroSeconds.load()); + record->SchedulerClock->Add(now.MicroSeconds() - group.get()->LastNowRecalc.MicroSeconds()); + record->Vtime->Add(delta); + record->EntitiesWeight->Set(v.Next()->EntitiesWeight); + record->Limit->Add(FromDuration(now - group.get()->LastNowRecalc) * group.get()->Capacity); + record->Weight->Set(group.get()->Capacity); } } + v.Publish(); +} + +void TComputeScheduler::AdvanceTime(TMonotonic now) { + Impl->WeightsUpdater.UpdateAll(); for (size_t i = 0; i < Impl->Records.size(); ++i) { - auto& v = Impl->Records[i]->MutableStats; - { - auto group = v.Current(); - if (group.get()->LastNowRecalc > now) { - continue; - } - double delta = 0; - - auto tracked = Impl->Records[i]->TrackedMicroSeconds.load(); - v.Next()->MaxLimitDeviation = Impl->SmoothPeriod.MicroSeconds() * v.Next()->Weight; - v.Next()->LastNowRecalc = now; - v.Next()->TrackedBefore = - Max( - tracked - FromDuration(Impl->ForgetInteval) * group.get()->Weight, - Min(group.get()->Limit(now) - group.get()->MaxLimitDeviation, tracked)); - - if (!group.get()->Disabled && group.get()->EntitiesWeight > MinEntitiesWeight) { - delta = FromDuration(now - group.get()->LastNowRecalc) * group.get()->Weight / group.get()->EntitiesWeight; - v.Next()->MaxDeviation = (FromDuration(Impl->SmoothPeriod) * v.Next()->Weight) / v.Next()->EntitiesWeight; - } + Impl->AdvanceTime(now, Impl->Records[i].get()); + } + Impl->CollectGroups(); + if (Impl->Counters) { + Impl->Counters->SchedulerGroupsCount->Set(Impl->Records.size()); + Impl->Counters->SchedulerValuesCount->Set(Impl->WeightsUpdater.ValuesCount()); + } +} - if (Impl->VtimeCounters.size() > i && Impl->VtimeCounters[i]) { - Impl->SchedulerLimitUs[i]->Set(group.get()->Limit(now)); - Impl->SchedulerTrackedUs[i]->Set(Impl->Records[i]->TrackedMicroSeconds.load()); - Impl->SchedulerClock[i]->Add(now.MicroSeconds() - group.get()->LastNowRecalc.MicroSeconds()); - Impl->VtimeCounters[i]->Add(delta); - Impl->EntitiesWeightCounters[i]->Set(v.Next()->EntitiesWeight); - Impl->LimitCounters[i]->Add(FromDuration(now - group.get()->LastNowRecalc) * group.get()->Weight); - Impl->WeightCounters[i]->Set(group.get()->Weight); - } - } - v.Publish(); +void TComputeScheduler::Deregister(TSchedulerEntityHandle& self, TMonotonic now) { + for (auto group : (*self).Groups) { + auto* next = group->MutableStats.Next(); + next->EntitiesWeight -= (*self).Weight; + Impl->AdvanceTime(now, group); } } -void TComputeScheduler::Deregister(TSchedulerEntity& self, TMonotonic now) { - auto* group = self.Group->MutableStats.Next(); - group->EntitiesWeight -= self.Weight; +ui64 TComputeScheduler::MakePerQueryGroup(TMonotonic now, double share, TString baseGroup) { + auto baseId = Impl->GroupId.at(baseGroup); + auto perQueryShare = Impl->WeightsUpdater.FindOrAddParameter({baseGroup, TImpl::PerQueryShare}, share); - Impl->AssignWeights(); - AdvanceTime(now); + Impl->CreateGroup(MakeHolder(Impl->Records[baseId]->Share.Get(), perQueryShare), now); + ui64 res = Impl->Records.size() - 1; + Impl->AdvanceTime(now, Impl->Records[res].get()); + return res; } void TSchedulerEntityHandle::TrackTime(TDuration time, TMonotonic now) { @@ -356,35 +687,42 @@ void TComputeScheduler::SetForgetInterval(TDuration period) { } bool TComputeScheduler::Disabled(TString group) { - auto ptr = Impl->PoolId.FindPtr(group); + auto ptr = Impl->GroupId.FindPtr(group); return !ptr || Impl->Records[*ptr]->MutableStats.Current().get()->Disabled; } -bool TComputeScheduler::Disable(TString group, TMonotonic now) { - auto ptr = Impl->PoolId.FindPtr(group); - if (Impl->Records[*ptr]->MutableStats.Current().get()->Weight > MinEntitiesWeight) { - return false; +void TComputeScheduler::Disable(TString group, TMonotonic now) { + auto ptr = Impl->GroupId.FindPtr(group); + // if ptr == 0 it's already disabled + if (ptr) { + Impl->Records[*ptr]->MutableStats.Next()->Disabled = true; + Impl->AdvanceTime(now, Impl->Records[*ptr].get()); } - Impl->Records[*ptr]->MutableStats.Next()->Disabled = true; - AdvanceTime(now); - return true; } -void TComputeScheduler::UpdateMaxShare(TString group, double share, TMonotonic now) { - auto ptr = Impl->PoolId.FindPtr(group); +void TComputeScheduler::UpdateGroupShare(TString group, double share, TMonotonic now) { + auto ptr = Impl->GroupId.FindPtr(group); + + auto* shareValue = Impl->WeightsUpdater.FindOrAddParameter({group, TImpl::Share}, share); + shareValue->SetValue(share); if (!ptr) { - Impl->CreateGroup(group, share); + auto cap = MakeHolder(&Impl->SumCores, shareValue); + Impl->CreateGroup(std::move(cap), now, group); } else { auto& record = Impl->Records[*ptr]; - record->Share = share; + record->MutableStats.Next()->Disabled = false; + Impl->AdvanceTime(now, record.get()); } - Impl->AssignWeights(); - AdvanceTime(now); +} + +void TComputeScheduler::UpdatePerQueryShare(TString group, double share, TMonotonic) { + auto ptr = Impl->WeightsUpdater.FindOrAddParameter({group, TImpl::PerQueryShare}, share); + ptr->SetValue(share); } void TComputeScheduler::SetCapacity(ui64 cores) { - Impl->SumCores = cores; + Impl->SumCores.SetValue(cores); } ::NMonitoring::TDynamicCounters::TCounterPtr TComputeScheduler::GetGroupUsageCounter(TString group) const { @@ -462,13 +800,12 @@ class TSchedulerActor : public TActorBootstrapped { void Handle(TEvSchedulerDeregister::TPtr& ev) { if (ev->Get()->SchedulerEntity) { - Opts.Scheduler->Deregister(*ev->Get()->SchedulerEntity, TlsActivationContext->Monotonic()); + Opts.Scheduler->Deregister(ev->Get()->SchedulerEntity, TlsActivationContext->Monotonic()); } } void Handle(TEvSchedulerNewPool::TPtr& ev) { Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(ev->Get()->Database, ev->Get()->Pool)); - Opts.Scheduler->UpdateMaxShare(ev->Get()->Pool, ev->Get()->MaxShare, TlsActivationContext->Monotonic()); } void Handle(TEvPingPool::TPtr& ev) { @@ -477,11 +814,21 @@ class TSchedulerActor : public TActorBootstrapped { void Handle(NWorkload::TEvUpdatePoolInfo::TPtr& ev) { if (ev->Get()->Config.has_value()) { - Opts.Scheduler->UpdateMaxShare(ev->Get()->PoolId, ev->Get()->Config->TotalCpuLimitPercentPerNode / 100.0, TlsActivationContext->Monotonic()); - } else { - if (!Opts.Scheduler->Disable(ev->Get()->PoolId, TlsActivationContext->Monotonic())) { - Schedule(Opts.ActivePoolPollingTimeout.ToDeadLine(), new TEvPingPool(ev->Get()->Database, ev->Get()->PoolId)); + auto totalShare = ev->Get()->Config->TotalCpuLimitPercentPerNode / 100.0; + auto queryShare = ev->Get()->Config->QueryCpuLimitPercentPerNode / 100.0; + + if (totalShare <= 0 && queryShare > 0) { + totalShare = 1; + } + + if (queryShare <= 0) { + queryShare = 1; } + + Opts.Scheduler->UpdateGroupShare(ev->Get()->PoolId, totalShare, TlsActivationContext->Monotonic()); + Opts.Scheduler->UpdatePerQueryShare(ev->Get()->PoolId, queryShare, TlsActivationContext->Monotonic()); + } else { + Opts.Scheduler->Disable(ev->Get()->PoolId, TlsActivationContext->Monotonic()); } } diff --git a/ydb/core/kqp/runtime/kqp_compute_scheduler.h b/ydb/core/kqp/runtime/kqp_compute_scheduler.h index fd66610fbd33..37bb1b2008a0 100644 --- a/ydb/core/kqp/runtime/kqp_compute_scheduler.h +++ b/ydb/core/kqp/runtime/kqp_compute_scheduler.h @@ -61,22 +61,27 @@ class TComputeScheduler { void ReportCounters(TIntrusivePtr); - void UpdateMaxShare(TString, double, TMonotonic now); void SetCapacity(ui64 cores); + void UpdateGroupShare(TString name, double share, TMonotonic now); + void UpdatePerQueryShare(TString name, double share, TMonotonic now); + + ui64 MakePerQueryGroup(TMonotonic now, double share, TString baseGroup); + void AddToGroup(TMonotonic now, ui64, TSchedulerEntityHandle&); + void SetMaxDeviation(TDuration); void SetForgetInterval(TDuration); ::NMonitoring::TDynamicCounters::TCounterPtr GetGroupUsageCounter(TString group) const; - TSchedulerEntityHandle Enroll(TString group, double weight, TMonotonic now); + TSchedulerEntityHandle Enroll(TString group, i64 weight, TMonotonic now); void AdvanceTime(TMonotonic now); - void Deregister(TSchedulerEntity& self, TMonotonic now); + void Deregister(TSchedulerEntityHandle& self, TMonotonic now); bool Disabled(TString group); - bool Disable(TString group, TMonotonic now); + void Disable(TString group, TMonotonic now); private: struct TImpl; @@ -114,12 +119,10 @@ struct TEvSchedulerDeregister : public TEventLocal { TString Database; TString Pool; - double MaxShare; - TEvSchedulerNewPool(TString database, TString pool, double maxShare) + TEvSchedulerNewPool(TString database, TString pool) : Database(database) , Pool(pool) - , MaxShare(maxShare) { } }; diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto index 7399ec1cc164..2436d4b97b13 100644 --- a/ydb/core/protos/kqp.proto +++ b/ydb/core/protos/kqp.proto @@ -589,9 +589,11 @@ message TEvStartKqpTasksRequest { optional string SchedulerGroup = 9; optional double MemoryPoolPercent = 10 [default = 100]; optional string Database = 11; - optional double MaxCpuShare = 12; optional uint64 LockTxId = 13; optional uint32 LockNodeId = 14; + + optional double PoolMaxCpuShare = 12; + optional double QueryCpuShare = 16; } message TEvStartKqpTasksResponse { diff --git a/ydb/core/resource_pools/resource_pool_settings.cpp b/ydb/core/resource_pools/resource_pool_settings.cpp index 1ee8d077bbcd..e6d9d1dfedcc 100644 --- a/ydb/core/resource_pools/resource_pool_settings.cpp +++ b/ydb/core/resource_pools/resource_pool_settings.cpp @@ -57,7 +57,8 @@ std::unordered_map TPoolSettings::GetProperti {"queue_size", &QueueSize}, {"query_memory_limit_percent_per_node", &QueryMemoryLimitPercentPerNode}, {"database_load_cpu_threshold", &DatabaseLoadCpuThreshold}, - {"total_cpu_limit_percent_per_node", &TotalCpuLimitPercentPerNode} + {"total_cpu_limit_percent_per_node", &TotalCpuLimitPercentPerNode}, + {"query_cpu_limit_percent_per_node", &QueryCpuLimitPercentPerNode}, }; if (!restricted) { properties.insert({"query_cancel_after_seconds", &QueryCancelAfter}); diff --git a/ydb/core/resource_pools/resource_pool_settings.h b/ydb/core/resource_pools/resource_pool_settings.h index c0303a951bb2..4d2052be586d 100644 --- a/ydb/core/resource_pools/resource_pool_settings.h +++ b/ydb/core/resource_pools/resource_pool_settings.h @@ -45,6 +45,7 @@ struct TPoolSettings : public TSettingsBase { TPercent QueryMemoryLimitPercentPerNode = -1; // Percent from node memory capacity, -1 = disabled TPercent DatabaseLoadCpuThreshold = -1; // -1 = disabled TPercent TotalCpuLimitPercentPerNode = -1; // -1 = disabled + TPercent QueryCpuLimitPercentPerNode = -1; // -1 = disabled; }; } // namespace NKikimr::NResourcePool From c89ab17a438187d359f663294e391625de61250a Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Mon, 7 Oct 2024 12:36:39 +0300 Subject: [PATCH 5/5] YQ-3689 added database id for workload manager (#9768) --- ydb/core/kqp/common/events/workload_service.h | 28 ++-- .../kqp/common/kqp_user_request_context.cpp | 3 +- .../kqp/common/kqp_user_request_context.h | 1 + ydb/core/kqp/common/simple/query_id.cpp | 5 +- ydb/core/kqp/common/simple/query_id.h | 3 +- .../kqp/compile_service/kqp_compile_actor.cpp | 2 +- .../compile_service/kqp_compile_service.cpp | 2 +- ydb/core/kqp/executer_actor/kqp_planner.cpp | 1 + .../executer_actor/kqp_scheme_executer.cpp | 2 + .../kqp/executer_actor/ut/kqp_executer_ut.cpp | 2 +- .../resource_pool_classifier/checker.cpp | 24 +-- .../resource_pool_classifier/manager.cpp | 2 +- ydb/core/kqp/gateway/kqp_gateway.h | 3 +- ydb/core/kqp/gateway/kqp_ic_gateway.cpp | 20 ++- ydb/core/kqp/host/kqp_gateway_proxy.cpp | 1 + ydb/core/kqp/host/kqp_host.cpp | 1 + .../kqp/node_service/kqp_node_service.cpp | 2 +- .../kqp/provider/yql_kikimr_gateway_ut.cpp | 2 +- ydb/core/kqp/provider/yql_kikimr_provider.h | 9 ++ .../kqp_proxy_databases_cache.cpp | 17 +- .../kqp/proxy_service/kqp_proxy_service.cpp | 12 +- .../proxy_service/kqp_proxy_service_impl.h | 81 +++++----- .../kqp/runtime/kqp_compute_scheduler.cpp | 10 +- ydb/core/kqp/runtime/kqp_compute_scheduler.h | 6 +- .../kqp/session_actor/kqp_query_state.cpp | 12 +- ydb/core/kqp/session_actor/kqp_query_state.h | 1 + .../kqp/session_actor/kqp_session_actor.cpp | 5 +- .../kqp/session_actor/kqp_worker_actor.cpp | 2 +- ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp | 2 +- ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp | 43 +++++- ydb/core/kqp/workload_service/actors/actors.h | 8 +- .../actors/pool_handlers_acors.cpp | 56 +++---- .../workload_service/actors/scheme_actors.cpp | 42 ++--- ydb/core/kqp/workload_service/common/events.h | 36 ++--- .../kqp/workload_service/common/helpers.cpp | 24 +++ .../kqp/workload_service/common/helpers.h | 3 + .../workload_service/kqp_workload_service.cpp | 119 +++++++------- .../kqp_workload_service_impl.h | 25 +-- .../workload_service/tables/table_queries.cpp | 146 +++++++++--------- .../workload_service/tables/table_queries.h | 8 +- .../common/kqp_workload_service_ut_common.cpp | 11 +- .../common/kqp_workload_service_ut_common.h | 1 + .../ut/kqp_workload_service_tables_ut.cpp | 4 +- .../ut/kqp_workload_service_ut.cpp | 85 +++++++++- ydb/core/protos/kqp.proto | 1 + ydb/services/metadata/manager/abstract.h | 1 + ydb/tools/query_replay/query_compiler.cpp | 6 +- ydb/tools/query_replay_yt/query_compiler.cpp | 6 +- 48 files changed, 530 insertions(+), 356 deletions(-) diff --git a/ydb/core/kqp/common/events/workload_service.h b/ydb/core/kqp/common/events/workload_service.h index 9781b951e2d5..385e4c7eded8 100644 --- a/ydb/core/kqp/common/events/workload_service.h +++ b/ydb/core/kqp/common/events/workload_service.h @@ -14,24 +14,24 @@ namespace NKikimr::NKqp::NWorkload { struct TEvSubscribeOnPoolChanges : public NActors::TEventLocal { - TEvSubscribeOnPoolChanges(const TString& database, const TString& poolId) - : Database(database) + TEvSubscribeOnPoolChanges(const TString& databaseId, const TString& poolId) + : DatabaseId(databaseId) , PoolId(poolId) {} - const TString Database; + const TString DatabaseId; const TString PoolId; }; struct TEvPlaceRequestIntoPool : public NActors::TEventLocal { - TEvPlaceRequestIntoPool(const TString& database, const TString& sessionId, const TString& poolId, TIntrusiveConstPtr userToken) - : Database(database) + TEvPlaceRequestIntoPool(const TString& databaseId, const TString& sessionId, const TString& poolId, TIntrusiveConstPtr userToken) + : DatabaseId(databaseId) , SessionId(sessionId) , PoolId(poolId) , UserToken(userToken) {} - const TString Database; + const TString DatabaseId; const TString SessionId; TString PoolId; // Can be changed to default pool id TIntrusiveConstPtr UserToken; @@ -52,15 +52,15 @@ struct TEvContinueRequest : public NActors::TEventLocal { - TEvCleanupRequest(const TString& database, const TString& sessionId, const TString& poolId, TDuration duration, TDuration cpuConsumed) - : Database(database) + TEvCleanupRequest(const TString& databaseId, const TString& sessionId, const TString& poolId, TDuration duration, TDuration cpuConsumed) + : DatabaseId(databaseId) , SessionId(sessionId) , PoolId(poolId) , Duration(duration) , CpuConsumed(cpuConsumed) {} - const TString Database; + const TString DatabaseId; const TString SessionId; const TString PoolId; const TDuration Duration; @@ -78,23 +78,24 @@ struct TEvCleanupResponse : public NActors::TEventLocal { - TEvUpdatePoolInfo(const TString& database, const TString& poolId, const std::optional& config, const std::optional& securityObject) - : Database(database) + TEvUpdatePoolInfo(const TString& databaseId, const TString& poolId, const std::optional& config, const std::optional& securityObject) + : DatabaseId(databaseId) , PoolId(poolId) , Config(config) , SecurityObject(securityObject) {} - const TString Database; + const TString DatabaseId; const TString PoolId; const std::optional Config; const std::optional SecurityObject; }; struct TEvFetchDatabaseResponse : public NActors::TEventLocal { - TEvFetchDatabaseResponse(Ydb::StatusIds::StatusCode status, const TString& database, bool serverless, TPathId pathId, NYql::TIssues issues) + TEvFetchDatabaseResponse(Ydb::StatusIds::StatusCode status, const TString& database, const TString& databaseId, bool serverless, TPathId pathId, NYql::TIssues issues) : Status(status) , Database(database) + , DatabaseId(databaseId) , Serverless(serverless) , PathId(pathId) , Issues(std::move(issues)) @@ -102,6 +103,7 @@ struct TEvFetchDatabaseResponse : public NActors::TEventLocal& resultMap) { resultMap["TraceId"] = ctx.TraceId; resultMap["Database"] = ctx.Database; + resultMap["DatabaseId"] = ctx.DatabaseId; resultMap["SessionId"] = ctx.SessionId; resultMap["CurrentExecutionId"] = ctx.CurrentExecutionId; resultMap["CustomerSuppliedId"] = ctx.CustomerSuppliedId; diff --git a/ydb/core/kqp/common/kqp_user_request_context.h b/ydb/core/kqp/common/kqp_user_request_context.h index 1aa4a0574d67..9a17252fccf7 100644 --- a/ydb/core/kqp/common/kqp_user_request_context.h +++ b/ydb/core/kqp/common/kqp_user_request_context.h @@ -11,6 +11,7 @@ namespace NKikimr::NKqp { struct TUserRequestContext : public TAtomicRefCount { TString TraceId; TString Database; + TString DatabaseId; TString SessionId; TString CurrentExecutionId; TString CustomerSuppliedId; diff --git a/ydb/core/kqp/common/simple/query_id.cpp b/ydb/core/kqp/common/simple/query_id.cpp index b2066bbd8b88..241abd3cf82c 100644 --- a/ydb/core/kqp/common/simple/query_id.cpp +++ b/ydb/core/kqp/common/simple/query_id.cpp @@ -10,11 +10,12 @@ namespace NKikimr::NKqp { -TKqpQueryId::TKqpQueryId(const TString& cluster, const TString& database, const TString& text, +TKqpQueryId::TKqpQueryId(const TString& cluster, const TString& database, const TString& databaseId, const TString& text, const TKqpQuerySettings& settings, std::shared_ptr> queryParameterTypes, const TGUCSettings& gUCSettings) : Cluster(cluster) , Database(database) + , DatabaseId(databaseId) , Text(text) , Settings(settings) , QueryParameterTypes(queryParameterTypes) @@ -42,6 +43,7 @@ bool TKqpQueryId::IsSql() const { bool TKqpQueryId::operator==(const TKqpQueryId& other) const { if (!(Cluster == other.Cluster && Database == other.Database && + DatabaseId == other.DatabaseId && UserSid == other.UserSid && Text == other.Text && Settings == other.Settings && @@ -79,6 +81,7 @@ TString TKqpQueryId::SerializeToString() const { TStringBuilder result = TStringBuilder() << "{" << "Cluster: " << Cluster << ", " << "Database: " << Database << ", " + << "DatabaseId: " << DatabaseId << ", " << "UserSid: " << UserSid << ", " << "Text: " << EscapeC(Text) << ", " << "Settings: " << Settings.SerializeToString() << ", "; diff --git a/ydb/core/kqp/common/simple/query_id.h b/ydb/core/kqp/common/simple/query_id.h index 7dd7ef0dc5a7..120d524d9808 100644 --- a/ydb/core/kqp/common/simple/query_id.h +++ b/ydb/core/kqp/common/simple/query_id.h @@ -13,6 +13,7 @@ namespace NKikimr::NKqp { struct TKqpQueryId { TString Cluster; TString Database; + TString DatabaseId; TString UserSid; TString Text; TKqpQuerySettings Settings; @@ -21,7 +22,7 @@ struct TKqpQueryId { TGUCSettings GUCSettings; public: - TKqpQueryId(const TString& cluster, const TString& database, const TString& text, + TKqpQueryId(const TString& cluster, const TString& database, const TString& databaseId, const TString& text, const TKqpQuerySettings& settings, std::shared_ptr> queryParameterTypes, const TGUCSettings& gUCSettings); diff --git a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp index 4079fe093f3f..2a5c02ad776c 100644 --- a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp +++ b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp @@ -266,7 +266,7 @@ class TKqpCompileActor : public TActorBootstrapped { std::shared_ptr loader = std::make_shared( QueryId.Cluster, TlsActivationContext->ActorSystem(), Config, true, TempTablesState); - Gateway = CreateKikimrIcGateway(QueryId.Cluster, QueryId.Settings.QueryType, QueryId.Database, std::move(loader), + Gateway = CreateKikimrIcGateway(QueryId.Cluster, QueryId.Settings.QueryType, QueryId.Database, QueryId.DatabaseId, std::move(loader), ctx.ExecutorThread.ActorSystem, ctx.SelfID.NodeId(), counters, QueryServiceConfig); Gateway->SetToken(QueryId.Cluster, UserToken); diff --git a/ydb/core/kqp/compile_service/kqp_compile_service.cpp b/ydb/core/kqp/compile_service/kqp_compile_service.cpp index 4be278911ccf..bf921172acd2 100644 --- a/ydb/core/kqp/compile_service/kqp_compile_service.cpp +++ b/ydb/core/kqp/compile_service/kqp_compile_service.cpp @@ -153,7 +153,7 @@ class TKqpQueryCache { } } } - return TKqpQueryId{query.Cluster, query.Database, ast.Root->ToString(), query.Settings, astPgParams, query.GUCSettings}; + return TKqpQueryId{query.Cluster, query.Database, query.DatabaseId, ast.Root->ToString(), query.Settings, astPgParams, query.GUCSettings}; } TKqpCompileResult::TConstPtr FindByQuery(const TKqpQueryId& query, bool promote) { diff --git a/ydb/core/kqp/executer_actor/kqp_planner.cpp b/ydb/core/kqp/executer_actor/kqp_planner.cpp index d9eaf4619cf9..d35663cb9ae7 100644 --- a/ydb/core/kqp/executer_actor/kqp_planner.cpp +++ b/ydb/core/kqp/executer_actor/kqp_planner.cpp @@ -236,6 +236,7 @@ std::unique_ptr TKqpPlanner::SerializeReque request.SetSchedulerGroup(UserRequestContext->PoolId); request.SetDatabase(Database); + request.SetDatabaseId(UserRequestContext->DatabaseId); if (UserRequestContext->PoolConfig.has_value()) { request.SetMemoryPoolPercent(UserRequestContext->PoolConfig->QueryMemoryLimitPercentPerNode); request.SetPoolMaxCpuShare(UserRequestContext->PoolConfig->TotalCpuLimitPercentPerNode / 100.0); diff --git a/ydb/core/kqp/executer_actor/kqp_scheme_executer.cpp b/ydb/core/kqp/executer_actor/kqp_scheme_executer.cpp index 727e00a928dd..095e46f559c8 100644 --- a/ydb/core/kqp/executer_actor/kqp_scheme_executer.cpp +++ b/ydb/core/kqp/executer_actor/kqp_scheme_executer.cpp @@ -72,6 +72,7 @@ class TKqpSchemeExecuter : public TActorBootstrapped { , RequestType(requestType) , KqpTempTablesAgentActor(kqpTempTablesAgentActor) { + YQL_ENSURE(RequestContext); YQL_ENSURE(PhyTx); YQL_ENSURE(PhyTx->GetType() == NKqpProto::TKqpPhyTx::TYPE_SCHEME); @@ -402,6 +403,7 @@ class TKqpSchemeExecuter : public TActorBootstrapped { NMetadata::NModifications::IOperationsManager::TExternalModificationContext context; context.SetDatabase(Database); + context.SetDatabaseId(RequestContext->DatabaseId); context.SetActorSystem(actorSystem); if (UserToken) { context.SetUserToken(*UserToken); diff --git a/ydb/core/kqp/executer_actor/ut/kqp_executer_ut.cpp b/ydb/core/kqp/executer_actor/ut/kqp_executer_ut.cpp index 7e921e3c3578..937ceb64e1d7 100644 --- a/ydb/core/kqp/executer_actor/ut/kqp_executer_ut.cpp +++ b/ydb/core/kqp/executer_actor/ut/kqp_executer_ut.cpp @@ -42,7 +42,7 @@ NKqpProto::TKqpPhyTx BuildTxPlan(const TString& sql, TIntrusivePtr [[maybe_unused]] TIntrusivePtr MakeIcGateway(const TKikimrRunner& kikimr) { auto actorSystem = kikimr.GetTestServer().GetRuntime()->GetAnyNodeActorSystem(); - return CreateKikimrIcGateway(TString(DefaultKikimrClusterName), "/Root", TKqpGatewaySettings(), + return CreateKikimrIcGateway(TString(DefaultKikimrClusterName), "/Root", "/Root", TKqpGatewaySettings(), actorSystem, kikimr.GetTestServer().GetRuntime()->GetNodeId(0), TAlignedPagePoolCounters(), kikimr.GetTestServer().GetSettings().AppConfig->GetQueryServiceConfig()); } diff --git a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp index 7200999c7ca1..10f4d6b1f54c 100644 --- a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp +++ b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/checker.cpp @@ -48,13 +48,13 @@ class TRanksCheckerActor : public NKikimr::TQueryBase { using TBase = NKikimr::TQueryBase; public: - TRanksCheckerActor(const TString& database, const TString& sessionId, const TString& transactionId, const std::unordered_map& ranksToCheck) + TRanksCheckerActor(const TString& databaseId, const TString& sessionId, const TString& transactionId, const std::unordered_map& ranksToCheck) : TBase(NKikimrServices::KQP_GATEWAY, sessionId) - , Database(database) + , DatabaseId(databaseId) , RanksToCheck(ranksToCheck) { TxId = transactionId; - SetOperationInfo(__func__, Database); + SetOperationInfo(__func__, DatabaseId); } void OnRunQuery() override { @@ -62,13 +62,13 @@ class TRanksCheckerActor : public NKikimr::TQueryBase { TStringBuilder sql = TStringBuilder() << R"( -- TRanksCheckerActor::OnRunQuery - DECLARE $database AS Text; + DECLARE $database_id AS Text; )"; NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(CanonizePath(Database)) + .AddParam("$database_id") + .Utf8(CanonizePath(DatabaseId)) .Build(); if (!RanksToCheck.empty()) { @@ -79,7 +79,7 @@ class TRanksCheckerActor : public NKikimr::TQueryBase { SELECT rank, name FROM `)" << tablePath << R"(` - WHERE database = $database + WHERE database = $database_id AND rank IN $ranks; )"; @@ -97,7 +97,7 @@ class TRanksCheckerActor : public NKikimr::TQueryBase { MAX(rank) AS MaxRank, COUNT(*) AS NumberClassifiers FROM `)" << tablePath << R"(` - WHERE database = $database; + WHERE database = $database_id; )"; RunDataQuery(sql, ¶ms, TTxControl::ContinueTx()); @@ -149,7 +149,7 @@ class TRanksCheckerActor : public NKikimr::TQueryBase { } private: - const TString Database; + const TString DatabaseId; const std::unordered_map RanksToCheck; ui64 ExpectedResultSets = 1; @@ -235,8 +235,8 @@ class TResourcePoolClassifierPreparationActor : public TActorBootstrappedGetClassifierConfig(CanonizePath(object.GetDatabase()), object.GetName())) { - FailAndPassAway(TStringBuilder() << "Classifier with name " << object.GetName() << " not found in database " << object.GetDatabase()); + if (!snapshot->GetClassifierConfig(object.GetDatabase(), object.GetName())) { + FailAndPassAway(TStringBuilder() << "Classifier with name " << object.GetName() << " not found in database with id " << object.GetDatabase()); return; } } @@ -279,7 +279,7 @@ class TResourcePoolClassifierPreparationActor : public TActorBootstrapped>( - SelfId(), Context.GetExternalData().GetDatabase(), AlterContext.GetSessionId(), AlterContext.GetTransactionId(), ranksToNames + SelfId(), Context.GetExternalData().GetDatabaseId(), AlterContext.GetSessionId(), AlterContext.GetTransactionId(), ranksToNames )); } diff --git a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp index dc884cd029ba..5791ab5befa7 100644 --- a/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp +++ b/ydb/core/kqp/gateway/behaviour/resource_pool_classifier/manager.cpp @@ -13,7 +13,7 @@ using namespace NResourcePool; NMetadata::NInternal::TTableRecord GetResourcePoolClassifierRecord(const NYql::TObjectSettingsImpl& settings, const NMetadata::NModifications::IOperationsManager::TInternalModificationContext& context) { NMetadata::NInternal::TTableRecord result; - result.SetColumn(TResourcePoolClassifierConfig::TDecoder::Database, NMetadata::NInternal::TYDBValue::Utf8(CanonizePath(context.GetExternalData().GetDatabase()))); + result.SetColumn(TResourcePoolClassifierConfig::TDecoder::Database, NMetadata::NInternal::TYDBValue::Utf8(context.GetExternalData().GetDatabaseId())); result.SetColumn(TResourcePoolClassifierConfig::TDecoder::Name, NMetadata::NInternal::TYDBValue::Utf8(settings.GetObjectId())); return result; } diff --git a/ydb/core/kqp/gateway/kqp_gateway.h b/ydb/core/kqp/gateway/kqp_gateway.h index 9a46611da4fa..681f02a0c3e7 100644 --- a/ydb/core/kqp/gateway/kqp_gateway.h +++ b/ydb/core/kqp/gateway/kqp_gateway.h @@ -188,6 +188,7 @@ class IKqpGateway : public NYql::IKikimrGateway { public: virtual TString GetDatabase() = 0; + virtual TString GetDatabaseId() = 0; virtual bool GetDomainLoginOnly() = 0; virtual TMaybe GetDomainName() = 0; @@ -231,7 +232,7 @@ class IKqpGateway : public NYql::IKikimrGateway { const Ydb::Table::TransactionSettings& txSettings, const NActors::TActorId& target) = 0; }; -TIntrusivePtr CreateKikimrIcGateway(const TString& cluster, NKikimrKqp::EQueryType queryType, const TString& database, +TIntrusivePtr CreateKikimrIcGateway(const TString& cluster, NKikimrKqp::EQueryType queryType, const TString& database, const TString& databaseId, std::shared_ptr&& metadataLoader, NActors::TActorSystem* actorSystem, ui32 nodeId, TKqpRequestCounters::TPtr counters, const NKikimrConfig::TQueryServiceConfig& queryServiceConfig = NKikimrConfig::TQueryServiceConfig()); diff --git a/ydb/core/kqp/gateway/kqp_ic_gateway.cpp b/ydb/core/kqp/gateway/kqp_ic_gateway.cpp index fa565e66b210..1d219192ea7f 100644 --- a/ydb/core/kqp/gateway/kqp_ic_gateway.cpp +++ b/ydb/core/kqp/gateway/kqp_ic_gateway.cpp @@ -536,10 +536,11 @@ class TKqpSchemeExecuterRequestHandler: public TActorBootstrapped& requestType, const TString& database, - TIntrusiveConstPtr userToken, TPromise promise) + const TString& databaseId, TIntrusiveConstPtr userToken, TPromise promise) : PhyTx(std::move(phyTx)) , QueryType(queryType) , Database(database) + , DatabaseId(databaseId) , UserToken(std::move(userToken)) , Promise(promise) , RequestType(requestType) @@ -547,6 +548,7 @@ class TKqpSchemeExecuterRequestHandler: public TActorBootstrapped(); + ctx->DatabaseId = DatabaseId; IActor* actor = CreateKqpSchemeExecuter(PhyTx, QueryType, SelfId(), RequestType, Database, UserToken, false /* temporary */, TString() /* sessionId */, ctx); Register(actor); Become(&TThis::WaitState); @@ -578,6 +580,7 @@ class TKqpSchemeExecuterRequestHandler: public TActorBootstrapped UserToken; TPromise Promise; const TMaybe RequestType; @@ -700,11 +703,12 @@ class TKikimrIcGateway : public IKqpGateway { using TNavigate = NSchemeCache::TSchemeCacheNavigate; public: - TKikimrIcGateway(const TString& cluster, NKikimrKqp::EQueryType queryType, const TString& database, std::shared_ptr&& metadataLoader, + TKikimrIcGateway(const TString& cluster, NKikimrKqp::EQueryType queryType, const TString& database, const TString& databaseId, std::shared_ptr&& metadataLoader, TActorSystem* actorSystem, ui32 nodeId, TKqpRequestCounters::TPtr counters, const NKikimrConfig::TQueryServiceConfig& queryServiceConfig) : Cluster(cluster) , QueryType(queryType) , Database(database) + , DatabaseId(databaseId) , ActorSystem(actorSystem) , NodeId(nodeId) , Counters(counters) @@ -727,6 +731,10 @@ class TKikimrIcGateway : public IKqpGateway { return Database; } + TString GetDatabaseId() override { + return DatabaseId; + } + TMaybe GetSetting(const TString& cluster, const TString& name) override { Y_UNUSED(cluster); Y_UNUSED(name); @@ -1471,6 +1479,7 @@ class TKikimrIcGateway : public IKqpGateway { context.SetUserToken(*GetUserToken()); } context.SetDatabase(Owner.Database); + context.SetDatabaseId(Owner.DatabaseId); context.SetActorSystem(Owner.ActorSystem); return DoExecute(cBehaviour, settings, context).Apply([](const NThreading::TFuture& f) { if (f.HasValue() && !f.HasException() && f.GetValue().Ok()) { @@ -2222,7 +2231,7 @@ class TKikimrIcGateway : public IKqpGateway { TFuture SendSchemeExecuterRequest(const TString&, const TMaybe& requestType, const std::shared_ptr& phyTx) override { auto promise = NewPromise(); - IActor* requestHandler = new TKqpSchemeExecuterRequestHandler(phyTx, QueryType, requestType, Database, UserToken, promise); + IActor* requestHandler = new TKqpSchemeExecuterRequestHandler(phyTx, QueryType, requestType, Database, DatabaseId, UserToken, promise); RegisterActor(requestHandler); return promise.GetFuture(); } @@ -2302,6 +2311,7 @@ class TKikimrIcGateway : public IKqpGateway { TString Cluster; const NKikimrKqp::EQueryType QueryType; TString Database; + TString DatabaseId; TActorSystem* ActorSystem; ui32 NodeId; TKqpRequestCounters::TPtr Counters; @@ -2313,11 +2323,11 @@ class TKikimrIcGateway : public IKqpGateway { } // namespace -TIntrusivePtr CreateKikimrIcGateway(const TString& cluster, NKikimrKqp::EQueryType queryType, const TString& database, +TIntrusivePtr CreateKikimrIcGateway(const TString& cluster, NKikimrKqp::EQueryType queryType, const TString& database, const TString& databaseId, std::shared_ptr&& metadataLoader, TActorSystem* actorSystem, ui32 nodeId, TKqpRequestCounters::TPtr counters, const NKikimrConfig::TQueryServiceConfig& queryServiceConfig) { - return MakeIntrusive(cluster, queryType, database, std::move(metadataLoader), actorSystem, nodeId, + return MakeIntrusive(cluster, queryType, database, databaseId, std::move(metadataLoader), actorSystem, nodeId, counters, queryServiceConfig); } diff --git a/ydb/core/kqp/host/kqp_gateway_proxy.cpp b/ydb/core/kqp/host/kqp_gateway_proxy.cpp index a5c89dad6a46..29986334cd8e 100644 --- a/ydb/core/kqp/host/kqp_gateway_proxy.cpp +++ b/ydb/core/kqp/host/kqp_gateway_proxy.cpp @@ -1262,6 +1262,7 @@ class TKqpGatewayProxy : public IKikimrGateway { NMetadata::NModifications::IOperationsManager::TExternalModificationContext context; context.SetDatabase(SessionCtx->GetDatabase()); + context.SetDatabaseId(SessionCtx->GetDatabaseId()); context.SetActorSystem(ActorSystem); if (SessionCtx->GetUserToken()) { context.SetUserToken(*SessionCtx->GetUserToken()); diff --git a/ydb/core/kqp/host/kqp_host.cpp b/ydb/core/kqp/host/kqp_host.cpp index a3ebd5ae0cb2..f2074d895c1b 100644 --- a/ydb/core/kqp/host/kqp_host.cpp +++ b/ydb/core/kqp/host/kqp_host.cpp @@ -1062,6 +1062,7 @@ class TKqpHost : public IKqpHost { } SessionCtx->SetDatabase(database); + SessionCtx->SetDatabaseId(Gateway->GetDatabaseId()); SessionCtx->SetCluster(cluster); if (tempTablesState) { SessionCtx->SetSessionId(tempTablesState->SessionId); diff --git a/ydb/core/kqp/node_service/kqp_node_service.cpp b/ydb/core/kqp/node_service/kqp_node_service.cpp index ac482ab9a1d3..4bdb4d5e4def 100644 --- a/ydb/core/kqp/node_service/kqp_node_service.cpp +++ b/ydb/core/kqp/node_service/kqp_node_service.cpp @@ -204,7 +204,7 @@ class TKqpNodeService : public TActorBootstrapped { } if (share > 0) { Scheduler->UpdateGroupShare(schedulerGroup, share, schedulerNow); - Send(SchedulerActorId, new TEvSchedulerNewPool(msg.GetDatabase(), schedulerGroup)); + Send(SchedulerActorId, new TEvSchedulerNewPool(msg.GetDatabaseId(), schedulerGroup)); } else { schedulerGroup = ""; } diff --git a/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp b/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp index 79777b77307f..a5422321bcf8 100644 --- a/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp +++ b/ydb/core/kqp/provider/yql_kikimr_gateway_ut.cpp @@ -74,7 +74,7 @@ TIntrusivePtr GetIcGateway(Tests::TServer& server) { counters->TxProxyMon = new NTxProxy::TTxProxyMon(server.GetRuntime()->GetAppData(0).Counters); std::shared_ptr loader = std::make_shared(TestCluster, server.GetRuntime()->GetAnyNodeActorSystem(),TIntrusivePtr(nullptr), false); - return CreateKikimrIcGateway(TestCluster, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, "/Root", std::move(loader), server.GetRuntime()->GetAnyNodeActorSystem(), + return CreateKikimrIcGateway(TestCluster, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, "/Root", "/Root", std::move(loader), server.GetRuntime()->GetAnyNodeActorSystem(), server.GetRuntime()->GetNodeId(0), counters, server.GetSettings().AppConfig->GetQueryServiceConfig()); } diff --git a/ydb/core/kqp/provider/yql_kikimr_provider.h b/ydb/core/kqp/provider/yql_kikimr_provider.h index 003dd61a6295..246fd20b3e99 100644 --- a/ydb/core/kqp/provider/yql_kikimr_provider.h +++ b/ydb/core/kqp/provider/yql_kikimr_provider.h @@ -512,6 +512,10 @@ class TKikimrSessionContext : public TThrRefBase { return Database; } + TString GetDatabaseId() const { + return DatabaseId; + } + const TString& GetSessionId() const { return SessionId; } @@ -524,6 +528,10 @@ class TKikimrSessionContext : public TThrRefBase { Database = database; } + void SetDatabaseId(const TString& databaseId) { + DatabaseId = databaseId; + } + void SetSessionId(const TString& sessionId) { SessionId = sessionId; } @@ -562,6 +570,7 @@ class TKikimrSessionContext : public TThrRefBase { TString UserName; TString Cluster; TString Database; + TString DatabaseId; TString SessionId; TKikimrConfiguration::TPtr Configuration; TIntrusivePtr TablesData; diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_databases_cache.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_databases_cache.cpp index ebac128f8f9e..4b949031032c 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_databases_cache.cpp +++ b/ydb/core/kqp/proxy_service/kqp_proxy_databases_cache.cpp @@ -94,7 +94,9 @@ class TDatabaseSubscriberActor : public TActor { } databaseStateIt->FetchRequestIsRunning = false; - UpdateDatabaseState(*databaseStateIt, ev->Get()->PathId, ev->Get()->Serverless); + databaseStateIt->LastUpdateTime = TInstant::Now(); + databaseStateIt->DatabaseId = ev->Get()->DatabaseId; + databaseStateIt->Serverless = ev->Get()->Serverless; SendSubscriberInfo(*databaseStateIt, ev->Get()->Status, ev->Get()->Issues); if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) { @@ -151,12 +153,6 @@ class TDatabaseSubscriberActor : public TActor { ) private: - static void UpdateDatabaseState(TDatabaseState& databaseState, TPathId pathId, bool serverless) { - databaseState.LastUpdateTime = TInstant::Now(); - databaseState.DatabaseId = (serverless ? TStringBuilder() << pathId.OwnerId << ":" << pathId.LocalPathId << ":" : TStringBuilder()) << databaseState.Database; - databaseState.Serverless = serverless; - } - void UnsubscribeFromSchemeCache(TDatabaseState& databaseState) const { if (databaseState.WatchKey) { Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvWatchRemove(databaseState.WatchKey)); @@ -198,13 +194,6 @@ TDatabasesCache::TDatabasesCache(TDuration idleTimeout) : IdleTimeout(idleTimeout) {} -const TString& TDatabasesCache::GetTenantName() { - if (!TenantName) { - TenantName = CanonizePath(AppData()->TenantName); - } - return TenantName; -} - void TDatabasesCache::UpdateDatabaseInfo(TEvKqp::TEvUpdateDatabaseInfo::TPtr& event, TActorContext actorContext) { auto it = DatabasesCache.find(event->Get()->Database); if (it == DatabasesCache.end()) { diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp index 126165cf434a..b21490755b6e 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp +++ b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp @@ -1602,19 +1602,19 @@ class TKqpProxyService : public TActorBootstrapped { bool TryFillPoolInfoFromCache(TEvKqp::TEvQueryRequest::TPtr& ev, ui64 requestId) { ResourcePoolsCache.UpdateFeatureFlags(FeatureFlags, ActorContext()); - const auto& database = ev->Get()->GetDatabase(); - if (!ResourcePoolsCache.ResourcePoolsEnabled(database)) { + const auto& databaseId = ev->Get()->GetDatabaseId(); + if (!ResourcePoolsCache.ResourcePoolsEnabled(databaseId)) { ev->Get()->SetPoolId(""); return true; } const auto& userToken = ev->Get()->GetUserToken(); if (!ev->Get()->GetPoolId()) { - ev->Get()->SetPoolId(ResourcePoolsCache.GetPoolId(database, userToken, ActorContext())); + ev->Get()->SetPoolId(ResourcePoolsCache.GetPoolId(databaseId, userToken, ActorContext())); } const auto& poolId = ev->Get()->GetPoolId(); - const auto& poolInfo = ResourcePoolsCache.GetPoolInfo(database, poolId, ActorContext()); + const auto& poolInfo = ResourcePoolsCache.GetPoolInfo(databaseId, poolId, ActorContext()); if (!poolInfo) { return true; } @@ -1873,12 +1873,12 @@ class TKqpProxyService : public TActorBootstrapped { } void Handle(NWorkload::TEvUpdatePoolInfo::TPtr& ev) { - ResourcePoolsCache.UpdatePoolInfo(ev->Get()->Database, ev->Get()->PoolId, ev->Get()->Config, ev->Get()->SecurityObject, ActorContext()); + ResourcePoolsCache.UpdatePoolInfo(ev->Get()->DatabaseId, ev->Get()->PoolId, ev->Get()->Config, ev->Get()->SecurityObject, ActorContext()); } void Handle(TEvKqp::TEvUpdateDatabaseInfo::TPtr& ev) { if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) { - ResourcePoolsCache.UpdateDatabaseInfo(ev->Get()->Database, ev->Get()->Serverless); + ResourcePoolsCache.UpdateDatabaseInfo(ev->Get()->DatabaseId, ev->Get()->Serverless); } DatabasesCache.UpdateDatabaseInfo(ev, ActorContext()); } diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h b/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h index 323458630d15..94f8f0526dff 100644 --- a/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h +++ b/ydb/core/kqp/proxy_service/kqp_proxy_service_impl.h @@ -432,9 +432,9 @@ class TResourcePoolsCache { }; struct TDatabaseInfo { - std::unordered_map ResourcePoolsClassifiers = {}; - std::map RankToClassifierInfo = {}; - std::unordered_map> UserToResourcePool = {}; + std::unordered_map ResourcePoolsClassifiers = {}; // Classifier name to config + std::map RankToClassifierInfo = {}; // Classifier rank to config + std::unordered_map> UserToResourcePool = {}; // UserSID to (resource pool, classifier rank) bool Serverless = false; }; @@ -445,7 +445,7 @@ class TResourcePoolsCache { }; public: - bool ResourcePoolsEnabled(const TString& database) const { + bool ResourcePoolsEnabled(const TString& databaseId) const { if (!EnableResourcePools) { return false; } @@ -454,19 +454,19 @@ class TResourcePoolsCache { return true; } - const auto databaseInfo = GetDatabaseInfo(database); + const auto databaseInfo = GetDatabaseInfo(databaseId); return !databaseInfo || !databaseInfo->Serverless; } - TString GetPoolId(const TString& database, const TIntrusiveConstPtr& userToken, TActorContext actorContext) { + TString GetPoolId(const TString& databaseId, const TIntrusiveConstPtr& userToken, TActorContext actorContext) { if (!userToken || userToken->GetUserSID().empty()) { return NResourcePool::DEFAULT_POOL_ID; } - TDatabaseInfo& databaseInfo = *GetOrCreateDatabaseInfo(database); - auto [resultPoolId, resultRank] = GetPoolIdFromClassifiers(database, userToken->GetUserSID(), databaseInfo, userToken, actorContext); + TDatabaseInfo& databaseInfo = *GetOrCreateDatabaseInfo(databaseId); + auto [resultPoolId, resultRank] = GetPoolIdFromClassifiers(databaseId, userToken->GetUserSID(), databaseInfo, userToken, actorContext); for (const auto& userSID : userToken->GetGroupSIDs()) { - const auto& [poolId, rank] = GetPoolIdFromClassifiers(database, userSID, databaseInfo, userToken, actorContext); + const auto& [poolId, rank] = GetPoolIdFromClassifiers(databaseId, userSID, databaseInfo, userToken, actorContext); if (poolId && (!resultPoolId || resultRank > rank)) { resultPoolId = poolId; resultRank = rank; @@ -476,10 +476,10 @@ class TResourcePoolsCache { return resultPoolId ? resultPoolId : NResourcePool::DEFAULT_POOL_ID; } - std::optional GetPoolInfo(const TString& database, const TString& poolId, TActorContext actorContext) const { - auto it = PoolsCache.find(GetPoolKey(database, poolId)); + std::optional GetPoolInfo(const TString& databaseId, const TString& poolId, TActorContext actorContext) const { + auto it = PoolsCache.find(GetPoolKey(databaseId, poolId)); if (it == PoolsCache.end()) { - actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(database, poolId)); + actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(databaseId, poolId)); return std::nullopt; } return it->second; @@ -491,14 +491,14 @@ class TResourcePoolsCache { UpdateResourcePoolClassifiersSubscription(actorContext); } - void UpdateDatabaseInfo(const TString& database, bool serverless) { - GetOrCreateDatabaseInfo(database)->Serverless = serverless; + void UpdateDatabaseInfo(const TString& databaseId, bool serverless) { + GetOrCreateDatabaseInfo(databaseId)->Serverless = serverless; } - void UpdatePoolInfo(const TString& database, const TString& poolId, const std::optional& config, const std::optional& securityObject, TActorContext actorContext) { + void UpdatePoolInfo(const TString& databaseId, const TString& poolId, const std::optional& config, const std::optional& securityObject, TActorContext actorContext) { bool clearClassifierCache = false; - const TString& poolKey = GetPoolKey(database, poolId); + const TString& poolKey = GetPoolKey(databaseId, poolId); if (!config) { auto it = PoolsCache.find(poolKey); if (it == PoolsCache.end()) { @@ -511,7 +511,7 @@ class TResourcePoolsCache { } else { // Refresh pool subscription it->second.Expired = true; - actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(database, poolId)); + actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(databaseId, poolId)); } } else { auto& poolInfo = PoolsCache[poolKey]; @@ -522,16 +522,16 @@ class TResourcePoolsCache { } if (clearClassifierCache) { - GetOrCreateDatabaseInfo(database)->UserToResourcePool.clear(); + GetOrCreateDatabaseInfo(databaseId)->UserToResourcePool.clear(); } } void UpdateResourcePoolClassifiersInfo(const TResourcePoolClassifierSnapshot* snapsot, TActorContext actorContext) { auto resourcePoolClassifierConfigs = snapsot->GetResourcePoolClassifierConfigs(); - for (auto& [database, databaseInfo] : DatabasesCache) { - auto it = resourcePoolClassifierConfigs.find(database); + for (auto& [databaseId, databaseInfo] : DatabasesCache) { + auto it = resourcePoolClassifierConfigs.find(databaseId); if (it != resourcePoolClassifierConfigs.end()) { - UpdateDatabaseResourcePoolClassifiers(database, databaseInfo, std::move(it->second), actorContext); + UpdateDatabaseResourcePoolClassifiers(databaseId, databaseInfo, std::move(it->second), actorContext); resourcePoolClassifierConfigs.erase(it); } else if (!databaseInfo.ResourcePoolsClassifiers.empty()) { databaseInfo.ResourcePoolsClassifiers.clear(); @@ -539,8 +539,8 @@ class TResourcePoolsCache { databaseInfo.UserToResourcePool.clear(); } } - for (auto& [database, configsMap] : resourcePoolClassifierConfigs) { - UpdateDatabaseResourcePoolClassifiers(database, *GetOrCreateDatabaseInfo(database), std::move(configsMap), actorContext); + for (auto& [databaseId, configsMap] : resourcePoolClassifierConfigs) { + UpdateDatabaseResourcePoolClassifiers(databaseId, *GetOrCreateDatabaseInfo(databaseId), std::move(configsMap), actorContext); } } @@ -567,7 +567,7 @@ class TResourcePoolsCache { } } - void UpdateDatabaseResourcePoolClassifiers(const TString& database, TDatabaseInfo& databaseInfo, std::unordered_map&& configsMap, TActorContext actorContext) { + void UpdateDatabaseResourcePoolClassifiers(const TString& databaseId, TDatabaseInfo& databaseInfo, std::unordered_map&& configsMap, TActorContext actorContext) { if (databaseInfo.ResourcePoolsClassifiers == configsMap) { return; } @@ -579,12 +579,12 @@ class TResourcePoolsCache { const auto& classifierSettings = classifier.GetClassifierSettings(); databaseInfo.RankToClassifierInfo.insert({classifier.GetRank(), TClassifierInfo(classifierSettings)}); if (!PoolsCache.contains(classifierSettings.ResourcePool)) { - actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(database, classifierSettings.ResourcePool)); + actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(databaseId, classifierSettings.ResourcePool)); } } } - std::pair GetPoolIdFromClassifiers(const TString& database, const TString& userSID, TDatabaseInfo& databaseInfo, const TIntrusiveConstPtr& userToken, TActorContext actorContext) const { + std::pair GetPoolIdFromClassifiers(const TString& databaseId, const TString& userSID, TDatabaseInfo& databaseInfo, const TIntrusiveConstPtr& userToken, TActorContext actorContext) const { auto& usersMap = databaseInfo.UserToResourcePool; if (const auto it = usersMap.find(userSID); it != usersMap.end()) { return it->second; @@ -597,9 +597,9 @@ class TResourcePoolsCache { continue; } - auto it = PoolsCache.find(GetPoolKey(database, classifier.PoolId)); + auto it = PoolsCache.find(GetPoolKey(databaseId, classifier.PoolId)); if (it == PoolsCache.end()) { - actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(database, classifier.PoolId)); + actorContext.Send(MakeKqpWorkloadServiceId(actorContext.SelfID.NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(databaseId, classifier.PoolId)); continue; } @@ -616,21 +616,20 @@ class TResourcePoolsCache { return {poolId, rank}; } - TDatabaseInfo* GetOrCreateDatabaseInfo(const TString& database) { - const TString& path = CanonizePath(database); - if (const auto it = DatabasesCache.find(path); it != DatabasesCache.end()) { + TDatabaseInfo* GetOrCreateDatabaseInfo(const TString& databaseId) { + if (const auto it = DatabasesCache.find(databaseId); it != DatabasesCache.end()) { return &it->second; } - return &DatabasesCache.insert({path, TDatabaseInfo{}}).first->second; + return &DatabasesCache.insert({databaseId, TDatabaseInfo{}}).first->second; } - const TDatabaseInfo* GetDatabaseInfo(const TString& database) const { - const auto it = DatabasesCache.find(CanonizePath(database)); + const TDatabaseInfo* GetDatabaseInfo(const TString& databaseId) const { + const auto it = DatabasesCache.find(databaseId); return it != DatabasesCache.end() ? &it->second : nullptr; } - static TString GetPoolKey(const TString& database, const TString& poolId) { - return CanonizePath(TStringBuilder() << database << "/" << poolId); + static TString GetPoolKey(const TString& databaseId, const TString& poolId) { + return TStringBuilder() << databaseId << "/" << poolId; } private: @@ -660,13 +659,10 @@ class TDatabasesCache { template bool SetDatabaseIdOrDefer(TEvent& event, i32 requestType, TActorContext actorContext) { - if (!event->Get()->GetDatabaseId().empty()) { - return true; - } - const auto& database = CanonizePath(event->Get()->GetDatabase()); - if (database.empty() || database == GetTenantName()) { - event->Get()->SetDatabaseId(GetTenantName()); + const auto& tenantName = CanonizePath(AppData()->TenantName); + if (database.empty() || database == tenantName) { + event->Get()->SetDatabaseId(tenantName); return true; } @@ -690,7 +686,6 @@ class TDatabasesCache { void StopSubscriberActor(TActorContext actorContext) const; private: - const TString& GetTenantName(); void SubscribeOnDatabase(const TString& database, TActorContext actorContext); void PingDatabaseSubscription(const TString& database, TActorContext actorContext) const; diff --git a/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp b/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp index db3a9b844400..a29f837a7564 100644 --- a/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp +++ b/ydb/core/kqp/runtime/kqp_compute_scheduler.cpp @@ -734,11 +734,11 @@ ::NMonitoring::TDynamicCounters::TCounterPtr TComputeScheduler::GetGroupUsageCou struct TEvPingPool : public TEventLocal { - TString Database; + TString DatabaseId; TString Pool; - TEvPingPool(TString database, TString pool) - : Database(database) + TEvPingPool(TString databaseId, TString pool) + : DatabaseId(databaseId) , Pool(pool) { } @@ -805,11 +805,11 @@ class TSchedulerActor : public TActorBootstrapped { } void Handle(TEvSchedulerNewPool::TPtr& ev) { - Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(ev->Get()->Database, ev->Get()->Pool)); + Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(ev->Get()->DatabaseId, ev->Get()->Pool)); } void Handle(TEvPingPool::TPtr& ev) { - Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(ev->Get()->Database, ev->Get()->Pool)); + Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvSubscribeOnPoolChanges(ev->Get()->DatabaseId, ev->Get()->Pool)); } void Handle(NWorkload::TEvUpdatePoolInfo::TPtr& ev) { diff --git a/ydb/core/kqp/runtime/kqp_compute_scheduler.h b/ydb/core/kqp/runtime/kqp_compute_scheduler.h index 37bb1b2008a0..4c1f5f8c972f 100644 --- a/ydb/core/kqp/runtime/kqp_compute_scheduler.h +++ b/ydb/core/kqp/runtime/kqp_compute_scheduler.h @@ -117,11 +117,11 @@ struct TEvSchedulerDeregister : public TEventLocal { - TString Database; + TString DatabaseId; TString Pool; - TEvSchedulerNewPool(TString database, TString pool) - : Database(database) + TEvSchedulerNewPool(TString databaseId, TString pool) + : DatabaseId(databaseId) , Pool(pool) { } diff --git a/ydb/core/kqp/session_actor/kqp_query_state.cpp b/ydb/core/kqp/session_actor/kqp_query_state.cpp index f45f214b044f..984957857bd3 100644 --- a/ydb/core/kqp/session_actor/kqp_query_state.cpp +++ b/ydb/core/kqp/session_actor/kqp_query_state.cpp @@ -199,12 +199,12 @@ std::unique_ptr TKqpQueryState::BuildCompileRequest(s TGUCSettings gUCSettings = gUCSettingsPtr ? *gUCSettingsPtr : TGUCSettings(); switch (GetAction()) { case NKikimrKqp::QUERY_ACTION_EXECUTE: - query = TKqpQueryId(Cluster, Database, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); + query = TKqpQueryId(Cluster, Database, UserRequestContext->DatabaseId, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); keepInCache = GetQueryKeepInCache() && query->IsSql(); break; case NKikimrKqp::QUERY_ACTION_PREPARE: - query = TKqpQueryId(Cluster, Database, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); + query = TKqpQueryId(Cluster, Database, UserRequestContext->DatabaseId, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); keepInCache = query->IsSql(); break; @@ -214,7 +214,7 @@ std::unique_ptr TKqpQueryState::BuildCompileRequest(s break; case NKikimrKqp::QUERY_ACTION_EXPLAIN: - query = TKqpQueryId(Cluster, Database, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); + query = TKqpQueryId(Cluster, Database, UserRequestContext->DatabaseId, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); keepInCache = false; break; @@ -255,11 +255,11 @@ std::unique_ptr TKqpQueryState::BuildReCompileReque switch (GetAction()) { case NKikimrKqp::QUERY_ACTION_EXPLAIN: case NKikimrKqp::QUERY_ACTION_EXECUTE: - query = TKqpQueryId(Cluster, Database, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); + query = TKqpQueryId(Cluster, Database, UserRequestContext->DatabaseId, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); break; case NKikimrKqp::QUERY_ACTION_PREPARE: - query = TKqpQueryId(Cluster, Database, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); + query = TKqpQueryId(Cluster, Database, UserRequestContext->DatabaseId, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); break; case NKikimrKqp::QUERY_ACTION_EXECUTE_PREPARED: @@ -299,7 +299,7 @@ std::unique_ptr TKqpQueryState::BuildCompileSplittedR switch (GetAction()) { case NKikimrKqp::QUERY_ACTION_EXECUTE: - query = TKqpQueryId(Cluster, Database, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); + query = TKqpQueryId(Cluster, Database, UserRequestContext->DatabaseId, GetQuery(), settings, GetQueryParameterTypes(), gUCSettings); break; default: YQL_ENSURE(false); diff --git a/ydb/core/kqp/session_actor/kqp_query_state.h b/ydb/core/kqp/session_actor/kqp_query_state.h index 8d7ec00ec70c..94ec2e2b8aa0 100644 --- a/ydb/core/kqp/session_actor/kqp_query_state.h +++ b/ydb/core/kqp/session_actor/kqp_query_state.h @@ -88,6 +88,7 @@ class TKqpQueryState : public TNonCopyable { } UserRequestContext->PoolId = RequestEv->GetPoolId(); UserRequestContext->PoolConfig = RequestEv->GetPoolConfig(); + UserRequestContext->DatabaseId = RequestEv->GetDatabaseId(); } // the monotonously growing counter, the ordinal number of the query, diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp index fe5f31be2fc1..9ea406597e75 100644 --- a/ydb/core/kqp/session_actor/kqp_session_actor.cpp +++ b/ydb/core/kqp/session_actor/kqp_session_actor.cpp @@ -249,7 +249,7 @@ class TKqpSessionActor : public TActorBootstrapped { } Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvPlaceRequestIntoPool( - QueryState->Database, + QueryState->UserRequestContext->DatabaseId, SessionId, QueryState->UserRequestContext->PoolId, QueryState->UserToken @@ -415,6 +415,7 @@ class TKqpSessionActor : public TActorBootstrapped { << " text: " << QueryState->GetQuery() << " rpcActor: " << QueryState->RequestActorId << " database: " << QueryState->GetDatabase() + << " databaseId: " << QueryState->UserRequestContext->DatabaseId << " pool id: " << QueryState->UserRequestContext->PoolId ); @@ -2100,7 +2101,7 @@ class TKqpSessionActor : public TActorBootstrapped { const auto& stats = QueryState->QueryStats; auto event = std::make_unique( - QueryState->Database, SessionId, QueryState->UserRequestContext->PoolId, + QueryState->UserRequestContext->DatabaseId, SessionId, QueryState->UserRequestContext->PoolId, TDuration::MicroSeconds(stats.DurationUs), TDuration::MicroSeconds(stats.WorkerCpuTimeUs) ); diff --git a/ydb/core/kqp/session_actor/kqp_worker_actor.cpp b/ydb/core/kqp/session_actor/kqp_worker_actor.cpp index 176179a04b2c..a18f69b7f347 100644 --- a/ydb/core/kqp/session_actor/kqp_worker_actor.cpp +++ b/ydb/core/kqp/session_actor/kqp_worker_actor.cpp @@ -182,7 +182,7 @@ class TKqpWorkerActor : public TActorBootstrapped { std::shared_ptr loader = std::make_shared( Settings.Cluster, TlsActivationContext->ActorSystem(), Config, false, nullptr); - Gateway = CreateKikimrIcGateway(Settings.Cluster, QueryState->RequestEv->GetType(), Settings.Database, std::move(loader), + Gateway = CreateKikimrIcGateway(Settings.Cluster, QueryState->RequestEv->GetType(), Settings.Database, QueryState->RequestEv->GetDatabaseId(), std::move(loader), ctx.ExecutorThread.ActorSystem, ctx.SelfID.NodeId(), RequestCounters, QueryServiceConfig); Config->FeatureFlags = AppData(ctx)->FeatureFlags; diff --git a/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp b/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp index c36fda70cb2b..38de83491cbb 100644 --- a/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp +++ b/ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp @@ -35,7 +35,7 @@ TIntrusivePtr GetIcGateway(Tests::TServer& server) { counters->Counters = new TKqpCounters(server.GetRuntime()->GetAppData(0).Counters); counters->TxProxyMon = new NTxProxy::TTxProxyMon(server.GetRuntime()->GetAppData(0).Counters); std::shared_ptr loader = std::make_shared(TestCluster, server.GetRuntime()->GetAnyNodeActorSystem(),TIntrusivePtr(nullptr),false); - return NKqp::CreateKikimrIcGateway(TestCluster, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, "/Root", std::move(loader), server.GetRuntime()->GetAnyNodeActorSystem(), + return NKqp::CreateKikimrIcGateway(TestCluster, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, "/Root", "/Root", std::move(loader), server.GetRuntime()->GetAnyNodeActorSystem(), server.GetRuntime()->GetNodeId(0), counters, server.GetSettings().AppConfig->GetQueryServiceConfig()); } diff --git a/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp b/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp index f169f6f961b1..55d1c6b16a05 100644 --- a/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp +++ b/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -6618,7 +6619,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { // DROP RESOURCE POOL CLASSIFIER checkQuery("DROP RESOURCE POOL CLASSIFIER MyResourcePoolClassifier;", EStatus::GENERIC_ERROR, - "Classifier with name MyResourcePoolClassifier not found in database /Root"); + "Classifier with name MyResourcePoolClassifier not found in database with id /Root"); } Y_UNIT_TEST(DisableResourcePoolClassifiersOnServerless) { @@ -6799,16 +6800,19 @@ Y_UNIT_TEST_SUITE(KqpScheme) { UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "The rank could not be set automatically, the maximum rank of the resource pool classifier is too high: 9223372036854775807"); } - TString FetchResourcePoolClassifiers(TKikimrRunner& kikimr) { - auto& runtime = *kikimr.GetTestServer().GetRuntime(); - const TActorId edgeActor = runtime.AllocateEdgeActor(); - runtime.Send(NMetadata::NProvider::MakeServiceId(runtime.GetNodeId()), edgeActor, new NMetadata::NProvider::TEvAskSnapshot(std::make_shared())); + TString FetchResourcePoolClassifiers(TTestActorRuntime& runtime, ui32 nodeIndex) { + const TActorId edgeActor = runtime.AllocateEdgeActor(nodeIndex); + runtime.Send(NMetadata::NProvider::MakeServiceId(runtime.GetNodeId(nodeIndex)), edgeActor, new NMetadata::NProvider::TEvAskSnapshot(std::make_shared()), nodeIndex); const auto response = runtime.GrabEdgeEvent(edgeActor); UNIT_ASSERT(response); return response->Get()->GetSnapshotAs()->SerializeToString(); } + TString FetchResourcePoolClassifiers(TKikimrRunner& kikimr) { + return FetchResourcePoolClassifiers(*kikimr.GetTestServer().GetRuntime(), 0); + } + Y_UNIT_TEST(CreateResourcePoolClassifier) { NKikimrConfig::TAppConfig config; config.MutableFeatureFlags()->SetEnableResourcePools(true); @@ -6842,6 +6846,31 @@ Y_UNIT_TEST_SUITE(KqpScheme) { UNIT_ASSERT_VALUES_EQUAL(FetchResourcePoolClassifiers(kikimr), "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"member_name\":\"test@user\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"},{\"rank\":1020,\"name\":\"AnotherResourcePoolClassifier\",\"config\":{\"member_name\":\"another@user\",\"resource_pool\":\"test_pool\"},\"database\":\"\\/Root\"}]}"); } + Y_UNIT_TEST(CreateResourcePoolClassifierOnServerless) { + auto ydb = NWorkload::TYdbSetupSettings() + .CreateSampleTenants(true) + .EnableResourcePoolsOnServerless(true) + .Create(); + + const auto& serverlessTenant = ydb->GetSettings().GetServerlessTenantName(); + NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(R"( + CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH ( + RANK=20, + RESOURCE_POOL="test_pool" + );)", + NWorkload::TQueryRunnerSettings() + .PoolId("") + .Database(serverlessTenant) + .NodeIndex(1) + )); + + const auto pathId = ydb->FetchDatabase(serverlessTenant)->Get()->PathId; + UNIT_ASSERT_VALUES_EQUAL( + FetchResourcePoolClassifiers(*ydb->GetRuntime(), 1), + TStringBuilder() << "{\"resource_pool_classifiers\":[{\"rank\":20,\"name\":\"MyResourcePoolClassifier\",\"config\":{\"resource_pool\":\"test_pool\"},\"database\":\"" << pathId.OwnerId << ":" << pathId.LocalPathId << ":\\/Root\\/test-serverless\"}]}" + ); + } + Y_UNIT_TEST(DoubleCreateResourcePoolClassifier) { NKikimrConfig::TAppConfig config; config.MutableFeatureFlags()->SetEnableResourcePools(true); @@ -6951,7 +6980,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { )"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString()); - UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Classifier with name MyResourcePoolClassifier not found in database /Root"); + UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Classifier with name MyResourcePoolClassifier not found in database with id /Root"); } Y_UNIT_TEST(DropResourcePoolClassifier) { @@ -6998,7 +7027,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) { auto query = "DROP RESOURCE POOL CLASSIFIER MyResourcePoolClassifier;"; auto result = session.ExecuteSchemeQuery(query).GetValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString()); - UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Classifier with name MyResourcePoolClassifier not found in database /Root"); + UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "Classifier with name MyResourcePoolClassifier not found in database with id /Root"); } Y_UNIT_TEST(DisableMetadataObjectsOnServerless) { diff --git a/ydb/core/kqp/workload_service/actors/actors.h b/ydb/core/kqp/workload_service/actors/actors.h index 6e55f7553f55..d21706f3df7b 100644 --- a/ydb/core/kqp/workload_service/actors/actors.h +++ b/ydb/core/kqp/workload_service/actors/actors.h @@ -6,16 +6,16 @@ namespace NKikimr::NKqp::NWorkload { // Pool state holder -NActors::IActor* CreatePoolHandlerActor(const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters); +NActors::IActor* CreatePoolHandlerActor(const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters); // Fetch pool and create default pool if needed NActors::IActor* CreatePoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaultPoolExists); // Fetch and create pool in scheme shard -NActors::IActor* CreatePoolFetcherActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr userToken); -NActors::IActor* CreatePoolCreatorActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr userToken, NACLibProto::TDiffACL diffAcl); +NActors::IActor* CreatePoolFetcherActor(const NActors::TActorId& replyActorId, const TString& databaseId, const TString& poolId, TIntrusiveConstPtr userToken); +NActors::IActor* CreatePoolCreatorActor(const NActors::TActorId& replyActorId, const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr userToken, NACLibProto::TDiffACL diffAcl); -// Checks that database is serverless +// Checks that database is serverless and return database id NActors::IActor* CreateDatabaseFetcherActor(const NActors::TActorId& replyActorId, const TString& database, TIntrusiveConstPtr userToken = nullptr, NACLib::EAccessRights checkAccess = NACLib::EAccessRights::NoAccess); // Cpu load fetcher actor diff --git a/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp b/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp index 77512a4d529f..63fb1adbe068 100644 --- a/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp +++ b/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp @@ -46,9 +46,9 @@ class TPoolHandlerActorBase : public TActor { NMonitoring::TDynamicCounters::TCounterPtr QueueSizeLimit; NMonitoring::TDynamicCounters::TCounterPtr LoadCpuThreshold; - TCommonCounters(NMonitoring::TDynamicCounterPtr counters, const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig) + TCommonCounters(NMonitoring::TDynamicCounterPtr counters, const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig) : CountersRoot(counters) - , CountersSubgroup(counters->GetSubgroup("pool", CanonizePath(TStringBuilder() << database << "/" << poolId))) + , CountersSubgroup(counters->GetSubgroup("pool", TStringBuilder() << databaseId << "/" << poolId)) { Register(); UpdateConfigCounters(poolConfig); @@ -125,10 +125,10 @@ class TPoolHandlerActorBase : public TActor { }; public: - TPoolHandlerActorBase(void (TDerived::* requestFunc)(TAutoPtr& ev), const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters) + TPoolHandlerActorBase(void (TDerived::* requestFunc)(TAutoPtr& ev), const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters) : TBase(requestFunc) - , Counters(counters, database, poolId, poolConfig) - , Database(database) + , Counters(counters, databaseId, poolId, poolConfig) + , DatabaseId(databaseId) , PoolId(poolId) , QueueSizeLimit(GetMaxQueueSize(poolConfig)) , InFlightLimit(GetMaxInFlight(poolConfig)) @@ -161,7 +161,7 @@ class TPoolHandlerActorBase : public TActor { } SendPoolInfoUpdate(std::nullopt, std::nullopt, Subscribers); - this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvStopPoolHandlerResponse(Database, PoolId)); + this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvStopPoolHandlerResponse(DatabaseId, PoolId)); Counters.OnCleanup(ResetCountersOnStrop); @@ -187,7 +187,7 @@ class TPoolHandlerActorBase : public TActor { void Handle(TEvPrivate::TEvResolvePoolResponse::TPtr& ev) { auto event = std::move(ev->Get()->Event); const TString& sessionId = event->Get()->SessionId; - this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPlaceRequestIntoPoolResponse(Database, PoolId, sessionId)); + this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPlaceRequestIntoPoolResponse(DatabaseId, PoolId, sessionId)); const TActorId& workerActorId = event->Sender; if (!InFlightLimit) { @@ -360,7 +360,7 @@ class TPoolHandlerActorBase : public TActor { void SendPoolInfoUpdate(const std::optional& config, const std::optional& securityObject, const std::unordered_set& subscribers) const { for (const auto& subscriber : subscribers) { - this->Send(subscriber, new TEvUpdatePoolInfo(Database, PoolId, config, securityObject)); + this->Send(subscriber, new TEvUpdatePoolInfo(DatabaseId, PoolId, config, securityObject)); } } @@ -390,7 +390,7 @@ class TPoolHandlerActorBase : public TActor { void RemoveRequest(TRequest* request) { auto event = std::make_unique( - Database, PoolId, request->Duration, request->CpuConsumed, request->UsedCpuQuota + DatabaseId, PoolId, request->Duration, request->CpuConsumed, request->UsedCpuQuota ); this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), event.release()); @@ -424,7 +424,7 @@ class TPoolHandlerActorBase : public TActor { } TString LogPrefix() const { - return TStringBuilder() << "[TPoolHandlerActorBase] ActorId: " << this->SelfId() << ", Database: " << Database << ", PoolId: " << PoolId << ", "; + return TStringBuilder() << "[TPoolHandlerActorBase] ActorId: " << this->SelfId() << ", DatabaseId: " << DatabaseId << ", PoolId: " << PoolId << ", "; } private: @@ -482,8 +482,8 @@ class TPoolHandlerActorBase : public TActor { RefreshState(true); if (ShouldResign()) { - const TActorId& newHandler = this->RegisterWithSameMailbox(CreatePoolHandlerActor(Database, PoolId, poolConfig, Counters.CountersRoot)); - this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvResignPoolHandler(Database, PoolId, newHandler)); + const TActorId& newHandler = this->RegisterWithSameMailbox(CreatePoolHandlerActor(DatabaseId, PoolId, poolConfig, Counters.CountersRoot)); + this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvResignPoolHandler(DatabaseId, PoolId, newHandler)); } } @@ -501,7 +501,7 @@ class TPoolHandlerActorBase : public TActor { TCommonCounters Counters; // Configuration - const TString Database; + const TString DatabaseId; const TString PoolId; ui64 QueueSizeLimit = std::numeric_limits::max(); ui64 InFlightLimit = std::numeric_limits::max(); @@ -527,8 +527,8 @@ class TUnlimitedPoolHandlerActor : public TPoolHandlerActorBase; public: - TUnlimitedPoolHandlerActor(const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters) - : TBase(&TBase::StateFuncBase, database, poolId, poolConfig, counters) + TUnlimitedPoolHandlerActor(const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters) + : TBase(&TBase::StateFuncBase, databaseId, poolId, poolConfig, counters) { Y_ENSURE(!ShouldResign()); } @@ -591,8 +591,8 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBaseInc(); if (!PreparingFinished) { - this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPrepareTablesRequest(Database, PoolId)); + this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPrepareTablesRequest(DatabaseId, PoolId)); } RefreshState(); @@ -679,7 +679,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBaseRegister(CreateRefreshPoolStateActor(this->SelfId(), Database, PoolId, LEASE_DURATION, Counters.CountersSubgroup)); + this->Register(CreateRefreshPoolStateActor(this->SelfId(), DatabaseId, PoolId, LEASE_DURATION, Counters.CountersSubgroup)); } } @@ -813,7 +813,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBaseRegister(CreateStartRequestActor(this->SelfId(), Database, PoolId, sessionId, LEASE_DURATION, Counters.CountersSubgroup)); + this->Register(CreateStartRequestActor(this->SelfId(), DatabaseId, PoolId, sessionId, LEASE_DURATION, Counters.CountersSubgroup)); GetRequest(sessionId)->CleanupRequired = true; } break; @@ -846,7 +846,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase(); event->Record.SetPoolId(PoolId); - event->Record.SetDatabase(Database); + event->Record.SetDatabase(DatabaseId); this->Send(MakeKqpWorkloadServiceId(nodeId), std::move(event)); RefreshState(); return; @@ -911,7 +911,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBaseRegister(CreateStartRequestActor(this->SelfId(), Database, PoolId, sessionId, LEASE_DURATION, Counters.CountersSubgroup)); + this->Register(CreateStartRequestActor(this->SelfId(), DatabaseId, PoolId, sessionId, LEASE_DURATION, Counters.CountersSubgroup)); GetRequest(sessionId)->CleanupRequired = true; } } @@ -928,7 +928,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBaseRegister(CreateStartRequestActor(this->SelfId(), Database, PoolId, std::nullopt, LEASE_DURATION, Counters.CountersSubgroup)); + this->Register(CreateStartRequestActor(this->SelfId(), DatabaseId, PoolId, std::nullopt, LEASE_DURATION, Counters.CountersSubgroup)); } } } @@ -943,7 +943,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBaseRegister(CreateDelayRequestActor(this->SelfId(), Database, PoolId, sessionId, request->StartTime, GetWaitDeadline(request->StartTime), LEASE_DURATION, Counters.CountersSubgroup)); + this->Register(CreateDelayRequestActor(this->SelfId(), DatabaseId, PoolId, sessionId, request->StartTime, GetWaitDeadline(request->StartTime), LEASE_DURATION, Counters.CountersSubgroup)); DelayedRequests.emplace_back(sessionId); request->CleanupRequired = true; } @@ -956,7 +956,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBaseRegister(CreateCleanupRequestsActor(this->SelfId(), Database, PoolId, FinishedRequests, Counters.CountersSubgroup)); + this->Register(CreateCleanupRequestsActor(this->SelfId(), DatabaseId, PoolId, FinishedRequests, Counters.CountersSubgroup)); FinishedRequests.clear(); FifoCounters.FinishingRequestsCount->Set(0); } @@ -1054,11 +1054,11 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase { void StartPoolFetchRequest() const { LOG_D("Start pool fetching"); - Register(CreatePoolFetcherActor(SelfId(), Event->Get()->Database, Event->Get()->PoolId, Event->Get()->UserToken)); + Register(CreatePoolFetcherActor(SelfId(), Event->Get()->DatabaseId, Event->Get()->PoolId, Event->Get()->UserToken)); } void Handle(TEvPrivate::TEvFetchPoolResponse::TPtr& ev) { @@ -74,13 +74,13 @@ class TPoolResolverActor : public TActorBootstrapped { diffAcl.AddAccess(NACLib::EAccessType::Allow, useAccess, BUILTIN_ACL_ROOT); auto token = MakeIntrusive(BUILTIN_ACL_METADATA, TVector{}); - Register(CreatePoolCreatorActor(SelfId(), Event->Get()->Database, Event->Get()->PoolId, NResourcePool::TPoolSettings(), token, diffAcl)); + Register(CreatePoolCreatorActor(SelfId(), Event->Get()->DatabaseId, Event->Get()->PoolId, NResourcePool::TPoolSettings(), token, diffAcl)); } void Handle(TEvPrivate::TEvCreatePoolResponse::TPtr& ev) { if (ev->Get()->Status != Ydb::StatusIds::SUCCESS) { LOG_E("Failed to create default pool " << ev->Get()->Status << ", issues: " << ev->Get()->Issues.ToOneLineString()); - Reply(ev->Get()->Status, GroupIssues(ev->Get()->Issues, "Failed to create default pool")); + Reply(ev->Get()->Status, GroupIssues(ev->Get()->Issues, TStringBuilder() << "Failed to create default pool in database " << Event->Get()->DatabaseId)); return; } @@ -96,7 +96,7 @@ class TPoolResolverActor : public TActorBootstrapped { private: TString LogPrefix() const { - return TStringBuilder() << "[TPoolResolverActor] ActorId: " << SelfId() << ", Database: " << Event->Get()->Database << ", PoolId: " << Event->Get()->PoolId << ", SessionId: " << Event->Get()->SessionId << ", "; + return TStringBuilder() << "[TPoolResolverActor] ActorId: " << SelfId() << ", DatabaseId: " << Event->Get()->DatabaseId << ", PoolId: " << Event->Get()->PoolId << ", SessionId: " << Event->Get()->SessionId << ", "; } void Reply(NResourcePool::TPoolSettings poolConfig, TPathId pathId) { @@ -122,9 +122,9 @@ class TPoolResolverActor : public TActorBootstrapped { class TPoolFetcherActor : public TSchemeActorBase { public: - TPoolFetcherActor(const TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr userToken) + TPoolFetcherActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, TIntrusiveConstPtr userToken) : ReplyActorId(replyActorId) - , Database(database) + , DatabaseId(databaseId) , PoolId(poolId) , UserToken(userToken) {} @@ -180,7 +180,7 @@ class TPoolFetcherActor : public TSchemeActorBase { LOG_D("Start pool fetching"); auto event = NTableCreator::BuildSchemeCacheNavigateRequest( {{".metadata/workload_manager/pools", PoolId}}, - Database ? Database : AppData()->TenantName, + DatabaseIdToDatabase(DatabaseId), UserToken ); event->ResultSet[0].Access |= NACLib::SelectRow; @@ -193,7 +193,7 @@ class TPoolFetcherActor : public TSchemeActorBase { } TString LogPrefix() const override { - return TStringBuilder() << "[TPoolFetcherActor] ActorId: " << SelfId() << ", Database: " << Database << ", PoolId: " << PoolId << ", "; + return TStringBuilder() << "[TPoolFetcherActor] ActorId: " << SelfId() << ", DatabaseId: " << DatabaseId << ", PoolId: " << PoolId << ", "; } private: @@ -221,13 +221,13 @@ class TPoolFetcherActor : public TSchemeActorBase { } Issues.AddIssues(std::move(issues)); - Send(ReplyActorId, new TEvPrivate::TEvFetchPoolResponse(status, Database, PoolId, PoolConfig, PathIdFromPathId(PathId), std::move(Issues))); + Send(ReplyActorId, new TEvPrivate::TEvFetchPoolResponse(status, DatabaseId, PoolId, PoolConfig, PathIdFromPathId(PathId), std::move(Issues))); PassAway(); } private: const TActorId ReplyActorId; - const TString Database; + const TString DatabaseId; const TString PoolId; const TIntrusiveConstPtr UserToken; @@ -240,9 +240,9 @@ class TPoolCreatorActor : public TSchemeActorBase { using TBase = TSchemeActorBase; public: - TPoolCreatorActor(const TActorId& replyActorId, const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr userToken, NACLibProto::TDiffACL diffAcl) + TPoolCreatorActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr userToken, NACLibProto::TDiffACL diffAcl) : ReplyActorId(replyActorId) - , Database(database) + , DatabaseId(databaseId) , PoolId(poolId) , UserToken(userToken) , DiffAcl(diffAcl) @@ -326,7 +326,7 @@ class TPoolCreatorActor : public TSchemeActorBase { auto event = std::make_unique(); auto& schemeTx = *event->Record.MutableTransaction()->MutableModifyScheme(); - schemeTx.SetWorkingDir(JoinPath({Database ? Database : AppData()->TenantName, ".metadata/workload_manager/pools"})); + schemeTx.SetWorkingDir(JoinPath({DatabaseIdToDatabase(DatabaseId), ".metadata/workload_manager/pools"})); schemeTx.SetOperationType(NKikimrSchemeOp::ESchemeOpCreateResourcePool); schemeTx.SetInternal(true); @@ -345,7 +345,7 @@ class TPoolCreatorActor : public TSchemeActorBase { } TString LogPrefix() const override { - return TStringBuilder() << "[TPoolCreatorActor] ActorId: " << SelfId() << ", Database: " << Database << ", PoolId: " << PoolId << ", "; + return TStringBuilder() << "[TPoolCreatorActor] ActorId: " << SelfId() << ", DatabaseId: " << DatabaseId << ", PoolId: " << PoolId << ", "; } private: @@ -432,7 +432,7 @@ class TPoolCreatorActor : public TSchemeActorBase { private: const TActorId ReplyActorId; - const TString Database; + const TString DatabaseId; const TString PoolId; const TIntrusiveConstPtr UserToken; const NACLibProto::TDiffACL DiffAcl; @@ -532,13 +532,13 @@ class TDatabaseFetcherActor : public TSchemeActorBase { void Reply(Ydb::StatusIds::StatusCode status, NYql::TIssues issues = {}) { if (status == Ydb::StatusIds::SUCCESS) { - LOG_D("Database info successfully fetched"); + LOG_D("Database info successfully fetched, serverless: " << Serverless); } else { LOG_W("Failed to fetch database info, " << status << ", issues: " << issues.ToOneLineString()); } Issues.AddIssues(std::move(issues)); - Send(ReplyActorId, new TEvFetchDatabaseResponse(status, Database, Serverless, PathId, std::move(Issues))); + Send(ReplyActorId, new TEvFetchDatabaseResponse(status, Database, CreateDatabaseId(Database, Serverless, PathId), Serverless, PathId, std::move(Issues))); PassAway(); } @@ -570,12 +570,12 @@ IActor* CreatePoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaul return new TPoolResolverActor(std::move(event), defaultPoolExists); } -IActor* CreatePoolFetcherActor(const TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr userToken) { - return new TPoolFetcherActor(replyActorId, database, poolId, userToken); +IActor* CreatePoolFetcherActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, TIntrusiveConstPtr userToken) { + return new TPoolFetcherActor(replyActorId, databaseId, poolId, userToken); } -IActor* CreatePoolCreatorActor(const TActorId& replyActorId, const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr userToken, NACLibProto::TDiffACL diffAcl) { - return new TPoolCreatorActor(replyActorId, database, poolId, poolConfig, userToken, diffAcl); +IActor* CreatePoolCreatorActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr userToken, NACLibProto::TDiffACL diffAcl) { + return new TPoolCreatorActor(replyActorId, databaseId, poolId, poolConfig, userToken, diffAcl); } IActor* CreateDatabaseFetcherActor(const TActorId& replyActorId, const TString& database, TIntrusiveConstPtr userToken, NACLib::EAccessRights checkAccess) { diff --git a/ydb/core/kqp/workload_service/common/events.h b/ydb/core/kqp/workload_service/common/events.h index 57e332ac66b3..911dccd36f0e 100644 --- a/ydb/core/kqp/workload_service/common/events.h +++ b/ydb/core/kqp/workload_service/common/events.h @@ -74,9 +74,9 @@ struct TEvPrivate { }; struct TEvFetchPoolResponse : public NActors::TEventLocal { - TEvFetchPoolResponse(Ydb::StatusIds::StatusCode status, const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TPathId pathId, NYql::TIssues issues) + TEvFetchPoolResponse(Ydb::StatusIds::StatusCode status, const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TPathId pathId, NYql::TIssues issues) : Status(status) - , Database(database) + , DatabaseId(databaseId) , PoolId(poolId) , PoolConfig(poolConfig) , PathId(pathId) @@ -84,7 +84,7 @@ struct TEvPrivate { {} const Ydb::StatusIds::StatusCode Status; - const TString Database; + const TString DatabaseId; const TString PoolId; const NResourcePool::TPoolSettings PoolConfig; const TPathId PathId; @@ -102,37 +102,37 @@ struct TEvPrivate { }; struct TEvPrepareTablesRequest : public NActors::TEventLocal { - TEvPrepareTablesRequest(const TString& database, const TString& poolId) - : Database(database) + TEvPrepareTablesRequest(const TString& databaseId, const TString& poolId) + : DatabaseId(databaseId) , PoolId(poolId) {} - const TString Database; + const TString DatabaseId; const TString PoolId; }; struct TEvPlaceRequestIntoPoolResponse : public NActors::TEventLocal { - TEvPlaceRequestIntoPoolResponse(const TString& database, const TString& poolId, const TString& sessionId) - : Database(database) + TEvPlaceRequestIntoPoolResponse(const TString& databaseId, const TString& poolId, const TString& sessionId) + : DatabaseId(databaseId) , PoolId(poolId) , SessionId(sessionId) {} - const TString Database; + const TString DatabaseId; const TString PoolId; const TString SessionId; }; struct TEvFinishRequestInPool : public NActors::TEventLocal { - TEvFinishRequestInPool(const TString& database, const TString& poolId, TDuration duration, TDuration cpuConsumed, bool adjustCpuQuota) - : Database(database) + TEvFinishRequestInPool(const TString& databaseId, const TString& poolId, TDuration duration, TDuration cpuConsumed, bool adjustCpuQuota) + : DatabaseId(databaseId) , PoolId(poolId) , Duration(duration) , CpuConsumed(cpuConsumed) , AdjustCpuQuota(adjustCpuQuota) {} - const TString Database; + const TString DatabaseId; const TString PoolId; const TDuration Duration; const TDuration CpuConsumed; @@ -140,13 +140,13 @@ struct TEvPrivate { }; struct TEvResignPoolHandler : public NActors::TEventLocal { - TEvResignPoolHandler(const TString& database, const TString& poolId, const TActorId& newHandler) - : Database(database) + TEvResignPoolHandler(const TString& databaseId, const TString& poolId, const TActorId& newHandler) + : DatabaseId(databaseId) , PoolId(poolId) , NewHandler(newHandler) {} - const TString Database; + const TString DatabaseId; const TString PoolId; const TActorId NewHandler; }; @@ -160,12 +160,12 @@ struct TEvPrivate { }; struct TEvStopPoolHandlerResponse : public NActors::TEventLocal { - TEvStopPoolHandlerResponse(const TString& database, const TString& poolId) - : Database(database) + TEvStopPoolHandlerResponse(const TString& databaseId, const TString& poolId) + : DatabaseId(databaseId) , PoolId(poolId) {} - const TString Database; + const TString DatabaseId; const TString PoolId; }; diff --git a/ydb/core/kqp/workload_service/common/helpers.cpp b/ydb/core/kqp/workload_service/common/helpers.cpp index 79fd3fff5c69..b9ad774c6f84 100644 --- a/ydb/core/kqp/workload_service/common/helpers.cpp +++ b/ydb/core/kqp/workload_service/common/helpers.cpp @@ -1,8 +1,32 @@ #include "helpers.h" +#include +#include + namespace NKikimr::NKqp::NWorkload { +TString CreateDatabaseId(const TString& database, bool serverless, TPathId pathId) { + TString databasePath = CanonizePath(database); + TString tennantPath = CanonizePath(AppData()->TenantName); + if (databasePath.empty() || databasePath == tennantPath) { + return tennantPath; + } + + if (serverless) { + databasePath = TStringBuilder() << pathId.OwnerId << ":" << pathId.LocalPathId << ":" << databasePath; + } + return databasePath; +} + +TString DatabaseIdToDatabase(TStringBuf databaseId) { + TStringBuf id; + TStringBuf database; + return databaseId.TrySplit("/", id, database) + ? CanonizePath(TString(database)) // Serverless + : CanonizePath(TString(databaseId)); // Dedicated +} + NYql::TIssues GroupIssues(const NYql::TIssues& issues, const TString& message) { NYql::TIssue rootIssue(message); for (const NYql::TIssue& issue : issues) { diff --git a/ydb/core/kqp/workload_service/common/helpers.h b/ydb/core/kqp/workload_service/common/helpers.h index 163b2d765ed1..a4c933b92356 100644 --- a/ydb/core/kqp/workload_service/common/helpers.h +++ b/ydb/core/kqp/workload_service/common/helpers.h @@ -99,6 +99,9 @@ class TSchemeActorBase : public NActors::TActorBootstrapped { }; +TString CreateDatabaseId(const TString& database, bool serverless, TPathId pathId); +TString DatabaseIdToDatabase(TStringBuf databaseId); + NYql::TIssues GroupIssues(const NYql::TIssues& issues, const TString& message); void ParsePoolSettings(const NKikimrSchemeOp::TResourcePoolDescription& description, NResourcePool::TPoolSettings& poolConfig); diff --git a/ydb/core/kqp/workload_service/kqp_workload_service.cpp b/ydb/core/kqp/workload_service/kqp_workload_service.cpp index 9bcfe49f4a4e..af4c0536b45f 100644 --- a/ydb/core/kqp/workload_service/kqp_workload_service.cpp +++ b/ydb/core/kqp/workload_service/kqp_workload_service.cpp @@ -143,15 +143,15 @@ class TKqpWorkloadService : public TActorBootstrapped { } void Handle(TEvSubscribeOnPoolChanges::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; if (!EnabledResourcePools) { - Send(ev->Sender, new TEvUpdatePoolInfo(database, poolId, std::nullopt, std::nullopt)); + Send(ev->Sender, new TEvUpdatePoolInfo(databaseId, poolId, std::nullopt, std::nullopt)); return; } - LOG_D("Recieved subscription request, Database: " << database << ", PoolId: " << poolId); - GetOrCreateDatabaseState(database)->DoSubscribeRequest(std::move(ev)); + LOG_D("Recieved subscription request, DatabaseId: " << databaseId << ", PoolId: " << poolId); + GetOrCreateDatabaseState(databaseId)->DoSubscribeRequest(std::move(ev)); } void Handle(TEvPlaceRequestIntoPool::TPtr& ev) { @@ -161,28 +161,28 @@ class TKqpWorkloadService : public TActorBootstrapped { return; } - const TString& database = ev->Get()->Database; - LOG_D("Recieved new request from " << workerActorId << ", Database: " << database << ", PoolId: " << ev->Get()->PoolId << ", SessionId: " << ev->Get()->SessionId); - GetOrCreateDatabaseState(database)->DoPlaceRequest(std::move(ev)); + const TString& databaseId = ev->Get()->DatabaseId; + LOG_D("Recieved new request from " << workerActorId << ", DatabaseId: " << databaseId << ", PoolId: " << ev->Get()->PoolId << ", SessionId: " << ev->Get()->SessionId); + GetOrCreateDatabaseState(databaseId)->DoPlaceRequest(std::move(ev)); } void Handle(TEvCleanupRequest::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; const TString& sessionId = ev->Get()->SessionId; - if (GetOrCreateDatabaseState(database)->PendingSessionIds.contains(sessionId)) { - LOG_D("Finished request with worker actor " << ev->Sender << ", wait for place request, Database: " << database << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId); - GetOrCreateDatabaseState(database)->PendingCancelRequests[sessionId].emplace_back(std::move(ev)); + if (GetOrCreateDatabaseState(databaseId)->PendingSessionIds.contains(sessionId)) { + LOG_D("Finished request with worker actor " << ev->Sender << ", wait for place request, DatabaseId: " << databaseId << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId); + GetOrCreateDatabaseState(databaseId)->PendingCancelRequests[sessionId].emplace_back(std::move(ev)); return; } - auto poolState = GetPoolState(database, poolId); + auto poolState = GetPoolState(databaseId, poolId); if (!poolState) { ReplyCleanupError(ev->Sender, Ydb::StatusIds::NOT_FOUND, TStringBuilder() << "Pool " << poolId << " not found"); return; } - LOG_D("Finished request with worker actor " << ev->Sender << ", Database: " << database << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId); + LOG_D("Finished request with worker actor " << ev->Sender << ", DatabaseId: " << databaseId << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId); poolState->DoCleanupRequest(std::move(ev)); } @@ -232,28 +232,33 @@ class TKqpWorkloadService : public TActorBootstrapped { private: void Handle(TEvFetchDatabaseResponse::TPtr& ev) { - GetOrCreateDatabaseState(ev->Get()->Database)->UpdateDatabaseInfo(ev); + if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) { + LOG_D("Successfully fetched database info, DatabaseId: " << ev->Get()->DatabaseId << ", Serverless: " << ev->Get()->Serverless); + } else { + LOG_D("Failed to fetch database info, DatabaseId: " << ev->Get()->DatabaseId << ", Status: " << ev->Get()->Status << ", Issues: " << ev->Get()->Issues.ToOneLineString()); + } + GetOrCreateDatabaseState(ev->Get()->DatabaseId)->UpdateDatabaseInfo(ev); } void Handle(TEvPrivate::TEvFetchPoolResponse::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; TActorId poolHandler; if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) { - LOG_D("Successfully fetched pool " << poolId << ", Database: " << database); - poolHandler = GetOrCreatePoolState(database, poolId, ev->Get()->PoolConfig)->PoolHandler; + LOG_D("Successfully fetched pool " << poolId << ", DatabaseId: " << databaseId); + poolHandler = GetOrCreatePoolState(databaseId, poolId, ev->Get()->PoolConfig)->PoolHandler; } else { - LOG_W("Failed to fetch pool " << poolId << ", Database: " << database << ", status: " << ev->Get()->Status << ", issues: " << ev->Get()->Issues.ToOneLineString()); + LOG_W("Failed to fetch pool " << poolId << ", DatabaseId: " << databaseId << ", status: " << ev->Get()->Status << ", issues: " << ev->Get()->Issues.ToOneLineString()); } - GetOrCreateDatabaseState(database)->UpdatePoolInfo(ev, poolHandler); + GetOrCreateDatabaseState(databaseId)->UpdatePoolInfo(ev, poolHandler); } void Handle(TEvPrivate::TEvResolvePoolResponse::TPtr& ev) { const auto& event = ev->Get()->Event; - const TString& database = event->Get()->Database; - auto databaseState = GetOrCreateDatabaseState(database); + const TString& databaseId = event->Get()->DatabaseId; + auto databaseState = GetOrCreateDatabaseState(databaseId); if (ev->Get()->DefaultPoolCreated) { databaseState->HasDefaultPool = true; } @@ -267,21 +272,21 @@ class TKqpWorkloadService : public TActorBootstrapped { return; } - LOG_D("Successfully fetched pool " << poolId << ", Database: " << database << ", SessionId: " << event->Get()->SessionId); + LOG_D("Successfully fetched pool " << poolId << ", DatabaseId: " << databaseId << ", SessionId: " << event->Get()->SessionId); - auto poolState = GetOrCreatePoolState(database, poolId, ev->Get()->PoolConfig); + auto poolState = GetOrCreatePoolState(databaseId, poolId, ev->Get()->PoolConfig); poolState->PendingRequests.emplace(std::move(ev)); poolState->StartPlaceRequest(); } void Handle(TEvPrivate::TEvPlaceRequestIntoPoolResponse::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; const TString& sessionId = ev->Get()->SessionId; - LOG_T("Request placed into pool, Database: " << database << ", PoolId: " << poolId << ", SessionId: " << sessionId); + LOG_T("Request placed into pool, DatabaseId: " << databaseId << ", PoolId: " << poolId << ", SessionId: " << sessionId); - auto poolState = GetPoolState(database, poolId); - GetOrCreateDatabaseState(database)->RemovePendingSession(sessionId, [this, poolState](TEvCleanupRequest::TPtr event) { + auto poolState = GetPoolState(databaseId, poolId); + GetOrCreateDatabaseState(databaseId)->RemovePendingSession(sessionId, [this, poolState](TEvCleanupRequest::TPtr event) { if (poolState) { poolState->DoCleanupRequest(std::move(event)); } else { @@ -302,11 +307,11 @@ class TKqpWorkloadService : public TActorBootstrapped { void Handle(TEvPrivate::TEvRefreshPoolState::TPtr& ev) { const auto& event = ev->Get()->Record; - const TString& database = event.GetDatabase(); + const TString& databaseId = event.GetDatabase(); const TString& poolId = event.GetPoolId(); - LOG_T("Got remote refresh request, Database: " << database << ", PoolId: " << poolId << ", NodeId: " << ev->Sender.NodeId()); + LOG_T("Got remote refresh request, DatabaseId: " << databaseId << ", PoolId: " << poolId << ", NodeId: " << ev->Sender.NodeId()); - if (auto poolState = GetPoolState(database, poolId)) { + if (auto poolState = GetPoolState(databaseId, poolId)) { Send(ev->Forward(poolState->PoolHandler)); } } @@ -321,11 +326,11 @@ class TKqpWorkloadService : public TActorBootstrapped { } void Handle(TEvPrivate::TEvFinishRequestInPool::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; - LOG_T("Request finished in pool, Database: " << database << ", PoolId: " << poolId << ", Duration: " << ev->Get()->Duration << ", CpuConsumed: " << ev->Get()->CpuConsumed << ", AdjustCpuQuota: " << ev->Get()->AdjustCpuQuota); + LOG_T("Request finished in pool, DatabaseId: " << databaseId << ", PoolId: " << poolId << ", Duration: " << ev->Get()->Duration << ", CpuConsumed: " << ev->Get()->CpuConsumed << ", AdjustCpuQuota: " << ev->Get()->AdjustCpuQuota); - if (auto poolState = GetPoolState(database, poolId)) { + if (auto poolState = GetPoolState(databaseId, poolId)) { poolState->OnRequestFinished(); } if (ev->Get()->AdjustCpuQuota) { @@ -335,11 +340,11 @@ class TKqpWorkloadService : public TActorBootstrapped { } void Handle(TEvPrivate::TEvPrepareTablesRequest::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; - LOG_T("Got create teables request, Database: " << database << ", PoolId: " << poolId); + LOG_T("Got create teables request, DatabaseId: " << databaseId << ", PoolId: " << poolId); - auto poolState = GetPoolState(database, poolId); + auto poolState = GetPoolState(databaseId, poolId); if (!poolState) { return; } @@ -348,7 +353,7 @@ class TKqpWorkloadService : public TActorBootstrapped { Send(poolState->PoolHandler, new TEvPrivate::TEvTablesCreationFinished(true, {})); } else { poolState->WaitingInitialization = true; - PendingHandlers.emplace(GetPoolKey(database, poolId)); + PendingHandlers.emplace(GetPoolKey(databaseId, poolId)); PrepareWorkloadServiceTables(); } } @@ -397,11 +402,11 @@ class TKqpWorkloadService : public TActorBootstrapped { } void Handle(TEvPrivate::TEvResignPoolHandler::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; - LOG_T("Got resign request, Database: " << database << ", PoolId: " << poolId); + LOG_T("Got resign request, DatabaseId: " << databaseId << ", PoolId: " << poolId); - if (auto poolState = GetPoolState(database, poolId)) { + if (auto poolState = GetPoolState(databaseId, poolId)) { if (poolState->NewPoolHandler) { Send(*poolState->NewPoolHandler, new TEvPrivate::TEvStopPoolHandler(false)); } @@ -411,12 +416,12 @@ class TKqpWorkloadService : public TActorBootstrapped { } void Handle(TEvPrivate::TEvStopPoolHandlerResponse::TPtr& ev) { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; const TString& poolId = ev->Get()->PoolId; - LOG_T("Got stop pool handler response, Database: " << database << ", PoolId: " << poolId); + LOG_T("Got stop pool handler response, DatabaseId: " << databaseId << ", PoolId: " << poolId); Counters.ActivePools->Dec(); - if (auto poolState = GetPoolState(database, poolId)) { + if (auto poolState = GetPoolState(databaseId, poolId)) { poolState->PreviousPoolHandlers.erase(ev->Sender); } } @@ -544,24 +549,24 @@ class TKqpWorkloadService : public TActorBootstrapped { Send(replyActorId, new TEvCleanupResponse(status, {NYql::TIssue(message)})); } - TDatabaseState* GetOrCreateDatabaseState(TString database) { - database = CanonizePath(database); - auto databaseIt = DatabaseToState.find(database); + TDatabaseState* GetOrCreateDatabaseState(TString databaseId) { + auto databaseIt = DatabaseToState.find(databaseId); if (databaseIt != DatabaseToState.end()) { return &databaseIt->second; } - return &DatabaseToState.insert({database, TDatabaseState{.ActorContext = ActorContext(), .EnabledResourcePoolsOnServerless = EnabledResourcePoolsOnServerless}}).first->second; + LOG_I("Creating new database state for id " << databaseId); + return &DatabaseToState.insert({databaseId, TDatabaseState{.ActorContext = ActorContext(), .EnabledResourcePoolsOnServerless = EnabledResourcePoolsOnServerless}}).first->second; } - TPoolState* GetOrCreatePoolState(const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig) { - const auto& poolKey = GetPoolKey(database, poolId); + TPoolState* GetOrCreatePoolState(const TString& databaseId, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig) { + const auto& poolKey = GetPoolKey(databaseId, poolId); if (auto poolState = GetPoolState(poolKey)) { return poolState; } LOG_I("Creating new handler for pool " << poolKey); - const auto poolHandler = Register(CreatePoolHandlerActor(database, poolId, poolConfig, EnableResourcePoolsCounters ? Counters.Counters : MakeIntrusive())); + const auto poolHandler = Register(CreatePoolHandlerActor(databaseId, poolId, poolConfig, EnableResourcePoolsCounters ? Counters.Counters : MakeIntrusive())); const auto poolState = &PoolIdToState.insert({poolKey, TPoolState{.PoolHandler = poolHandler, .ActorContext = ActorContext()}}).first->second; Counters.ActivePools->Inc(); @@ -570,8 +575,8 @@ class TKqpWorkloadService : public TActorBootstrapped { return poolState; } - TPoolState* GetPoolState(const TString& database, const TString& poolId) { - return GetPoolState(GetPoolKey(database, poolId)); + TPoolState* GetPoolState(const TString& databaseId, const TString& poolId) { + return GetPoolState(GetPoolKey(databaseId, poolId)); } TPoolState* GetPoolState(const TString& key) { @@ -582,8 +587,8 @@ class TKqpWorkloadService : public TActorBootstrapped { return nullptr; } - static TString GetPoolKey(const TString& database, const TString& poolId) { - return CanonizePath(TStringBuilder() << database << "/" << poolId); + static TString GetPoolKey(const TString& databaseId, const TString& poolId) { + return CanonizePath(TStringBuilder() << databaseId << "/" << poolId); } TString LogPrefix() const { @@ -599,10 +604,10 @@ class TKqpWorkloadService : public TActorBootstrapped { bool ServiceInitialized = false; bool IdleChecksStarted = false; ETablesCreationStatus TablesCreationStatus = ETablesCreationStatus::Cleanup; - std::unordered_set PendingHandlers; + std::unordered_set PendingHandlers; // DatabaseID/PoolID - std::unordered_map DatabaseToState; - std::unordered_map PoolIdToState; + std::unordered_map DatabaseToState; // DatabaseID to state + std::unordered_map PoolIdToState; // DatabaseID/PoolID to state std::unique_ptr CpuQuotaManager; ui32 NodeCount = 0; }; diff --git a/ydb/core/kqp/workload_service/kqp_workload_service_impl.h b/ydb/core/kqp/workload_service/kqp_workload_service_impl.h index e9e292d81dfc..5fd22dbb032b 100644 --- a/ydb/core/kqp/workload_service/kqp_workload_service_impl.h +++ b/ydb/core/kqp/workload_service/kqp_workload_service_impl.h @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -20,8 +21,8 @@ struct TDatabaseState { std::vector PendingRequersts = {}; std::unordered_set PendingSessionIds = {}; - std::unordered_map> PendingCancelRequests = {}; - std::unordered_map> PendingSubscriptions = {}; + std::unordered_map> PendingCancelRequests = {}; // Session ID to requests + std::unordered_map> PendingSubscriptions = {}; // Pool ID to subscribers bool HasDefaultPool = false; bool Serverless = false; bool DatabaseUnsupported = false; @@ -32,23 +33,23 @@ struct TDatabaseState { const TString& poolId = ev->Get()->PoolId; auto& subscribers = PendingSubscriptions[poolId]; if (subscribers.empty()) { - ActorContext.Register(CreatePoolFetcherActor(ActorContext.SelfID, ev->Get()->Database, poolId, nullptr)); + ActorContext.Register(CreatePoolFetcherActor(ActorContext.SelfID, ev->Get()->DatabaseId, poolId, nullptr)); } subscribers.emplace(ev->Sender); } void DoPlaceRequest(TEvPlaceRequestIntoPool::TPtr ev) { - TString database = ev->Get()->Database; + TString databaseId = ev->Get()->DatabaseId; PendingSessionIds.emplace(ev->Get()->SessionId); PendingRequersts.emplace_back(std::move(ev)); if (!EnabledResourcePoolsOnServerless && (TInstant::Now() - LastUpdateTime) > IDLE_DURATION) { - ActorContext.Register(CreateDatabaseFetcherActor(ActorContext.SelfID, database)); + ActorContext.Register(CreateDatabaseFetcherActor(ActorContext.SelfID, DatabaseIdToDatabase(databaseId))); } else if (!DatabaseUnsupported) { StartPendingRequests(); } else { - ReplyContinueError(Ydb::StatusIds::UNSUPPORTED, {NYql::TIssue(TStringBuilder() << "Unsupported database: " << database)}); + ReplyContinueError(Ydb::StatusIds::UNSUPPORTED, {NYql::TIssue(TStringBuilder() << "Unsupported database: " << databaseId)}); } } @@ -62,9 +63,9 @@ struct TDatabaseState { if (ev->Get()->Status == Ydb::StatusIds::SUCCESS && poolHandler) { ActorContext.Send(poolHandler, new TEvPrivate::TEvUpdatePoolSubscription(ev->Get()->PathId, subscribers)); } else { - const TString& database = ev->Get()->Database; + const TString& databaseId = ev->Get()->DatabaseId; for (const auto& subscriber : subscribers) { - ActorContext.Send(subscriber, new TEvUpdatePoolInfo(database, poolId, std::nullopt, std::nullopt)); + ActorContext.Send(subscriber, new TEvUpdatePoolInfo(databaseId, poolId, std::nullopt, std::nullopt)); } } subscribers.clear(); @@ -77,6 +78,10 @@ struct TDatabaseState { return; } + if (Serverless != ev->Get()->Serverless) { + ActorContext.Send(MakeKqpProxyID(ActorContext.SelfID.NodeId()), new TEvKqp::TEvUpdateDatabaseInfo(ev->Get()->Database, ev->Get()->DatabaseId, ev->Get()->Serverless)); + } + LastUpdateTime = TInstant::Now(); Serverless = ev->Get()->Serverless; StartPendingRequests(); @@ -159,8 +164,8 @@ struct TPoolState { void DoCleanupRequest(TEvCleanupRequest::TPtr event) { for (const auto& poolHandler : PreviousPoolHandlers) { ActorContext.Send(poolHandler, new TEvCleanupRequest( - event->Get()->Database, event->Get()->SessionId, event->Get()->PoolId, - event->Get()->Duration, event->Get()->CpuConsumed + event->Get()->DatabaseId, event->Get()->SessionId, + event->Get()->PoolId, event->Get()->Duration, event->Get()->CpuConsumed )); } ActorContext.Send(event->Forward(PoolHandler)); diff --git a/ydb/core/kqp/workload_service/tables/table_queries.cpp b/ydb/core/kqp/workload_service/tables/table_queries.cpp index ba6457288c94..2a18cadaaa6e 100644 --- a/ydb/core/kqp/workload_service/tables/table_queries.cpp +++ b/ydb/core/kqp/workload_service/tables/table_queries.cpp @@ -28,17 +28,17 @@ class TQueryBase : public NKikimr::TQueryBase { SetOperationInfo(operationName, traceId, counters); } - TQueryBase(const TString& operationName, const TString& traceId, const TString& database, const TString& sessionId, NMonitoring::TDynamicCounterPtr counters) - : TQueryBase(operationName, ComposeTraceId(traceId, database, sessionId), counters) + TQueryBase(const TString& operationName, const TString& traceId, const TString& databaseId, const TString& sessionId, NMonitoring::TDynamicCounterPtr counters) + : TQueryBase(operationName, ComposeTraceId(traceId, databaseId, sessionId), counters) {} - void UpdateLogInfo(const TString& traceId, const TString& database, const TString& sessionId) { - SetOperationInfo(OperationName, ComposeTraceId(traceId, database, sessionId), nullptr); + void UpdateLogInfo(const TString& traceId, const TString& databaseId, const TString& sessionId) { + SetOperationInfo(OperationName, ComposeTraceId(traceId, databaseId, sessionId), nullptr); } private: - static TString ComposeTraceId(const TString& traceId, const TString& database, const TString& sessionId) { - return TStringBuilder() << traceId << ", RequestDatabase: " << database << ", RequestSessionId: " << sessionId; + static TString ComposeTraceId(const TString& traceId, const TString& databaseId, const TString& sessionId) { + return TStringBuilder() << traceId << ", RequestDatabase: " << databaseId << ", RequestSessionId: " << sessionId; } }; @@ -306,9 +306,9 @@ class TCleanupTablesActor : public TSchemeActorBase { class TRefreshPoolStateQuery : public TQueryBase { public: - TRefreshPoolStateQuery(const TString& database, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) - : TQueryBase(__func__, poolId, database, "", counters) - , Database(database) + TRefreshPoolStateQuery(const TString& databaseId, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) + : TQueryBase(__func__, poolId, databaseId, "", counters) + , DatabaseId(databaseId) , PoolId(poolId) , LeaseDuration(leaseDuration) {} @@ -321,14 +321,14 @@ class TRefreshPoolStateQuery : public TQueryBase { TString sql = TStringBuilder() << R"( -- TRefreshPoolStateQuery::OnRunQuery - DECLARE $database AS Text; + DECLARE $database_id AS Text; DECLARE $pool_id AS Text; DECLARE $node_id AS Uint32; DECLARE $lease_duration AS Interval; UPDATE `)" << TTablesCreator::GetDelayedRequestsPath() << R"(` SET lease_deadline = CurrentUtcTimestamp() + $lease_duration - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND node_id = $node_id AND (wait_deadline IS NULL OR wait_deadline >= CurrentUtcTimestamp()) @@ -336,7 +336,7 @@ class TRefreshPoolStateQuery : public TQueryBase { UPDATE `)" << TTablesCreator::GetRunningRequestsPath() << R"(` SET lease_deadline = CurrentUtcTimestamp() + $lease_duration - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND node_id = $node_id AND lease_deadline >= CurrentUtcTimestamp(); @@ -344,8 +344,8 @@ class TRefreshPoolStateQuery : public TQueryBase { NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(Database) + .AddParam("$database_id") + .Utf8(DatabaseId) .Build() .AddParam("$pool_id") .Utf8(PoolId) @@ -364,27 +364,27 @@ class TRefreshPoolStateQuery : public TQueryBase { void OnLeaseUpdated() { TString sql = TStringBuilder() << R"( -- TRefreshPoolStateQuery::OnLeaseUpdated - DECLARE $database AS Text; + DECLARE $database_id AS Text; DECLARE $pool_id AS Text; SELECT COUNT(*) AS delayed_requests FROM `)" << TTablesCreator::GetDelayedRequestsPath() << R"(` - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND (wait_deadline IS NULL OR wait_deadline >= CurrentUtcTimestamp()) AND lease_deadline >= CurrentUtcTimestamp(); SELECT COUNT(*) AS running_requests FROM `)" << TTablesCreator::GetRunningRequestsPath() << R"(` - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND lease_deadline >= CurrentUtcTimestamp(); )"; NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(Database) + .AddParam("$database_id") + .Utf8(DatabaseId) .Build() .AddParam("$pool_id") .Utf8(PoolId) @@ -428,7 +428,7 @@ class TRefreshPoolStateQuery : public TQueryBase { } private: - const TString Database; + const TString DatabaseId; const TString PoolId; const TDuration LeaseDuration; @@ -438,9 +438,9 @@ class TRefreshPoolStateQuery : public TQueryBase { class TDelayRequestQuery : public TQueryBase { public: - TDelayRequestQuery(const TString& database, const TString& poolId, const TString& sessionId, TInstant startTime, TMaybe waitDeadline, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) - : TQueryBase(__func__, poolId, database, sessionId, counters) - , Database(database) + TDelayRequestQuery(const TString& databaseId, const TString& poolId, const TString& sessionId, TInstant startTime, TMaybe waitDeadline, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) + : TQueryBase(__func__, poolId, databaseId, sessionId, counters) + , DatabaseId(databaseId) , PoolId(poolId) , SessionId(sessionId) , StartTime(startTime) @@ -451,7 +451,7 @@ class TDelayRequestQuery : public TQueryBase { void OnRunQuery() override { TString sql = TStringBuilder() << R"( -- TDelayRequestQuery::OnRunQuery - DECLARE $database AS Text; + DECLARE $database_id AS Text; DECLARE $pool_id AS Text; DECLARE $start_time AS Timestamp; DECLARE $session_id AS Text; @@ -462,15 +462,15 @@ class TDelayRequestQuery : public TQueryBase { UPSERT INTO `)" << TTablesCreator::GetDelayedRequestsPath() << R"(` (database, pool_id, start_time, session_id, node_id, wait_deadline, lease_deadline) VALUES ( - $database, $pool_id, $start_time, $session_id, $node_id, $wait_deadline, + $database_id, $pool_id, $start_time, $session_id, $node_id, $wait_deadline, CurrentUtcTimestamp() + $lease_duration ); )"; NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(Database) + .AddParam("$database_id") + .Utf8(DatabaseId) .Build() .AddParam("$pool_id") .Utf8(PoolId) @@ -503,7 +503,7 @@ class TDelayRequestQuery : public TQueryBase { } private: - const TString Database; + const TString DatabaseId; const TString PoolId; const TString SessionId; const TInstant StartTime; @@ -514,9 +514,9 @@ class TDelayRequestQuery : public TQueryBase { class TStartFirstDelayedRequestQuery : public TQueryBase { public: - TStartFirstDelayedRequestQuery(const TString& database, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) - : TQueryBase(__func__, poolId, database, "", counters) - , Database(database) + TStartFirstDelayedRequestQuery(const TString& databaseId, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) + : TQueryBase(__func__, poolId, databaseId, "", counters) + , DatabaseId(databaseId) , PoolId(poolId) , LeaseDuration(leaseDuration) {} @@ -526,12 +526,12 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { TString sql = TStringBuilder() << R"( -- TStartFirstDelayedRequestQuery::OnRunQuery - DECLARE $database AS Text; + DECLARE $database_id AS Text; DECLARE $pool_id AS Text; SELECT database, pool_id, start_time, session_id, node_id FROM `)" << TTablesCreator::GetDelayedRequestsPath() << R"(` - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND (wait_deadline IS NULL OR wait_deadline >= CurrentUtcTimestamp()) AND lease_deadline >= CurrentUtcTimestamp() @@ -541,8 +541,8 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(Database) + .AddParam("$database_id") + .Utf8(DatabaseId) .Build() .AddParam("$pool_id") .Utf8(PoolId) @@ -583,7 +583,7 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { } RequestSessionId = *sessionId; - UpdateLogInfo(PoolId, Database, RequestSessionId); + UpdateLogInfo(PoolId, DatabaseId, RequestSessionId); TMaybe startTime = result.ColumnParser("start_time").GetOptionalTimestamp(); if (!startTime) { @@ -598,7 +598,7 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { void StartQueuedRequest() { TString sql = TStringBuilder() << R"( -- TStartFirstDelayedRequestQuery::StartQueuedRequest - DECLARE $database AS Text; + DECLARE $database_id AS Text; DECLARE $pool_id AS Text; DECLARE $start_time AS Timestamp; DECLARE $session_id AS Text; @@ -606,7 +606,7 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { DECLARE $lease_duration AS Interval; DELETE FROM `)" << TTablesCreator::GetDelayedRequestsPath() << R"(` - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND node_id = $node_id AND start_time = $start_time @@ -615,15 +615,15 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { UPSERT INTO `)" << TTablesCreator::GetRunningRequestsPath() << R"(` (database, pool_id, session_id, node_id, lease_deadline) VALUES ( - $database, $pool_id, $session_id, $node_id, + $database_id, $pool_id, $session_id, $node_id, CurrentUtcTimestamp() + $lease_duration ); )"; NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(Database) + .AddParam("$database_id") + .Utf8(DatabaseId) .Build() .AddParam("$pool_id") .Utf8(PoolId) @@ -654,7 +654,7 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { } private: - const TString Database; + const TString DatabaseId; const TString PoolId; const TDuration LeaseDuration; @@ -665,9 +665,9 @@ class TStartFirstDelayedRequestQuery : public TQueryBase { class TStartRequestQuery : public TQueryBase { public: - TStartRequestQuery(const TString& database, const TString& poolId, const TString& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) - : TQueryBase(__func__, poolId, database, sessionId, counters) - , Database(database) + TStartRequestQuery(const TString& databaseId, const TString& poolId, const TString& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) + : TQueryBase(__func__, poolId, databaseId, sessionId, counters) + , DatabaseId(databaseId) , PoolId(poolId) , SessionId(sessionId) , LeaseDuration(leaseDuration) @@ -676,7 +676,7 @@ class TStartRequestQuery : public TQueryBase { void OnRunQuery() override { TString sql = TStringBuilder() << R"( -- TStartRequestQuery::OnRunQuery - DECLARE $database AS Text; + DECLARE $database_id AS Text; DECLARE $pool_id AS Text; DECLARE $session_id AS Text; DECLARE $node_id AS Uint32; @@ -685,15 +685,15 @@ class TStartRequestQuery : public TQueryBase { UPSERT INTO `)" << TTablesCreator::GetRunningRequestsPath() << R"(` (database, pool_id, session_id, node_id, lease_deadline) VALUES ( - $database, $pool_id, $session_id, $node_id, + $database_id, $pool_id, $session_id, $node_id, CurrentUtcTimestamp() + $lease_duration ); )"; NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(Database) + .AddParam("$database_id") + .Utf8(DatabaseId) .Build() .AddParam("$pool_id") .Utf8(PoolId) @@ -720,7 +720,7 @@ class TStartRequestQuery : public TQueryBase { } private: - const TString Database; + const TString DatabaseId; const TString PoolId; const TString SessionId; const TDuration LeaseDuration; @@ -731,9 +731,9 @@ class TStartRequestActor : public TActorBootstrapped { using TStartRequestRetryQuery = TQueryRetryActor; public: - TStartRequestActor(const TActorId& replyActorId, const TString& database, const TString& poolId, const std::optional& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) + TStartRequestActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, const std::optional& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) : ReplyActorId(replyActorId) - , Database(database) + , DatabaseId(databaseId) , PoolId(poolId) , SessionId(sessionId) , LeaseDuration(leaseDuration) @@ -744,9 +744,9 @@ class TStartRequestActor : public TActorBootstrapped { Become(&TStartRequestActor::StateFunc); if (!SessionId) { - Register(new TStartFirstDelayedRequestRetryQuery(SelfId(), Database, PoolId, LeaseDuration, Counters)); + Register(new TStartFirstDelayedRequestRetryQuery(SelfId(), DatabaseId, PoolId, LeaseDuration, Counters)); } else { - Register(new TStartRequestRetryQuery(SelfId(), Database, PoolId, *SessionId, LeaseDuration, Counters)); + Register(new TStartRequestRetryQuery(SelfId(), DatabaseId, PoolId, *SessionId, LeaseDuration, Counters)); } } @@ -760,7 +760,7 @@ class TStartRequestActor : public TActorBootstrapped { private: const TActorId ReplyActorId; - const TString Database; + const TString DatabaseId; const TString PoolId; const std::optional SessionId; const TDuration LeaseDuration; @@ -770,9 +770,9 @@ class TStartRequestActor : public TActorBootstrapped { class TCleanupRequestsQuery : public TQueryBase { public: - TCleanupRequestsQuery(const TString& database, const TString& poolId, const std::vector& sessionIds, NMonitoring::TDynamicCounterPtr counters) - : TQueryBase(__func__, poolId, database, "", counters) - , Database(database) + TCleanupRequestsQuery(const TString& databaseId, const TString& poolId, const std::vector& sessionIds, NMonitoring::TDynamicCounterPtr counters) + : TQueryBase(__func__, poolId, databaseId, "", counters) + , DatabaseId(databaseId) , PoolId(poolId) , SessionIds(sessionIds) {} @@ -782,19 +782,19 @@ class TCleanupRequestsQuery : public TQueryBase { -- TCleanupRequestsQuery::OnRunQuery PRAGMA AnsiInForEmptyOrNullableItemsCollections; - DECLARE $database AS Text; + DECLARE $database_id AS Text; DECLARE $pool_id AS Text; DECLARE $node_id AS Uint32; DECLARE $session_ids AS List; DELETE FROM `)" << TTablesCreator::GetDelayedRequestsPath() << R"(` - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND node_id = $node_id AND session_id IN $session_ids; DELETE FROM `)" << TTablesCreator::GetRunningRequestsPath() << R"(` - WHERE database = $database + WHERE database = $database_id AND pool_id = $pool_id AND node_id = $node_id AND session_id IN $session_ids; @@ -802,8 +802,8 @@ class TCleanupRequestsQuery : public TQueryBase { NYdb::TParamsBuilder params; params - .AddParam("$database") - .Utf8(Database) + .AddParam("$database_id") + .Utf8(DatabaseId) .Build() .AddParam("$pool_id") .Utf8(PoolId) @@ -832,11 +832,11 @@ class TCleanupRequestsQuery : public TQueryBase { private: TString LogPrefix() const { - return TStringBuilder() << "[TCleanupRequestsQuery] ActorId: " << SelfId() << ", Database: " << Database << ", PoolId: " << PoolId << ", "; + return TStringBuilder() << "[TCleanupRequestsQuery] ActorId: " << SelfId() << ", DatabaseId: " << DatabaseId << ", PoolId: " << PoolId << ", "; } private: - const TString Database; + const TString DatabaseId; const TString PoolId; const std::vector SessionIds; }; @@ -851,20 +851,20 @@ IActor* CreateCleanupTablesActor() { return new TCleanupTablesActor(); } -IActor* CreateRefreshPoolStateActor(const TActorId& replyActorId, const TString& database, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) { - return new TQueryRetryActor(replyActorId, database, poolId, leaseDuration, counters); +IActor* CreateRefreshPoolStateActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) { + return new TQueryRetryActor(replyActorId, databaseId, poolId, leaseDuration, counters); } -IActor* CreateDelayRequestActor(const TActorId& replyActorId, const TString& database, const TString& poolId, const TString& sessionId, TInstant startTime, TMaybe waitDeadline, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) { - return new TQueryRetryActor, TDuration, NMonitoring::TDynamicCounterPtr>(replyActorId, database, poolId, sessionId, startTime, waitDeadline, leaseDuration, counters); +IActor* CreateDelayRequestActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, const TString& sessionId, TInstant startTime, TMaybe waitDeadline, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) { + return new TQueryRetryActor, TDuration, NMonitoring::TDynamicCounterPtr>(replyActorId, databaseId, poolId, sessionId, startTime, waitDeadline, leaseDuration, counters); } -IActor* CreateStartRequestActor(const TActorId& replyActorId, const TString& database, const TString& poolId, const std::optional& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) { - return new TStartRequestActor(replyActorId, database, poolId, sessionId, leaseDuration, counters); +IActor* CreateStartRequestActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, const std::optional& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters) { + return new TStartRequestActor(replyActorId, databaseId, poolId, sessionId, leaseDuration, counters); } -IActor* CreateCleanupRequestsActor(const TActorId& replyActorId, const TString& database, const TString& poolId, const std::vector& sessionIds, NMonitoring::TDynamicCounterPtr counters) { - return new TQueryRetryActor, NMonitoring::TDynamicCounterPtr>(replyActorId, database, poolId, sessionIds, counters); +IActor* CreateCleanupRequestsActor(const TActorId& replyActorId, const TString& databaseId, const TString& poolId, const std::vector& sessionIds, NMonitoring::TDynamicCounterPtr counters) { + return new TQueryRetryActor, NMonitoring::TDynamicCounterPtr>(replyActorId, databaseId, poolId, sessionIds, counters); } } // NKikimr::NKqp::NWorkload diff --git a/ydb/core/kqp/workload_service/tables/table_queries.h b/ydb/core/kqp/workload_service/tables/table_queries.h index f44afac3ba74..b7727bbff28a 100644 --- a/ydb/core/kqp/workload_service/tables/table_queries.h +++ b/ydb/core/kqp/workload_service/tables/table_queries.h @@ -12,11 +12,11 @@ NActors::IActor* CreateTablesCreator(); NActors::IActor* CreateCleanupTablesActor(); // Updates pool lease and returns pool description -NActors::IActor* CreateRefreshPoolStateActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters); +NActors::IActor* CreateRefreshPoolStateActor(const NActors::TActorId& replyActorId, const TString& databaseId, const TString& poolId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters); // Push / Start / Finish requests in pool -NActors::IActor* CreateDelayRequestActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, const TString& sessionId, TInstant startTime, TMaybe waitDeadline, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters); -NActors::IActor* CreateStartRequestActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, const std::optional& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters); -NActors::IActor* CreateCleanupRequestsActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, const std::vector& sessionIds, NMonitoring::TDynamicCounterPtr counters); +NActors::IActor* CreateDelayRequestActor(const NActors::TActorId& replyActorId, const TString& databaseId, const TString& poolId, const TString& sessionId, TInstant startTime, TMaybe waitDeadline, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters); +NActors::IActor* CreateStartRequestActor(const NActors::TActorId& replyActorId, const TString& databaseId, const TString& poolId, const std::optional& sessionId, TDuration leaseDuration, NMonitoring::TDynamicCounterPtr counters); +NActors::IActor* CreateCleanupRequestsActor(const NActors::TActorId& replyActorId, const TString& databaseId, const TString& poolId, const std::vector& sessionIds, NMonitoring::TDynamicCounterPtr counters); } // NKikimr::NKqp::NWorkload diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp index 43dbad70c660..7ff63e246426 100644 --- a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp +++ b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp @@ -452,7 +452,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup { TPoolStateDescription GetPoolDescription(TDuration leaseDuration = FUTURE_WAIT_TIMEOUT, const TString& poolId = "") const override { const auto& edgeActor = GetRuntime()->AllocateEdgeActor(); - GetRuntime()->Register(CreateRefreshPoolStateActor(edgeActor, Settings_.DomainName_, poolId ? poolId : Settings_.PoolId_, leaseDuration, GetRuntime()->GetAppData().Counters)); + GetRuntime()->Register(CreateRefreshPoolStateActor(edgeActor, CanonizePath(Settings_.DomainName_), poolId ? poolId : Settings_.PoolId_, leaseDuration, GetRuntime()->GetAppData().Counters)); auto response = GetRuntime()->GrabEdgeEvent(edgeActor, FUTURE_WAIT_TIMEOUT); UNIT_ASSERT_VALUES_EQUAL_C(response->Get()->Status, Ydb::StatusIds::SUCCESS, response->Get()->Issues.ToOneLineString()); @@ -498,6 +498,15 @@ class TWorkloadServiceYdbSetup : public IYdbSetup { } } + TEvFetchDatabaseResponse::TPtr FetchDatabase(const TString& database) const override { + const TActorId edgeActor = GetRuntime()->AllocateEdgeActor(); + GetRuntime()->Register(CreateDatabaseFetcherActor(edgeActor, database)); + const auto response = GetRuntime()->GrabEdgeEvent(edgeActor); + UNIT_ASSERT_C(response, "Got empty response from DatabaseFetcherActor"); + UNIT_ASSERT_VALUES_EQUAL_C(response->Get()->Status, Ydb::StatusIds::SUCCESS, response->Get()->Issues.ToOneLineString()); + return response; + } + // Coomon helpers TTestActorRuntime* GetRuntime() const override { return Server_->GetRuntime(); diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h index 8ff7fdb93244..0f4eb6c7cf6a 100644 --- a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h +++ b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h @@ -116,6 +116,7 @@ class IYdbSetup : public TThrRefBase { virtual void WaitPoolHandlersCount(i64 finalCount, std::optional initialCount = std::nullopt, TDuration timeout = FUTURE_WAIT_TIMEOUT) const = 0; virtual void StopWorkloadService(ui64 nodeIndex = 0) const = 0; virtual void ValidateWorkloadServiceCounters(bool checkTableCounters = true, const TString& poolId = "") const = 0; + virtual TEvFetchDatabaseResponse::TPtr FetchDatabase(const TString& database) const = 0; // Coomon helpers virtual TTestActorRuntime* GetRuntime() const = 0; diff --git a/ydb/core/kqp/workload_service/ut/kqp_workload_service_tables_ut.cpp b/ydb/core/kqp/workload_service/ut/kqp_workload_service_tables_ut.cpp index a004917a2a03..f5972e873f4e 100644 --- a/ydb/core/kqp/workload_service/ut/kqp_workload_service_tables_ut.cpp +++ b/ydb/core/kqp/workload_service/ut/kqp_workload_service_tables_ut.cpp @@ -27,7 +27,7 @@ void DelayRequest(TIntrusivePtr ydb, const TString& sessionId, TDurat auto runtime = ydb->GetRuntime(); const auto& edgeActor = runtime->AllocateEdgeActor(); - runtime->Register(CreateDelayRequestActor(edgeActor, settings.DomainName_, settings.PoolId_, sessionId, TInstant::Now(), Nothing(), leaseDuration, runtime->GetAppData().Counters)); + runtime->Register(CreateDelayRequestActor(edgeActor, CanonizePath(settings.DomainName_), settings.PoolId_, sessionId, TInstant::Now(), Nothing(), leaseDuration, runtime->GetAppData().Counters)); auto response = runtime->GrabEdgeEvent(edgeActor, FUTURE_WAIT_TIMEOUT); UNIT_ASSERT_VALUES_EQUAL_C(response->Get()->Status, Ydb::StatusIds::SUCCESS, response->Get()->Issues.ToOneLineString()); UNIT_ASSERT_VALUES_EQUAL(response->Get()->SessionId, sessionId); @@ -38,7 +38,7 @@ void StartRequest(TIntrusivePtr ydb, const TString& sessionId, TDurat auto runtime = ydb->GetRuntime(); const auto& edgeActor = runtime->AllocateEdgeActor(); - runtime->Register(CreateStartRequestActor(edgeActor, settings.DomainName_, settings.PoolId_, sessionId, leaseDuration, runtime->GetAppData().Counters)); + runtime->Register(CreateStartRequestActor(edgeActor, CanonizePath(settings.DomainName_), settings.PoolId_, sessionId, leaseDuration, runtime->GetAppData().Counters)); auto response = runtime->GrabEdgeEvent(edgeActor, FUTURE_WAIT_TIMEOUT); UNIT_ASSERT_VALUES_EQUAL_C(response->Get()->Status, Ydb::StatusIds::SUCCESS, response->Get()->Issues.ToOneLineString()); UNIT_ASSERT_VALUES_EQUAL(response->Get()->SessionId, sessionId); diff --git a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp index bb0ee347e6be..a020bd32ba6e 100644 --- a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp +++ b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp @@ -422,6 +422,56 @@ Y_UNIT_TEST_SUITE(ResourcePoolsDdl) { TSampleQueries::TSelect42::CheckResult(hangingRequest.GetResult()); } + Y_UNIT_TEST(TestCreateResourcePoolOnServerless) { + auto ydb = TYdbSetupSettings() + .CreateSampleTenants(true) + .EnableResourcePoolsOnServerless(true) + .Create(); + + const auto& serverlessTenant = ydb->GetSettings().GetServerlessTenantName(); + auto settings = TQueryRunnerSettings() + .PoolId("") + .Database(serverlessTenant) + .NodeIndex(1); + + const TString& poolId = "my_pool"; + TSampleQueries::CheckSuccess(ydb->ExecuteQuery(TStringBuilder() << R"( + CREATE RESOURCE POOL )" << poolId << R"( WITH ( + CONCURRENT_QUERY_LIMIT=1, + QUEUE_SIZE=0 + ); + )", settings)); + settings.PoolId(poolId); + + auto hangingRequest = ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query, settings.HangUpDuringExecution(true)); + ydb->WaitQueryExecution(hangingRequest); + + settings.HangUpDuringExecution(false); + + { // Rejected result + auto result = ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, settings.PoolId(poolId)); + UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::OVERLOADED, result.GetIssues().ToOneLineString()); + UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), TStringBuilder() << "Request was rejected, number of local pending requests is 1, number of global delayed/running requests is 1, sum of them is larger than allowed limit 0 (including concurrent query limit 1) for pool " << poolId); + } + + { // Check tables + auto result = ydb->ExecuteQuery(R"( + SELECT * FROM `.metadata/workload_manager/running_requests` + )", settings.PoolId(NResourcePool::DEFAULT_POOL_ID).Database(ydb->GetSettings().GetSharedTenantName())); + TSampleQueries::CheckSuccess(result); + + NYdb::TResultSetParser resultSet(result.GetResultSet(0)); + UNIT_ASSERT_C(resultSet.TryNextRow(), "Unexpected row count"); + + const auto& databaseId = resultSet.ColumnParser("database").GetOptionalUtf8(); + UNIT_ASSERT_C(databaseId, "Unexpected database response"); + UNIT_ASSERT_VALUES_EQUAL_C(*databaseId, ydb->FetchDatabase(serverlessTenant)->Get()->DatabaseId, "Unexpected database id"); + } + + ydb->ContinueQueryExecution(hangingRequest); + TSampleQueries::TSelect42::CheckResult(hangingRequest.GetResult()); + } + Y_UNIT_TEST(TestDefaultPoolRestrictions) { auto ydb = TYdbSetupSettings().Create(); @@ -615,22 +665,27 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { UNIT_ASSERT_STRING_CONTAINS(alterResult.GetIssues().ToOneLineString(), "You don't have access permissions for database Root"); } - void CreateSampleResourcePoolClassifier(TIntrusivePtr ydb, const TString& classifierId, const TString& userSID, const TString& poolId) { - ydb->ExecuteSchemeQuery(TStringBuilder() << R"( - GRANT ALL ON `/Root` TO `)" << userSID << R"(`; + void CreateSampleResourcePoolClassifier(TIntrusivePtr ydb, const TString& classifierId, const TQueryRunnerSettings& settings, const TString& poolId) { + TSampleQueries::CheckSuccess(ydb->ExecuteQuery(TStringBuilder() << R"( + GRANT ALL ON `)" << CanonizePath(settings.Database_ ? settings.Database_ : ydb->GetSettings().DomainName_) << R"(` TO `)" << settings.UserSID_ << R"(`; CREATE RESOURCE POOL )" << poolId << R"( WITH ( CONCURRENT_QUERY_LIMIT=0 ); CREATE RESOURCE POOL CLASSIFIER )" << classifierId << R"( WITH ( RESOURCE_POOL=")" << poolId << R"(", - MEMBER_NAME=")" << userSID << R"(" + MEMBER_NAME=")" << settings.UserSID_ << R"(" ); - )"); + )", TQueryRunnerSettings() + .UserSID(BUILTIN_ACL_METADATA) + .Database(settings.Database_) + .NodeIndex(settings.NodeIndex_) + .PoolId(NResourcePool::DEFAULT_POOL_ID) + )); } TString CreateSampleResourcePoolClassifier(TIntrusivePtr ydb, const TQueryRunnerSettings& settings, const TString& poolId) { const TString& classifierId = "my_pool_classifier"; - CreateSampleResourcePoolClassifier(ydb, classifierId, settings.UserSID_, poolId); + CreateSampleResourcePoolClassifier(ydb, classifierId, settings, poolId); return classifierId; } @@ -662,6 +717,24 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) { WaitForFail(ydb, settings, poolId); } + Y_UNIT_TEST(TestCreateResourcePoolClassifierOnServerless) { + auto ydb = TYdbSetupSettings() + .CreateSampleTenants(true) + .EnableResourcePoolsOnServerless(true) + .Create(); + + auto settings = TQueryRunnerSettings() + .PoolId("") + .UserSID("test@user") + .Database(ydb->GetSettings().GetServerlessTenantName()) + .NodeIndex(1); + + const TString& poolId = "my_pool"; + CreateSampleResourcePoolClassifier(ydb, settings, poolId); + + WaitForFail(ydb, settings, poolId); + } + Y_UNIT_TEST(TestAlterResourcePoolClassifier) { auto ydb = TYdbSetupSettings().Create(); diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto index 2436d4b97b13..18dd347cf667 100644 --- a/ydb/core/protos/kqp.proto +++ b/ydb/core/protos/kqp.proto @@ -589,6 +589,7 @@ message TEvStartKqpTasksRequest { optional string SchedulerGroup = 9; optional double MemoryPoolPercent = 10 [default = 100]; optional string Database = 11; + optional string DatabaseId = 17; optional uint64 LockTxId = 13; optional uint32 LockNodeId = 14; diff --git a/ydb/services/metadata/manager/abstract.h b/ydb/services/metadata/manager/abstract.h index 74a34997a1c9..65b6759628bc 100644 --- a/ydb/services/metadata/manager/abstract.h +++ b/ydb/services/metadata/manager/abstract.h @@ -76,6 +76,7 @@ class IOperationsManager { private: YDB_ACCESSOR_DEF(std::optional, UserToken); YDB_ACCESSOR_DEF(TString, Database); + YDB_ACCESSOR_DEF(TString, DatabaseId); using TActorSystemPtr = TActorSystem*; YDB_ACCESSOR_DEF(TActorSystemPtr, ActorSystem); YDB_ACCESSOR_DEF(NSQLTranslation::TTranslationSettings, TranslationSettings); diff --git a/ydb/tools/query_replay/query_compiler.cpp b/ydb/tools/query_replay/query_compiler.cpp index 8a2e417d5dcb..7c0c070e9e2b 100644 --- a/ydb/tools/query_replay/query_compiler.cpp +++ b/ydb/tools/query_replay/query_compiler.cpp @@ -254,9 +254,11 @@ class TReplayCompileActor: public TActorBootstrapped { } TKqpQuerySettings settings(queryType); + const auto& database = ReplayDetails["query_database"].GetStringSafe(); Query = std::make_unique( ReplayDetails["query_cluster"].GetStringSafe(), - ReplayDetails["query_database"].GetStringSafe(), + database, + database, queryText, settings, !queryParameterTypes.empty() @@ -288,7 +290,7 @@ class TReplayCompileActor: public TActorBootstrapped { counters->Counters = new TKqpCounters(c); counters->TxProxyMon = new NTxProxy::TTxProxyMon(c); - Gateway = CreateKikimrIcGateway(Query->Cluster, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, Query->Database, std::move(loader), + Gateway = CreateKikimrIcGateway(Query->Cluster, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, Query->Database, Query->DatabaseId, std::move(loader), TlsActivationContext->ExecutorThread.ActorSystem, SelfId().NodeId(), counters); auto federatedQuerySetup = std::make_optional({NYql::IHTTPGateway::Make(), nullptr, nullptr, nullptr, {}, {}, {}, nullptr, nullptr, {}}); diff --git a/ydb/tools/query_replay_yt/query_compiler.cpp b/ydb/tools/query_replay_yt/query_compiler.cpp index 50ed2095ee61..ba4176d8fc4e 100644 --- a/ydb/tools/query_replay_yt/query_compiler.cpp +++ b/ydb/tools/query_replay_yt/query_compiler.cpp @@ -589,9 +589,11 @@ class TReplayCompileActor: public TActorBootstrapped { } TKqpQuerySettings settings(queryType); + const auto& database = ReplayDetails["query_database"].GetStringSafe(); Query = std::make_unique( ReplayDetails["query_cluster"].GetStringSafe(), - ReplayDetails["query_database"].GetStringSafe(), + database, + database, queryText, settings, !queryParameterTypes.empty() @@ -621,7 +623,7 @@ class TReplayCompileActor: public TActorBootstrapped { counters->Counters = new TKqpCounters(c); counters->TxProxyMon = new NTxProxy::TTxProxyMon(c); - Gateway = CreateKikimrIcGateway(Query->Cluster, queryType, Query->Database, std::move(loader), + Gateway = CreateKikimrIcGateway(Query->Cluster, queryType, Query->Database, Query->DatabaseId, std::move(loader), TlsActivationContext->ExecutorThread.ActorSystem, SelfId().NodeId(), counters); auto federatedQuerySetup = std::make_optional({HttpGateway, nullptr, nullptr, nullptr, {}, {}, {}, nullptr, nullptr, {}}); KqpHost = CreateKqpHost(Gateway, Query->Cluster, Query->Database, Config, ModuleResolverState->ModuleResolver,