fix potential hang if no new reads are started in the new data batch (#26382)

Co-authored-by: Copilot <[email protected]>

Author: 2 people

ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp

@@ -305,8 +305,12 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
         ReadRowsCount += replyResultStats.ReadRowsCount;
         ReadBytesCount += replyResultStats.ReadBytesCount;
 
-        if (!StreamLookupWorker->IsOverloaded()) {
+        auto overloaded = StreamLookupWorker->IsOverloaded();
+        if (!overloaded.has_value()) {
             FetchInputRows();
+        } else {
+            CA_LOG_N("Pausing stream lookup because it's overloaded by reason: "
+                << overloaded.value_or("empty"));
         }
 
         if (Partitioning) {
@@ -316,9 +320,10 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
         const bool inputRowsFinished = LastFetchStatus == NUdf::EFetchStatus::Finish;
         const bool allReadsFinished = AllReadsFinished();
         const bool allRowsProcessed = StreamLookupWorker->AllRowsProcessed();
+        const bool hasPendingResults = StreamLookupWorker->HasPendingResults();
 
-        if (inputRowsFinished && allReadsFinished && !allRowsProcessed) {
-            // all reads are completed, but we have unprocessed rows
+        if (hasPendingResults) {
+            // has more results
             Send(ComputeActorId, new TEvNewAsyncInputDataArrived(InputIndex));
         }
 

ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp

@@ -303,6 +303,10 @@ class TKqpLookupRows : public TKqpStreamLookupWorker {
         return readRequests;
     }
 
+    bool HasPendingResults() final {
+        return !ReadResults.empty();
+    }
+
     void AddResult(TShardReadResult result) final {
         const auto& record = result.ReadResult->Get()->Record;
         YQL_ENSURE(record.GetStatus().GetCode() == Ydb::StatusIds::SUCCESS);
@@ -317,8 +321,8 @@ class TKqpLookupRows : public TKqpStreamLookupWorker {
         ReadResults.emplace_back(std::move(result));
     }
 
-    bool IsOverloaded() final {
-        return false;
+    std::optional<TString> IsOverloaded() final {
+        return std::nullopt;
     }
 
     TReadResultStats ReplyResult(NKikimr::NMiniKQL::TUnboxedValueBatch& batch, i64 freeSpace) final {
@@ -585,8 +589,20 @@ class TKqpJoinRows : public TKqpStreamLookupWorker {
         YQL_ENSURE(false);
     }
 
-    bool IsOverloaded() final {
-        return UnprocessedRows.size() >= MAX_IN_FLIGHT_LIMIT || PendingLeftRowsByKey.size() >= MAX_IN_FLIGHT_LIMIT || ResultRowsBySeqNo.size() >= MAX_IN_FLIGHT_LIMIT;
+    std::optional<TString> IsOverloaded() final {
+        if (UnprocessedRows.size() >= MAX_IN_FLIGHT_LIMIT ||
+            PendingLeftRowsByKey.size() >= MAX_IN_FLIGHT_LIMIT ||
+            ResultRowsBySeqNo.size() >= MAX_IN_FLIGHT_LIMIT)
+        {
+            TStringBuilder overloadDescriptor;
+            overloadDescriptor << "unprocessed rows: " << UnprocessedRows.size()
+                << ", pending left: " << PendingLeftRowsByKey.size()
+                << ", result rows: " << ResultRowsBySeqNo.size();
+
+            return TString(overloadDescriptor);
+        }
+
+        return std::nullopt;
     }
 
     std::vector<THolder<TEvDataShard::TEvRead>> RebuildRequest(const ui64& prevReadId, ui64& newReadId) final {
@@ -778,6 +794,16 @@ class TKqpJoinRows : public TKqpStreamLookupWorker {
         return requests;
     }
 
+    bool HasPendingResults() final {
+        auto nextSeqNo = KeepRowsOrder() ? ResultRowsBySeqNo.find(CurrentResultSeqNo) : ResultRowsBySeqNo.begin();
+
+        if (nextSeqNo != ResultRowsBySeqNo.end() && !nextSeqNo->second.Rows.empty()) {
+            return true;
+        }
+
+        return false;
+    }
+
     void AddResult(TShardReadResult result) final {
         const auto& record = result.ReadResult->Get()->Record;
         YQL_ENSURE(record.GetStatus().GetCode() == Ydb::StatusIds::SUCCESS);
@@ -880,7 +906,7 @@ class TKqpJoinRows : public TKqpStreamLookupWorker {
         batch.clear();
 
         auto getNextResult = [&]() {
-            if (!ShoulKeepRowsOrder()) {
+            if (!KeepRowsOrder()) {
                 return ResultRowsBySeqNo.begin();
             }
 
@@ -1063,12 +1089,12 @@ class TKqpJoinRows : public TKqpStreamLookupWorker {
         }
     };
 
-    bool ShoulKeepRowsOrder() const {
+    bool KeepRowsOrder() const {
         return Settings.KeepRowsOrder;
     }
 
     bool IsRowSeqNoValid(const ui64& seqNo) const {
-        if (!ShoulKeepRowsOrder()) {
+        if (!KeepRowsOrder()) {
             return true;
         }
 

ydb/core/kqp/runtime/kqp_stream_lookup_worker.h

@@ -78,8 +78,9 @@ class TKqpStreamLookupWorker {
     virtual TReadResultStats ReplyResult(NKikimr::NMiniKQL::TUnboxedValueBatch& batch, i64 freeSpace) = 0;
     virtual TReadResultStats ReadAllResult(std::function<void(TConstArrayRef<TCell>)> reader) = 0;
     virtual bool AllRowsProcessed() = 0;
+    virtual bool HasPendingResults() = 0;
     virtual void ResetRowsProcessing(ui64 readId) = 0;
-    virtual bool IsOverloaded() = 0;
+    virtual std::optional<TString> IsOverloaded() = 0;
 
 protected:
     const NMiniKQL::TTypeEnvironment& TypeEnv;