DataShard: remove cancelled requests from ProposeQueue (#44543)

4 files changed, 257 insertions, 75 deletions

diff --git a/ydb/core/tx/datashard/datashard.cpp b/ydb/core/tx/datashard/datashard.cpp
index 217a65cf0ee..9fd3f03e064 100644
--- a/ydb/core/tx/datashard/datashard.cpp
+++ b/ydb/core/tx/datashard/datashard.cpp
@@ -3515,73 +3515,71 @@ void TDataShard::Handle(TEvPrivate::TEvProgressTransaction::TPtr &ev, const TAct
     ExecuteProgressTx(ctx);
 }
 
+void TDataShard::SendCancelledProposeReply(const TProposeQueue::TItem& item, const TActorContext& ctx) {
+    TActorId target = item.Event->Sender;
+    ui64 cookie = item.Event->Cookie;
+    switch (item.Event->GetTypeRewrite()) {
+        case TEvDataShard::TEvProposeTransaction::EventType: {
+            auto* msg = item.Event->Get<TEvDataShard::TEvProposeTransaction>();
+            auto kind = msg->GetTxKind();
+            auto txId = msg->GetTxId();
+            auto result = new TEvDataShard::TEvProposeTransactionResult(
+                kind, TabletID(), txId,
+                NKikimrTxDataShard::TEvProposeTransactionResult::CANCELLED);
+            ctx.Send(target, result, 0, cookie);
+            break;
+        }
+        case NEvents::TDataEvents::TEvWrite::EventType: {
+            auto* msg = item.Event->Get<NEvents::TDataEvents::TEvWrite>();
+            auto result = NEvents::TDataEvents::TEvWriteResult::BuildError(TabletID(), msg->GetTxId(), NKikimrDataEvents::TEvWriteResult::STATUS_CANCELLED, "Canceled");
+            ctx.Send(target, result.release(), 0, cookie);
+            break;
+        }
+        default:
+            Y_ENSURE(false, "Unexpected event type " << item.Event->GetTypeRewrite());
+    }
+}
+
 void TDataShard::Handle(TEvPrivate::TEvDelayedProposeTransaction::TPtr &ev, const TActorContext &ctx) {
     Y_UNUSED(ev);
     IncCounter(COUNTER_PROPOSE_QUEUE_EV);
 
     if (ProposeQueue) {
+        // N.B. Cancelled items are removed immediately in Cancel() and never reach here.
         auto item = ProposeQueue.Dequeue();
         UpdateProposeQueueSize();
 
-        TDuration latency = TAppData::TimeProvider->Now() - item.ReceivedAt;
+        TDuration latency = TAppData::TimeProvider->Now() - item->ReceivedAt;
         IncCounter(COUNTER_PROPOSE_QUEUE_LATENCY, latency);
 
-        if (!item.Cancelled) {
-            // N.B. we don't call ProposeQueue.Reset(), tx will Ack() on its first Execute()
-
-            switch (item.Event->GetTypeRewrite()) {
-                case TEvDataShard::TEvProposeTransaction::EventType: {
-                    auto event = IEventHandle::Downcast<TEvDataShard::TEvProposeTransaction>(std::move(item.Event));
-                    NWilson::TSpan datashardTransactionSpan(TWilsonTablet::TabletTopLevel, std::move(event->TraceId), "Datashard.Transaction", NWilson::EFlags::AUTO_END);
-                    if (datashardTransactionSpan) {
-                        datashardTransactionSpan.Attribute("Shard", std::to_string(TabletID()));
-                    }
-
-                    auto userCtx = NACLib::TUserContextBuilder()
-                        .DeserializeFromEventHandle(*event.Get())
-                        .Build();
-                    Execute(new TTxProposeTransactionBase(this, std::move(event), item.ReceivedAt, item.TieBreakerIndex, /* delayed */ true, std::move(datashardTransactionSpan), userCtx), ctx);
-                    return;
-                }
-                case NEvents::TDataEvents::TEvWrite::EventType: {
-                    auto event = IEventHandle::Downcast<NEvents::TDataEvents::TEvWrite>(std::move(item.Event));
-                    NWilson::TSpan datashardTransactionSpan(TWilsonTablet::TabletTopLevel, std::move(event->TraceId), "Datashard.WriteTransaction", NWilson::EFlags::AUTO_END);
-                    if (datashardTransactionSpan) {
-                        datashardTransactionSpan.Attribute("Shard", std::to_string(TabletID()));
-                    }
-
-                    Execute(new TTxWrite(this, std::move(event), item.ReceivedAt, item.TieBreakerIndex, /* delayed */ true, std::move(datashardTransactionSpan)), ctx);
-                    return;
+        // N.B. we don't call ProposeQueue.Reset(), tx will Ack() on its first Execute()
+        switch (item->Event->GetTypeRewrite()) {
+            case TEvDataShard::TEvProposeTransaction::EventType: {
+                auto event = IEventHandle::Downcast<TEvDataShard::TEvProposeTransaction>(std::move(item->Event));
+                NWilson::TSpan datashardTransactionSpan(TWilsonTablet::TabletTopLevel, std::move(event->TraceId), "Datashard.Transaction", NWilson::EFlags::AUTO_END);
+                if (datashardTransactionSpan) {
+                    datashardTransactionSpan.Attribute("Shard", std::to_string(TabletID()));
                 }
-                default:
-                    Y_ENSURE(false, "Unexpected event type " << item.Event->GetTypeRewrite());
-            }
-        }
 
-        TActorId target = item.Event->Sender;
-        ui64 cookie = item.Event->Cookie;
-        switch (item.Event->GetTypeRewrite()) {
-            case TEvDataShard::TEvProposeTransaction::EventType: {
-                auto* msg = item.Event->Get<TEvDataShard::TEvProposeTransaction>();
-                auto kind = msg->GetTxKind();
-                auto txId = msg->GetTxId();
-                auto result = new TEvDataShard::TEvProposeTransactionResult(
-                    kind, TabletID(), txId,
-                    NKikimrTxDataShard::TEvProposeTransactionResult::CANCELLED);
-                ctx.Send(target, result, 0, cookie);
-                break;
+                auto userCtx = NACLib::TUserContextBuilder()
+                    .DeserializeFromEventHandle(*event.Get())
+                    .Build();
+                Execute(new TTxProposeTransactionBase(this, std::move(event), item->ReceivedAt, item->TieBreakerIndex, /* delayed */ true, std::move(datashardTransactionSpan), userCtx), ctx);
+                return;
             }
             case NEvents::TDataEvents::TEvWrite::EventType: {
-                auto* msg = item.Event->Get<NEvents::TDataEvents::TEvWrite>();
-                auto result = NEvents::TDataEvents::TEvWriteResult::BuildError(TabletID(), msg->GetTxId(), NKikimrDataEvents::TEvWriteResult::STATUS_CANCELLED, "Canceled");
-                ctx.Send(target, result.release(), 0, cookie);
-                break;
+                auto event = IEventHandle::Downcast<NEvents::TDataEvents::TEvWrite>(std::move(item->Event));
+                NWilson::TSpan datashardTransactionSpan(TWilsonTablet::TabletTopLevel, std::move(event->TraceId), "Datashard.WriteTransaction", NWilson::EFlags::AUTO_END);
+                if (datashardTransactionSpan) {
+                    datashardTransactionSpan.Attribute("Shard", std::to_string(TabletID()));
+                }
+
+                Execute(new TTxWrite(this, std::move(event), item->ReceivedAt, item->TieBreakerIndex, /* delayed */ true, std::move(datashardTransactionSpan)), ctx);
+                return;
             }
             default:
-                Y_ENSURE(false, "Unexpected event type " << item.Event->GetTypeRewrite());
+                Y_ENSURE(false, "Unexpected event type " << item->Event->GetTypeRewrite());
         }
-
-
     }
 
     // N.B. Ack directly since we didn't start any delayed transactions
diff --git a/ydb/core/tx/datashard/datashard__cancel_tx_proposal.cpp b/ydb/core/tx/datashard/datashard__cancel_tx_proposal.cpp
index 6ce62c15e19..7697265ed66 100644
--- a/ydb/core/tx/datashard/datashard__cancel_tx_proposal.cpp
+++ b/ydb/core/tx/datashard/datashard__cancel_tx_proposal.cpp
@@ -72,8 +72,12 @@ void TDataShard::Handle(TEvDataShard::TEvCancelTransactionProposal::TPtr &ev, co
     LOG_DEBUG_S(ctx, NKikimrServices::TX_DATASHARD, "Got TEvDataShard::TEvCancelTransactionProposal " << TabletID()
                 << " txId " <<  txId);
 
-    // Mark any queued proposals as cancelled
-    ProposeQueue.Cancel(txId);
+    // Immediately remove any queued proposals with this txId from the ProposeQueue,
+    // sending CANCELLED replies right away instead of waiting for their turn in the queue.
+    ProposeQueue.Cancel(txId, [this, &ctx](const TProposeQueue::TItem& item) {
+        SendCancelledProposeReply(item, ctx);
+    });
+    UpdateProposeQueueSize();
 
     // Cancel transactions that have already been proposed
     Execute(new TTxCancelTransactionProposal(this, txId), ctx);
diff --git a/ydb/core/tx/datashard/datashard_impl.h b/ydb/core/tx/datashard/datashard_impl.h
index c2bcc5513e1..95713ea11f8 100644
--- a/ydb/core/tx/datashard/datashard_impl.h
+++ b/ydb/core/tx/datashard/datashard_impl.h
@@ -2679,20 +2679,28 @@ private:
 
     class TProposeQueue : private TTxProgressIdempotentScalarQueue<TEvPrivate::TEvDelayedProposeTransaction> {
     public:
-        struct TItem : public TMoveOnly {
+        // TItem is heap-allocated and owned by TProposeQueue (via TIntrusiveList).
+        // Using an intrusive list allows O(1) removal from an arbitrary position,
+        // which is needed to physically remove cancelled items immediately in Cancel().
+        struct TItem : public TIntrusiveListItem<TItem> {
             TItem(TAutoPtr<IEventHandle>&& event, TInstant receivedAt, ui64 tieBreakerIndex)
                 : Event(std::move(event))
                 , ReceivedAt(receivedAt)
                 , TieBreakerIndex(tieBreakerIndex)
-                , Next(nullptr)
-                , Cancelled(false)
+                , NextForTxId(nullptr)
             { }
 
+            TItem(const TItem&) = delete;
+            TItem& operator=(const TItem&) = delete;
+            TItem(TItem&&) = delete;
+            TItem& operator=(TItem&&) = delete;
+
             TAutoPtr<IEventHandle> Event;
             TInstant ReceivedAt;
             ui64 TieBreakerIndex;
-            TItem* Next;
-            bool Cancelled;
+            // Linked list of items with the same txId within Items.
+            // N.B. there should almost always be exactly one propose per txId.
+            TItem* NextForTxId;
         };
 
         struct TItemList {
@@ -2700,14 +2708,23 @@ private:
             TItem* Last = nullptr;
         };
 
+        ~TProposeQueue() {
+            // Items are heap-allocated; destroy all of them.
+            while (!Items.Empty()) {
+                delete Items.PopFront();
+            }
+        }
+
         void Enqueue(TAutoPtr<IEventHandle> event, TInstant receivedAt, ui64 tieBreakerIndex, const TActorContext& ctx) {
-            TItem* item = &Items.emplace_back(std::move(event), receivedAt, tieBreakerIndex);
+            TItem* item = new TItem(std::move(event), receivedAt, tieBreakerIndex);
+            Items.PushBack(item);
+            ++Count;
 
             const ui64 txId = NEvWrite::TConvertor::GetTxId(item->Event);
 
             auto& links = TxIds[txId];
             if (Y_UNLIKELY(links.Last)) {
-                links.Last->Next = item;
+                links.Last->NextForTxId = item;
             } else {
                 links.First = item;
             }
@@ -2716,62 +2733,80 @@ private:
             Progress(ctx);
         }
 
-        TItem Dequeue() {
-            TItem* first = &Items.front();
+        // Removes the front item from the queue and transfers ownership to the caller.
+        THolder<TItem> Dequeue() {
+            Y_ENSURE(!Items.Empty());
+            TItem* first = Items.Front();
             const ui64 txId = NEvWrite::TConvertor::GetTxId(first->Event);
 
             auto it = TxIds.find(txId);
             Y_ENSURE(it != TxIds.end() && it->second.First == first,
-                "Consistency check: proposed txId " << txId << " in deque, but not in hashmap");
+                "Consistency check: proposed txId " << txId << " in list, but not in hashmap");
 
             // N.B. there should almost always be exactly one propose per txId
-            it->second.First = first->Next;
+            it->second.First = first->NextForTxId;
             if (Y_LIKELY(it->second.First == nullptr)) {
                 TxIds.erase(it);
             } else {
-                first->Next = nullptr;
+                first->NextForTxId = nullptr;
             }
 
-            TItem item = std::move(*first);
-            Items.pop_front();
-            return item;
+            first->Unlink();
+            --Count;
+            return THolder<TItem>(first);
         }
 
-        void Cancel(ui64 txId) {
+        // Physically removes all items with the given txId from the queue and calls
+        // onCancelled(item) for each one before deleting it.  Items are removed in
+        // O(1) per item (intrusive list unlink), so the total cost is O(k) where k
+        // is the number of items with this txId (almost always 1).
+        // The HasInFly slot is intentionally NOT released here: if the cancelled item
+        // was at the head of the queue the already-scheduled TEvDelayedProposeTransaction
+        // will arrive and find either the next non-cancelled item or an empty queue.
+        template <typename TOnCancelled>
+        void Cancel(ui64 txId, TOnCancelled&& onCancelled) {
             auto it = TxIds.find(txId);
             if (it != TxIds.end()) {
-                auto* item = it->second.First;
+                TItem* item = it->second.First;
                 while (item) {
-                    item->Cancelled = true;
-                    item = item->Next;
+                    TItem* next = item->NextForTxId;
+                    item->Unlink();
+                    --Count;
+                    THolder<TItem> holder(item);
+                    onCancelled(*holder);  // now Size() is accurate
+                    item = next;
                 }
+                TxIds.erase(it);
             }
         }
 
         void Ack(const TActorContext& ctx) {
             Reset(ctx);
-            if (Items) {
+            if (!Items.Empty()) {
                 Progress(ctx);
             }
         }
 
         explicit operator bool() const {
-            return bool(Items);
+            return !Items.Empty();
         }
 
         size_t Size() const {
-            return Items.size();
+            return Count;
         }
 
     private:
-        TDeque<TItem> Items;
+        TIntrusiveList<TItem> Items;
         THashMap<ui64, TItemList> TxIds;
+        size_t Count = 0;
     };
 
     TProposeQueue ProposeQueue;
     TVector<THolder<IEventHandle>> DelayedProposeQueue;
     TAsyncEvent DelayedProposeCoroutines;
 
+    void SendCancelledProposeReply(const TProposeQueue::TItem& item, const TActorContext& ctx);
+
     TActorId PersistentPipeCache;
     NTabletPipe::TClientRetryPolicy SchemeShardPipeRetryPolicy;
     TActorId SchemeShardPipe;   // For notifications about schema changes
diff --git a/ydb/core/tx/datashard/datashard_ut_write.cpp b/ydb/core/tx/datashard/datashard_ut_write.cpp
index 9bed7c639ce..e6152cce1be 100644
--- a/ydb/core/tx/datashard/datashard_ut_write.cpp
+++ b/ydb/core/tx/datashard/datashard_ut_write.cpp
@@ -1306,6 +1306,151 @@ Y_UNIT_TEST_SUITE(DataShardWrite) {
         }
     }
 
+    Y_UNIT_TEST(CancelMultipleImmediateTransactionsFromQueue) {
+        auto [runtime, server, sender] = TestCreateServer();
+
+        TShardedTableOptions opts;
+        const TString tableName = "table-1";
+        const auto [shards, tableId] = CreateShardedTable(server, sender, "/Root", tableName, opts);
+        const ui64 shard = shards[0];
+
+        TActorId shardActorId = ResolveTablet(runtime, shard, 0, false);
+
+        const ui64 txId1 = 101, txId2 = 102, txId3 = 103, txId4 = 104, txId5 = 105;
+
+        auto sendWrite = [&](ui64 txId, ui32 key, ui32 value) {
+            auto request = MakeWriteRequestOneKeyValue(
+                txId,
+                NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE,
+                NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPSERT,
+                tableId, opts.Columns_, key, value);
+            runtime.Send(new IEventHandle(shardActorId, sender, request.release()), 0, true);
+        };
+
+        // All 6 writes and both cancels are sent with viaActorSystem=true — they are
+        // queued in the shard mailbox but NOT dispatched yet.  Because the actor mailbox
+        // is FIFO, the shard will:
+        //   1. Enqueue tx1 → ProposeQueue sends TEvDelayedProposeTransaction to self (to
+        //      the END of the same mailbox, after all messages already in it)
+        //   2. Enqueue tx2 (first), tx2 (retry), tx3, tx4, tx5 (no new drive event —
+        //      idempotent scalar queue).  The two tx2 items are linked via NextForTxId.
+        //   3. Cancel tx2 → ProposeQueue.Cancel walks the NextForTxId chain, removes
+        //      BOTH tx2 items, sends 2 CANCELLED replies
+        //   4. Cancel tx3 → ProposeQueue.Cancel removes tx3, sends 1 CANCELLED reply
+        //   5. TEvDelayedProposeTransaction fires → drains remaining queue: tx1, tx4, tx5
+        // So we expect 6 total replies: 2×CANCELLED(tx2) + 1×CANCELLED(tx3) +
+        // 3×COMPLETED(tx1, tx4, tx5).
+
+        Cout << "========= Enqueue 6 immediate writes into ProposeQueue =========\n";
+        sendWrite(txId1, 1, 10);
+        sendWrite(txId2, 2, 20);         // first attempt
+        sendWrite(txId2, 2, 20);         // retry of the same txId — goes into NextForTxId chain
+        sendWrite(txId3, 3, 30);
+        sendWrite(txId4, 4, 40);
+        sendWrite(txId5, 5, 50);
+
+        Cout << "========= Send cancels for tx2 and tx3 (still not dispatched) =========\n";
+        {
+            auto cancel2 = std::make_unique<TEvDataShard::TEvCancelTransactionProposal>(txId2);
+            runtime.Send(new IEventHandle(shardActorId, sender, cancel2.release()), 0, true);
+            auto cancel3 = std::make_unique<TEvDataShard::TEvCancelTransactionProposal>(txId3);
+            runtime.Send(new IEventHandle(shardActorId, sender, cancel3.release()), 0, true);
+        }
+
+        Cout << "========= Collect 6 write results =========\n";
+        // Use GrabEdgeEventRethrow directly: WaitForWriteCompleted asserts a fixed status
+        // and would fail if it received STATUS_CANCELLED for tx2/tx3.
+        // Cancel(txId2) sends 2 replies (one per NextForTxId chain element), so we
+        // expect 6 results total, not 5.
+        THashMap<ui64, NKikimrDataEvents::TEvWriteResult::EStatus> results;
+        ui32 cancelledTx2Count = 0;
+        for (int i = 0; i < 6; ++i) {
+            auto ev = runtime.GrabEdgeEventRethrow<NEvents::TDataEvents::TEvWriteResult>(sender);
+            const auto& rec = ev->Get()->Record;
+            if (rec.GetTxId() == txId2 &&
+                rec.GetStatus() == NKikimrDataEvents::TEvWriteResult::STATUS_CANCELLED) {
+                ++cancelledTx2Count;
+            }
+            results[rec.GetTxId()] = rec.GetStatus();
+        }
+
+        // Both retries of tx2 must be cancelled
+        UNIT_ASSERT_VALUES_EQUAL(cancelledTx2Count, 2u);
+        UNIT_ASSERT_VALUES_EQUAL(results[txId1], NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+        UNIT_ASSERT_VALUES_EQUAL(results[txId2], NKikimrDataEvents::TEvWriteResult::STATUS_CANCELLED);
+        UNIT_ASSERT_VALUES_EQUAL(results[txId3], NKikimrDataEvents::TEvWriteResult::STATUS_CANCELLED);
+        UNIT_ASSERT_VALUES_EQUAL(results[txId4], NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+        UNIT_ASSERT_VALUES_EQUAL(results[txId5], NKikimrDataEvents::TEvWriteResult::STATUS_COMPLETED);
+
+        Cout << "========= Verify table state: only rows 1, 4, 5 =========\n";
+        {
+            auto tableData = ReadShardedTable(server, "/Root/table-1");
+            UNIT_ASSERT_VALUES_EQUAL(tableData,
+                "key = 1, value = 10\n"
+                "key = 4, value = 40\n"
+                "key = 5, value = 50\n");
+        }
+
+        Cout << "========= Verify shard still works correctly =========\n";
+        {
+            Upsert(runtime, sender, shard, tableId, opts.Columns_, 1, 106,
+                NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE);
+        }
+    }
+
+    // Companion test to CancelMultipleFromQueue that exercises the
+    // TEvProposeTransaction::EventType branch in SendCancelledProposeReply.
+    // Enqueues a TX_KIND_DATA|Immediate TEvProposeTransaction (old MiniKQL path)
+    // into ProposeQueue and cancels it before the queue is drained.
+    Y_UNIT_TEST(CancelImmediateProposeTransaction) {
+        auto [runtime, server, sender] = TestCreateServer();
+
+        TShardedTableOptions opts;
+        const TString tableName = "table-1";
+        const auto [shards, tableId] = CreateShardedTable(server, sender, "/Root", tableName, opts);
+        const ui64 shard = shards[0];
+        Y_UNUSED(shard);
+        Y_UNUSED(tableId);
+
+        TActorId shardActorId = ResolveTablet(runtime, shard, 0, false);
+
+        const ui64 txId = 201;
+
+        // Both messages are sent with viaActorSystem=true — they sit in the shard
+        // mailbox in FIFO order and are NOT dispatched until the runtime runs.
+        // ProposeQueue appends TEvDelayedProposeTransaction to the END of the
+        // mailbox when the first propose is enqueued, so the order is:
+        //   propose(txId) → cancel(txId) → TEvDelayedProposeTransaction
+        // The cancel fires before the drive event, so ProposeQueue.Cancel finds
+        // the item, removes it, and calls SendCancelledProposeReply through the
+        // TEvProposeTransaction::EventType branch, sending CANCELLED immediately.
+        // When TEvDelayedProposeTransaction fires the queue is already empty.
+        Cout << "========= Enqueue immediate TEvProposeTransaction into ProposeQueue =========\n";
+        {
+            auto request = std::make_unique<TEvDataShard::TEvProposeTransaction>(
+                NKikimrTxDataShard::TX_KIND_DATA,
+                sender,
+                txId,
+                /* txBody = */ TString{},
+                NDataShard::TTxFlags::Immediate);
+            runtime.Send(new IEventHandle(shardActorId, sender, request.release()), 0, true);
+        }
+
+        Cout << "========= Cancel while still in ProposeQueue =========\n";
+        {
+            auto cancel = std::make_unique<TEvDataShard::TEvCancelTransactionProposal>(txId);
+            runtime.Send(new IEventHandle(shardActorId, sender, cancel.release()), 0, true);
+        }
+
+        Cout << "========= Expect CANCELLED TEvProposeTransactionResult =========\n";
+        {
+            auto ev = runtime.GrabEdgeEventRethrow<TEvDataShard::TEvProposeTransactionResult>(sender);
+            UNIT_ASSERT_VALUES_EQUAL(ev->Get()->GetTxId(), txId);
+            UNIT_ASSERT_VALUES_EQUAL(ev->Get()->GetStatus(),
+                NKikimrTxDataShard::TEvProposeTransactionResult::CANCELLED);
+        }
+    }
+
     Y_UNIT_TEST_TWIN(UpsertPreparedManyTables, Volatile) {
         auto [runtime, server, sender] = TestCreateServer();