fix: concurrency conflict in MemPool.TryRemoveUnVerified (#3500)

* fix: concurrency conflict in MemPool.TryRemoveUnVerified * Remove method * Update src/Neo/Ledger/MemoryPool.cs * clean * Apply suggestions from code review * reformat --------- Co-authored-by: Shargon <shargon@gmail.com> Co-authored-by: Jimmy <jinghui@wayne.edu>
neo-project · Sep 30, 2024 · 6e74e33 · 6e74e33
1 parent 33d0913
commit 6e74e33
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 5 deletions.
diff --git a/src/Neo/Ledger/MemoryPool.cs b/src/Neo/Ledger/MemoryPool.cs
@@ -445,12 +445,20 @@ private void RemoveConflictsOfVerified(PoolItem item)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal bool TryRemoveUnVerified(UInt256 hash, [MaybeNullWhen(false)] out PoolItem? item)
         {
-            if (!_unverifiedTransactions.TryGetValue(hash, out item))
-                return false;
+            _txRwLock.EnterWriteLock();
+            try
+            {
+                if (!_unverifiedTransactions.TryGetValue(hash, out item))
+                    return false;
 
-            _unverifiedTransactions.Remove(hash);
-            _unverifiedSortedTransactions.Remove(item);
-            return true;
+                _unverifiedTransactions.Remove(hash);
+                _unverifiedSortedTransactions.Remove(item);
+                return true;
+            }
+            finally
+            {
+                _txRwLock.ExitWriteLock();
+            }
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]

diff --git a/tests/Neo.UnitTests/Ledger/UT_MemoryPool.cs b/tests/Neo.UnitTests/Ledger/UT_MemoryPool.cs
@@ -747,6 +747,25 @@ public void TestUpdatePoolForBlockPersisted()
             _unit.VerifiedCount.Should().Be(0);
         }
 
+        [TestMethod]
+        public void TestTryRemoveUnVerified()
+        {
+            AddTransactions(32);
+            _unit.SortedTxCount.Should().Be(32);
+
+            var txs = _unit.GetSortedVerifiedTransactions().ToArray();
+            _unit.InvalidateVerifiedTransactions();
+
+            _unit.SortedTxCount.Should().Be(0);
+
+            foreach (var tx in txs)
+            {
+                _unit.TryRemoveUnVerified(tx.Hash, out _);
+            }
+
+            _unit.UnVerifiedCount.Should().Be(0);
+        }
+
         public static StorageKey CreateStorageKey(int id, byte prefix, byte[] key = null)
         {
             byte[] buffer = GC.AllocateUninitializedArray<byte>(sizeof(byte) + (key?.Length ?? 0));