Refactor/minor worldstate cleanup (#6385)

* Remove address by storage * Make it clear what it it used for serving by hash * Fix test * TrieNodeRlpStore
NethermindEth · Dec 19, 2023 · 13b85a3 · 13b85a3
1 parent 482bc43
commit 13b85a3
Show file tree

Hide file tree

Showing 16 changed files with 45 additions and 53 deletions.
diff --git a/src/Nethermind/Nethermind.Init/Steps/InitializeNetwork.cs b/src/Nethermind/Nethermind.Init/Steps/InitializeNetwork.cs
@@ -166,7 +166,7 @@ private async Task Initialize(CancellationToken cancellationToken)
         _api.DisposeStack.Push(_api.Synchronizer);
 
         ISyncServer syncServer = _api.SyncServer = new SyncServer(
-            _api.TrieStore!.AsKeyValueStore(),
+            _api.TrieStore!.TrieNodeRlpStore,
             _api.DbProvider.CodeDb,
             _api.BlockTree,
             _api.ReceiptStorage!,

diff --git a/src/Nethermind/Nethermind.JsonRpc/Modules/Eth/EthRpcModule.cs b/src/Nethermind/Nethermind.JsonRpc/Modules/Eth/EthRpcModule.cs
@@ -176,13 +176,11 @@ public ResultWrapper<byte[]> eth_getStorageAt(Address address, UInt256 positionI
         }
 
         BlockHeader? header = searchResult.Object;
-        Account account = _stateReader.GetAccount(header!.StateRoot!, address);
-        if (account is null)
+        byte[] storage = _stateReader.GetStorage(header!.StateRoot!, address, positionIndex);
+        if (storage is null)
         {
             return ResultWrapper<byte[]>.Success(Array.Empty<byte>());
         }
-
-        byte[] storage = _stateReader.GetStorage(account.StorageRoot, positionIndex);
         return ResultWrapper<byte[]>.Success(storage!.PadLeft(32));
     }
 

diff --git a/src/Nethermind/Nethermind.State.Test/StateReaderTests.cs b/src/Nethermind/Nethermind.State.Test/StateReaderTests.cs
@@ -148,9 +148,7 @@ void CommitEverything()
 
             StateReader reader =
                 new(new TrieStore(stateDb, LimboLogs.Instance), Substitute.For<IDb>(), Logger);
-            Hash256 storageRoot = reader.GetStorageRoot(stateRoot0, _address1);
-            reader.GetStorage(storageRoot, storageCell.Index + 1).Should().BeEquivalentTo(new byte[] { 0 });
-            reader.GetStorage(Keccak.EmptyTreeHash, storageCell.Index + 1).Should().BeEquivalentTo(new byte[] { 0 });
+            reader.GetStorage(stateRoot0, _address1, storageCell.Index + 1).Should().BeEquivalentTo(new byte[] { 0 });
         }
 
         private Task StartTask(StateReader reader, Hash256 stateRoot, UInt256 value)
@@ -173,8 +171,7 @@ private Task StartStorageTask(StateReader reader, Hash256 stateRoot, StorageCell
                 {
                     for (int i = 0; i < 1000; i++)
                     {
-                        Hash256 storageRoot = reader.GetStorageRoot(stateRoot, storageCell.Address);
-                        byte[] result = reader.GetStorage(storageRoot, storageCell.Index);
+                        byte[] result = reader.GetStorage(stateRoot, storageCell.Address, storageCell.Index);
                         result.Should().BeEquivalentTo(value);
                     }
                 });
@@ -206,8 +203,7 @@ public async Task Get_storage()
             StateReader reader = new(
                 new TrieStore(dbProvider.StateDb, LimboLogs.Instance), dbProvider.CodeDb, Logger);
 
-            var account = reader.GetAccount(state.StateRoot, _address1);
-            var retrieved = reader.GetStorage(account.StorageRoot, storageCell.Index);
+            var retrieved = reader.GetStorage(state.StateRoot, _address1, storageCell.Index);
             retrieved.Should().BeEquivalentTo(initialValue);
 
             /* at this stage we set the value in storage to 1,2,3 at the tested storage cell */
@@ -230,7 +226,7 @@ It is a different stack of objects than the one that is used by the blockchain b
                We will try to retrieve the value by taking the state root from the processor.*/
 
             retrieved =
