Fix inacccurate memory metric due to concurrent writes (#7600)

NethermindEth · Oct 14, 2024 · 8e716e9 · 8e716e9
1 parent 6b4b485
commit 8e716e9
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 3 deletions.
diff --git a/src/Nethermind/Nethermind.Trie.Test/Pruning/TreeStoreTests.cs b/src/Nethermind/Nethermind.Trie.Test/Pruning/TreeStoreTests.cs
@@ -214,6 +214,32 @@ public void Memory_with_two_nodes_is_correct()
                 trieNode2.GetMemorySize(false) + ExpectedPerNodeKeyMemorySize);
         }
 
+        [Test]
+        public void Memory_with_concurrent_commits_is_correct()
+        {
+            using TrieStore fullTrieStore = CreateTrieStore(pruningStrategy: new TestPruningStrategy(true));
+
+            IScopedTrieStore trieStore = fullTrieStore.GetTrieStore(null);
+            PatriciaTree tree = new PatriciaTree(trieStore, LimboLogs.Instance);
+
+            Random rand = new Random(0);
+
+            Span<byte> key = stackalloc byte[32];
+            Span<byte> value = stackalloc byte[32];
+            for (int i = 0; i < 1000; i++)
+            {
+                rand.NextBytes(key);
+                rand.NextBytes(value);
+
+                tree.Set(key, value.ToArray());
+            }
+
+            tree.Commit(0);
+
+            fullTrieStore.MemoryUsedByDirtyCache.Should().Be(_scheme == INodeStorage.KeyScheme.Hash ? 591672L : 661820L);
+            fullTrieStore.CommittedNodesCount.Should().Be(1349);
+        }
+
         [Test]
         public void Memory_with_two_times_two_nodes_is_correct()
         {

diff --git a/src/Nethermind/Nethermind.Trie/Pruning/TrieStore.cs b/src/Nethermind/Nethermind.Trie/Pruning/TrieStore.cs
@@ -119,6 +119,11 @@ public long MemoryUsedByDirtyCache
             }
         }
 
+        public void IncrementMemoryUsedByDirtyCache(long nodeMemoryUsage)
+        {
+            Metrics.MemoryUsedByCache = Interlocked.Add(ref _memoryUsedByDirtyCache, nodeMemoryUsage);
+        }
+
         public int CommittedNodesCount
         {
             get => _committedNodesCount;
@@ -129,6 +134,11 @@ private set
             }
         }
 
+        private void IncrementCommittedNodesCount()
+        {
+            Metrics.CommittedNodesCount = Interlocked.Increment(ref _committedNodesCount);
+        }
+
         public int PersistedNodesCount
         {
             get => _persistedNodesCount;
@@ -139,6 +149,11 @@ private set
             }
         }
 
+        private void IncrementPersistedNodesCount()
+        {
+            Metrics.PersistedNodeCount = Interlocked.Increment(ref _persistedNodesCount);
+        }
+
         public int CachedNodesCount
         {
             get
@@ -178,7 +193,7 @@ private void CommitNode(long blockNumber, Hash256? address, ref TreePath path, i
                     PersistNode(address, path, node, blockNumber, writeFlags);
                 }
 
-                CommittedNodesCount++;
+                IncrementCommittedNodesCount();
             }
 
             [MethodImpl(MethodImplOptions.NoInlining)]
@@ -858,7 +873,7 @@ private void PersistNode(Hash256? address, in TreePath path, TrieNode currentNod
                 writeBatch.Set(address, path, currentNode.Keccak, currentNode.FullRlp, writeFlags);
                 currentNode.IsPersisted = true;
                 currentNode.LastSeen = Math.Max(blockNumber, currentNode.LastSeen);
-                PersistedNodesCount++;
+                IncrementPersistedNodesCount();
             }
             else
             {

diff --git a/src/Nethermind/Nethermind.Trie/Pruning/TrieStoreDirtyNodesCache.cs b/src/Nethermind/Nethermind.Trie/Pruning/TrieStoreDirtyNodesCache.cs
@@ -67,7 +67,7 @@ public void SaveInCache(in Key key, TrieNode node)
         if (TryAdd(key, node))
         {
             Metrics.CachedNodesCount = Interlocked.Increment(ref _count);
-            _trieStore.MemoryUsedByDirtyCache += node.GetMemorySize(false) + KeyMemoryUsage;
+            _trieStore.IncrementMemoryUsedByDirtyCache(node.GetMemorySize(false) + KeyMemoryUsage);
         }
     }