Consider inflight in full result writer (fix large memory consumption) YQL-17720

ydb/library/yql/providers/dq/actors/actors_ut.cpp

@@ -5,6 +5,7 @@
 #include <ydb/library/yql/providers/dq/actors/events.h>
 
 #include "result_receiver.h"
+#include "result_actor_base.h"
 
 using namespace NActors;
 using namespace NYql;
@@ -33,8 +34,8 @@ Y_UNIT_TEST(ReceiveStatus) {
     runtime.Initialize();
 
     auto sender = runtime.AllocateEdgeActor();
-    auto receiverId = runtime.Register(ResultReceiver().Release()); 
-    runtime.Send(new IEventHandle(receiverId, sender, new TEvReadyState(), 0, true)); 
+    auto receiverId = runtime.Register(ResultReceiver().Release());
+    runtime.Send(new IEventHandle(receiverId, sender, new TEvReadyState(), 0, true));
 }
 
 Y_UNIT_TEST(ReceiveError) {
@@ -49,5 +50,23 @@ Y_UNIT_TEST(ReceiveError) {
     UNIT_ASSERT_EQUAL(response->Record.GetStatusCode(), NYql::NDqProto::StatusIds::UNAVAILABLE);
 }
 
-} // Y_UNIT_TEST_SUITE(ResultReceiver) 
+Y_UNIT_TEST(WriteQueue) {
+    NYql::NDqs::NExecutionHelpers::TWriteQueue q;
+    UNIT_ASSERT(q.empty());
 
+    NYql::NDqs::NExecutionHelpers::TQueueItem item({}, ""); item.Size = 1000;
+    q.emplace(item);
+    UNIT_ASSERT_EQUAL(q.ByteSize, 1000);
+
+    item.Size = 11;
+    q.emplace(item);
+    UNIT_ASSERT_EQUAL(q.ByteSize, 1011);
+
+    q.pop();
+    UNIT_ASSERT_EQUAL(q.ByteSize, 11);
+
+    q.pop();
+    UNIT_ASSERT_EQUAL(q.ByteSize, 0);
+}
+
+} // Y_UNIT_TEST_SUITE(ResultReceiver)

ydb/library/yql/providers/dq/actors/result_actor_base.h

@@ -16,6 +16,64 @@
 
 namespace NYql::NDqs::NExecutionHelpers {
 
+struct TQueueItem {
+    TQueueItem(NDq::TDqSerializedBatch&& data, const TString& messageId)
+        : Data(std::move(data))
+        , MessageId(messageId)
+        , SentProcessedEvent(false)
+        , IsFinal(false)
+        , Size(Data.Size())
+        {
+        }
+
+    static TQueueItem Final() {
+        TQueueItem item({}, "FinalMessage");
+        item.SentProcessedEvent = true;
+        item.IsFinal = true;
+        return item;
+    }
+
+    NDq::TDqSerializedBatch Data;
+    const TString MessageId;
+    bool SentProcessedEvent = false;
+    bool IsFinal = false;
+    ui64 Size = 0;
+};
+
+struct TWriteQueue {
+    TQueue<TQueueItem> Queue;
+    ui64 ByteSize = 0;
+
+    template< class... Args >
+    decltype(auto) emplace( Args&&... args) {
+        Queue.emplace(std::forward<Args>(args)...);
+        ByteSize += Queue.back().Size;
+    }
+
+    auto& front() {
+        return Queue.front();
+    }
+
+    auto& back() {
+        return Queue.back();
+    }
+
+    auto pop() {
+        YQL_ENSURE(ByteSize >= Queue.front().Size);
+        ByteSize -= Queue.front().Size;
+        return Queue.pop();
+    }
+
+    auto empty() const {
+        return Queue.empty();
+    }
+
+    void clear() {
+        Queue.clear();
+        ByteSize = 0;
+    }
+};
+
     template <class TDerived>
     class TResultActorBase : public NYql::TSynchronizableRichActor<TDerived>, public NYql::TCounters {
     protected:
@@ -180,6 +238,10 @@ namespace NYql::NDqs::NExecutionHelpers {
             }
         }
 
+        ui64 InflightBytes() {
+            return WriteQueue.ByteSize;
+        }
+
     private:
         void OnQueryResult(TEvQueryResponse::TPtr& ev, const NActors::TActorContext&) {
             YQL_LOG_CTX_ROOT_SESSION_SCOPE(TraceId);
@@ -349,29 +411,6 @@ namespace NYql::NDqs::NExecutionHelpers {
             TBase::Send(FullResultWriterID, std::move(req));
         }
 
-    private:
-        struct TQueueItem {
-            TQueueItem(NDq::TDqSerializedBatch&& data, const TString& messageId)
-                : Data(std::move(data))
-                , MessageId(messageId)
-                , SentProcessedEvent(false)
-                , IsFinal(false)
-            {
-            }
-
-            static TQueueItem Final() {
-                TQueueItem item({}, "FinalMessage");
-                item.SentProcessedEvent = true;
-                item.IsFinal = true;
-                return item;
-            }
-
-            NDq::TDqSerializedBatch Data;
-            const TString MessageId;
-            bool SentProcessedEvent = false;
-            bool IsFinal = false;
-        };
-
     protected:
         const NActors::TActorId ExecuterID;
         const TString TraceId;
@@ -383,7 +422,7 @@ namespace NYql::NDqs::NExecutionHelpers {
         const bool FullResultTableEnabled;
         const NActors::TActorId GraphExecutionEventsId;
         const bool Discard;
-        TQueue<TQueueItem> WriteQueue;
+        TWriteQueue WriteQueue;
         ui64 SizeLimit;
         TMaybe<ui64> RowsLimit;
         ui64 Rows;

ydb/library/yql/providers/dq/actors/result_receiver.cpp

@@ -103,7 +103,7 @@ class TResultReceiver: public NYql::NDqs::NExecutionHelpers::TResultActorBase<TR
         auto req = MakeHolder<NDq::TEvDqCompute::TEvChannelDataAck>();
         req->Record.SetChannelId(message->Get()->Record.GetChannelData().GetChannelId());
         req->Record.SetSeqNo(message->Get()->Record.GetSeqNo());
-        req->Record.SetFreeSpace(256_MB);
+        req->Record.SetFreeSpace((i64)256_MB - (i64)InflightBytes());
         req->Record.SetFinish(EarlyFinish);  // set if premature finish started (when response limit reached and FullResultTable not enabled)
 
         Send(message->Sender, req.Release());