Only recalculate storage hashes once per block

src/Nethermind/Nethermind.Consensus/Processing/BlockProcessor.cs

@@ -249,7 +249,7 @@ protected virtual TxReceipt[] ProcessBlock(
         _withdrawalProcessor.ProcessWithdrawals(block, spec);
         ReceiptsTracer.EndBlockTrace();
 
-        _stateProvider.Commit(spec);
+        _stateProvider.PostCommit(spec);
 
         if (ShouldComputeStateRoot(block.Header))
         {

src/Nethermind/Nethermind.State/IWorldState.cs

@@ -104,4 +104,5 @@ public interface IWorldState : IJournal<Snapshot>, IReadOnlyStateProvider
     void Commit(IReleaseSpec releaseSpec, IWorldStateTracer? traver, bool isGenesis = false);
 
     void CommitTree(long blockNumber);
+    void PostCommit(IReleaseSpec releaseSpec);
 }

src/Nethermind/Nethermind.State/PartialStorageProviderBase.cs

@@ -150,6 +150,14 @@ public void Commit()
             Commit(NullStateTracer.Instance);
         }
 
+        public void PostCommit()
+        {
+            Commit(NullStateTracer.Instance);
+            PostCommitStorages();
+        }
+
+        protected virtual void PostCommitStorages() { }
+
         protected readonly struct ChangeTrace
         {
             public ChangeTrace(byte[]? before, byte[]? after)

src/Nethermind/Nethermind.State/PersistentStorageProvider.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading.Tasks;
 using Nethermind.Core;
 using Nethermind.Core.Collections;
 using Nethermind.Core.Crypto;
@@ -25,6 +26,7 @@ internal class PersistentStorageProvider : PartialStorageProviderBase
         private readonly ILogManager? _logManager;
         internal readonly IStorageTreeFactory _storageTreeFactory;
         private readonly ResettableDictionary<AddressAsKey, StorageTree> _storages = new();
+        private readonly HashSet<Address> _toUpdateRoots = new();
 
         /// <summary>
         /// EIP-1283
@@ -53,6 +55,7 @@ public override void Reset()
             _storages.Reset();
             _originalValues.Clear();
             _committedThisRound.Clear();
+            _toUpdateRoots.Clear();
         }
 
         /// <summary>
@@ -179,16 +182,17 @@ protected override void CommitCore(IStorageTracer tracer)
                 }
             }
 
-            // TODO: it seems that we are unnecessarily recalculating root hashes all the time in storage?
             foreach (Address address in toUpdateRoots)
             {
                 // since the accounts could be empty accounts that are removing (EIP-158)
                 if (_stateProvider.AccountExists(address))
                 {
-                    Hash256 root = RecalculateRootHash(address);
-
-                    // _logger.Warn($"Recalculating storage root {address}->{root} ({toUpdateRoots.Count})");
-                    _stateProvider.UpdateStorageRoot(address, root);
+                    _toUpdateRoots.Add(address);
+                }
+                else
+                {
+                    _toUpdateRoots.Remove(address);
+                    _storages.Remove(address);
                 }
             }
 
@@ -202,6 +206,75 @@ protected override void CommitCore(IStorageTracer tracer)
             }
         }
 
+        protected override void PostCommitStorages()
+        {
+            if (_toUpdateRoots.Count == 0)
+            {
+                return;
+            }
+
+            // Is overhead of parallel foreach worth it?
+            if (_toUpdateRoots.Count <= 4)
+            {
+                UpdateRootHashesSingleThread();
+            }
+            else
+            {
+                UpdateRootHashesMultiThread();
+            }
+
+            void UpdateRootHashesSingleThread()
+            {
+                foreach (KeyValuePair<AddressAsKey, StorageTree> kvp in _storages)
+                {
+                    if (!_toUpdateRoots.Contains(kvp.Key))
+                    {
+                        // Remove the storage tree if it was only read and not updated
+                        _storages.Remove(kvp.Key);
+                        continue;
+                    }
+
+                    StorageTree storageTree = kvp.Value;
+                    storageTree.UpdateRootHash();
+                    _stateProvider.UpdateStorageRoot(address: kvp.Key, storageTree.RootHash);
+                }
+
+                _toUpdateRoots.Clear();
+            }
+
+            void UpdateRootHashesMultiThread()
+            {
+                // We can recalculate the roots in parallel as they are all independent tries
+                Parallel.ForEach(_storages, kvp =>
+                {
+                    if (!_toUpdateRoots.Contains(kvp.Key))
+                    {
+                        // Wasn't updated don't recalculate; we don't remove here
+                        // in parallel foreach as it would be unsafe on non-concurrent dictionary
+                        return;
+                    }
+                    StorageTree storageTree = kvp.Value;
+                    storageTree.UpdateRootHash();
+                });
+
+                // Update the storage roots in the main thread non in parallel
+                foreach (KeyValuePair<AddressAsKey, StorageTree> kvp in _storages)
+                {
+                    if (!_toUpdateRoots.Contains(kvp.Key))
+                    {
+                        // Remove the storage tree if it was only read and not updated
+                        _storages.Remove(kvp.Key);
+                        continue;
+                    }
+
+                    // Update the storage root for the Account
+                    _stateProvider.UpdateStorageRoot(address: kvp.Key, kvp.Value.RootHash);
+                }
+
+                _toUpdateRoots.Clear();
+            }
+        }
+
         private void SaveToTree(HashSet<Address> toUpdateRoots, Change change)
         {
             if (_originalValues.TryGetValue(change.StorageCell, out byte[] initialValue) &&
@@ -223,14 +296,11 @@ private void SaveToTree(HashSet<Address> toUpdateRoots, Change change)
         /// <param name="blockNumber">Current block number</param>
         public void CommitTrees(long blockNumber)
         {
-            // _logger.Warn($"Storage block commit {blockNumber}");
             foreach (KeyValuePair<AddressAsKey, StorageTree> storage in _storages)
             {
                 storage.Value.Commit(blockNumber);
             }
 
-            // TODO: maybe I could update storage roots only now?
-
             // only needed here as there is no control over cached storage size otherwise
             _storages.Reset();
         }
@@ -305,6 +375,7 @@ public override void ClearStorage(Address address)
             // by means of CREATE 2 - notice that the cached trie may carry information about items that were not
             // touched in this block, hence were not zeroed above
             // TODO: how does it work with pruning?
+            _toUpdateRoots.Add(address);
             _storages[address] = new StorageTree(_trieStore.GetTrieStore(address.ToAccountPath), Keccak.EmptyTreeHash, _logManager);
         }
 

src/Nethermind/Nethermind.State/WorldState.cs

@@ -205,6 +205,12 @@ public void Commit(IReleaseSpec releaseSpec, IWorldStateTracer tracer, bool isGe
             _transientStorageProvider.Commit(tracer);
             _stateProvider.Commit(releaseSpec, tracer, isGenesis);
         }
+        public void PostCommit(IReleaseSpec releaseSpec)
+        {
+            _persistentStorageProvider.PostCommit();
+            _transientStorageProvider.PostCommit();
+            _stateProvider.Commit(releaseSpec, isGenesis: false);
+        }
 
         public Snapshot TakeSnapshot(bool newTransactionStart = false)
         {