Fix planning of volatile transactions with uncommitted effects (#2505)

Author: snaury

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -2006,23 +2006,68 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                                     TTopicTabletTxs& topicTxs) {
         TDatashardTxs datashardTxs;
 
-        std::vector<ui64> affectedShardsSet;
-        affectedShardsSet.reserve(datashardTasks.size());
-
         for (auto& [shardId, tasks]: datashardTasks) {
             auto [it, success] = datashardTxs.emplace(
                 shardId,
                 TasksGraph.GetMeta().Allocate<NKikimrTxDataShard::TKqpTransaction>());
 
             YQL_ENSURE(success, "unexpected duplicates in datashard transactions");
-            affectedShardsSet.emplace_back(shardId);
             NKikimrTxDataShard::TKqpTransaction* dsTxs = it->second;
             dsTxs->MutableTasks()->Reserve(tasks.size());
             for (auto& task: tasks) {
                 dsTxs->AddTasks()->Swap(task);
             }
         }
 
+        // Note: when locks map is present it will be mutated to avoid copying data
+        auto& locksMap = Request.DataShardLocks;
+        if (!locksMap.empty()) {
+            YQL_ENSURE(Request.LocksOp == ELocksOp::Commit || Request.LocksOp == ELocksOp::Rollback);
+        }
+
+        // Materialize (possibly empty) txs for all shards with locks (either commit or rollback)
+        for (auto& [shardId, locksList] : locksMap) {
+            YQL_ENSURE(!locksList.empty(), "unexpected empty locks list in DataShardLocks");
+
+            auto it = datashardTxs.find(shardId);
+            if (it == datashardTxs.end()) {
+                auto [emplaced, success] = datashardTxs.emplace(
+                    shardId,
+                    TasksGraph.GetMeta().Allocate<NKikimrTxDataShard::TKqpTransaction>());
+
+                YQL_ENSURE(success, "unexpected failure to emplace a datashard transaction");
+                it = emplaced;
+            }
+
+            NKikimrTxDataShard::TKqpTransaction* tx = it->second;
+            switch (Request.LocksOp) {
+                case ELocksOp::Commit:
+                    tx->MutableLocks()->SetOp(NKikimrDataEvents::TKqpLocks::Commit);
+                    break;
+                case ELocksOp::Rollback:
+                    tx->MutableLocks()->SetOp(NKikimrDataEvents::TKqpLocks::Rollback);
+                    break;
+                case ELocksOp::Unspecified:
+                    break;
+            }
+
+            // Move lock descriptions to the datashard tx
+            auto* protoLocks = tx->MutableLocks()->MutableLocks();
+            protoLocks->Reserve(locksList.size());
+            bool hasWrites = false;
+            for (auto& lock : locksList) {
+                hasWrites = hasWrites || lock.GetHasWrites();
+                protoLocks->Add(std::move(lock));
+            }
+            locksList.clear();
+
+            // When locks with writes are committed this commits accumulated effects
+            if (Request.LocksOp == ELocksOp::Commit && hasWrites) {
+                ShardsWithEffects.insert(shardId);
+                YQL_ENSURE(!ReadOnlyTx);
+            }
+        }
+
         Request.TopicOperations.BuildTopicTxs(topicTxs);
 
         const bool needRollback = Request.LocksOp == ELocksOp::Rollback;
@@ -2042,7 +2087,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
             // TODO: add support in the future
             topicTxs.empty() &&
             // We only want to use volatile transactions for multiple shards
-            (affectedShardsSet.size() + topicTxs.size()) > 1 &&
+            (datashardTxs.size() + topicTxs.size()) > 1 &&
             // We cannot use volatile transactions with persistent channels
             // Note: currently persistent channels are never used
             !HasPersistentChannels);
@@ -2055,30 +2100,29 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
             // Transactions with topics must always use generic readsets
             !topicTxs.empty());
 
