Merge b8d7260 into dbf5fb8

Author: abyss7

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -2608,6 +2608,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 this->Become(&TThis::WaitShutdownState);
                 LOG_I("Waiting for shutdown of " << Planner->GetPendingComputeTasks().size() << " tasks and "
                     << Planner->GetPendingComputeActors().size() << " compute actors");
+                // TODO(ilezhankin): the CA awaiting timeout should be configurable.
                 TActivationContext::Schedule(TDuration::Seconds(10), new IEventHandle(SelfId(), SelfId(), new TEvents::TEvPoison));
             }
         } else {
@@ -2643,17 +2644,10 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
     }
 
     void HandleShutdown(TEvDqCompute::TEvState::TPtr& ev) {
-        if (ev->Get()->Record.GetState() == NDqProto::COMPUTE_STATE_FAILURE) {
-            YQL_ENSURE(Planner);
-
-            TActorId actor = ev->Sender;
-            ui64 taskId = ev->Get()->Record.GetTaskId();
+        HandleComputeState(ev);
 
-            Planner->CompletedCA(taskId, actor);
-
-            if (Planner->GetPendingComputeTasks().empty() && Planner->GetPendingComputeActors().empty()) {
-                PassAway();
-            }
+        if (Planner->GetPendingComputeTasks().empty() && Planner->GetPendingComputeActors().empty()) {
+            PassAway();
         }
     }
 

ydb/core/kqp/executer_actor/kqp_executer_impl.h

@@ -380,7 +380,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         this->Send(channelComputeActorId, ackEv.Release(), /* TODO: undelivery */ 0, /* cookie */ channelId);
     }
 
-    void HandleComputeStats(NYql::NDq::TEvDqCompute::TEvState::TPtr& ev) {
+    bool HandleComputeState(NYql::NDq::TEvDqCompute::TEvState::TPtr& ev) {
         TActorId computeActor = ev->Sender;
         auto& state = ev->Get()->Record;
         ui64 taskId = state.GetTaskId();
@@ -409,27 +409,13 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         }
 
         YQL_ENSURE(Planner);
-        bool populateChannels = Planner->AcknowledgeCA(taskId, computeActor, &state);
+        bool ack = Planner->AcknowledgeCA(taskId, computeActor, &state);
 
         switch (state.GetState()) {
-            case NYql::NDqProto::COMPUTE_STATE_UNKNOWN: {
-                YQL_ENSURE(false, "unexpected state from " << computeActor << ", task: " << taskId);
-                return;
-            }
-
-            case NYql::NDqProto::COMPUTE_STATE_EXECUTING: {
-                if (populateChannels) {
-                    auto& task = TasksGraph.GetTask(taskId);
-                    THashMap<TActorId, THashSet<ui64>> updates;
-                    CollectTaskChannelsUpdates(task, updates);
-                    PropagateChannelsUpdates(updates);
-                }
-                break;
-            }
-
             case NYql::NDqProto::COMPUTE_STATE_FAILURE:
             case NYql::NDqProto::COMPUTE_STATE_FINISHED: {
                 ExtraData[computeActor].Swap(state.MutableExtraData());
+
                 if (Stats) {
                     Stats->AddComputeActorStats(
                         computeActor.NodeId(),
@@ -440,9 +426,41 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
 
                 LastTaskId = taskId;
                 LastComputeActorId = computeActor.ToString();
-                YQL_ENSURE(Planner);
+
                 Planner->CompletedCA(taskId, computeActor);
             }
+            default:
+                ; // ignore all other states.
+        }
+
+        return ack;
+    }
+
+    void HandleComputeStats(NYql::NDq::TEvDqCompute::TEvState::TPtr& ev) {
+        TActorId computeActor = ev->Sender;
+        auto& state = ev->Get()->Record;
+        ui64 taskId = state.GetTaskId();
+
+        bool populateChannels = HandleComputeState(ev);
+
+        switch (state.GetState()) {
+            case NYql::NDqProto::COMPUTE_STATE_UNKNOWN: {
+                YQL_ENSURE(false, "unexpected state from " << computeActor << ", task: " << taskId);
+                return;
+            }
+
+            case NYql::NDqProto::COMPUTE_STATE_EXECUTING: {
+                if (populateChannels) {
+                    auto& task = TasksGraph.GetTask(taskId);
+                    THashMap<TActorId, THashSet<ui64>> updates;
+                    CollectTaskChannelsUpdates(task, updates);
+                    PropagateChannelsUpdates(updates);
+                }
+                break;
+            }
+
+            default:
+                ; // ignore all other states.
         }
 
         if (state.GetState() == NYql::NDqProto::COMPUTE_STATE_FAILURE) {
@@ -1854,6 +1872,9 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
 
     void PassAway() override {
         YQL_ENSURE(AlreadyReplied && ResponseEv);
+
+        // Actualize stats with the last stats from terminated CAs, but keep the status.
+        FillResponseStats(ResponseEv->Record.GetResponse().GetStatus());
         this->Send(Target, ResponseEv.release());
 
         for (auto channelPair: ResultChannelProxies) {