Merge 16ceb2b into a169e9b

GrigoriyPA · web-flow · commit 43937fbe1542 · 2024-07-29T10:04:20.000Z
diff --git a/ydb/core/kqp/common/events/query.h b/ydb/core/kqp/common/events/query.h
@@ -1,4 +1,5 @@
 #pragma once
+#include <ydb/core/resource_pools/resource_pool_settings.h>
 #include <ydb/core/protos/kqp.pb.h>
 #include <ydb/core/kqp/common/simple/kqp_event_ids.h>
 #include <ydb/core/kqp/common/kqp_user_request_context.h>
@@ -342,6 +343,14 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
         return Record.GetRequest().GetPoolId();
     }
 
+    void SetPoolConfig(const NResourcePool::TPoolSettings& config) {
+        PoolConfig = config;
+    }
+
+    std::optional<NResourcePool::TPoolSettings> GetPoolConfig() const {
+        return PoolConfig;
+    }
+
     mutable NKikimrKqp::TEvQueryRequest Record;
 
 private:
@@ -370,6 +379,7 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
     TDuration CancelAfter;
     TIntrusivePtr<TUserRequestContext> UserRequestContext;
     TDuration ProgressStatsPeriod;
+    std::optional<NResourcePool::TPoolSettings> PoolConfig;
 };
 
 struct TEvDataQueryStreamPart: public TEventPB<TEvDataQueryStreamPart,
diff --git a/ydb/core/kqp/common/events/workload_service.h b/ydb/core/kqp/common/events/workload_service.h
@@ -24,6 +24,7 @@ struct TEvPlaceRequestIntoPool : public NActors::TEventLocal<TEvPlaceRequestInto
     const TString SessionId;
     TString PoolId;  // Can be changed to default pool id
     TIntrusiveConstPtr<NACLib::TUserToken> UserToken;
+    bool FromCache = false;
 };
 
 struct TEvContinueRequest : public NActors::TEventLocal<TEvContinueRequest, TKqpWorkloadServiceEvents::EvContinueRequest> {
@@ -66,4 +67,18 @@ struct TEvCleanupResponse : public NActors::TEventLocal<TEvCleanupResponse, TKqp
     const NYql::TIssues Issues;
 };
 
+struct TEvUpdatePoolInfo : public NActors::TEventLocal<TEvUpdatePoolInfo, TKqpWorkloadServiceEvents::EvUpdatePoolInfo> {
+    TEvUpdatePoolInfo(const TString& database, const TString& poolId, const std::optional<NResourcePool::TPoolSettings>& config, const std::optional<NACLib::TSecurityObject>& securityObject)
+        : Database(database)
+        , PoolId(poolId)
+        , Config(config)
+        , SecurityObject(securityObject)
+    {}
+
+    const TString Database;
+    const TString PoolId;
+    const std::optional<NResourcePool::TPoolSettings> Config;
+    const std::optional<NACLib::TSecurityObject> SecurityObject;
+};
+
 }  // NKikimr::NKqp::NWorkload
diff --git a/ydb/core/kqp/common/kqp_user_request_context.h b/ydb/core/kqp/common/kqp_user_request_context.h
@@ -15,7 +15,7 @@ namespace NKikimr::NKqp {
         TString CurrentExecutionId;
         TString CustomerSuppliedId;
         TString PoolId;
-        NResourcePool::TPoolSettings PoolConfig;
+        std::optional<NResourcePool::TPoolSettings> PoolConfig;
 
         TUserRequestContext() = default;
 
diff --git a/ydb/core/kqp/common/simple/kqp_event_ids.h b/ydb/core/kqp/common/simple/kqp_event_ids.h
@@ -174,6 +174,7 @@ struct TKqpWorkloadServiceEvents {
         EvContinueRequest,
         EvCleanupRequest,
         EvCleanupResponse,
+        EvUpdatePoolInfo,
     };
 };
 
diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
@@ -11,6 +11,7 @@
 #include <ydb/core/cms/console/console.h>
 #include <ydb/core/kqp/counters/kqp_counters.h>
 #include <ydb/core/kqp/common/events/script_executions.h>
+#include <ydb/core/kqp/common/events/workload_service.h>
 #include <ydb/core/kqp/common/kqp_lwtrace_probes.h>
 #include <ydb/core/kqp/common/kqp_timeouts.h>
 #include <ydb/core/kqp/compile_service/kqp_compile_service.h>
@@ -696,11 +697,8 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
             LocalSessions->AttachQueryText(sessionInfo, ev->Get()->GetQuery());
         }
 
-        if (!FeatureFlags.GetEnableResourcePools()) {
-            ev->Get()->SetPoolId("");
-        } else if (!ev->Get()->GetPoolId()) {
-            // TODO: do not use default pool if there is no limits
-            ev->Get()->SetPoolId(NResourcePool::DEFAULT_POOL_ID);
+        if (!TryGetPoolInfoFromCache(ev, requestId)) {
+            return;
         }
 
         TActorId targetId;
@@ -1352,6 +1350,7 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
             hFunc(TEvInterconnect::TEvNodeDisconnected, Handle);
             hFunc(TEvKqp::TEvListSessionsRequest, Handle);
             hFunc(TEvKqp::TEvListProxyNodesRequest, Handle);
+            hFunc(NWorkload::TEvUpdatePoolInfo, Handle);
         default:
             Y_ABORT("TKqpProxyService: unexpected event type: %" PRIx32 " event: %s",
                 ev->GetTypeRewrite(), ev->ToString().data());
@@ -1559,6 +1558,43 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
         }
     }
 
+    bool TryGetPoolInfoFromCache(TEvKqp::TEvQueryRequest::TPtr& ev, ui64 requestId) {
+        if (!FeatureFlags.GetEnableResourcePools()) {
+            ev->Get()->SetPoolId("");
+            return true;
+        }
+
+        if (!ev->Get()->GetPoolId()) {
+            ev->Get()->SetPoolId(NResourcePool::DEFAULT_POOL_ID);
+        }
+
+        const auto& poolId = ev->Get()->GetPoolId();
+        const auto& poolInfo = ResourcePoolsCache.GetPoolInfo(ev->Get()->GetDatabase(), poolId);
+        if (!poolInfo) {
+            return true;
+        }
+
+        const auto& securityObject = poolInfo->SecurityObject;
+        const auto& userToken = ev->Get()->GetUserToken();
+        if (securityObject && userToken && !userToken->GetSerializedToken().empty()) {
+            if (!securityObject->CheckAccess(NACLib::EAccessRights::DescribeSchema, *userToken)) {
+                ReplyProcessError(Ydb::StatusIds::NOT_FOUND, TStringBuilder() << "Resource pool " << poolId << " not found or you don't have access permissions", requestId);
+                return false;
+            }
+            if (!securityObject->CheckAccess(NACLib::EAccessRights::SelectRow, *userToken)) {
+                ReplyProcessError(Ydb::StatusIds::UNAUTHORIZED, TStringBuilder() << "You don't have access permissions for resource pool " << poolId, requestId);
+                return false;
+            }
+        }
+
+        const auto& poolConfig = poolInfo->Config;
+        if (poolConfig.ConcurrentQueryLimit == -1 && poolConfig.DatabaseLoadCpuThreshold < 0.0 && !poolConfig.QueryCancelAfter) {
+            ev->Get()->SetPoolConfig(poolConfig);
+        }
+
+        return true;
+    }
+
     void UpdateYqlLogLevels() {
         const auto& kqpYqlName = NKikimrServices::EServiceKikimr_Name(NKikimrServices::KQP_YQL);
         for (auto &entry : LogConfig.GetEntry()) {
@@ -1744,6 +1780,10 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
         Send(ev->Sender, result.release(), 0, ev->Cookie);
     }
 
+    void Handle(NWorkload::TEvUpdatePoolInfo::TPtr& ev) {
+        ResourcePoolsCache.UpdatePoolInfo(ev->Get()->Database, ev->Get()->PoolId, ev->Get()->Config, ev->Get()->SecurityObject);
+    }
+
 private:
     NKikimrConfig::TLogConfig LogConfig;
     NKikimrConfig::TTableServiceConfig TableServiceConfig;
@@ -1805,6 +1845,8 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
     std::deque<TDelayedEvent> DelayedEventsQueue;
     bool IsLookupByRmScheduled = false;
     TActorId KqpTempTablesAgentActor;
+
+    TResourcePoolsCache ResourcePoolsCache;
 };
 
 } // 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
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ydb/core/base/appdata.h>
+#include <ydb/core/base/path.h>
 #include <ydb/core/kqp/common/kqp.h>
 #include <ydb/core/kqp/counters/kqp_counters.h>
 #include <ydb/core/kqp/rm_service/kqp_rm_service.h>
