Fix flow control in run script actor (#11241)

Author: UgnineSirdis

ydb/core/grpc_services/query/rpc_execute_query.cpp

@@ -28,6 +28,25 @@ 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>();
+        resp->Record.SetSeqNo(*LastSeqNo);
+        resp->Record.SetFreeSpace(AckedFreeSpaceBytes);
+        resp->Record.SetChannelId(ChannelId);
+
+        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,
@@ -292,28 +311,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) {
@@ -334,19 +341,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& ctx) {
@@ -489,7 +492,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
     NKikimrKqp::EQueryAction QueryAction;
     TRpcFlowControlState FlowControl_;
     TMap<ui64, TProducerState> StreamChannels_;
-    
+
     NWilson::TSpan Span_;
 };
 

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);
     }