Merge 78b6755 into b1f9ae2

Author: GrigoriyPA

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,

ydb/core/kqp/common/events/workload_service.h

@@ -66,4 +66,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

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;
 

ydb/core/kqp/common/simple/kqp_event_ids.h

@@ -174,6 +174,7 @@ struct TKqpWorkloadServiceEvents {
         EvContinueRequest,
         EvCleanupRequest,
         EvCleanupResponse,
+        EvUpdatePoolInfo,
     };
 };
 

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>
@@ -698,11 +699,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 (!TryFillPoolInfoFromCache(ev, requestId)) {
+            return;
         }
 
         TActorId targetId;
@@ -1354,6 +1352,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());
@@ -1561,6 +1560,43 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
         }
     }
 
+    bool TryFillPoolInfoFromCache(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 (!NWorkload::IsWorkloadServiceRequired(poolConfig)) {
+            ev->Get()->SetPoolConfig(poolConfig);
+        }
+
+        return true;
+    }
+
     void UpdateYqlLogLevels() {
         const auto& kqpYqlName = NKikimrServices::EServiceKikimr_Name(NKikimrServices::KQP_YQL);
         for (auto &entry : LogConfig.GetEntry()) {
@@ -1746,6 +1782,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;
@@ -1807,6 +1847,8 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
     std::deque<TDelayedEvent> DelayedEventsQueue;
     bool IsLookupByRmScheduled = false;
     TActorId KqpTempTablesAgentActor;
+
+    TResourcePoolsCache ResourcePoolsCache;
 };
 
 } // namespace

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

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,

ydb/core/kqp/session_actor/kqp_session_actor.cpp

@@ -242,6 +242,12 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
     }
 
     void PassRequestToResourcePool() {
+        if (QueryState->UserRequestContext->PoolConfig) {
+            LOG_D("request placed into pool from cache: " << QueryState->UserRequestContext->PoolId);
+            CompileQuery();
+            return;
+        }
+
         Send(MakeKqpWorkloadServiceId(SelfId().NodeId()), new NWorkload::TEvPlaceRequestIntoPool(
             QueryState->Database,
             SessionId,

ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp

@@ -148,7 +148,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
 
         // Schemeboard events
         hFunc(TEvTxProxySchemeCache::TEvWatchNotifyUpdated, Handle);
-        IgnoreFunc(TEvTxProxySchemeCache::TEvWatchNotifyDeleted);
+        hFunc(TEvTxProxySchemeCache::TEvWatchNotifyDeleted, Handle);
         IgnoreFunc(TEvTxProxySchemeCache::TEvWatchNotifyUnavailable);
     )
 
@@ -157,6 +157,8 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
             this->Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvWatchRemove(0));
         }
 
+        SendPoolInfoUpdate(std::nullopt, std::nullopt);
+
         Counters.OnCleanup();
 
         TBase::PassAway();
@@ -261,6 +263,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:
@@ -326,6 +346,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;

ydb/core/kqp/workload_service/kqp_workload_service.cpp

@@ -412,7 +412,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() {
@@ -466,7 +466,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 {
@@ -538,6 +538,10 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
 
 }  // anonymous namespace
 
+bool IsWorkloadServiceRequired(const NResourcePool::TPoolSettings& config) {
+    return config.ConcurrentQueryLimit != -1 || config.DatabaseLoadCpuThreshold >= 0.0 || config.QueryCancelAfter;
+}
+
 }  // namespace NWorkload
 
 IActor* CreateKqpWorkloadService(NMonitoring::TDynamicCounterPtr counters) {