Properly handle all CA states and stats

Author: abyss7

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -2647,26 +2647,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;
-            auto& state = ev->Get()->Record;
-            ui64 taskId = state.GetTaskId();
-
-            if (Stats) {
-                Stats->AddComputeActorStats(
-                    actor.NodeId(),
-                    std::move(*state.MutableStats()),
-                    TDuration::MilliSeconds(AggregationSettings.GetCollectLongTasksStatsTimeoutMs())
-                );
-            }
-
-            Planner->CompletedCA(taskId, actor);
+        HandleComputeState(ev);
 
-            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) {