the race between TEvLockStatus and TEvProposeTransaction

ydb/core/kqp/session_actor/kqp_query_state.cpp

@@ -323,8 +323,14 @@ void TKqpQueryState::AddOffsetsToTransaction() {
     const auto& operations = GetTopicOperations();
 
     TMaybe<TString> consumer;
-    if (operations.HasConsumer())
+    if (operations.HasConsumer()) {
         consumer = operations.GetConsumer();
+    }
+
+    TMaybe<ui32> supportivePartition;
+    if (operations.HasSupportivePartition()) {
+        supportivePartition = operations.GetSupportivePartition();
+    }
 
     TopicOperations = NTopic::TTopicOperations();
 
@@ -334,7 +340,7 @@ void TKqpQueryState::AddOffsetsToTransaction() {
 
         for (auto& partition : topic.partitions()) {
             if (partition.partition_offsets().empty()) {
-                TopicOperations.AddOperation(path, partition.partition_id());
+                TopicOperations.AddOperation(path, partition.partition_id(), supportivePartition);
             } else {
                 for (auto& range : partition.partition_offsets()) {
                     YQL_ENSURE(consumer.Defined());

ydb/core/kqp/topics/kqp_topics.cpp

@@ -8,6 +8,17 @@
 
 namespace NKikimr::NKqp::NTopic {
 
+static void UpdateSupportivePartition(TMaybe<ui32>& lhs, const TMaybe<ui32>& rhs)
+{
+    if (lhs) {
+        if ((rhs != Nothing()) && (rhs != lhs)) {
+            lhs = Max<ui32>();
+        }
+    } else {
+        lhs = rhs;
+    }
+}
+
 //
 // TConsumerOperations
 //
@@ -78,7 +89,8 @@ void TTopicPartitionOperations::AddOperation(const TString& topic, ui32 partitio
     Operations_[consumer].AddOperation(consumer, range);
 }
 
-void TTopicPartitionOperations::AddOperation(const TString& topic, ui32 partition)
+void TTopicPartitionOperations::AddOperation(const TString& topic, ui32 partition,
+                                             TMaybe<ui32> supportivePartition)
 {
     Y_ABORT_UNLESS(Topic_.Empty() || Topic_ == topic);
     Y_ABORT_UNLESS(Partition_.Empty() || Partition_ == partition);
@@ -88,6 +100,8 @@ void TTopicPartitionOperations::AddOperation(const TString& topic, ui32 partitio
         Partition_ = partition;
     }
 
+    UpdateSupportivePartition(SupportivePartition_, supportivePartition);
+
     HasWriteOperations_ = true;
 }
 
@@ -112,6 +126,9 @@ void TTopicPartitionOperations::BuildTopicTxs(THashMap<ui64, NKikimrPQ::TDataTra
         NKikimrPQ::TPartitionOperation* o = tx.MutableOperations()->Add();
         o->SetPartitionId(*Partition_);
         o->SetPath(*Topic_);
+        if (SupportivePartition_.Defined()) {
+            o->SetSupportivePartition(*SupportivePartition_);
+        }
     }
 }
 
@@ -127,6 +144,8 @@ void TTopicPartitionOperations::Merge(const TTopicPartitionOperations& rhs)
         TabletId_ = rhs.TabletId_;
     }
 
+    UpdateSupportivePartition(SupportivePartition_, rhs.SupportivePartition_);
+
     for (auto& [key, value] : rhs.Operations_) {
         Operations_[key].Merge(value);
     }
@@ -240,10 +259,11 @@ void TTopicOperations::AddOperation(const TString& topic, ui32 partition,
     HasReadOperations_ = true;
 }
 
-void TTopicOperations::AddOperation(const TString& topic, ui32 partition)
+void TTopicOperations::AddOperation(const TString& topic, ui32 partition,
+                                    TMaybe<ui32> supportivePartition)
 {
     TTopicPartition key{topic, partition};
-    Operations_[key].AddOperation(topic, partition);
+    Operations_[key].AddOperation(topic, partition, supportivePartition);
     HasWriteOperations_ = true;
 }
 

ydb/core/kqp/topics/kqp_topics.h

@@ -49,7 +49,8 @@ class TTopicPartitionOperations {
     void AddOperation(const TString& topic, ui32 partition,
                       const TString& consumer,
                       const Ydb::Topic::OffsetsRange& range);
-    void AddOperation(const TString& topic, ui32 partition);
+    void AddOperation(const TString& topic, ui32 partition,
+                      TMaybe<ui32> supportivePartition);
 
     void BuildTopicTxs(THashMap<ui64, NKikimrPQ::TDataTransaction> &txs);
 
@@ -67,6 +68,7 @@ class TTopicPartitionOperations {
     THashMap<TString, TConsumerOperations> Operations_;
     bool HasWriteOperations_ = false;
     TMaybe<ui64> TabletId_;
+    TMaybe<ui32> SupportivePartition_;
 };
 
 struct TTopicPartition {
@@ -98,7 +100,8 @@ class TTopicOperations {
     void AddOperation(const TString& topic, ui32 partition,
                       const TString& consumer,
                       const Ydb::Topic::OffsetsRange& range);
-    void AddOperation(const TString& topic, ui32 partition);
+    void AddOperation(const TString& topic, ui32 partition,
+                      TMaybe<ui32> supportivePartition);
 
     void FillSchemeCacheNavigate(NSchemeCache::TSchemeCacheNavigate& navigate,
                                  TMaybe<TString> consumer);

ydb/core/persqueue/partition.cpp

@@ -3042,6 +3042,8 @@ void TPartition::Handle(TEvPQ::TEvCheckPartitionStatusRequest::TPtr& ev, const T
 
 void TPartition::HandleOnInit(TEvPQ::TEvDeletePartition::TPtr& ev, const TActorContext&)
 {
+    Y_ABORT_UNLESS(IsSupportive());
+
     PendingEvents.emplace_back(ev->ReleaseBase().Release());
 }
 

ydb/core/persqueue/partition_id.h

@@ -29,14 +29,15 @@ class TPartitionId {
 
     size_t GetHash() const
     {
-        return MultiHash(OriginalPartitionId, WriteId);
+        return MultiHash(MultiHash(OriginalPartitionId, WriteId), InternalPartitionId);
     }
 
     bool IsEqual(const TPartitionId& rhs) const
     {
         return
             (OriginalPartitionId == rhs.OriginalPartitionId) &&
-            (WriteId == rhs.WriteId);
+            (WriteId == rhs.WriteId) &&
+            (InternalPartitionId == rhs.InternalPartitionId);
     }
 
     void ToStream(IOutputStream& s) const

ydb/core/persqueue/partition_write.cpp

@@ -41,6 +41,9 @@ void TPartition::ReplyOwnerOk(const TActorContext& ctx, const ui64 dst, const TS
     r->SetOwnerCookie(cookie);
     r->SetStatus(PartitionConfig ? PartitionConfig->GetStatus() : NKikimrPQ::ETopicPartitionStatus::Active);
     r->SetSeqNo(seqNo);
+    if (IsSupportive()) {
+        r->SetSupportivePartition(Partition.InternalPartitionId);
+    }
 
     ctx.Send(Tablet, response.Release());
 }

ydb/core/persqueue/pq_impl.cpp

@@ -2001,7 +2001,8 @@ void TPersQueue::HandleWriteRequest(const ui64 responseCookie, const TActorId& p
                     "Tablet " << TabletID() <<
                     " Write in transaction." <<
                     " Partition: " << req.GetPartition() <<
-                    ", WriteId: " << req.GetWriteId());
+                    ", WriteId: " << req.GetWriteId() <<
+                    ", NeedSupportivePartition: " << req.GetNeedSupportivePartition());
     }
 
     for (ui32 i = 0; i < req.CmdWriteSize(); ++i) {
@@ -2198,7 +2199,8 @@ void TPersQueue::HandleReserveBytesRequest(const ui64 responseCookie, const TAct
                     "Tablet " << TabletID() <<
                     " Reserve bytes in transaction." <<
                     " Partition: " << req.GetPartition() <<
-                    ", WriteId: " << req.GetWriteId());
+                    ", WriteId: " << req.GetWriteId() <<
+                    ", NeedSupportivePartition: " << req.GetNeedSupportivePartition());
     }
 
     InitResponseBuilder(responseCookie, 1, COUNTER_LATENCY_PQ_RESERVE_BYTES);
@@ -2574,6 +2576,8 @@ void TPersQueue::HandleWriteRequestForSupportivePartition(const ui64 responseCoo
                                                           const NKikimrClient::TPersQueuePartitionRequest& req,
                                                           const TActorContext& ctx)
 {
+    Y_ABORT_UNLESS(req.HasWriteId());
+
     const TPartitionInfo& partition = GetPartitionInfo(req);
     const TActorId& actorId = partition.Actor;
 
@@ -2648,6 +2652,14 @@ void TPersQueue::HandleEventForSupportivePartition(const ui64 responseCookie,
                                                    sender);
         }
     } else {
+        if (!req.GetNeedSupportivePartition()) {
+            ReplyError(ctx,
+                       responseCookie,
+                       NPersQueue::NErrorCode::PRECONDITION_FAILED,
+                       "lost messages");
+            return;
+        }
+
         //
         // этап 1:
         // - создать запись в TxWrites
@@ -2665,6 +2677,7 @@ void TPersQueue::HandleEventForSupportivePartition(const ui64 responseCookie,
         TPartitionId partitionId(originalPartitionId, writeId, NextSupportivePartitionId++);
 
         writeInfo.Partitions.emplace(originalPartitionId, partitionId);
+        TxWritesChanged = true;
         AddSupportivePartition(partitionId);
 
         Y_ABORT_UNLESS(Partitions.contains(partitionId));
@@ -3115,6 +3128,38 @@ void TPersQueue::Handle(TEvPersQueue::TEvProposeTransaction::TPtr& ev, const TAc
 
 }
 
+bool TPersQueue::CheckTxWriteOperation(const NKikimrPQ::TPartitionOperation& operation,
+                                       ui64 writeId) const
+{
+    TPartitionId partitionId(operation.GetPartitionId(),
+                             writeId,
+                             operation.GetSupportivePartition());
+    return Partitions.contains(partitionId);
+}
+
+bool TPersQueue::CheckTxWriteOperations(const NKikimrPQ::TDataTransaction& txBody) const
+{
+    if (!txBody.HasWriteId()) {
+        return true;
+    }
+
+    ui64 writeId = txBody.GetWriteId();
+
+    for (auto& operation : txBody.GetOperations()) {
+        auto isWrite = [](const NKikimrPQ::TPartitionOperation& o) {
+            return !o.HasBegin();
+        };
+
+        if (isWrite(operation)) {
+            if (!CheckTxWriteOperation(operation, writeId)) {
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
 void TPersQueue::HandleDataTransaction(TAutoPtr<TEvPersQueue::TEvProposeTransaction> ev,
                                        const TActorContext& ctx)
 {
@@ -3123,23 +3168,6 @@ void TPersQueue::HandleDataTransaction(TAutoPtr<TEvPersQueue::TEvProposeTransact
     Y_ABORT_UNLESS(event.HasData());
     const NKikimrPQ::TDataTransaction& txBody = event.GetData();
 
-    for (auto& operation : txBody.GetOperations()) {
-        Y_ABORT_UNLESS(!operation.HasPath() || (operation.GetPath() == TopicPath));
-
-        bool isWriteOperation = !operation.HasBegin();
-
-        LOG_DEBUG_S(ctx, NKikimrServices::PERSQUEUE,
-                    "Tablet " << TabletID() <<
-                    " tx=" << event.GetTxId() <<
-                    ", write_id=" << txBody.GetWriteId() <<
-                    ", path=" << operation.GetPath() <<
-                    ", partition=" << operation.GetPartitionId() <<
-                    ", consumer=" << operation.GetConsumer() <<
-                    ", begin=" << operation.GetBegin() <<
-                    ", end=" << operation.GetEnd() <<
-                    ", is_write=" << isWriteOperation);
-    }
-
     if (TabletState != NKikimrPQ::ENormal) {
         SendProposeTransactionAbort(ActorIdFromProto(event.GetSourceActor()),
                                     event.GetTxId(),
@@ -3158,6 +3186,13 @@ void TPersQueue::HandleDataTransaction(TAutoPtr<TEvPersQueue::TEvProposeTransact
         return;
     }
 
+    if (!CheckTxWriteOperations(txBody)) {
+        SendProposeTransactionAbort(ActorIdFromProto(event.GetSourceActor()),
+                                    event.GetTxId(),
+                                    ctx);
+        return;
+    }
+
     TMaybe<TPartitionId> partitionId = FindPartitionId(txBody);
     if (!partitionId.Defined()) {
         SendProposeTransactionAbort(ActorIdFromProto(event.GetSourceActor()),
@@ -3392,7 +3427,8 @@ void TPersQueue::BeginWriteTxs(const TActorContext& ctx)
         CanProcessPlanStepQueue() ||
         CanProcessWriteTxs() ||
         CanProcessDeleteTxs() ||
-        CanProcessTxWrites()
+        CanProcessTxWrites() ||
+        TxWritesChanged
         ;
     if (!canProcess) {
         return;
@@ -3441,6 +3477,8 @@ void TPersQueue::EndWriteTxs(const NKikimrClient::TResponse& resp,
         return;
     }
 
+    TxWritesChanged = false;
+
     SendReplies(ctx);
     CheckChangedTxStates(ctx);
     CreateSupportivePartitionActors(ctx);
@@ -4473,6 +4511,7 @@ void TPersQueue::Handle(TEvPQ::TEvDeletePartitionDone::TPtr& ev, const TActorCon
         }
         TxWrites.erase(writeId);
     }
+    TxWritesChanged = true;
 
     TryWriteTxs(ctx);
 }

ydb/core/persqueue/pq_impl.h

@@ -205,6 +205,7 @@ class TPersQueue : public NKeyValue::TKeyValueFlat {
     };
 
     THashMap<ui64, TTxWriteInfo> TxWrites;
+    bool TxWritesChanged = false;
     ui32 NextSupportivePartitionId = 100'000;
 
     TActorId CacheActor;
@@ -494,6 +495,10 @@ class TPersQueue : public NKeyValue::TKeyValueFlat {
 
     void BeginDeleteTx(const TDistributedTransaction& tx);
     void BeginDeletePartitions(TTxWriteInfo& writeInfo);
+
+    bool CheckTxWriteOperation(const NKikimrPQ::TPartitionOperation& operation,
+                               ui64 writeId) const;
+    bool CheckTxWriteOperations(const NKikimrPQ::TDataTransaction& txBody) const;
 };
 
 

ydb/core/persqueue/ut/pqtablet_ut.cpp

@@ -74,6 +74,7 @@ struct TGetOwnershipRequestParams {
     TMaybe<ui32> Partition;
     TMaybe<ui64> MsgNo;
     TMaybe<ui64> WriteId;
+    TMaybe<bool> NeedSupportivePartition;
     TMaybe<TString> Owner; // o
     TMaybe<ui64> Cookie;
 };
@@ -550,6 +551,9 @@ std::unique_ptr<TEvPersQueue::TEvRequest> TPQTabletFixture::MakeGetOwnershipRequ
     if (params.WriteId.Defined()) {
         request->SetWriteId(*params.WriteId);
     }
+    if (params.NeedSupportivePartition.Defined()) {
+        request->SetNeedSupportivePartition(*params.NeedSupportivePartition);
+    }
     if (params.Cookie.Defined()) {
         request->SetCookie(*params.Cookie);
     }
@@ -1281,6 +1285,7 @@ Y_UNIT_TEST_F(ProposeTx_Command_After_Propose, TPQTabletFixture)
 
     SyncGetOwnership({.Partition=partitionId,
                      .WriteId=writeId,
+                     .NeedSupportivePartition=true,
                      .Owner="-=[ 0wn3r ]=-",
                      .Cookie=4},
                      {.Cookie=4,