[Stable-25-1-2] Flush only insert before distributed commit (#20586)

Author: nikvas0

ydb/core/kqp/runtime/kqp_write_actor.cpp

@@ -333,6 +333,9 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {
             std::move(columnsMetadata),
             std::move(writeIndexes),
             priority);
+
+        // At current time only insert operation can fail.
+        NeedToFlushBeforeCommit |= (operationType == NKikimrDataEvents::TEvWrite::TOperation::OPERATION_INSERT);
         CA_LOG_D("Open: token=" << token);
         return token;
     }
@@ -365,6 +368,7 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {
     void CleanupClosedTokens() {
         YQL_ENSURE(ShardedWriteController);
         ShardedWriteController->CleanupClosedTokens();
+        NeedToFlushBeforeCommit = false;
     }
 
     void SetParentTraceId(NWilson::TTraceId traceId) {
@@ -1301,6 +1305,10 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {
         Stats.AffectedPartitions.clear();
     }
 
+    bool FlushBeforeCommit() const {
+        return NeedToFlushBeforeCommit;
+    }
+
 private:
     NActors::TActorId PipeCacheId = NKikimr::MakePipePerNodeCacheID(false);
 
@@ -1329,6 +1337,7 @@ class TKqpTableWriteActor : public TActorBootstrapped<TKqpTableWriteActor> {
 
     IKqpTransactionManagerPtr TxManager;
     bool Closed = false;
+    bool NeedToFlushBeforeCommit = false;
     EMode Mode = EMode::WRITE;
     THashMap<ui64, TInstant> SendTime;
 
@@ -1659,7 +1668,6 @@ class TKqpDirectWriteActor : public TActorBootstrapped<TKqpDirectWriteActor>, pu
     std::optional<TKqpTableWriteActor::TWriteToken> WriteToken;
 
     bool Closed = false;
-
     bool WaitingForTableActor = false;
 
     NWilson::TSpan DirectWriteActorSpan;
@@ -1877,6 +1885,23 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
         Process();
     }
 
+    bool NeedToFlush() {
+        const bool outOfMemory = GetTotalFreeSpace() <= 0;
+        const bool needToFlush = outOfMemory
+            || State == EState::FLUSHING
+            || State == EState::PREPARING
+            || State == EState::COMMITTING
+            || State == EState::ROLLINGBACK;
+        return needToFlush;
+    }
+
+    bool NeedToFlushActor(const TKqpTableWriteActor* actor) {
+        return NeedToFlush()
+            && (State != EState::FLUSHING
+                || !TxId // Flush between queries
+                || actor->FlushBeforeCommit()); // Flush before commit
+    }
+
     bool Process() {
         ProcessRequestQueue();
         if (!ProcessWrite()) {
@@ -1887,7 +1912,9 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
         if (State == EState::FLUSHING) {
             bool isEmpty = true;
             for (auto& [_, info] : WriteInfos) {
-                isEmpty = isEmpty && info.WriteTableActor->IsReady() && info.WriteTableActor->IsEmpty();
+                if (NeedToFlushActor(info.WriteTableActor)) {
+                    isEmpty = isEmpty && info.WriteTableActor->IsReady() && info.WriteTableActor->IsEmpty();
+                }
             }
             if (isEmpty) {
                 OnFlushed();
@@ -1942,14 +1969,7 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
     }
 
     bool ProcessWrite() {
-        const bool outOfMemory = GetTotalFreeSpace() <= 0;
-        const bool needToFlush = outOfMemory
-            || State == EState::FLUSHING
-            || State == EState::PREPARING
-            || State == EState::COMMITTING
-            || State == EState::ROLLINGBACK;
-
-        if (!EnableStreamWrite && outOfMemory) {
+        if (!EnableStreamWrite && GetTotalFreeSpace() <= 0) {
             ReplyErrorAndDie(
                 NYql::NDqProto::StatusIds::PRECONDITION_FAILED,
                 NYql::TIssuesIds::KIKIMR_PRECONDITION_FAILED,
@@ -1959,10 +1979,10 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
             return false;
         }
 
-        if (needToFlush) {
+        if (NeedToFlush()) {
             CA_LOG_D("Flush data");
             for (auto& [_, info] : WriteInfos) {
-                if (info.WriteTableActor->IsReady()) {
+                if (info.WriteTableActor->IsReady() && NeedToFlushActor(info.WriteTableActor)) {
                     if (!info.WriteTableActor->FlushToShards()) {
                         return false;
                     }
@@ -1992,13 +2012,13 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
 
         CA_LOG_D("Start prepare for distributed commit");
         YQL_ENSURE(State == EState::WRITING);
-        YQL_ENSURE(!NeedToFlushBeforeCommit);
         State = EState::PREPARING;
         for (auto& [_, queue] : DataQueues) {
             YQL_ENSURE(queue.empty());
         }
         YQL_ENSURE(TxId);
         for (auto& [_, info] : WriteInfos) {
+            AFL_ENSURE(!info.WriteTableActor->FlushBeforeCommit());
             info.WriteTableActor->SetPrepare(*TxId);
         }
         Close();
@@ -2452,7 +2472,13 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
         } else {
             AFL_ENSURE(ev->Get()->TxId);
             TxId = ev->Get()->TxId;
-            if (NeedToFlushBeforeCommit) {
+
+            bool needToFlushBeforeCommit = false;
+            for (auto& [_, info] : WriteInfos) {
+                needToFlushBeforeCommit |= info.WriteTableActor->FlushBeforeCommit();
+            }
+
+            if (needToFlushBeforeCommit) {
                 Flush(std::move(ev->TraceId));
             } else {
                 TxManager->StartPrepare();
@@ -2837,8 +2863,19 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
         UpdateTracingState("Write", BufferWriteActorSpan.GetTraceId());
         OnOperationFinished(Counters->BufferActorFlushLatencyHistogram);
         State = EState::WRITING;
-        AFL_ENSURE(!TxId || NeedToFlushBeforeCommit); // TxId => NeedToFlushBeforeCommit
-        NeedToFlushBeforeCommit = false;
+
+        for (auto& [_, info] : WriteInfos) {
+            AFL_ENSURE(TxId || info.WriteTableActor->IsEmpty());
+            if (info.WriteTableActor->IsEmpty()) {
+                info.WriteTableActor->CleanupClosedTokens();
+            }
+            if (!TxId) {
+                info.WriteTableActor->Unlink();
+            }
+
+            AFL_ENSURE(!info.WriteTableActor->FlushBeforeCommit());
+        }
+
         if (TxId) {
             TxManager->StartPrepare();
             Prepare(std::nullopt);
@@ -2972,7 +3009,6 @@ class TKqpBufferWriteActor :public TActorBootstrapped<TKqpBufferWriteActor>, pub
 
     EState State;
     bool HasError = false;
-    bool NeedToFlushBeforeCommit = false;
     THashMap<TPathId, std::queue<TBufferWriteMessage>> DataQueues;
 
     struct TAckMessage {

ydb/core/kqp/runtime/kqp_write_table.cpp

@@ -1381,7 +1381,6 @@ class TShardedWriteController : public IShardedWriteController {
     }
 
     void CleanupClosedTokens() override {
-        AFL_ENSURE(IsEmpty());
         for (auto it = WriteInfos.begin(); it != WriteInfos.end();) {
             if (it->second.Closed) {
                 AFL_ENSURE(it->second.Serializer->IsFinished());