-                reader.GetStorage(processorStateProvider.GetStorageRoot(storageCell.Address), storageCell.Index);
+                reader.GetStorage(processorStateProvider.StateRoot, storageCell.Address, storageCell.Index);
             retrieved.Should().BeEquivalentTo(newValue);
 
             /* If it failed then it means that the blockchain bridge cached the previous call value */

diff --git a/src/Nethermind/Nethermind.State/IStateReader.cs b/src/Nethermind/Nethermind.State/IStateReader.cs
@@ -12,7 +12,7 @@ public interface IStateReader
     {
         Account? GetAccount(Hash256 stateRoot, Address address);
 
-        byte[]? GetStorage(Hash256 storageRoot, in UInt256 index);
+        byte[]? GetStorage(Hash256 stateRoot, Address address, in UInt256 index);
 
         byte[]? GetCode(Hash256 codeHash);
 

diff --git a/src/Nethermind/Nethermind.State/StateReader.cs b/src/Nethermind/Nethermind.State/StateReader.cs
@@ -33,8 +33,12 @@ public StateReader(ITrieStore? trieStore, IKeyValueStore? codeDb, ILogManager? l
             return GetState(stateRoot, address);
         }
 
-        public byte[] GetStorage(Hash256 storageRoot, in UInt256 index)
+        public byte[] GetStorage(Hash256 stateRoot, Address address, in UInt256 index)
         {
+            Account? account = GetAccount(stateRoot, address);
+            if (account == null) return null;
+
+            Hash256 storageRoot = account.StorageRoot;
             if (storageRoot == Keccak.EmptyTreeHash)
             {
                 return new byte[] { 0 };

diff --git a/src/Nethermind/Nethermind.Synchronization.Test/OldStyleFullSynchronizerTests.cs b/src/Nethermind/Nethermind.Synchronization.Test/OldStyleFullSynchronizerTests.cs
@@ -85,7 +85,7 @@ public async Task Setup()
                 stateReader,
                 LimboLogs.Instance);
             _syncServer = new SyncServer(
-                trieStore.AsKeyValueStore(),
+                trieStore.TrieNodeRlpStore,
                 _codeDb,
                 _blockTree,
                 _receiptStorage,

diff --git a/src/Nethermind/Nethermind.Synchronization.Test/SyncServerTests.cs b/src/Nethermind/Nethermind.Synchronization.Test/SyncServerTests.cs
@@ -645,7 +645,7 @@ public void GetNodeData_returns_cached_trie_nodes()
             MemDb stateDb = new();
             TrieStore trieStore = new(stateDb, Prune.WhenCacheReaches(10.MB()), NoPersistence.Instance, LimboLogs.Instance);
             ctx.SyncServer = new SyncServer(
-                trieStore.AsKeyValueStore(),
+                trieStore.TrieNodeRlpStore,
                 new MemDb(),
                 localBlockTree,
                 NullReceiptStorage.Instance,

diff --git a/src/Nethermind/Nethermind.Synchronization.Test/SyncThreadTests.cs b/src/Nethermind/Nethermind.Synchronization.Test/SyncThreadTests.cs
@@ -367,7 +367,7 @@ private SyncTestContext CreateSyncManager(int index)
 
             ISyncModeSelector selector = synchronizer.SyncModeSelector;
             SyncServer syncServer = new(
-                trieStore.AsKeyValueStore(),
+                trieStore.TrieNodeRlpStore,
                 codeDb,
                 tree,
                 receiptStorage,

diff --git a/src/Nethermind/Nethermind.Synchronization.Test/SynchronizerTests.cs b/src/Nethermind/Nethermind.Synchronization.Test/SynchronizerTests.cs
@@ -401,7 +401,7 @@ ISyncConfig GetSyncConfig() =>
                 }
 
                 SyncServer = new SyncServer(
-                    trieStore.AsKeyValueStore(),
+                    trieStore.TrieNodeRlpStore,
                     codeDb,
                     BlockTree,
                     NullReceiptStorage.Instance,

diff --git a/src/Nethermind/Nethermind.Synchronization.Test/Trie/HealingTreeTests.cs b/src/Nethermind/Nethermind.Synchronization.Test/Trie/HealingTreeTests.cs
@@ -83,7 +83,7 @@ private void recovery_works<T>(
             k => throw new MissingTrieNodeException("", new TrieNodeException("", _key), path, 1),
             k => new TrieNode(NodeType.Leaf) { Key = path });
         TestMemDb db = new();
-        trieStore.AsKeyValueStore().Returns(db);
+        trieStore.TrieNodeRlpStore.Returns(db);
 
         ITrieNodeRecovery<GetTrieNodesRequest> recovery = Substitute.For<ITrieNodeRecovery<GetTrieNodesRequest>>();
         recovery.CanRecover.Returns(isMainThread);
@@ -95,7 +95,7 @@ private void recovery_works<T>(
         if (isMainThread && successfullyRecovered)
         {
             action.Should().NotThrow();
-            db.KeyWasWritten(kvp => Bytes.AreEqual(kvp.Item1, ValueKeccak.Compute(_rlp).Bytes) && Bytes.AreEqual(kvp.Item2, _rlp));
+            trieStore.Received().Set(ValueKeccak.Compute(_rlp), _rlp);
         }
         else
         {

diff --git a/src/Nethermind/Nethermind.Synchronization/Trie/HealingStateTree.cs b/src/Nethermind/Nethermind.Synchronization/Trie/HealingStateTree.cs
@@ -82,7 +82,7 @@ private bool Recover(in ValueHash256 rlpHash, ReadOnlySpan<byte> pathPart, Hash2
             byte[]? rlp = _recovery.Recover(rlpHash, request).GetAwaiter().GetResult();
             if (rlp is not null)
             {
-                TrieStore.AsKeyValueStore().Set(rlpHash.Bytes, rlp);
+                TrieStore.Set(rlpHash, rlp);
                 return true;
             }
         }

diff --git a/src/Nethermind/Nethermind.Synchronization/Trie/HealingStorageTree.cs b/src/Nethermind/Nethermind.Synchronization/Trie/HealingStorageTree.cs
@@ -81,7 +81,7 @@ private bool Recover(in ValueHash256 rlpHash, ReadOnlySpan<byte> pathPart)
             byte[]? rlp = _recovery.Recover(rlpHash, request).GetAwaiter().GetResult();
             if (rlp is not null)
             {
-                TrieStore.AsKeyValueStore().Set(rlpHash.Bytes, rlp);
+                TrieStore.Set(rlpHash, rlp);
                 return true;
             }
         }

diff --git a/src/Nethermind/Nethermind.Trie/Pruning/ITrieStore.cs b/src/Nethermind/Nethermind.Trie/Pruning/ITrieStore.cs
@@ -19,6 +19,10 @@ public interface ITrieStore : ITrieNodeResolver, IDisposable
 
         event EventHandler<ReorgBoundaryReached>? ReorgBoundaryReached;
 
-        IKeyValueStore AsKeyValueStore();
+        // Used for serving via hash
+        IReadOnlyKeyValueStore TrieNodeRlpStore { get; }
+
+        // Used by healing
+        void Set(in ValueHash256 hash, byte[] rlp);
     }
 }
diff --git a/src/Nethermind/Nethermind.Trie/Pruning/NullTrieStore.cs b/src/Nethermind/Nethermind.Trie/Pruning/NullTrieStore.cs
@@ -17,8 +17,6 @@ public void CommitNode(long blockNumber, NodeCommitInfo nodeCommitInfo, WriteFla
 
         public void FinishBlockCommit(TrieType trieType, long blockNumber, TrieNode? root, WriteFlags flags = WriteFlags.None) { }
 
-        public static void HackPersistOnShutdown() { }
-
         public IReadOnlyTrieStore AsReadOnly(IKeyValueStore keyValueStore) => this;
 
         public event EventHandler<ReorgBoundaryReached> ReorgBoundaryReached
@@ -27,7 +25,7 @@ public event EventHandler<ReorgBoundaryReached> ReorgBoundaryReached
             remove { }
         }
 
-        public IKeyValueStore AsKeyValueStore() => null!;
+        public IReadOnlyKeyValueStore TrieNodeRlpStore => null!;
 
         public TrieNode FindCachedOrUnknown(Hash256 hash) => new(NodeType.Unknown, hash);
 
@@ -37,6 +35,8 @@ public event EventHandler<ReorgBoundaryReached> ReorgBoundaryReached
 
         public void Dispose() { }
 
-        public static byte[]? Get(ReadOnlySpan<byte> key, ReadFlags flags = ReadFlags.None) => null;
+        public void Set(in ValueHash256 hash, byte[] rlp)
+        {
+        }
     }
 }
diff --git a/src/Nethermind/Nethermind.Trie/Pruning/ReadOnlyTrieStore.cs b/src/Nethermind/Nethermind.Trie/Pruning/ReadOnlyTrieStore.cs
@@ -14,13 +14,13 @@ public class ReadOnlyTrieStore : IReadOnlyTrieStore
     {
         private readonly TrieStore _trieStore;
         private readonly IKeyValueStore? _readOnlyStore;
-        private readonly ReadOnlyValueStore _publicStore;
+        private readonly IReadOnlyKeyValueStore _publicStore;
 
         public ReadOnlyTrieStore(TrieStore trieStore, IKeyValueStore? readOnlyStore)
         {
             _trieStore = trieStore ?? throw new ArgumentNullException(nameof(trieStore));
             _readOnlyStore = readOnlyStore;
-            _publicStore = new ReadOnlyValueStore(_trieStore.AsKeyValueStore());
+            _publicStore = _trieStore.TrieNodeRlpStore;
         }
 
         public TrieNode FindCachedOrUnknown(Hash256 hash) =>
@@ -45,24 +45,12 @@ public event EventHandler<ReorgBoundaryReached> ReorgBoundaryReached
             remove { }
         }
 
-        public IKeyValueStore AsKeyValueStore() => _publicStore;
+        public IReadOnlyKeyValueStore TrieNodeRlpStore => _publicStore;
 
-        public void Dispose() { }
-
-        public static void Set(ReadOnlySpan<byte> key, byte[]? value, WriteFlags flags = WriteFlags.None) { }
-
-        private class ReadOnlyValueStore : IKeyValueStore
+        public void Set(in ValueHash256 hash, byte[] rlp)
         {
-            private readonly IKeyValueStore _keyValueStore;
-
-            public ReadOnlyValueStore(IKeyValueStore keyValueStore)
-            {
-                _keyValueStore = keyValueStore;
-            }
-
-            public byte[]? Get(ReadOnlySpan<byte> key, ReadFlags flags = ReadFlags.None) => _keyValueStore.Get(key, flags);
-
-            public void Set(ReadOnlySpan<byte> key, byte[]? value, WriteFlags flags = WriteFlags.None) { }
         }
+
+        public void Dispose() { }
     }
 }