@@ -415,4 +416,40 @@ class TLocalSessionsRegistry {
     }
 };
 
+class TResourcePoolsCache {
+    struct TPoolInfo {
+        NResourcePool::TPoolSettings Config;
+        std::optional<NACLib::TSecurityObject> SecurityObject;
+    };
+
+public:
+    std::optional<TPoolInfo> GetPoolInfo(const TString& database, const TString& poolId) const {
+        auto it = PoolsCache.find(GetPoolKey(database, poolId));
+        if (it == PoolsCache.end()) {
+            return std::nullopt;
+        }
+        return it->second;
+    }
+
+    void UpdatePoolInfo(const TString& database, const TString& poolId, const std::optional<NResourcePool::TPoolSettings>& config, const std::optional<NACLib::TSecurityObject>& securityObject) {
+        const TString& poolKey = GetPoolKey(database, poolId);
+        if (!config) {
+            PoolsCache.erase(poolKey);
+            return;
+        }
+
+        auto& poolInfo = PoolsCache[poolKey];
+        poolInfo.Config = *config;
+        poolInfo.SecurityObject = securityObject;
+    }
+
+private:
+    static TString GetPoolKey(const TString& database, const TString& poolId) {
+        return CanonizePath(TStringBuilder() << database << "/" << poolId);
+    }
+
+private:
+    std::unordered_map<TString, TPoolInfo> PoolsCache;
+};
+
 }  // namespace NKikimr::NKqp
diff --git a/ydb/core/kqp/session_actor/kqp_query_state.h b/ydb/core/kqp/session_actor/kqp_query_state.h
@@ -85,6 +85,7 @@ class TKqpQueryState : public TNonCopyable {
             UserRequestContext = MakeIntrusive<TUserRequestContext>(RequestEv->GetTraceId(), Database, sessionId);
         }
         UserRequestContext->PoolId = RequestEv->GetPoolId();
+        UserRequestContext->PoolConfig = RequestEv->GetPoolConfig();
     }
 
     // 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
@@ -242,14 +242,20 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
     }
 
     void PassRequestToResourcePool() {
-        Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvPlaceRequestIntoPool(
-            QueryState->Database,
-            SessionId,
-            QueryState->UserRequestContext->PoolId,
-            QueryState->UserToken
-        ), IEventHandle::FlagTrackDelivery);
+        const TString& poolId = QueryState->UserRequestContext->PoolId;
+        auto event = std::make_unique<NWorkload::TEvPlaceRequestIntoPool>(QueryState->Database, SessionId, poolId, QueryState->UserToken);
 
