One more fix

Author: abyss7

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -2601,6 +2601,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
     void Shutdown() override {
         if (Planner) {
             if (Planner->GetPendingComputeTasks().empty() && Planner->GetPendingComputeActors().empty()) {
+                LOG_I("Shutdown immediately - nothing to wait");
                 PassAway();
             } else {
                 this->Become(&TThis::WaitShutdownState);
@@ -2641,15 +2642,17 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
     }
 
     void HandleShutdown(TEvDqCompute::TEvState::TPtr& ev) {
-        YQL_ENSURE(Planner);
+        if (ev->Get()->Record.GetState() == NDqProto::COMPUTE_STATE_FAILURE) {
+            YQL_ENSURE(Planner);
 
-        TActorId actor = ev->Sender;
-        ui64 taskId = ev->Get()->Record.GetTaskId();
+            TActorId actor = ev->Sender;
+            ui64 taskId = ev->Get()->Record.GetTaskId();
 
-        Planner->CompletedCA(taskId, actor);
+            Planner->CompletedCA(taskId, actor);
 
-        if (Planner->GetPendingComputeTasks().empty() && Planner->GetPendingComputeActors().empty()) {
-            PassAway();
+            if (Planner->GetPendingComputeTasks().empty() && Planner->GetPendingComputeActors().empty()) {
+                PassAway();
+            }
         }
     }
 

ydb/core/kqp/ut/query/kqp_limits_ut.cpp

@@ -860,14 +860,17 @@ Y_UNIT_TEST_SUITE(KqpLimits) {
         auto& runtime = *kikimr.GetTestServer().GetRuntime();
         runtime.SetObserverFunc([&](TAutoPtr<IEventHandle>& ev) {
             if (ev->GetTypeRewrite() == NYql::NDq::TEvDqCompute::TEvState::EventType) {
+                ++totalEvState;
                 if (!firstEvState) {
                     executerId = ev->Recipient;
                     ev = new IEventHandle(ev->Recipient, ev->Sender,
                             new NKikimr::NKqp::TEvKqp::TEvAbortExecution(NYql::NDqProto::StatusIds::UNSPECIFIED, NYql::TIssues()));
                     firstEvState = true;
                 }
-                ++totalEvState;
+            } else if (ev->GetTypeRewrite() == NKikimr::NKqp::TEvKqpExecuter::TEvTxResponse::EventType && ev->Sender == executerId) {
+                UNIT_ASSERT_C(totalEvState == actorCount*2, "Executer sent response before waiting for CAs");
             }
+
             return TTestActorRuntime::EEventAction::PROCESS;
         });
 
@@ -912,21 +915,21 @@ Y_UNIT_TEST_SUITE(KqpLimits) {
         auto& runtime = *kikimr.GetTestServer().GetRuntime();
         runtime.SetObserverFunc([&](TAutoPtr<IEventHandle>& ev) {
             if (ev->GetTypeRewrite() == NYql::NDq::TEvDqCompute::TEvState::EventType) {
+                ++totalEvState;
                 if (!firstEvState) {
                     executerId = ev->Recipient;
                     ev = new IEventHandle(ev->Recipient, ev->Sender,
                             new NKikimr::NKqp::TEvKqp::TEvAbortExecution(NYql::NDqProto::StatusIds::UNSPECIFIED, NYql::TIssues()));
                     firstEvState = true;
-                }
-                ++totalEvState;
-
-                if (totalEvState == actorCount*2) {
+                } else {
                     return TTestActorRuntime::EEventAction::DROP;
                 }
             } else if (ev->GetTypeRewrite() == TEvents::TEvPoison::EventType && totalEvState == actorCount*2 &&
                 ev->Sender == executerId && ev->Recipient == executerId)
             {
                 timeoutPoison = true;
+            } else if (ev->GetTypeRewrite() == NKikimr::NKqp::TEvKqpExecuter::TEvTxResponse::EventType && ev->Sender == executerId) {
+                UNIT_ASSERT_C(timeoutPoison, "Executer sent response before waiting for CAs");
             }
 
             return TTestActorRuntime::EEventAction::PROCESS;