From fec14ab89a942b868482ee378516100d13973566 Mon Sep 17 00:00:00 2001
From: erikzhang <erik@neo.org>
Date: Tue, 26 Nov 2019 19:00:34 +0800
Subject: [PATCH] Use ConcurrentDictionary

---
 src/neo/Persistence/Memory/Snapshot.cs | 20 +++++++-------------
 src/neo/Persistence/Memory/Store.cs    | 22 ++++++----------------
 2 files changed, 13 insertions(+), 29 deletions(-)

diff --git a/src/neo/Persistence/Memory/Snapshot.cs b/src/neo/Persistence/Memory/Snapshot.cs
index 0d1b4ad98a..82455ddea8 100644
--- a/src/neo/Persistence/Memory/Snapshot.cs
+++ b/src/neo/Persistence/Memory/Snapshot.cs
@@ -1,40 +1,34 @@
 using Neo.IO;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
-using System.Threading;
 
 namespace Neo.Persistence.Memory
 {
     internal class Snapshot : ISnapshot
     {
-        private readonly ReaderWriterLockSlim readerWriterLock;
-        private readonly Dictionary<byte[], byte[]>[] innerData;
+        private readonly ConcurrentDictionary<byte[], byte[]>[] innerData;
         private readonly ImmutableDictionary<byte[], byte[]>[] immutableData;
-        private readonly Dictionary<byte[], byte[]>[] writeBatch;
+        private readonly ConcurrentDictionary<byte[], byte[]>[] writeBatch;
 
-        public Snapshot(Dictionary<byte[], byte[]>[] innerData, ReaderWriterLockSlim readerWriterLock)
+        public Snapshot(ConcurrentDictionary<byte[], byte[]>[] innerData)
         {
-            this.readerWriterLock = readerWriterLock;
             this.innerData = innerData;
-            readerWriterLock.EnterReadLock();
             this.immutableData = innerData.Select(p => p.ToImmutableDictionary(ByteArrayEqualityComparer.Default)).ToArray();
-            readerWriterLock.ExitReadLock();
-            this.writeBatch = new Dictionary<byte[], byte[]>[innerData.Length];
+            this.writeBatch = new ConcurrentDictionary<byte[], byte[]>[innerData.Length];
             for (int i = 0; i < writeBatch.Length; i++)
-                writeBatch[i] = new Dictionary<byte[], byte[]>(ByteArrayEqualityComparer.Default);
+                writeBatch[i] = new ConcurrentDictionary<byte[], byte[]>(ByteArrayEqualityComparer.Default);
         }
 
         public void Commit()
         {
-            readerWriterLock.EnterWriteLock();
             for (int i = 0; i < writeBatch.Length; i++)
                 foreach (var pair in writeBatch[i])
                     if (pair.Value is null)
-                        innerData[i].Remove(pair.Key);
+                        innerData[i].TryRemove(pair.Key, out _);
                     else
                         innerData[i][pair.Key] = pair.Value;
-            readerWriterLock.ExitWriteLock();
         }
 
         public void Delete(byte table, byte[] key)
diff --git a/src/neo/Persistence/Memory/Store.cs b/src/neo/Persistence/Memory/Store.cs
index 39944b903d..5e2d11eeb9 100644
--- a/src/neo/Persistence/Memory/Store.cs
+++ b/src/neo/Persistence/Memory/Store.cs
@@ -1,33 +1,29 @@
 using Neo.IO;
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
-using System.Threading;
 
 namespace Neo.Persistence.Memory
 {
     public class Store : IStore
     {
-        private readonly ReaderWriterLockSlim readerWriterLock = new ReaderWriterLockSlim();
-        private readonly Dictionary<byte[], byte[]>[] innerData;
+        private readonly ConcurrentDictionary<byte[], byte[]>[] innerData;
 
         public Store()
         {
-            innerData = new Dictionary<byte[], byte[]>[256];
+            innerData = new ConcurrentDictionary<byte[], byte[]>[256];
             for (int i = 0; i < innerData.Length; i++)
-                innerData[i] = new Dictionary<byte[], byte[]>(ByteArrayEqualityComparer.Default);
+                innerData[i] = new ConcurrentDictionary<byte[], byte[]>(ByteArrayEqualityComparer.Default);
         }
 
         public void Delete(byte table, byte[] key)
         {
-            readerWriterLock.EnterWriteLock();
-            innerData[table].Remove(key.EnsureNotNull());
-            readerWriterLock.ExitWriteLock();
+            innerData[table].TryRemove(key.EnsureNotNull(), out _);
         }
 
         public void Dispose()
         {
-            readerWriterLock.Dispose();
         }
 
         public IEnumerable<(byte[] Key, byte[] Value)> Find(byte table, byte[] prefix)
@@ -36,15 +32,13 @@ public IEnumerable<(byte[] Key, byte[] Value)> Find(byte table, byte[] prefix)
             if (prefix?.Length > 0)
                 records = records.Where(p => p.Key.Length >= prefix.Length && p.Key.Take(prefix.Length).SequenceEqual(prefix));
             records = records.OrderBy(p => p.Key, ByteArrayComparer.Default);
-            readerWriterLock.EnterReadLock();
             foreach (var pair in records)
                 yield return (pair.Key, pair.Value);
-            readerWriterLock.ExitReadLock();
         }
 
         public ISnapshot GetSnapshot()
         {
-            return new Snapshot(innerData, readerWriterLock);
+            return new Snapshot(innerData);
         }
 
         public void Put(byte table, byte[] key, byte[] value)
@@ -54,16 +48,12 @@ public void Put(byte table, byte[] key, byte[] value)
 
         public void PutSync(byte table, byte[] key, byte[] value)
         {
-            readerWriterLock.EnterWriteLock();
             innerData[table][key.EnsureNotNull()] = value;
-            readerWriterLock.ExitWriteLock();
         }
 
         public byte[] TryGet(byte table, byte[] key)
         {
-            readerWriterLock.EnterReadLock();
             innerData[table].TryGetValue(key.EnsureNotNull(), out byte[] value);
-            readerWriterLock.ExitReadLock();
             return value;
         }
     }