DataShard: prioritize cancellation processing for reads

ydb/core/tablet_flat/flat_executor.cpp

@@ -224,10 +224,10 @@ void TExecutor::RecreatePageCollectionsCache() noexcept
 
     if (TransactionWaitPads) {
         for (auto &xpair : TransactionWaitPads) {
+            xpair.second->WaitingSpan.EndOk();
             TSeat* seat = xpair.second->Seat;
             Y_ABORT_UNLESS(seat->State == ESeatState::Waiting);
             seat->State = ESeatState::None;
-            xpair.second->WaitingSpan.EndOk();
 
             if (seat->Cancelled) {
                 FinishCancellation(seat, false);
@@ -627,11 +627,11 @@ void TExecutor::ActivateWaitingTransactions(TPrivatePageCache::TPage::TWaitQueue
         bool cancelled = false;
         while (TPrivatePageCacheWaitPad *waitPad = waitPadsQueue->Pop()) {
             if (auto it = TransactionWaitPads.find(waitPad); it != TransactionWaitPads.end()) {
+                it->second->WaitingSpan.EndOk();
                 TSeat* seat = it->second->Seat;
                 Y_ABORT_UNLESS(seat->State == ESeatState::Waiting);
                 seat->State = ESeatState::None;
-                seat->WaitingSpan.EndOk();
-                TransactionWaitPads.erase(waitPad);
+                TransactionWaitPads.erase(it);
                 if (seat->Cancelled) {
                     FinishCancellation(seat, false);
                     cancelled = true;

ydb/core/tx/datashard/datashard__read_iterator.cpp

@@ -2399,6 +2399,7 @@ class TDataShard::TTxReadViaPipeline : public NTabletFlatExecutor::TTransactionB
                         return true;
                     }
                     auto& state = *readIt->second;
+                    state.EnqueuedLocalTxId = 0;
                     ReplyError(
                         Ydb::StatusIds::INTERNAL_ERROR,
                         TStringBuilder() << "Failed to sync follower: " << errMessage
@@ -2417,6 +2418,7 @@ class TDataShard::TTxReadViaPipeline : public NTabletFlatExecutor::TTransactionB
                 const auto& record = request->Record;
 
                 Y_ABORT_UNLESS(state.State == TReadIteratorState::EState::Init);
+                state.EnqueuedLocalTxId = 0;
 
                 bool setUsingSnapshotFlag = false;
 
@@ -3221,7 +3223,9 @@ void TDataShard::Handle(TEvDataShard::TEvRead::TPtr& ev, const TActorContext& ct
 
     SetCounter(COUNTER_READ_ITERATORS_COUNT, ReadIterators.size());
 
-    Executor()->Execute(new TTxReadViaPipeline(this, localReadId, request->ReadSpan.GetTraceId()), ctx);
+    // Note: we may have a read cancellation already in the mailbox, so we
+    // enqueue a low-priority transaction.
+    state.EnqueuedLocalTxId = Executor()->EnqueueLowPriority(new TTxReadViaPipeline(this, localReadId, request->ReadSpan.GetTraceId()));
 }
 
 void TDataShard::Handle(TEvDataShard::TEvReadContinue::TPtr& ev, const TActorContext& ctx) {
@@ -3454,6 +3458,9 @@ void TDataShard::DeleteReadIterator(TReadIteratorsMap::iterator it) {
     if (state.IsExhausted()) {
         DecCounter(COUNTER_READ_ITERATORS_EXHAUSTED_COUNT);
     }
+    if (state.EnqueuedLocalTxId) {
+        Executor()->CancelTransaction(state.EnqueuedLocalTxId);
+    }
     ReadIteratorsByLocalReadId.erase(state.LocalReadId);
     ReadIterators.erase(it);
     SetCounter(COUNTER_READ_ITERATORS_COUNT, ReadIterators.size());

ydb/core/tx/datashard/datashard_ut_read_iterator.cpp

@@ -4969,4 +4969,41 @@ Y_UNIT_TEST_SUITE(DataShardReadIteratorBatchMode) {
 
 }
 
+Y_UNIT_TEST_SUITE(DataShardReadIteratorFastCancel) {
+
+    Y_UNIT_TEST(ShouldProcessFastCancel) {
+        TTestHelper helper;
+
+        auto request1 = helper.GetBaseReadRequest("table-1", 1);
+        AddRangeQuery<ui32>(*request1, { 1 }, true, { 1 }, true);
+        auto readResult1 = helper.SendRead("table-1", request1.release());
+        CheckResult(helper.Tables.at("table-1").UserTable, *readResult1, {
+            {1, 1, 1, 100},
+        });
+        UNIT_ASSERT(readResult1->Record.GetFinished());
+
+        auto& runtime = *helper.Server->GetRuntime();
+        auto sender = runtime.AllocateEdgeActor();
+        auto shardActor = ResolveTablet(runtime, helper.Tables.at("table-1").TabletId);
+
+        auto request2 = helper.GetBaseReadRequest("table-1", 2);
+        AddRangeQuery<ui32>(*request2, { 1 }, true, { 1 }, true);
+        auto request3 = helper.GetBaseReadRequest("table-1", 3);
+        AddRangeQuery<ui32>(*request3, { 1 }, true, { 1 }, true);
+        auto cancel2 = std::make_unique<TEvDataShard::TEvReadCancel>();
+        cancel2->Record.SetReadId(2);
+        auto cancel3 = std::make_unique<TEvDataShard::TEvReadCancel>();
+        cancel3->Record.SetReadId(3);
+
+        runtime.Send(new IEventHandle(shardActor, sender, request2.release()), 0, true);
+        runtime.Send(new IEventHandle(shardActor, sender, request3.release()), 0, true);
+        runtime.Send(new IEventHandle(shardActor, sender, cancel2.release()), 0, true);
+        runtime.Send(new IEventHandle(shardActor, sender, cancel3.release()), 0, true);
+
+        auto ev = runtime.GrabEdgeEvent<TEvDataShard::TEvReadResult>(sender, TDuration::Seconds(1));
+        UNIT_ASSERT_C(!ev, "Unexpected response received (should have been cancelled)");
+    }
+
+}
+
 } // namespace NKikimr

ydb/core/tx/datashard/datashard_ut_trace.cpp

@@ -359,6 +359,7 @@ Y_UNIT_TEST_SUITE(TDataShardTrace) {
                         Repeat(
                             ExpectedSpan("Datashard.Read",
                                 ExpectedSpan("Tablet.Transaction",
+                                    ExpectedSpan("Tablet.Transaction.Enqueued"),
                                     ExpectedSpan("Tablet.Transaction.Execute",
                                         Repeat("Datashard.Unit", 3)),
                                     // No extra page fault with btree index (root is in meta)
@@ -443,6 +444,7 @@ Y_UNIT_TEST_SUITE(TDataShardTrace) {
                         Repeat(
                             ExpectedSpan("Datashard.Read",
                                 ExpectedSpan("Tablet.Transaction",
+                                    ExpectedSpan("Tablet.Transaction.Enqueued"),
                                     ExpectedSpan("Tablet.Transaction.Execute",
                                         Repeat("Datashard.Unit", 4)),
                                     "Tablet.Transaction.Complete"),

ydb/core/tx/datashard/read_iterator.h

@@ -218,6 +218,9 @@ struct TReadIteratorState {
     TActorId ScanActorId;
     // temporary storage for forwarded events until scan has started
     std::vector<std::unique_ptr<IEventHandle>> ScanPendingEvents;
+
+    // May be used to cancel enqueued transactions
+    ui64 EnqueuedLocalTxId = 0;
 };
 
 using TReadIteratorsMap = THashMap<TReadIteratorId, TReadIteratorState, TReadIteratorId::THash>;