Fix use-after-free in CommittingOps tracking (#8712)

Author: snaury

ydb/core/tx/datashard/datashard_pipeline.cpp

@@ -2274,11 +2274,15 @@ void TPipeline::AddCommittingOp(const TOperation::TPtr& op) {
     if (!Self->IsMvccEnabled() || op->IsReadOnly())
         return;
 
+    Y_VERIFY_S(!op->GetCommittingOpsVersion(),
+        "Trying to AddCommittingOp " << *op << " more than once");
+
     TRowVersion version = Self->GetReadWriteVersions(op.Get()).WriteVersion;
     if (op->IsImmediate())
         CommittingOps.Add(op->GetTxId(), version);
     else
         CommittingOps.Add(version);
+    op->SetCommittingOpsVersion(version);
 }
 
 void TPipeline::RemoveCommittingOp(const TRowVersion& version) {
@@ -2288,13 +2292,13 @@ void TPipeline::RemoveCommittingOp(const TRowVersion& version) {
 }
 
 void TPipeline::RemoveCommittingOp(const TOperation::TPtr& op) {
-    if (!Self->IsMvccEnabled() || op->IsReadOnly())
-        return;
-
-    if (op->IsImmediate())
-        CommittingOps.Remove(op->GetTxId());
-    else
-        CommittingOps.Remove(TRowVersion(op->GetStep(), op->GetTxId()));
+    if (const auto& version = op->GetCommittingOpsVersion()) {
+        if (op->IsImmediate())
+            CommittingOps.Remove(op->GetTxId(), *version);
+        else
+            CommittingOps.Remove(*version);
+        op->ResetCommittingOpsVersion();
+    }
 }
 
 bool TPipeline::WaitCompletion(const TOperation::TPtr& op) const {

ydb/core/tx/datashard/datashard_pipeline.h

@@ -424,11 +424,13 @@ class TPipeline : TNonCopyable {
             ui64 Step;
             ui64 TxId;
             mutable ui32 Counter;
+            mutable ui32 TxCounter;
 
             TItem(const TRowVersion& from)
                 : Step(from.Step)
                 , TxId(from.TxId)
                 , Counter(1u)
+                , TxCounter(0u)
             {}
 
             friend constexpr bool operator<(const TItem& a, const TItem& b) {
@@ -442,6 +444,7 @@ class TPipeline : TNonCopyable {
 
         using TItemsSet = TSet<TItem>;
         using TTxIdMap = THashMap<ui64, TItemsSet::iterator>;
+
     public:
         inline void Add(ui64 txId, TRowVersion version) {
             auto res = ItemsSet.emplace(version);
@@ -450,6 +453,7 @@ class TPipeline : TNonCopyable {
             auto res2 = TxIdMap.emplace(txId, res.first);
             Y_VERIFY_S(res2.second, "Unexpected duplicate immediate tx " << txId
                     << " committing at " << version);
+            res.first->TxCounter += 1;
         }
 
         inline void Add(TRowVersion version) {
@@ -458,17 +462,29 @@ class TPipeline : TNonCopyable {
                 res.first->Counter += 1;
         }
 
-        inline void Remove(ui64 txId) {
-            if (auto it = TxIdMap.find(txId); it != TxIdMap.end()) {
-                if (--it->second->Counter == 0)
-                    ItemsSet.erase(it->second);
-                TxIdMap.erase(it);
-            }
+        inline void Remove(ui64 txId, TRowVersion version) {
+            auto it = TxIdMap.find(txId);
+            Y_VERIFY_S(it != TxIdMap.end(), "Removing immediate tx " << txId << " " << version
+                    << " does not match a previous Add");
+            Y_VERIFY_S(TRowVersion(it->second->Step, it->second->TxId) == version, "Removing immediate tx " << txId << " " << version
+                    << " does not match a previous Add " << TRowVersion(it->second->Step, it->second->TxId));
+            Y_VERIFY_S(it->second->TxCounter > 0, "Removing immediate tx " << txId << " " << version
+                    << " with a mismatching TxCounter");
+            --it->second->TxCounter;
+            if (--it->second->Counter == 0)
+                ItemsSet.erase(it->second);
+            TxIdMap.erase(it);
         }
 
         inline void Remove(TRowVersion version) {
-            if (auto it = ItemsSet.find(version); it != ItemsSet.end() && --it->Counter == 0)
+            auto it = ItemsSet.find(version);
+            Y_VERIFY_S(it != ItemsSet.end(), "Removing version " << version
+                    << " does not match a previous Add");
+            if (--it->Counter == 0) {
+                Y_VERIFY_S(it->TxCounter == 0, "Removing version " << version
+                    << " while TxCounter has active references, possible Add/Remove mismatch");
                 ItemsSet.erase(it);
+            }
         }
 
         inline bool HasOpsBelow(TRowVersion upperBound) const {