-        if (auto locksMap = Request.DataShardLocks;
-            !locksMap.empty() ||
-            VolatileTx ||
+        if (!locksMap.empty() || VolatileTx ||
             Request.TopicOperations.HasReadOperations())
         {
             YQL_ENSURE(Request.LocksOp == ELocksOp::Commit || Request.LocksOp == ELocksOp::Rollback || VolatileTx);
 
             bool needCommit = Request.LocksOp == ELocksOp::Commit || VolatileTx;
 
-            auto locksOp = needCommit
-                ? NKikimrDataEvents::TKqpLocks::Commit
-                : NKikimrDataEvents::TKqpLocks::Rollback;
-
             absl::flat_hash_set<ui64> sendingShardsSet;
             absl::flat_hash_set<ui64> receivingShardsSet;
 
             // Gather shards that need to send/receive readsets (shards with effects)
             if (needCommit) {
-                for (auto& shardId: affectedShardsSet) {
+                for (auto& [shardId, tx] : datashardTxs) {
+                    if (tx->HasLocks()) {
+                        // Locks may be broken so shards with locks need to send readsets
+                        sendingShardsSet.insert(shardId);
+                    }
                     if (ShardsWithEffects.contains(shardId)) {
                         // Volatile transactions may abort effects, so they send readsets
                         if (VolatileTx) {
                             sendingShardsSet.insert(shardId);
                         }
+                        // Effects are only applied when all locks are valid
                         receivingShardsSet.insert(shardId);
                     }
                 }
@@ -2093,44 +2137,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 }
             }
 
-            // Gather locks that need to be committed or erased
-            for (auto& [shardId, locksList] : locksMap) {
-                NKikimrTxDataShard::TKqpTransaction* tx = nullptr;
-                auto it = datashardTxs.find(shardId);
-                if (it != datashardTxs.end()) {
-                    tx = it->second;
-                } else {
-                    auto [eIt, success] = datashardTxs.emplace(
-                        shardId,
-                        TasksGraph.GetMeta().Allocate<NKikimrTxDataShard::TKqpTransaction>());
-                    tx = eIt->second;
-                }
-
-                tx->MutableLocks()->SetOp(locksOp);
-
-                if (!locksList.empty()) {
-                    auto* protoLocks = tx->MutableLocks()->MutableLocks();
-                    protoLocks->Reserve(locksList.size());
-                    bool hasWrites = false;
-                    for (auto& lock : locksList) {
-                        hasWrites = hasWrites || lock.GetHasWrites();
-                        protoLocks->Add()->Swap(&lock);
-                    }
-
-                    if (needCommit) {
-                        // We also send the result on commit
-                        sendingShardsSet.insert(shardId);
-
-                        if (hasWrites) {
-                            // Tx with uncommitted changes can be aborted due to conflicts,
-                            // so shards with write locks should receive readsets
-                            receivingShardsSet.insert(shardId);
-                            YQL_ENSURE(!ReadOnlyTx);
-                        }
-                    }
-                }
-            }
-
+            // Encode sending/receiving shards in tx bodies
             if (needCommit) {
                 NProtoBuf::RepeatedField<ui64> sendingShards(sendingShardsSet.begin(), sendingShardsSet.end());
                 NProtoBuf::RepeatedField<ui64> receivingShards(receivingShardsSet.begin(), receivingShardsSet.end());
@@ -2139,23 +2146,13 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 std::sort(receivingShards.begin(), receivingShards.end());
 
                 for (auto& [shardId, shardTx] : datashardTxs) {
-                    shardTx->MutableLocks()->SetOp(locksOp);
+                    shardTx->MutableLocks()->SetOp(NKikimrDataEvents::TKqpLocks::Commit);
                     *shardTx->MutableLocks()->MutableSendingShards() = sendingShards;
                     *shardTx->MutableLocks()->MutableReceivingShards() = receivingShards;
                 }
 
                 for (auto& [_, tx] : topicTxs) {
-                    switch (locksOp) {
-                    case NKikimrDataEvents::TKqpLocks::Commit:
-                        tx.SetOp(NKikimrPQ::TDataTransaction::Commit);
-                        break;
-                    case NKikimrDataEvents::TKqpLocks::Rollback:
-                        tx.SetOp(NKikimrPQ::TDataTransaction::Rollback);
-                        break;
-                    case NKikimrDataEvents::TKqpLocks::Unspecified:
-                        break;
-                    }
-
+                    tx.SetOp(NKikimrPQ::TDataTransaction::Commit);
                     *tx.MutableSendingShards() = sendingShards;
                     *tx.MutableReceivingShards() = receivingShards;
                 }

ydb/core/tx/datashard/datashard_ut_order.cpp

@@ -4279,6 +4279,7 @@ Y_UNIT_TEST(UncommittedReadSetAck) {
 
     bool capturePlanSteps = true;
     TVector<THolder<IEventHandle>> capturedPlanSteps;
+    TVector<ui64> capturedPlanTxIds;
     THashSet<ui64> passReadSetTxIds;
     ui64 observedReadSets = 0;
     TVector<THolder<IEventHandle>> capturedReadSets;
@@ -4294,6 +4295,12 @@ Y_UNIT_TEST(UncommittedReadSetAck) {
             case TEvTxProcessing::TEvPlanStep::EventType: {
                 if (nodeIndex == 1 && ev->GetRecipientRewrite() == table3actor && capturePlanSteps) {
                     Cerr << "... captured plan step for table-3" << Endl;
+                    auto* msg = ev->Get<TEvTxProcessing::TEvPlanStep>();
+                    for (const auto& tx : msg->Record.GetTransactions()) {
+                        ui64 txId = tx.GetTxId();
+                        capturedPlanTxIds.push_back(txId);
+                        Cerr << "... captured plan step tx " << txId << " for table-3" << Endl;
+                    }
                     capturedPlanSteps.emplace_back(ev.Release());
                     return TTestActorRuntime::EEventAction::DROP;
                 }
@@ -4303,6 +4310,12 @@ Y_UNIT_TEST(UncommittedReadSetAck) {
                 if (nodeIndex == 1 && ev->GetRecipientRewrite() == table3actor) {
                     auto* msg = ev->Get<TEvTxProcessing::TEvReadSet>();
                     ui64 txId = msg->Record.GetTxId();
+                    if ((msg->Record.GetFlags() & NKikimrTx::TEvReadSet::FLAG_EXPECT_READSET) &&
+                        (msg->Record.GetFlags() & NKikimrTx::TEvReadSet::FLAG_NO_DATA))
+                    {
+                        Cerr << "... passing expectation for txid# " << txId << Endl;
+                        break;
+                    }
                     ++observedReadSets;
                     if (!passReadSetTxIds.contains(txId)) {
                         Cerr << "... readset for txid# " << txId << " was blocked" << Endl;
@@ -4353,20 +4366,11 @@ Y_UNIT_TEST(UncommittedReadSetAck) {
         }
     };
 
-    waitFor([&]{ return capturedPlanSteps.size() > 0; }, "plan step");
-    UNIT_ASSERT_VALUES_EQUAL(capturedPlanSteps.size(), 1u);
-    ui64 realTxId1, realTxId2;
-    {
-        auto* msg = capturedPlanSteps[0]->Get<TEvTxProcessing::TEvPlanStep>();
-        TVector<ui64> realTxIds;
-        for (const auto& tx : msg->Record.GetTransactions()) {
-            realTxIds.emplace_back(tx.GetTxId());
-        }
-        UNIT_ASSERT_VALUES_EQUAL(realTxIds.size(), 2u);
-        std::sort(realTxIds.begin(), realTxIds.end());
-        realTxId1 = realTxIds.at(0);
-        realTxId2 = realTxIds.at(1);
-    }
+    waitFor([&]{ return capturedPlanTxIds.size() >= 2; }, "captured transactions");
+    UNIT_ASSERT_C(capturedPlanTxIds.size(), 2u);
+    std::sort(capturedPlanTxIds.begin(), capturedPlanTxIds.end());
+    ui64 realTxId1 = capturedPlanTxIds.at(0);
+    ui64 realTxId2 = capturedPlanTxIds.at(1);
 
     // Unblock and resend the plan step message
     capturePlanSteps = false;
@@ -4375,7 +4379,7 @@ Y_UNIT_TEST(UncommittedReadSetAck) {
     }
     capturedPlanSteps.clear();
 
-    // Wait until there are 2 readset messages
+    // Wait until there are 2 readset messages (with data)
     waitFor([&]{ return capturedReadSets.size() >= 2; }, "initial readsets");
     SimulateSleep(runtime, TDuration::MilliSeconds(5));
 

ydb/core/tx/datashard/datashard_ut_read_iterator.cpp

@@ -348,7 +348,7 @@ struct TTestHelper {
         auto &runtime = *Server->GetRuntime();
         Sender = runtime.AllocateEdgeActor();
 
-        runtime.SetLogPriority(NKikimrServices::TX_DATASHARD, NLog::PRI_NOTICE);
+        runtime.SetLogPriority(NKikimrServices::TX_DATASHARD, NLog::PRI_TRACE);
         runtime.SetLogPriority(NKikimrServices::TX_PROXY, NLog::PRI_INFO);
 
         InitRoot(Server, Sender);
@@ -818,7 +818,11 @@ struct TTestHelper {
                     break;
                 }
                 case TEvTxProcessing::EvReadSet: {
-                    if (dropRS) {
+                    auto* msg = event->Get<TEvTxProcessing::TEvReadSet>();
+                    auto flags = msg->Record.GetFlags();
+                    auto isExpect = flags & NKikimrTx::TEvReadSet::FLAG_EXPECT_READSET;
+                    auto isNoData = flags & NKikimrTx::TEvReadSet::FLAG_NO_DATA;
+                    if (dropRS && !(isExpect && isNoData)) {
                         result.ReadSets.push_back(std::move(event));
                         return TTestActorRuntime::EEventAction::DROP;
                     }
@@ -852,7 +856,10 @@ struct TTestHelper {
             )"));
         }
 
-        waitFor([&]{ return result.ReadSets.size() == 1; }, "intercepted RS");
+        const bool usesVolatileTxs = runtime.GetAppData(0).FeatureFlags.GetEnableDataShardVolatileTransactions();
+        const size_t expectedReadSets = 1 + (finalUpserts && usesVolatileTxs ? 2 : 0);
+
+        waitFor([&]{ return result.ReadSets.size() == expectedReadSets; }, "intercepted RS");
 
         // restore original observer (note we used lambda function and stack variables)
         Server->GetRuntime()->SetObserverFunc(prevObserverFunc);
@@ -2576,7 +2583,9 @@ Y_UNIT_TEST_SUITE(DataShardReadIterator) {
         TPortManager pm;
         TServerSettings serverSettings(pm.GetPort(2134));
         serverSettings.SetDomainName("Root")
-            .SetUseRealThreads(false);
+            .SetUseRealThreads(false)
+            // Blocked volatile transactions block reads, disable
+            .SetEnableDataShardVolatileTransactions(false);
 
         const ui64 shardCount = 1;
         TTestHelper helper(serverSettings, shardCount);
@@ -3600,7 +3609,7 @@ Y_UNIT_TEST_SUITE(DataShardReadIteratorPageFaults) {
         auto& runtime = *server->GetRuntime();
         auto sender = runtime.AllocateEdgeActor();
 
-        runtime.SetLogPriority(NKikimrServices::TX_DATASHARD, NLog::PRI_NOTICE);
+        runtime.SetLogPriority(NKikimrServices::TX_DATASHARD, NLog::PRI_TRACE);
         runtime.SetLogPriority(NKikimrServices::TX_PROXY, NLog::PRI_INFO);
         // runtime.SetLogPriority(NKikimrServices::TABLET_EXECUTOR, NLog::PRI_DEBUG);
 

ydb/core/tx/datashard/datashard_ut_rs.cpp

@@ -48,7 +48,10 @@ struct IsReadSet {
     {
         if (ev.GetTypeRewrite() == TEvTxProcessing::EvReadSet) {
             auto &rec = ev.Get<TEvTxProcessing::TEvReadSet>()->Record;
-            if (rec.GetTabletSource() == Source && rec.GetTabletDest() == Dest) {
+            bool isExpectation = (
+                (rec.GetFlags() & NKikimrTx::TEvReadSet::FLAG_EXPECT_READSET) &&
+                (rec.GetFlags() & NKikimrTx::TEvReadSet::FLAG_NO_DATA));
+            if (rec.GetTabletSource() == Source && rec.GetTabletDest() == Dest && !isExpectation) {
                 return true;
             }
         }
@@ -64,7 +67,9 @@ Y_UNIT_TEST_SUITE(TDataShardRSTest) {
         TPortManager pm;
         TServerSettings serverSettings(pm.GetPort(2134));
         serverSettings.SetDomainName("Root")
-            .SetUseRealThreads(false);
+            .SetUseRealThreads(false)
+            // Volatile transactions avoid storing readsets in InReadSets table
+            .SetEnableDataShardVolatileTransactions(false);
 
         Tests::TServer::TPtr server = new TServer(serverSettings);
         auto &runtime = *server->GetRuntime();
@@ -250,7 +255,13 @@ Y_UNIT_TEST_SUITE(TDataShardRSTest) {
         TPortManager pm;
         TServerSettings serverSettings(pm.GetPort(2134));
         serverSettings.SetDomainName("Root")
-            .SetUseRealThreads(false);
+            .SetUseRealThreads(false)
+            // This test expects rs acks to be delayed during one of restarts,
+            // which doesn't happen with volatile transactions. With volatile
+            // transactions both upserts have already executed, one of them is
+            // just waiting for confirmation before making changes visible.
+            // Since acks are not delayed they are just gone when dropped.
+            .SetEnableDataShardVolatileTransactions(false);
 
         Tests::TServer::TPtr server = new TServer(serverSettings);
         auto &runtime = *server->GetRuntime();
@@ -278,7 +289,10 @@ Y_UNIT_TEST_SUITE(TDataShardRSTest) {
         auto captureRS = [shard1,shard3](TAutoPtr<IEventHandle> &event) -> auto {
             if (event->GetTypeRewrite() == TEvTxProcessing::EvReadSet) {
                 auto &rec = event->Get<TEvTxProcessing::TEvReadSet>()->Record;
-                if (rec.GetTabletSource() == shard1) {
+                bool isExpectation = (
+                    (rec.GetFlags() & NKikimrTx::TEvReadSet::FLAG_EXPECT_READSET) &&
+                    (rec.GetFlags() & NKikimrTx::TEvReadSet::FLAG_NO_DATA));
+                if (rec.GetTabletSource() == shard1 && !isExpectation) {
                     return TTestActorRuntime::EEventAction::DROP;
                 }
             } else if (event->GetTypeRewrite() == TEvTxProcessing::EvReadSetAck) {
@@ -359,6 +373,9 @@ Y_UNIT_TEST_SUITE(TDataShardRSTest) {
             switch (ev->GetTypeRewrite()) {
                 case TEvTxProcessing::TEvReadSet::EventType: {
                     auto* msg = ev->Get<TEvTxProcessing::TEvReadSet>();
+                    if (msg->Record.GetFlags() & NKikimrTx::TEvReadSet::FLAG_NO_DATA) {
+                        break;
+                    }
                     NKikimrTx::TReadSetData genericData;
                     Y_ABORT_UNLESS(genericData.ParseFromString(msg->Record.GetReadSet()));
                     Cerr << "... generic readset: " << genericData.DebugString() << Endl;
@@ -419,6 +436,13 @@ Y_UNIT_TEST_SUITE(TDataShardRSTest) {
             switch (ev->GetTypeRewrite()) {
                 case TEvTxProcessing::TEvReadSet::EventType: {
                     auto* msg = ev->Get<TEvTxProcessing::TEvReadSet>();
+                    if (msg->Record.GetFlags() & NKikimrTx::TEvReadSet::FLAG_NO_DATA) {
+                        if (!(msg->Record.GetFlags() & NKikimrTx::TEvReadSet::FLAG_EXPECT_READSET)) {
+                            Cerr << "... nodata readset" << Endl;
+                            ++readSets;
+                        }
+                        break;
+                    }
                     NKikimrTx::TReadSetData genericData;
                     Y_ABORT_UNLESS(genericData.ParseFromString(msg->Record.GetReadSet()));
                     Cerr << "... generic readset: " << genericData.DebugString() << Endl;