Merge result changes 24 4

ydb/core/grpc_services/query/rpc_execute_query.cpp

@@ -11,6 +11,8 @@
 #include <ydb/public/api/protos/ydb_query.pb.h>
 
 #include <ydb/library/actors/core/actor_bootstrapped.h>
+#include <ydb/library/wilson_ids/wilson.h>
+
 
 namespace NKikimr::NGRpcService {
 
@@ -25,6 +27,23 @@ struct TProducerState {
     TMaybe<ui64> LastSeqNo;
     i64 AckedFreeSpaceBytes = 0;
     TActorId ActorId;
+    ui64 ChannelId = 0;
+
+    void SendAck(const NActors::TActorIdentity& actor) const {
+        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(*LastSeqNo, ChannelId);
+        resp->Record.SetFreeSpace(AckedFreeSpaceBytes);
+
+        actor.Send(ActorId, resp.Release());
+    }
+
+    bool ResumeIfStopped(const NActors::TActorIdentity& actor, i64 freeSpaceBytes) {
+        if (LastSeqNo && AckedFreeSpaceBytes <= 0) {
+            AckedFreeSpaceBytes = freeSpaceBytes;
+            SendAck(actor);
+            return true;
+        }
+        return false;
+    }
 };
 
 bool FillTxSettings(const Ydb::Query::TransactionSettings& from, Ydb::Table::TransactionSettings& to,
@@ -163,7 +182,9 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
 
     TExecuteQueryRPC(TEvExecuteQueryRequest* request, ui64 inflightLimitBytes)
         : Request_(request)
-        , FlowControl_(inflightLimitBytes) {}
+        , FlowControl_(inflightLimitBytes)
+        , Span_(TWilsonGrpc::RequestActor, request->GetWilsonTraceId(),
+                "RequestProxy.RpcOperationRequestActor", NWilson::EFlags::AUTO_END) {}
 
     void Bootstrap(const TActorContext &ctx) {
         this->Become(&TExecuteQueryRPC::StateWork);
@@ -255,7 +276,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
             settings,
             req->pool_id());
 
-        if (!ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release())) {
+        if (!ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release(), 0, 0, Span_.GetTraceId())) {
             NYql::TIssues issues;
             issues.AddIssue(MakeIssue(NKikimrIssues::TIssuesIds::DEFAULT_ERROR, "Internal error"));
             ReplyFinishStream(Ydb::StatusIds::INTERNAL_ERROR, std::move(issues));
@@ -286,28 +307,16 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
         }
 
         const i64 freeSpaceBytes = FlowControl_.FreeSpaceBytes();
-
-        for (auto& pair : StreamChannels_) {
-            const auto& channelId = pair.first;
-            auto& channel = pair.second;
-
-            if (freeSpaceBytes > 0 && channel.LastSeqNo && channel.AckedFreeSpaceBytes <= 0) {
-                LOG_DEBUG_S(ctx, NKikimrServices::RPC_REQUEST, this->SelfId() << "Resume execution, "
-                    << ", channel: " << channelId
-                    << ", seqNo: " << channel.LastSeqNo
-                    << ", freeSpace: " << freeSpaceBytes);
-
-                auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
-                resp->Record.SetSeqNo(*channel.LastSeqNo);
-                resp->Record.SetFreeSpace(freeSpaceBytes);
-                resp->Record.SetChannelId(channelId);
-
-                ctx.Send(channel.ActorId, resp.Release());
-
-                channel.AckedFreeSpaceBytes = freeSpaceBytes;
+        if (freeSpaceBytes > 0) {
+            for (auto& [channelId, channel] : StreamChannels_) {
+                if (channel.ResumeIfStopped(SelfId(), freeSpaceBytes)) {
+                    LOG_DEBUG_S(ctx, NKikimrServices::RPC_REQUEST, this->SelfId() << "Resume execution, "
+                        << ", channel: " << channelId
+                        << ", seqNo: " << channel.LastSeqNo
+                        << ", freeSpace: " << freeSpaceBytes);
+                }
             }
         }
-
     }
 
     void Handle(NKqp::TEvKqpExecuter::TEvStreamData::TPtr& ev, const TActorContext& ctx) {
@@ -328,19 +337,15 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
         channel.ActorId = ev->Sender;
         channel.LastSeqNo = ev->Get()->Record.GetSeqNo();
         channel.AckedFreeSpaceBytes = freeSpaceBytes;
+        channel.ChannelId = ev->Get()->Record.GetChannelId();
 
         LOG_DEBUG_S(ctx, NKikimrServices::RPC_REQUEST, this->SelfId() << "Send stream data ack"
             << ", seqNo: " << ev->Get()->Record.GetSeqNo()
             << ", freeSpace: " << freeSpaceBytes
             << ", to: " << ev->Sender
             << ", queue: " << FlowControl_.QueueSize());
 
-        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
-        resp->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
-        resp->Record.SetFreeSpace(freeSpaceBytes);
-        resp->Record.SetChannelId(ev->Get()->Record.GetChannelId());
-
-        ctx.Send(channel.ActorId, resp.Release());
+        channel.SendAck(SelfId());
     }
 
     void Handle(NKqp::TEvKqp::TEvQueryResponse::TPtr& ev, const TActorContext&) {
@@ -415,6 +420,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
     void ReplySerializedAndFinishStream(Ydb::StatusIds::StatusCode status, TString&& buf) {
         const auto finishStreamFlag = NYdbGrpc::IRequestContextBase::EStreamCtrl::FINISH;
         Request_->SendSerializedResult(std::move(buf), status, finishStreamFlag);
+        NWilson::EndSpanWithStatus(Span_, status);
         this->PassAway();
     }
 
@@ -453,6 +459,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
         } else {
             Request_->FinishStream(status);
         }
+        NWilson::EndSpanWithStatus(Span_, status);
         this->PassAway();
     }
 
@@ -479,6 +486,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
     NKikimrKqp::EQueryAction QueryAction;
     TRpcFlowControlState FlowControl_;
     TMap<ui64, TProducerState> StreamChannels_;
+
+    NWilson::TSpan Span_;
 };
 
 } // namespace