-        Become(&TKqpSessionActor::ExecuteState);
+        TEventFlags flags = 0;
+        if (QueryState->UserRequestContext->PoolConfig) {
+            LOG_D("request placed into pool from cache: " << poolId);
+            event->FromCache = true;
+            CompileQuery();
+        } else {
+            flags |= IEventHandle::FlagTrackDelivery;
+            Become(&TKqpSessionActor::ExecuteState);
+        }
+
+        Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), event.release(), flags);
     }
 
     void ForwardRequest(TEvKqp::TEvQueryRequest::TPtr& ev) {
diff --git a/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp b/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp
@@ -141,7 +141,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
 
         // Schemeboard events
         hFunc(TEvTxProxySchemeCache::TEvWatchNotifyUpdated, Handle);
-        IgnoreFunc(TEvTxProxySchemeCache::TEvWatchNotifyDeleted);
+        hFunc(TEvTxProxySchemeCache::TEvWatchNotifyDeleted, Handle);
         IgnoreFunc(TEvTxProxySchemeCache::TEvWatchNotifyUnavailable);
     )
 
@@ -150,6 +150,8 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
             this->Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvWatchRemove(0));
         }
 
+        SendPoolInfoUpdate(std::nullopt, std::nullopt);
+
         Counters.OnCleanup();
 
         TBase::PassAway();
@@ -254,6 +256,24 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
         NResourcePool::TPoolSettings poolConfig;
         ParsePoolSettings(result->GetPathDescription().GetResourcePoolDescription(), poolConfig);
         UpdatePoolConfig(poolConfig);
+
+        const auto& pathDescription = result->GetPathDescription().GetSelf();
+        NACLib::TSecurityObject object(pathDescription.GetOwner(), false);
+        if (object.MutableACL()->ParseFromString(pathDescription.GetEffectiveACL())) {
+            SendPoolInfoUpdate(poolConfig, object);
+        } else {
+            SendPoolInfoUpdate(poolConfig, std::nullopt);
+        }
+    }
+
+    void Handle(TEvTxProxySchemeCache::TEvWatchNotifyDeleted::TPtr& ev) {
+        if (ev->Get()->Key != WatchKey) {
+            // Skip old paths watch notifications
+            return;
+        }
+
+        LOG_D("Got delete notification");
+        SendPoolInfoUpdate(std::nullopt, std::nullopt);
     }
 
 public:
@@ -318,6 +338,10 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
         RemoveRequest(request);
     }
 
+    void SendPoolInfoUpdate(const std::optional<NResourcePool::TPoolSettings>& config, const std::optional<NACLib::TSecurityObject>& securityObject) const {
+        this->Send(MakeKqpProxyID(this->SelfId().NodeId()), new TEvUpdatePoolInfo(Database, PoolId, config, securityObject));
+    }
+
 protected:
     virtual bool ShouldResign() const = 0;
     virtual void OnScheduleRequest(TRequest* request) = 0;
diff --git a/ydb/core/kqp/workload_service/kqp_workload_service.cpp b/ydb/core/kqp/workload_service/kqp_workload_service.cpp
@@ -132,7 +132,18 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
             return;
         }
 
-        LOG_D("Recieved new request from " << workerActorId << ", Database: " << ev->Get()->Database << ", PoolId: " << ev->Get()->PoolId << ", SessionId: " << ev->Get()->SessionId);
+        const TString& database = ev->Get()->Database;
+        const TString& poolId = ev->Get()->PoolId;
+
+        LOG_D("Recieved new request from " << workerActorId << ", Database: " << database << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId);
+        if (ev->Get()->FromCache) {
+            if (auto poolState = GetPoolState(database, poolId)) {
+                // Update LastUpdateTime due to pool idle checks 
+                poolState->LastUpdateTime = TInstant::Now();
+            }
+            return;
+        }
+
         bool hasDefaultPool = DatabasesWithDefaultPool.contains(CanonizePath(ev->Get()->Database));
         Register(CreatePoolResolverActor(std::move(ev), hasDefaultPool, EnabledResourcePoolsOnServerless));
     }
@@ -392,7 +403,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
         }
         IdleChecksStarted = true;
 
