Better handling of active requests inside node state (#28825)

Author: abyss7

ydb/core/kqp/compute_actor/kqp_compute_actor.cpp

@@ -152,19 +152,17 @@ void TShardsScanningPolicy::FillRequestScanFeatures(const NKikimrTxDataShard::TK
 }
 
 TConclusionStatus TCPULimits::DeserializeFromProto(const NKikimrKqp::TEvStartKqpTasksRequest& config) {
+    const static auto maxThreadsCount = TActivationContext::ActorSystem()->GetPoolMaxThreadsCount(TActivationContext::AsActorContext().SelfID.PoolID());
     const auto share = config.GetPoolMaxCpuShare();
     if (share <= 0 || 1 < share) {
         return TConclusionStatus::Fail("cpu share have to be in (0, 1] interval");
     }
-    NActors::TExecutorPoolStats poolStats;
-    TVector<NActors::TExecutorThreadStats> threadsStats;
-    TActivationContext::ActorSystem()->GetPoolStats(TActivationContext::AsActorContext().SelfID.PoolID(), poolStats, threadsStats);
-    CPUGroupThreadsLimit = Max<ui64>(poolStats.MaxThreadCount, 1) * share;
+    CPUGroupThreadsLimit = Max<ui64>(1, maxThreadsCount) * share;
     CPUGroupName = config.GetPoolId();
     return TConclusionStatus::Success();
 }
 
-}
+} // namespace NKqp
 } // namespace NKikimr
 
 namespace NKikimr::NKqp {

ydb/core/kqp/compute_actor/kqp_compute_actor.h

@@ -56,9 +56,8 @@ class TCPULimits {
 };
 
 IActor* CreateKqpComputeActor(const TActorId& executerId, ui64 txId, NYql::NDqProto::TDqTask* task,
-    NYql::NDq::IDqAsyncIoFactory::TPtr asyncIoFactory,
-    const NYql::NDq::TComputeRuntimeSettings& settings, const NYql::NDq::TComputeMemoryLimits& memoryLimits,
-    NWilson::TTraceId traceId,
+    NYql::NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, const NYql::NDq::TComputeRuntimeSettings& settings,
+    const NYql::NDq::TComputeMemoryLimits& memoryLimits, NWilson::TTraceId traceId,
     TIntrusivePtr<NActors::TProtoArenaHolder> arena,
     const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup, const TGUCSettings::TPtr& GUCSettings,
     NScheduler::TSchedulableActorOptions schedulableOptions,

ydb/core/kqp/compute_actor/kqp_compute_actor_factory.cpp

@@ -2,6 +2,7 @@
 #include "kqp_compute_actor.h"
 
 #include <ydb/core/kqp/common/kqp_resolve.h>
+#include <ydb/core/kqp/node_service/kqp_node_state.h>
 #include <ydb/core/kqp/rm_service/kqp_resource_estimation.h>
 
 namespace NKikimr::NKqp::NComputeActor {
@@ -11,24 +12,17 @@ struct TMemoryQuotaManager : public NYql::NDq::TGuaranteeQuotaManager {
 
     TMemoryQuotaManager(std::shared_ptr<NRm::IKqpResourceManager> resourceManager
         , NRm::EKqpMemoryPool memoryPool
-        , std::shared_ptr<IKqpNodeState> state
         , TIntrusivePtr<NRm::TTxState> tx
         , TIntrusivePtr<NRm::TTaskState> task
         , ui64 limit)
     : NYql::NDq::TGuaranteeQuotaManager(limit, limit)
     , ResourceManager(std::move(resourceManager))
     , MemoryPool(memoryPool)
-    , State(std::move(state))
     , Tx(std::move(tx))
     , Task(std::move(task))
-    {
-    }
+    {}
 
     ~TMemoryQuotaManager() override {
-        if (State) {
-            State->OnTaskTerminate(Tx->TxId, Task->TaskId, Success);
-        }
-
         ResourceManager->FreeResources(Tx, Task);
     }
 
@@ -71,7 +65,6 @@ struct TMemoryQuotaManager : public NYql::NDq::TGuaranteeQuotaManager {
 
     std::shared_ptr<NRm::IKqpResourceManager> ResourceManager;
     NRm::EKqpMemoryPool MemoryPool;
-    std::shared_ptr<IKqpNodeState> State;
     TIntrusivePtr<NRm::TTxState> Tx;
     TIntrusivePtr<NRm::TTaskState> Task;
     bool Success = true;
@@ -169,7 +162,6 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
         memoryLimits.MemoryQuotaManager = std::make_shared<TMemoryQuotaManager>(
             ResourceManager_,
             args.MemoryPool,
-            std::move(args.State),
             std::move(args.TxInfo),
             std::move(task),
             limit);
@@ -193,12 +185,14 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
         }
 
         NYql::NDq::IMemoryQuotaManager::TWeakPtr memoryQuotaManager = memoryLimits.MemoryQuotaManager;
-        runtimeSettings.TerminateHandler = [memoryQuotaManager]
+        runtimeSettings.TerminateHandler = [memoryQuotaManager, state=args.State, txId=args.TxId, taskId=args.Task->GetId()]
             (bool success, const NYql::TIssues& issues) {
-                auto manager = memoryQuotaManager.lock();
-                if (manager) {
+                if (auto manager = memoryQuotaManager.lock()) {
                     static_cast<TMemoryQuotaManager*>(manager.get())->TerminateHandler(success, issues);
                 }
+                if (state) {
+                    state->OnTaskFinished(txId, taskId, success);
+                }
             };
 
         NKikimrTxDataShard::TKqpTransaction::TScanTaskMeta meta;
@@ -224,8 +218,7 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
             YQL_ENSURE(args.ComputesByStages);
             auto& info = args.ComputesByStages->UpsertTaskWithScan(*args.Task, meta);
             IActor* computeActor = CreateKqpScanComputeActor(
-                args.ExecuterId, args.TxId,
-                args.Task, AsyncIoFactory, runtimeSettings, memoryLimits,
+                args.ExecuterId, args.TxId, args.Task, AsyncIoFactory, runtimeSettings, memoryLimits,
                 std::move(args.TraceId), std::move(args.Arena),
                 std::move(schedulableOptions), args.BlockTrackingMode);
             TActorId result = TlsActivationContext->Register(computeActor);
@@ -236,8 +229,10 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
             if (!args.SerializedGUCSettings.empty()) {
                 GUCSettings = std::make_shared<TGUCSettings>(args.SerializedGUCSettings);
             }
-            IActor* computeActor = ::NKikimr::NKqp::CreateKqpComputeActor(args.ExecuterId, args.TxId, args.Task, AsyncIoFactory,
-                runtimeSettings, memoryLimits, std::move(args.TraceId), std::move(args.Arena), FederatedQuerySetup, GUCSettings,
+            IActor* computeActor = NKqp::CreateKqpComputeActor(
+                args.ExecuterId, args.TxId, args.Task, AsyncIoFactory, runtimeSettings, memoryLimits,
+                std::move(args.TraceId), std::move(args.Arena),
+                FederatedQuerySetup, GUCSettings,
                 std::move(schedulableOptions), args.BlockTrackingMode, std::move(args.UserToken), args.Database);
             return args.ShareMailbox ? TlsActivationContext->AsActorContext().RegisterWithSameMailbox(computeActor) :
                 TlsActivationContext->AsActorContext().Register(computeActor);

ydb/core/kqp/compute_actor/kqp_compute_actor_factory.h

@@ -11,6 +11,7 @@
 
 namespace NKikimr::NKqp {
     struct TKqpFederatedQuerySetup;
+    class TNodeState;
 }
 
 namespace NKikimr::NKqp::NComputeActor {
@@ -92,13 +93,6 @@ class TComputeStagesWithScan {
     }
 };
 
-struct IKqpNodeState {
-    virtual ~IKqpNodeState() = default;
-
-    virtual void OnTaskTerminate(ui64 txId, ui64 taskId, bool success) = 0;
-};
-
-
 struct IKqpNodeComputeActorFactory {
     virtual ~IKqpNodeComputeActorFactory() = default;
 
@@ -127,7 +121,7 @@ struct IKqpNodeComputeActorFactory {
         const NKikimrConfig::TTableServiceConfig::EBlockTrackingMode BlockTrackingMode;
 
         TComputeStagesWithScan* ComputesByStages = nullptr;
-        std::shared_ptr<IKqpNodeState> State = nullptr;
+        std::shared_ptr<TNodeState> State = nullptr;
         TIntrusiveConstPtr<NACLib::TUserToken> UserToken;
         TString Database;
 

ydb/core/kqp/compute_actor/kqp_scan_compute_actor.cpp

@@ -174,7 +174,6 @@ void TKqpScanComputeActor::Handle(TEvScanExchange::TEvTerminateFromFetcher::TPtr
     ALS_DEBUG(NKikimrServices::KQP_COMPUTE) << "TEvTerminateFromFetcher: " << ev->Sender << "/" << SelfId();
     TBase::InternalError(ev->Get()->GetStatusCode(), ev->Get()->GetIssues());
     State = ev->Get()->GetState();
-    DoTerminateImpl();
 }
 
 void TKqpScanComputeActor::Handle(TEvScanExchange::TEvSendData::TPtr& ev) {

ydb/core/kqp/compute_actor/kqp_scan_compute_actor.h

@@ -65,10 +65,6 @@ class TKqpScanComputeActor: public NScheduler::TSchedulableComputeActorBase<TKqp
     const EBlockTrackingMode BlockTrackingMode;
 
 public:
-    ~TKqpScanComputeActor() override {
-        DoTerminateImpl();
-    }
-
     static constexpr NKikimrServices::TActivity::EType ActorActivityType() {
         return NKikimrServices::TActivity::KQP_SCAN_COMPUTE_ACTOR;
     }