ydb/core/grpc_services/rpc_stream_execute_scan_query.cpp

@@ -268,12 +268,10 @@ class TStreamExecuteScanQueryRPC : public TActorBootstrapped<TStreamExecuteScanQ
                 << ", freeSpace: " << freeSpaceBytes
                 << ", to: " << ExecuterActorId_);
 
-            auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
-            resp->Record.SetSeqNo(*LastSeqNo_);
+            // scan query has single result set, so it's ok to put zero as channelId here.
+            auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(*LastSeqNo_, 0);
             resp->Record.SetFreeSpace(freeSpaceBytes);
-
             ctx.Send(ExecuterActorId_, resp.Release());
-
             AckedFreeSpaceBytes_ = freeSpaceBytes;
         }
     }
@@ -326,6 +324,7 @@ class TStreamExecuteScanQueryRPC : public TActorBootstrapped<TStreamExecuteScanQ
     }
 
     void Handle(NKqp::TEvKqp::TEvAbortExecution::TPtr& ev, const TActorContext& ctx) {
+
         auto& record = ev->Get()->Record;
         NYql::TIssues issues = ev->Get()->GetIssues();
 
@@ -361,8 +360,7 @@ class TStreamExecuteScanQueryRPC : public TActorBootstrapped<TStreamExecuteScanQ
             << ", to: " << ev->Sender
             << ", queue: " << FlowControl_.QueueSize());
 
-        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
-        resp->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
+        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(ev->Get()->Record.GetSeqNo(), ev->Get()->Record.GetChannelId());
         resp->Record.SetFreeSpace(freeSpaceBytes);
 
         ctx.Send(ev->Sender, resp.Release());

ydb/core/grpc_services/rpc_stream_execute_yql_script.cpp

@@ -279,8 +279,7 @@ class TStreamExecuteYqlScriptRPC
             << ", to: " << ev->Sender
             << ", queue: " << FlowControl_.QueueSize());
 
-        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
-        resp->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
+        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(ev->Get()->Record.GetSeqNo(), ev->Get()->Record.GetChannelId());
         resp->Record.SetFreeSpace(freeSpaceBytes);
 
         ctx.Send(ev->Sender, resp.Release());
@@ -320,8 +319,7 @@ class TStreamExecuteYqlScriptRPC
                     << ", freeSpace: " << freeSpaceBytes
                     << ", to: " << GatewayRequestHandlerActorId_);
 
