Fixed errors with transactions in the PQ tablet

ydb/core/persqueue/events/internal.h

@@ -1060,6 +1060,7 @@ struct TEvPQ {
     };
 
     struct TEvGetWriteInfoRequest : public TEventLocal<TEvGetWriteInfoRequest, EvGetWriteInfoRequest> {
+        TActorId OriginalPartition;
     };
 
     struct TEvGetWriteInfoResponse : public TEventLocal<TEvGetWriteInfoResponse, EvGetWriteInfoResponse> {

ydb/core/persqueue/partition.cpp

@@ -981,6 +981,7 @@ void TPartition::HandleOnInit(TEvPQ::TEvGetWriteInfoRequest::TPtr& ev, const TAc
 
     Y_ABORT_UNLESS(IsSupportive());
 
+    ev->Get()->OriginalPartition = ev->Sender;
     PendingEvents.emplace_back(ev->ReleaseBase().Release());
 }
 
@@ -1097,11 +1098,16 @@ void TPartition::Handle(TEvPQ::TEvTxRollback::TPtr& ev, const TActorContext& ctx
 
 void TPartition::Handle(TEvPQ::TEvGetWriteInfoRequest::TPtr& ev, const TActorContext& ctx) {
     PQ_LOG_D("Handle TEvPQ::TEvGetWriteInfoRequest");
+    TActorId originalPartition = ev->Get()->OriginalPartition;
+    if (!originalPartition) {
+        // original message
+        originalPartition = ev->Sender;
+    }
     if (ClosedInternalPartition || WaitingForPreviousBlobQuota() || (CurrentStateFunc() != &TThis::StateIdle)) {
         PQ_LOG_D("Send TEvPQ::TEvGetWriteInfoError");
         auto* response = new TEvPQ::TEvGetWriteInfoError(Partition.InternalPartitionId,
                                                          "Write info requested while writes are not complete");
-        ctx.Send(ev->Sender, response);
+        ctx.Send(originalPartition, response);
         ClosedInternalPartition = true;
         return;
     }
@@ -1125,7 +1131,7 @@ void TPartition::Handle(TEvPQ::TEvGetWriteInfoRequest::TPtr& ev, const TActorCon
     response->InputLags = std::move(SupportivePartitionTimeLag);
 
     PQ_LOG_D("Send TEvPQ::TEvGetWriteInfoResponse");
-    ctx.Send(ev->Sender, response);
+    ctx.Send(originalPartition, response);
 }
 
 void TPartition::WriteInfoResponseHandler(

ydb/core/persqueue/pq_impl.cpp

@@ -1148,11 +1148,10 @@ void TPersQueue::ReadConfig(const NKikimrClient::TKeyValueResponse::TReadResult&
             }
 
             Txs.emplace(tx.GetTxId(), tx);
+            SetTxInFlyCounter();
 
             if (tx.HasStep()) {
-                if (std::make_pair(tx.GetStep(), tx.GetTxId()) >= std::make_pair(ExecStep, ExecTxId)) {
-                    PlannedTxs.emplace_back(tx.GetStep(), tx.GetTxId());
-                }
+                PlannedTxs.emplace_back(tx.GetStep(), tx.GetTxId());
             }
 
             if (tx.HasWriteId()) {
@@ -1164,6 +1163,8 @@ void TPersQueue::ReadConfig(const NKikimrClient::TKeyValueResponse::TReadResult&
 
     EndInitTransactions();
     EndReadConfig(ctx);
+
+    SetTxCounters();
 }
 
 void TPersQueue::EndReadConfig(const TActorContext& ctx)
@@ -3269,6 +3270,7 @@ void TPersQueue::HandleWakeup(const TActorContext& ctx) {
     }
     MeteringSink.MayFlush(ctx.Now());
     DeleteExpiredTransactions(ctx);
+    SetTxCounters();
     ctx.Schedule(TDuration::Seconds(5), new TEvents::TEvWakeup());
 }
 
@@ -3289,6 +3291,33 @@ void TPersQueue::DeleteExpiredTransactions(const TActorContext& ctx)
     TryWriteTxs(ctx);
 }
 
+void TPersQueue::SetTxCounters()
+{
+    SetTxCompleteLagCounter();
+    SetTxInFlyCounter();
+}
+
+void TPersQueue::SetTxCompleteLagCounter()
+{
+    ui64 lag = 0;
+
+    if (!TxQueue.empty()) {
+        ui64 firstTxStep = TxQueue.front().first;
+        ui64 currentStep = TAppData::TimeProvider->Now().MilliSeconds();
+
+        if (currentStep > firstTxStep) {
+            lag = currentStep - firstTxStep;
+        }
+    }
+
+    Counters->Simple()[COUNTER_PQ_TABLET_TX_COMPLETE_LAG] = lag;
+}
+
+void TPersQueue::SetTxInFlyCounter()
+{
+    Counters->Simple()[COUNTER_PQ_TABLET_TX_IN_FLY] = Txs.size();
+}
+
 void TPersQueue::Handle(TEvPersQueue::TEvCancelTransactionProposal::TPtr& ev, const TActorContext& ctx)
 {
     PQ_LOG_D("Handle TEvPersQueue::TEvCancelTransactionProposal");
@@ -3763,6 +3792,7 @@ void TPersQueue::ProcessProposeTransactionQueue(const TActorContext& ctx)
 
         const NKikimrPQ::TEvProposeTransaction& event = front->GetRecord();
         TDistributedTransaction& tx = Txs[event.GetTxId()];
+        SetTxInFlyCounter();
 
         switch (tx.State) {
         case NKikimrPQ::TTransaction::UNKNOWN:
@@ -3844,6 +3874,7 @@ void TPersQueue::ProcessPlanStepQueue(const TActorContext& ctx)
                         TryExecuteTxs(ctx, tx);
 
                         TxQueue.emplace_back(step, txId);
+                        SetTxCompleteLagCounter();
                     } else {
                         LOG_WARN_S(ctx, NKikimrServices::PERSQUEUE,
                                    "Tablet " << TabletID() <<
@@ -4580,6 +4611,11 @@ void TPersQueue::CheckTxState(const TActorContext& ctx,
         Y_ABORT_UNLESS(tx.TxId == TxsOrder[tx.State].front(),
                        "PQ %" PRIu64 ", TxId %" PRIu64 ", FrontTxId %" PRIu64,
                        TabletID(), tx.TxId, TxsOrder[tx.State].front());
+        Y_ABORT_UNLESS(tx.TxId == TxQueue.front().second,
+                       "PQ %" PRIu64 ", TxId %" PRIu64 ", FrontTxId %" PRIu64,
+                       TabletID(), tx.TxId, TxQueue.front().second);
+        TxQueue.pop_front();
+        SetTxCompleteLagCounter();
 
         SendEvReadSetAckToSenders(ctx, tx);
 
@@ -4599,13 +4635,10 @@ void TPersQueue::CheckTxState(const TActorContext& ctx,
 
     case NKikimrPQ::TTransaction::DELETING:
         // The PQ tablet has persisted its state. Now she can delete the transaction and take the next one.
-        if (!TxQueue.empty() && (TxQueue.front().second == tx.TxId)) {
-            TxQueue.pop_front();
-        }
-
         DeleteWriteId(tx.WriteId);
         PQ_LOG_D("delete TxId " << tx.TxId);
         Txs.erase(tx.TxId);
+        SetTxInFlyCounter();
 
         // If this was the last transaction, then you need to send responses to messages about changes
         // in the status of the PQ tablet (if they came)

ydb/core/persqueue/pq_impl.h

@@ -440,6 +440,10 @@ class TPersQueue : public NKeyValue::TKeyValueFlat {
     void DeleteExpiredTransactions(const TActorContext& ctx);
     void Handle(TEvPersQueue::TEvCancelTransactionProposal::TPtr& ev, const TActorContext& ctx);
 
+    void SetTxCounters();
+    void SetTxCompleteLagCounter();
+    void SetTxInFlyCounter();
+
     bool CanProcessProposeTransactionQueue() const;
     bool CanProcessPlanStepQueue() const;
     bool CanProcessWriteTxs() const;

ydb/core/persqueue/ut/pqtablet_mock.cpp

@@ -94,7 +94,10 @@ void TPQTabletMock::Handle(TEvTxProcessing::TEvReadSet::TPtr& ev, const TActorCo
 {
     Y_UNUSED(ctx);
 
-    ReadSet = ev->Get()->Record;
+    const auto& record = ev->Get()->Record;
+
+    ReadSet = record;
+    ReadSets[std::make_pair(record.GetStep(), record.GetTxId())].push_back(record);
 }
 
 void TPQTabletMock::Handle(TEvTxProcessing::TEvReadSetAck::TPtr& ev, const TActorContext& ctx)

ydb/core/persqueue/ut/pqtablet_mock.h

@@ -32,6 +32,8 @@ class TPQTabletMock : public TActor<TPQTabletMock>, public NTabletFlatExecutor::
     TMaybe<NKikimrTx::TEvReadSet> ReadSet;
     TMaybe<NKikimrTx::TEvReadSetAck> ReadSetAck;
 
+    THashMap<std::pair<ui64, ui64>, TVector<NKikimrTx::TEvReadSet>> ReadSets;
+
 private:
     struct TEvPQTablet {
         enum EEv {

ydb/core/persqueue/ut/pqtablet_ut.cpp

@@ -139,6 +139,7 @@ class TPQTabletFixture : public NUnitTest::TBaseFixture {
         TMaybe<ui64> Target;
         TMaybe<NKikimrTx::TReadSetData::EDecision> Decision;
         TMaybe<ui64> Producer;
+        TMaybe<size_t> Count;
     };
 
     struct TReadSetAckMatcher {
@@ -191,6 +192,7 @@ class TPQTabletFixture : public NUnitTest::TBaseFixture {
     void WaitPlanStepAccepted(const TPlanStepAcceptedMatcher& matcher = {});
 
     void WaitReadSet(NHelpers::TPQTabletMock& tablet, const TReadSetMatcher& matcher);
+    void WaitReadSetEx(NHelpers::TPQTabletMock& tablet, const TReadSetMatcher& matcher);
     void SendReadSet(const TReadSetParams& params);
 
     void WaitReadSetAck(NHelpers::TPQTabletMock& tablet, const TReadSetAckMatcher& matcher);
@@ -236,6 +238,9 @@ class TPQTabletFixture : public NUnitTest::TBaseFixture {
     void StartPQCalcPredicateObserver(size_t& received);
     void WaitForPQCalcPredicate(size_t& received, size_t expected);
 
+    void WaitForTxState(ui64 txId, NKikimrPQ::TTransaction::EState state);
+    void WaitForExecStep(ui64 step);
+
     //
     // TODO(abcdef): для тестирования повторных вызовов нужны примитивы Send+Wait
     //
@@ -456,6 +461,15 @@ void TPQTabletFixture::WaitReadSet(NHelpers::TPQTabletMock& tablet, const TReadS
     }
 }
 
+void TPQTabletFixture::WaitReadSetEx(NHelpers::TPQTabletMock& tablet, const TReadSetMatcher& matcher)
+{
+    TDispatchOptions options;
+    options.CustomFinalCondition = [&]() {
+        return tablet.ReadSets[std::make_pair(*matcher.Step, *matcher.TxId)].size() >= *matcher.Count;
+    };
+    UNIT_ASSERT(Ctx->Runtime->DispatchEvents(options));
+}
+
 void TPQTabletFixture::SendReadSet(const TReadSetParams& params)
 {
     NKikimrTx::TReadSetData payload;
@@ -952,6 +966,70 @@ void TPQTabletFixture::WaitForPQCalcPredicate(size_t& received, size_t expected)
     UNIT_ASSERT(Ctx->Runtime->DispatchEvents(options));
 }
 
+void TPQTabletFixture::WaitForTxState(ui64 txId, NKikimrPQ::TTransaction::EState state)
+{
+    const TString key = GetTxKey(txId);
+
+    while (true) {
+        auto request = std::make_unique<TEvKeyValue::TEvRequest>();
+        request->Record.SetCookie(12345);
+        auto cmd = request->Record.AddCmdReadRange();
+        auto range = cmd->MutableRange();
+        range->SetFrom(key);
+        range->SetIncludeFrom(true);
+        range->SetTo(key);
+        range->SetIncludeTo(true);
+        cmd->SetIncludeData(true);
+        SendToPipe(Ctx->Edge, request.release());
+
+        auto response = Ctx->Runtime->GrabEdgeEvent<TEvKeyValue::TEvResponse>();
+        UNIT_ASSERT_VALUES_EQUAL(response->Record.GetStatus(), NMsgBusProxy::MSTATUS_OK);
+        const auto& result = response->Record.GetReadRangeResult(0);
+        UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), static_cast<ui32>(NKikimrProto::OK));
+        const auto& pair = result.GetPair(0);
+
+        NKikimrPQ::TTransaction tx;
+        Y_ABORT_UNLESS(tx.ParseFromString(pair.GetValue()));
+
+        if (tx.GetState() == state) {
+            return;
+        }
+    }
+
+    UNIT_FAIL("transaction " << txId << " has not entered the " << state << " state");
+}
+
+void TPQTabletFixture::WaitForExecStep(ui64 step)
+{
+    while (true) {
+        auto request = std::make_unique<TEvKeyValue::TEvRequest>();
+        request->Record.SetCookie(12345);
+        auto cmd = request->Record.AddCmdReadRange();
+        auto range = cmd->MutableRange();
+        range->SetFrom("_txinfo");
+        range->SetIncludeFrom(true);
+        range->SetTo("_txinfo");
+        range->SetIncludeTo(true);
+        cmd->SetIncludeData(true);
+        SendToPipe(Ctx->Edge, request.release());
+
+        auto response = Ctx->Runtime->GrabEdgeEvent<TEvKeyValue::TEvResponse>();
+        UNIT_ASSERT_VALUES_EQUAL(response->Record.GetStatus(), NMsgBusProxy::MSTATUS_OK);
+        const auto& result = response->Record.GetReadRangeResult(0);
+        UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), static_cast<ui32>(NKikimrProto::OK));
+        const auto& pair = result.GetPair(0);
+
+        NKikimrPQ::TTabletTxInfo txInfo;
+        Y_ABORT_UNLESS(txInfo.ParseFromString(pair.GetValue()));
+
+        if (txInfo.GetExecStep() == step) {
+            return;
+        }
+    }
+
+    UNIT_FAIL("expected execution step " << step);
+}
+
 Y_UNIT_TEST_F(Parallel_Transactions_1, TPQTabletFixture)
 {
     TestParallelTransactions("consumer", "consumer");
@@ -1730,6 +1808,62 @@ Y_UNIT_TEST_F(Huge_ProposeTransacton, TPQTabletFixture)
     WaitPlanStepAccepted({.Step=100});
 }
 
+Y_UNIT_TEST_F(After_Restarting_The_Tablet_Sends_A_TEvReadSet_For_Transactions_In_The_EXECUTED_State, TPQTabletFixture)
+{
+    const ui64 txId_1 = 67890;
+    const ui64 txId_2 = txId_1 + 1;
+    const ui64 mockTabletId = 22222;
+
+    NHelpers::TPQTabletMock* tablet = CreatePQTabletMock(mockTabletId);
+    PQTabletPrepare({.partitions=1}, {}, *Ctx);
+
+    // 1st tx
+    SendProposeTransactionRequest({.TxId=txId_1,
+                                  .Senders={mockTabletId}, .Receivers={mockTabletId},
+                                  .TxOps={
+                                  {.Partition=0, .Consumer="user", .Begin=0, .End=0, .Path="/topic"},
+                                  }});
+    WaitProposeTransactionResponse({.TxId=txId_1,
+                                   .Status=NKikimrPQ::TEvProposeTransactionResult::PREPARED});
+
+    SendPlanStep({.Step=100, .TxIds={txId_1}});
+
+    WaitForCalcPredicateResult();
+
+    tablet->SendReadSet(*Ctx->Runtime, {.Step=100, .TxId=txId_1, .Target=Ctx->TabletId, .Decision=NKikimrTx::TReadSetData::DECISION_COMMIT});
+
+    WaitProposeTransactionResponse({.TxId=txId_1,
+                                   .Status=NKikimrPQ::TEvProposeTransactionResult::COMPLETE});
+
+    WaitForTxState(txId_1, NKikimrPQ::TTransaction::EXECUTED);
+
+    tablet->ReadSet = Nothing();
+
+    // 2nd tx
+    SendProposeTransactionRequest({.TxId=txId_2,
+                                  .Senders={mockTabletId}, .Receivers={mockTabletId},
+                                  .TxOps={
+                                  {.Partition=0, .Consumer="user", .Begin=0, .End=0, .Path="/topic"},
+                                  }});
+    WaitProposeTransactionResponse({.TxId=txId_2,
+                                   .Status=NKikimrPQ::TEvProposeTransactionResult::PREPARED});
+
+    SendPlanStep({.Step=110, .TxIds={txId_2}});
+
+    WaitForCalcPredicateResult();
+
+    WaitReadSetEx(*tablet, {.Step=110, .TxId=txId_2, .Decision=NKikimrTx::TReadSetData::DECISION_COMMIT, .Count=1});
+
+    // the PQ tablet has moved a step forward
+    WaitForExecStep(110);
+
+    // restart PQ tablet
+    PQTabletRestart(*Ctx);
+
+    // the PQ tablet should send a TEvReadSet for the executed transaction
+    WaitReadSetEx(*tablet, {.Step=100, .TxId=txId_1, .Decision=NKikimrTx::TReadSetData::DECISION_COMMIT, .Count=2});
+}
+
 }
 
 }

ydb/core/protos/counters_pq.proto

@@ -68,6 +68,8 @@ enum ESimpleCounters {
     COUNTER_PQ_TABLET_OPENED_PIPES = 5                               [(CounterOpts) = {Name: "OpenedPipes"}];
     COUNTER_PQ_TABLET_INFLIGHT = 6                                   [(CounterOpts) = {Name: "RequestsInflight"}];
 
+    COUNTER_PQ_TABLET_TX_COMPLETE_LAG = 7                            [(CounterOpts) = {Name: "TxCompleteLag"}];
+    COUNTER_PQ_TABLET_TX_IN_FLY = 8                                  [(CounterOpts) = {Name: "TxInFly"}];
 }
 
 enum EPercentileCounters {