diff --git a/src/Nethermind/Nethermind.Trie/Pruning/TrieStore.cs b/src/Nethermind/Nethermind.Trie/Pruning/TrieStore.cs
@@ -812,7 +812,7 @@ void PersistNode(TrieNode n)
             });
         }
 
-        private byte[]? Get(ReadOnlySpan<byte> key, ReadFlags flags = ReadFlags.None)
+        private byte[]? GetByHash(ReadOnlySpan<byte> key, ReadFlags flags = ReadFlags.None)
         {
             return _pruningStrategy.PruningEnabled
                    && _dirtyNodes.AllNodes.TryGetValue(new ValueHash256(key), out TrieNode? trieNode)
@@ -823,9 +823,14 @@ void PersistNode(TrieNode n)
                 : _keyValueStore.Get(key, flags);
         }
 
-        public IKeyValueStore AsKeyValueStore() => _publicStore;
+        public IReadOnlyKeyValueStore TrieNodeRlpStore => _publicStore;
 
-        private class TrieKeyValueStore : IKeyValueStore
+        public void Set(in ValueHash256 hash, byte[] rlp)
+        {
+            _keyValueStore.Set(hash, rlp);
+        }
+
+        private class TrieKeyValueStore : IReadOnlyKeyValueStore
         {
             private readonly TrieStore _trieStore;
 
@@ -834,10 +839,7 @@ public TrieKeyValueStore(TrieStore trieStore)
                 _trieStore = trieStore;
             }
 
-            public byte[]? Get(ReadOnlySpan<byte> key, ReadFlags flags = ReadFlags.None) => _trieStore.Get(key, flags);
-
-            public void Set(ReadOnlySpan<byte> key, byte[]? value, WriteFlags flags = WriteFlags.None)
-                => _trieStore._keyValueStore.Set(key, value, flags);
+            public byte[]? Get(ReadOnlySpan<byte> key, ReadFlags flags = ReadFlags.None) => _trieStore.GetByHash(key, flags);
         }
     }
 }