-                auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
-                resp->Record.SetSeqNo(*LastSeqNo_);
+                auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(*LastSeqNo_, 0);
                 resp->Record.SetFreeSpace(freeSpaceBytes);
 
                 ctx.Send(GatewayRequestHandlerActorId_, resp.Release());

ydb/core/kqp/executer_actor/kqp_executer.h

@@ -70,7 +70,19 @@ struct TEvKqpExecuter {
         TKqpExecuterEvents::EvStreamData> {};
 
     struct TEvStreamDataAck : public TEventPB<TEvStreamDataAck, NKikimrKqp::TEvExecuterStreamDataAck,
-        TKqpExecuterEvents::EvStreamDataAck> {};
+        TKqpExecuterEvents::EvStreamDataAck>
+    {
+        friend class TEventPBBase;
+        explicit TEvStreamDataAck(ui64 seqno, ui64 channelId)
+        {
+            Record.SetSeqNo(seqno);
+            Record.SetChannelId(channelId);
+        }
+
+    private:
+        // using a little hack to hide default empty constructor
+        TEvStreamDataAck() = default;
+    };
 
     struct TEvStreamProfile : public TEventPB<TEvStreamProfile, NKikimrKqp::TEvExecuterStreamProfile,
         TKqpExecuterEvents::EvStreamProfile> {};

ydb/core/kqp/gateway/kqp_ic_gateway.cpp

@@ -202,9 +202,8 @@ class TKqpScanQueryRequestHandler : public TRequestHandlerBase<
             ResultSet.set_truncated(true);
         }
 
-        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
+        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(ev->Get()->Record.GetSeqNo(), ev->Get()->Record.GetChannelId());
         resp->Record.SetEnough(truncated);
-        resp->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
         resp->Record.SetFreeSpace(ResultSetBytesLimit);
         ctx.Send(ev->Sender, resp.Release());
     }
@@ -497,8 +496,7 @@ class TKqpGenericQueryRequestHandler: public TRequestHandlerBase<
             }
         }
 
-        auto response = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
-        response->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
+        auto response = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(ev->Get()->Record.GetSeqNo(), ev->Get()->Record.GetChannelId());
         response->Record.SetFreeSpace(SizeLimit && SizeLimit < std::numeric_limits<i64>::max() ? SizeLimit : std::numeric_limits<i64>::max());
         Send(ev->Sender, response.Release());
     }

ydb/core/kqp/proxy_service/kqp_script_executions.cpp

@@ -1032,7 +1032,7 @@ class TForgetScriptExecutionOperationActor : public TActorBootstrapped<TForgetSc
     void Reply(Ydb::StatusIds::StatusCode status, NYql::TIssues issues = {}) {
         if (!ExecutionEntryExists && status == Ydb::StatusIds::SUCCESS) {
             status = Ydb::StatusIds::NOT_FOUND;
-            issues.AddIssue("No such execution");   
+            issues.AddIssue("No such execution");
         }
 
         if (status == Ydb::StatusIds::SUCCESS) {
@@ -1796,26 +1796,26 @@ class TSaveScriptExecutionResultQuery : public TQueryBase {
             .AddParam("$items");
 
         param
-                .BeginList();
+            .BeginList();
 
         auto row = FirstRow;
         for (const auto& rowValue : ResultSet.rows()) {
             auto rowValueSerialized = rowValue.SerializeAsString();
             SavedSize += rowValueSerialized.size();
             param
-                    .AddListItem()
-                    .BeginStruct()
-                        .AddMember("row_id")
-                            .Int64(row++)
-                        .AddMember("result_set")
-                            .String(std::move(rowValueSerialized))
-                        .AddMember("accumulated_size")
-                            .Int64(AccumulatedSize + SavedSize)
-                    .EndStruct();
+                .AddListItem()
+                .BeginStruct()
+                    .AddMember("row_id")
+                        .Int64(row++)
+                    .AddMember("result_set")
+                        .String(std::move(rowValueSerialized))
+                    .AddMember("accumulated_size")
+                        .Int64(AccumulatedSize + SavedSize)
+                .EndStruct();
         }
         param
-                .EndList()
-                .Build();
+            .EndList()
+            .Build();
 
         RunDataQuery(sql, &params);
     }