Refactor write actor buffers (#19300)

Author: nikvas0

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -266,10 +266,6 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
 
         if (TxManager) {
             TxManager->SetHasSnapshot(GetSnapshot().IsValid());
-
-            for (const ui64& shardId : TxManager->GetShards()) {
-                Stats->AffectedShards.insert(shardId);
-            }
         }
 
         if (!BufferActorId || (ReadOnlyTx && Request.LocksOp != ELocksOp::Rollback)) {
@@ -392,6 +388,12 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
             }
         }
 
+        if (TxManager) {
+            for (const ui64& shardId : TxManager->GetShards()) {
+                Stats->AffectedShards.insert(shardId);
+            }
+        }
+
         auto resultSize = ResponseEv->GetByteSize();
         if (resultSize > (int)ReplySizeLimit) {
             TString message;

ydb/core/kqp/runtime/kqp_write_actor.cpp

@@ -1495,6 +1495,7 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu
             WriteTableActor->Write(WriteToken, Batcher->Build());
             if (Closed) {
                 WriteTableActor->Close(WriteToken);
+                WriteTableActor->FlushBuffers();
                 WriteTableActor->Close();
             }
         } catch (const TMemoryLimitExceededException&) {
@@ -1536,6 +1537,10 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu
                 return;
             }
 
+            if (!Closed && outOfMemory) {
+                WriteTableActor->FlushBuffers();
+            }
+
             if (Closed || outOfMemory) {
                 if (!WriteTableActor->FlushToShards()) {
                     return;
@@ -2025,6 +2030,7 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
             bool flushFailed = false;
             ForEachWriteActor([&](TKqpTableWriteActor* actor, const TActorId) {
                 if (!flushFailed && actor->IsReady()) {
+                    actor->FlushBuffers();
                     if (!actor->FlushToShards()) {
                         flushFailed = true;
                     }

ydb/core/kqp/runtime/kqp_write_table.cpp

@@ -1385,8 +1385,6 @@ class TShardedWriteController : public IShardedWriteController {
             for (const auto& [token, writeInfo] : WriteInfos) {
                 if (writeInfo.Closed) {
                     Close(token);
-                } else {
-                    FlushSerializer(token, GetMemory() >= Settings.MemoryLimitTotal);
                 }
             }
         }
@@ -1442,21 +1440,13 @@ class TShardedWriteController : public IShardedWriteController {
             data->AttachAlloc(Alloc);
         }
         info.Serializer->AddData(std::move(data));
-
-        if (info.Metadata.Priority == 0) {
-            FlushSerializer(token, GetMemory() >= Settings.MemoryLimitTotal);
-        }
     }
 
     void Close(TWriteToken token) override {
         auto& info = WriteInfos.at(token);
         AFL_ENSURE(info.Serializer);
         info.Closed = true;
         info.Serializer->Close();
-        if (info.Metadata.Priority == 0) {
-            FlushSerializer(token, true);
-            AFL_ENSURE(info.Serializer->IsFinished());
-        }
     }
 
     void CleanupClosedTokens() override {
@@ -1474,13 +1464,8 @@ class TShardedWriteController : public IShardedWriteController {
     void FlushBuffers() override {
         TVector<TWriteToken> writeTokensFoFlush;
         for (const auto& [token, writeInfo] : WriteInfos) {
-            AFL_ENSURE(writeInfo.Closed);
-            if (writeInfo.Metadata.Priority != 0) {
-                if (!writeInfo.Serializer->IsFinished()) {
-                    writeTokensFoFlush.push_back(token);
-                }
-            } else {
-                AFL_ENSURE(writeInfo.Serializer->IsFinished());
+             if ((writeInfo.Metadata.Priority == 0 || writeInfo.Closed) && !writeInfo.Serializer->IsFinished()) {
+                writeTokensFoFlush.push_back(token);
             }
         }
 
@@ -1495,7 +1480,11 @@ class TShardedWriteController : public IShardedWriteController {
 
         for (const TWriteToken token : writeTokensFoFlush) {
             FlushSerializer(token, true);
-            AFL_ENSURE(WriteInfos.at(token).Serializer->IsFinished());
+            const auto& writeInfo = WriteInfos.at(token);
+            if (writeInfo.Metadata.Priority != 0) {
+                AFL_ENSURE(writeInfo.Closed);
+                AFL_ENSURE(writeInfo.Serializer->IsFinished());
+            }
         }
     }