-        Schedule(IDLE_DURATION / 2, new TEvents::TEvWakeup());
+        Schedule(IDLE_DURATION / 2, new TEvents::TEvWakeup(static_cast<ui64>(EWakeUp::IdleCheck)));
     }
 
     void RunIdleCheck() {
@@ -446,7 +457,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
     }
 
     void ScheduleNodeInfoRequest() const {
-        Schedule(IDLE_DURATION * 2, new TEvents::TEvWakeup(static_cast<ui64>(EWakeUp::StartCpuLoadRequest)));
+        Schedule(IDLE_DURATION * 2, new TEvents::TEvWakeup(static_cast<ui64>(EWakeUp::StartNodeInfoRequest)));
     }
 
     void RunNodeInfoRequest() const {
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
@@ -260,6 +260,9 @@ Y_UNIT_TEST_SUITE(KqpWorkloadService) {
         TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, TQueryRunnerSettings().PoolId(NResourcePool::DEFAULT_POOL_ID)));
 
         ydb->WaitPoolHandlersCount(0, 2, TDuration::Seconds(95));
+
+        // Check pool creation after cleanup
+        TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(TSampleQueries::TSelect42::Query));
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -85,6 +85,7 @@ class TKqpQueryState : public TNonCopyable {`
`85`	`85`	`UserRequestContext = MakeIntrusive<TUserRequestContext>(RequestEv->GetTraceId(), Database, sessionId);`
`86`	`86`	`}`
`87`	`87`	`UserRequestContext->PoolId = RequestEv->GetPoolId();`
	`88`	`+ UserRequestContext->PoolConfig = RequestEv->GetPoolConfig();`
`88`	`89`	`}`
`89`	`90`
`90`	`91`	`// the monotonously growing counter, the ordinal number of the query,`
Original file line number	Diff line number	Diff line change
`@@ -132,7 +132,18 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {`
`132`	`132`	`return;`
`133`	`133`	`}`
`134`	`134`
`135`		`- LOG_D("Recieved new request from " << workerActorId << ", Database: " << ev->Get()->Database << ", PoolId: " << ev->Get()->PoolId << ", SessionId: " << ev->Get()->SessionId);`
	`135`	`+ const TString& database = ev->Get()->Database;`
	`136`	`+ const TString& poolId = ev->Get()->PoolId;`
	`137`	`+`
	`138`	`+ LOG_D("Recieved new request from " << workerActorId << ", Database: " << database << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId);`
	`139`	`+ if (ev->Get()->FromCache) {`
	`140`	`+ if (auto poolState = GetPoolState(database, poolId)) {`
	`141`	`+ // Update LastUpdateTime due to pool idle checks`
	`142`	`+ poolState->LastUpdateTime = TInstant::Now();`
	`143`	`+ }`
	`144`	`+ return;`
	`145`	`+ }`
	`146`	`+`
`136`	`147`	`bool hasDefaultPool = DatabasesWithDefaultPool.contains(CanonizePath(ev->Get()->Database));`
`137`	`148`	`Register(CreatePoolResolverActor(std::move(ev), hasDefaultPool, EnabledResourcePoolsOnServerless));`
`138`	`149`	`}`
`@@ -392,7 +403,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {`
`392`	`403`	`}`
`393`	`404`	`IdleChecksStarted = true;`
`394`	`405`
`395`		`- Schedule(IDLE_DURATION / 2, new TEvents::TEvWakeup());`
	`406`	`+ Schedule(IDLE_DURATION / 2, new TEvents::TEvWakeup(static_cast<ui64>(EWakeUp::IdleCheck)));`
`396`	`407`	`}`
`397`	`408`
`398`	`409`	`void RunIdleCheck() {`
`@@ -446,7 +457,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {`
`446`	`457`	`}`
`447`	`458`
`448`	`459`	`void ScheduleNodeInfoRequest() const {`
`449`		`- Schedule(IDLE_DURATION * 2, new TEvents::TEvWakeup(static_cast<ui64>(EWakeUp::StartCpuLoadRequest)));`
	`460`	`+ Schedule(IDLE_DURATION * 2, new TEvents::TEvWakeup(static_cast<ui64>(EWakeUp::StartNodeInfoRequest)));`
`450`	`461`	`}`
`451`	`462`
`452`	`463`	`void RunNodeInfoRequest() const {`