Fix review issues

Author: CyberROFL

ydb/core/tx/datashard/datashard.cpp

@@ -864,19 +864,36 @@ void TDataShard::PersistChangeRecord(NIceDb::TNiceDb& db, const TChangeRecord& r
                 TSchemaSnapshotKey(res.first->second.TableId, res.first->second.SchemaVersion));
         }
 
-        db.GetDatabase().OnRollback([this, order = record.GetOrder()] {
-            auto it = ChangesQueue.find(order);
-            Y_VERIFY_S(it != ChangesQueue.end(), "Cannot find change record: " << order);
+        const auto key = TCommittingChangeRecordsKey::FromRecord(record);
+        if (!CommittingChangeRecords.contains(key)) {
+            db.GetDatabase().OnCommit([this, key] {
+                auto it = CommittingChangeRecords.find(key);
+                Y_ABORT_UNLESS(it != CommittingChangeRecords.end());
+                CommittingChangeRecords.erase(it);
+            });
+            db.GetDatabase().OnRollback([this, key] {
+                auto it = CommittingChangeRecords.find(key);
+                Y_ABORT_UNLESS(it != CommittingChangeRecords.end());
+
+                for (const auto order : it->second) {
+                    auto cIt = ChangesQueue.find(order);
+                    Y_VERIFY_S(cIt != ChangesQueue.end(), "Cannot find change record: " << order);
+
+                    if (cIt->second.SchemaSnapshotAcquired) {
+                        const auto snapshotKey = TSchemaSnapshotKey(cIt->second.TableId, cIt->second.SchemaVersion);
+                        if (const auto last = SchemaSnapshotManager.ReleaseReference(snapshotKey)) {
+                            ScheduleRemoveSchemaSnapshot(snapshotKey);
+                        }
+                    }
 
-            if (it->second.SchemaSnapshotAcquired) {
-                const auto snapshotKey = TSchemaSnapshotKey(it->second.TableId, it->second.SchemaVersion);
-                if (const auto last = SchemaSnapshotManager.ReleaseReference(snapshotKey)) {
-                    ScheduleRemoveSchemaSnapshot(snapshotKey);
+                    ChangesQueue.erase(cIt);
                 }
-            }
 
-            ChangesQueue.erase(it);
-        });
+                CommittingChangeRecords.erase(it);
+            });
+        }
+
+        CommittingChangeRecords[key].push_back(record.GetOrder());
     } else {
         auto& state = LockChangeRecords[lockId];
         Y_ABORT_UNLESS(state.Changes.empty() || state.Changes.back().LockOffset < record.GetLockOffset(),

ydb/core/tx/datashard/datashard_impl.h

@@ -2861,6 +2861,20 @@ class TDataShard
     bool OutChangeSenderSuspended = false;
     THolder<IChangeRecordSerializer> ChangeRecordDebugSerializer;
 
+    struct TCommittingChangeRecordsKey: std::tuple<ui64, ui64, ui64> {
+        explicit TCommittingChangeRecordsKey(ui64 group, ui64 step, ui64 txId)
+            : std::tuple<ui64, ui64, ui64>(group, step, txId)
+        {}
+
+        static TCommittingChangeRecordsKey FromRecord(const TChangeRecord& record) {
+            return TCommittingChangeRecordsKey(record.GetGroup(), record.GetStep(), record.GetTxId());
+        }
+
+        explicit operator size_t() const noexcept {
+            return THash<std::tuple<ui64, ui64, ui64>>()(*this);
+        }
+    };
+
     struct TUncommittedLockChangeRecords {
         TVector<IDataShardChangeCollector::TChange> Changes;
         size_t PersistentCount = 0;
@@ -2876,6 +2890,7 @@ class TDataShard
         size_t Count = 0;
     };
 
+    THashMap<TCommittingChangeRecordsKey, TVector<ui64>> CommittingChangeRecords;
     THashMap<ui64, TUncommittedLockChangeRecords> LockChangeRecords; // ui64 is lock id
     THashMap<ui64, TCommittedLockChangeRecords> CommittedLockChangeRecords; // ui64 is lock id
     TVector<ui64> PendingLockChangeRecordsToRemove;