Fix participant nodes count for sinks (#16137)

Author: nikvas0

.github/config/muted_ya.txt

@@ -27,8 +27,6 @@ ydb/core/kqp/ut/query KqpAnalyze.AnalyzeTable+ColumnStore
 ydb/core/kqp/ut/query KqpAnalyze.AnalyzeTable-ColumnStore
 ydb/core/kqp/ut/query KqpLimits.OutOfSpaceYQLUpsertFail+useSink
 ydb/core/kqp/ut/query KqpLimits.QSReplySizeEnsureMemoryLimits+useSink
-ydb/core/kqp/ut/query KqpStats.OneShardLocalExec+UseSink
-ydb/core/kqp/ut/query KqpStats.OneShardNonLocalExec+UseSink
 ydb/core/kqp/ut/query KqpStats.SysViewClientLost
 ydb/core/kqp/ut/scheme KqpOlapScheme.TenThousandColumns
 ydb/core/kqp/ut/scheme KqpScheme.AlterAsyncReplication

ydb/core/kqp/common/kqp_tx_manager.cpp

@@ -162,6 +162,14 @@ class TKqpTransactionManager : public IKqpTransactionManager {
         return nullptr;
     }
 
+    void AddParticipantNode(const ui32 nodeId) override {
+        ParticipantNodes.insert(nodeId);
+    }
+
+    const THashSet<ui32>& GetParticipantNodes() const override {
+        return ParticipantNodes;
+    }
+
     void SetTopicOperations(NTopic::TTopicOperations&& topicOperations) override {
         TopicOperations = std::move(topicOperations);
     }
@@ -506,6 +514,8 @@ class TKqpTransactionManager : public IKqpTransactionManager {
     THashMap<ui64, TShardInfo> ShardsInfo;
     std::unordered_set<TString> TablePathes;
 
+    THashSet<ui32> ParticipantNodes;
+
     THashMap<TTableId, std::shared_ptr<const TVector<TKeyDesc::TPartitionInfo>>> TablePartitioning;
 
     bool AllowVolatile = false;

ydb/core/kqp/common/kqp_tx_manager.h

@@ -71,6 +71,9 @@ class IKqpTransactionManager {
     virtual void BuildTopicTxs(NTopic::TTopicOperationTransactions& txs) = 0;
     virtual bool HasTopics() const = 0;
 
+    virtual void AddParticipantNode(const ui32 nodeId) = 0;
+    virtual const THashSet<ui32>& GetParticipantNodes() const = 0;
+
     virtual bool IsTxPrepared() const = 0;
     virtual bool IsTxFinished() const = 0;
 

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -1277,6 +1277,10 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
             Stats->AddDatashardStats(std::move(*res->Record.MutableTxStats()));
         }
 
+        if (TxManager) {
+            TxManager->AddParticipantNode(ev->Sender.NodeId());
+        }
+
         switch (ev->Get()->GetStatus()) {
             case NKikimrDataEvents::TEvWriteResult::STATUS_UNSPECIFIED: {
                 YQL_ENSURE(false);

ydb/core/kqp/executer_actor/kqp_executer_impl.h

@@ -265,6 +265,9 @@ class TKqpExecuterBase : public TActor<TDerived> {
         for (auto& [shardId, nodeId] : ShardIdToNodeId) {
             ShardsOnNode[nodeId].push_back(shardId);
             ParticipantNodes.emplace(nodeId);
+            if (TxManager) {
+                TxManager->AddParticipantNode(nodeId);
+            }
         }
 
         if (IsDebugLogEnabled()) {

ydb/core/kqp/runtime/kqp_write_actor.cpp

@@ -548,6 +548,8 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {
             << ", Cookie=" << ev->Cookie);
         UpdateStats(ev->Get()->Record.GetTxStats());
 
+        TxManager->AddParticipantNode(ev->Sender.NodeId());
+
         switch (ev->Get()->GetStatus()) {
         case NKikimrDataEvents::TEvWriteResult::STATUS_UNSPECIFIED: {
             CA_LOG_E("Got UNSPECIFIED for table `"
@@ -2349,6 +2351,8 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
             return (tableInfo.Pathes.size() == 1 ? "Table: " : "Tables: ")  + builder;
         };
 
+        TxManager->AddParticipantNode(ev->Sender.NodeId());
+
         // TODO: get rid of copy-paste
         switch (ev->Get()->GetStatus()) {
         case NKikimrDataEvents::TEvWriteResult::STATUS_UNSPECIFIED: {

ydb/core/kqp/session_actor/kqp_query_state.h

@@ -176,12 +176,16 @@ class TKqpQueryState : public TNonCopyable {
         if (RequestEv->GetRequestCtx() == nullptr) {
             return false;
         }
-        if (ParticipantNodes.size() == 1) {
+        if (IsSingleNodeExecution()) {
             return *ParticipantNodes.begin() == nodeId;
         }
         return false;
     }
 
+    bool IsSingleNodeExecution() const {
+        return ParticipantNodes.size() == 1;
+    }
+
     NKikimrKqp::EQueryAction GetAction() const {
         return QueryAction;
     }

ydb/core/kqp/session_actor/kqp_session_actor.cpp

@@ -1627,8 +1627,12 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             QueryState->QueryStats.Executions.back().Swap(executerResults.MutableStats());
         }
 
-        for (auto nodeId : ev->ParticipantNodes) {
-            QueryState->ParticipantNodes.emplace(nodeId);
+        if (QueryState->TxCtx->TxManager) {
+            QueryState->ParticipantNodes = QueryState->TxCtx->TxManager->GetParticipantNodes();
+        } else {
+            for (auto nodeId : ev->ParticipantNodes) {
+                QueryState->ParticipantNodes.emplace(nodeId);
+            }
         }
 
         if (response->GetStatus() != Ydb::StatusIds::SUCCESS) {
@@ -2377,7 +2381,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             QueryState->PoolHandlerActor = Nothing();
         }
 
-        if (QueryState && QueryState->ParticipantNodes.size() == 1) {
+        if (QueryState && QueryState->IsSingleNodeExecution()) {
             Counters->TotalSingleNodeReqCount->Inc();
             if (!QueryState->IsLocalExecution(SelfId().NodeId())) {
                 Counters->NonLocalSingleNodeReqCount->Inc();

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

@@ -685,6 +685,8 @@ Y_UNIT_TEST_TWIN(OneShardLocalExec, UseSink) {
     auto session = db.CreateSession().GetValueSync().GetSession();
 
     TKqpCounters counters(kikimr.GetTestServer().GetRuntime()->GetAppData().Counters);
+
+    UNIT_ASSERT_VALUES_EQUAL(counters.TotalSingleNodeReqCount->Val(), 1);
     {
         auto result = session.ExecuteDataQuery(R"(
             SELECT * FROM `/Root/KeyValue` WHERE Key = 1;