perf(trie): share cached nodes across read-only walks (CloneForReadOnly 19.2 GB -> 1.4 MB)

src/Nethermind/Nethermind.Consensus/Stateless/WitnessCapturingTrieStore.cs

@@ -16,18 +16,19 @@ namespace Nethermind.Consensus.Stateless;
 /// Delegates all logic to base store except for writing trie nodes (readonly!)
 /// Adds logic for capturing trie nodes accessed during execution and state root recomputation.
 /// </remarks>
-public class WitnessCapturingTrieStore(IReadOnlyTrieStore baseStore) : ITrieStore
+public class WitnessCapturingTrieStore(IReadOnlyTrieStore baseStore) : ITrieStore, IScopedReadOnlyTraversalProvider
 {
+    private readonly IReadOnlyTrieStore _baseStore = baseStore;
     private readonly ConcurrentDictionary<Hash256AsKey, byte[]> _rlpCollector = new();
 
     public IEnumerable<byte[]> TouchedNodesRlp => _rlpCollector.Select(static kvp => kvp.Value);
 
-    public void Dispose() => baseStore.Dispose();
+    public void Dispose() => _baseStore.Dispose();
 
     public TrieNode FindCachedOrUnknown(Hash256? address, in TreePath path, Hash256 hash)
     {
-        TrieNode node = baseStore.FindCachedOrUnknown(address, in path, hash);
-        if (node.NodeType != NodeType.Unknown) _rlpCollector.TryAdd(node.Keccak, node.FullRlp.ToArray());
+        TrieNode node = _baseStore.FindCachedOrUnknown(address, in path, hash);
+        CaptureNode(node);
         return node;
     }
 
@@ -37,21 +38,59 @@ public TrieNode FindCachedOrUnknown(Hash256? address, in TreePath path, Hash256
 
     public byte[]? TryLoadRlp(Hash256? address, in TreePath path, Hash256 hash, ReadFlags flags = ReadFlags.None)
     {
-        byte[]? rlp = baseStore.TryLoadRlp(address, in path, hash, flags);
-        if (rlp is not null) _rlpCollector.TryAdd(hash, rlp);
+        byte[]? rlp = _baseStore.TryLoadRlp(address, in path, hash, flags);
+        CaptureRlp(hash, rlp);
         return rlp;
     }
 
-    public bool HasRoot(Hash256 stateRoot) => baseStore.HasRoot(stateRoot);
+    public bool HasRoot(Hash256 stateRoot) => _baseStore.HasRoot(stateRoot);
 
-    public IDisposable BeginScope(BlockHeader? baseBlock) => baseStore.BeginScope(baseBlock);
+    public IDisposable BeginScope(BlockHeader? baseBlock) => _baseStore.BeginScope(baseBlock);
 
     public IScopedTrieStore GetTrieStore(Hash256? address) => new ScopedTrieStore(this, address);
 
-    public INodeStorage.KeyScheme Scheme => baseStore.Scheme;
+    public INodeStorage.KeyScheme Scheme => _baseStore.Scheme;
 
     public IBlockCommitter BeginBlockCommit(long blockNumber) => NullCommitter.Instance;
 
     // WitnessCapturingTrieStore is read-only, so we return a no-op committer that doesn't persist any trie nodes
     public ICommitter BeginCommit(Hash256? address, TrieNode? root, WriteFlags writeFlags) => NullCommitter.Instance;
+
+    private void CaptureNode(TrieNode node)
+    {
+        if (node.NodeType != NodeType.Unknown)
+        {
+            _rlpCollector.TryAdd(node.Keccak, node.FullRlp.ToArray());
+        }
+    }
+
+    private void CaptureRlp(Hash256 hash, byte[]? rlp)
+    {
+        if (rlp is not null)
+        {
+            _rlpCollector.TryAdd(hash, rlp);
+        }
+    }
+
+    public ITrieNodeResolver? GetReadOnlyTraversalResolver(Hash256? address) =>
+        _baseStore.GetTrieStore(address) is ITrieNodeResolverSource source
+            && source.GetReadOnlyTraversalResolver() is { } readOnlyResolver
+                ? new WitnessCapturingReadOnlyTraversalResolver(this, address, readOnlyResolver)
+                : null;
+
+    private sealed class WitnessCapturingReadOnlyTraversalResolver(
+        WitnessCapturingTrieStore fullTrieStore,
+        Hash256? address,
+        ITrieNodeResolver inner) : ReadOnlyTraversalResolverBase(fullTrieStore, address)
+    {
+        public override TrieNode FindCachedOrUnknown(in TreePath path, Hash256 hash)
+        {
+            TrieNode node = inner.FindCachedOrUnknown(path, hash);
+            fullTrieStore.CaptureNode(node);
+            return node;
+        }
+
+        protected override ITrieNodeResolver WithAddress(Hash256? address1) =>
+            new WitnessCapturingReadOnlyTraversalResolver(fullTrieStore, address1, inner.GetStorageTrieNodeResolver(address1));
+    }
 }

src/Nethermind/Nethermind.Core/Buffers/CappedArray.cs

@@ -25,11 +25,18 @@ namespace Nethermind.Core.Buffers;
     public static object EmptyBoxed { get; } = _empty;
 
     private readonly T[]? _array;
+    private readonly int _offset;
     private readonly int _length;
 
     public CappedArray(T[]? array, int length)
+        : this(array, 0, length)
+    {
+    }
+
+    public CappedArray(T[]? array, int offset, int length)
     {
         _array = array;
+        _offset = offset;
         _length = length;
     }
 
@@ -38,8 +45,15 @@ public CappedArray(T[]? array)
         if (array is not null)
         {
             _array = array;
+            _offset = 0;
             _length = array.Length;
         }
+        else
+        {
+            _array = null;
+            _offset = 0;
+            _length = 0;
+        }
     }
 
     public static implicit operator ReadOnlySpan<T>(in CappedArray<T> array) => array.AsSpan();
@@ -51,46 +65,49 @@ public T this[int index]
         get
         {
             T[] array = _array!;
-            if (index >= _length || (uint)index >= (uint)array.Length)
+            int arrayIndex = _offset + index;
+            if (index >= _length || (uint)arrayIndex >= (uint)array.Length)
             {
                 ThrowArgumentOutOfRangeException();
             }
 
-            return array[index];
+            return array[arrayIndex];
         }
         set
         {
             T[] array = _array!;
-            if (index >= _length || (uint)index >= (uint)array.Length)
+            int arrayIndex = _offset + index;
+            if (index >= _length || (uint)arrayIndex >= (uint)array.Length)
             {
                 ThrowArgumentOutOfRangeException();
             }
 
-            array[index] = value;
+            array[arrayIndex] = value;
         }
     }
 
     [DoesNotReturn, StackTraceHidden]
     private static void ThrowArgumentOutOfRangeException() => throw new ArgumentOutOfRangeException();
 
     public int Length => _length;
+    public int Offset => _offset;
     public int UnderlyingLength => _array?.Length ?? 0;
     public T[]? UnderlyingArray => _array;
-    public bool IsUncapped => _length == _array?.Length;
+    public bool IsUncapped => _offset == 0 && _length == _array?.Length;
     public bool IsNull => _array is null;
     public bool IsNotNull => _array is not null;
     public bool IsNullOrEmpty => _length == 0;
     public bool IsNotNullOrEmpty => _length > 0;
-    public Span<T> AsSpan() => _array.AsSpan(0, _length);
-    public Span<T> AsSpan(int start, int length) => _array.AsSpan(start, length);
+    public Span<T> AsSpan() => _array.AsSpan(_offset, _length);
+    public Span<T> AsSpan(int start, int length) => _array.AsSpan(_offset + start, length);
 
     public T[]? ToArray()
     {
         T[]? array = _array;
 
         if (array is null) return null;
         if (array.Length == 0) return [];
-        if (_length == array.Length) return array;
+        if (_offset == 0 && _length == array.Length) return array;
         return AsSpan().ToArray();
     }
 
@@ -99,5 +116,5 @@ public T this[int index]
         : base.ToString();
 
     public ArraySegment<T> AsArraySegment() => AsArraySegment(0, _length);
-    public ArraySegment<T> AsArraySegment(int start, int length) => new(_array!, start, length);
+    public ArraySegment<T> AsArraySegment(int start, int length) => new(_array!, _offset + start, length);
 }

src/Nethermind/Nethermind.Synchronization/SnapSync/SnapUpperBoundAdapter.cs

@@ -14,7 +14,7 @@ namespace Nethermind.Synchronization.SnapSync;
 /// UpperBound. This is to prevent double writes on partitioned snap ranges.
 /// </summary>
 /// <param name="baseTrieStore"></param>
-public class SnapUpperBoundAdapter(IScopedTrieStore baseTrieStore) : IScopedTrieStore
+public class SnapUpperBoundAdapter(IScopedTrieStore baseTrieStore) : IScopedTrieStore, ITrieNodeResolverSource
 {
     public ValueHash256 UpperBound = ValueKeccak.MaxValue;
 
@@ -30,6 +30,16 @@ public class SnapUpperBoundAdapter(IScopedTrieStore baseTrieStore) : IScopedTrie
 
     public ICommitter BeginCommit(TrieNode? root, WriteFlags writeFlags = WriteFlags.None) => new BoundedSnapCommitter(baseTrieStore.BeginCommit(root, writeFlags), UpperBound);
 
+    public ITrieNodeResolver? GetReadOnlyTraversalResolver()
+    {
+        if (baseTrieStore is ITrieNodeResolverSource source)
+        {
+            return source.GetReadOnlyTraversalResolver();
+        }
+
+        return null;
+    }
+
     private sealed class BoundedSnapCommitter(ICommitter baseCommitter, ValueHash256 subtreeLimit) : ICommitter
     {
         public void Dispose() => baseCommitter.Dispose();

src/Nethermind/Nethermind.Trie.Test/OverlayTrieStoreTests.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Linq;
 using FluentAssertions;
+using Nethermind.Core;
 using Nethermind.Core.Crypto;
 using Nethermind.Core.Test;
 using Nethermind.Core.Test.Builders;
@@ -75,4 +76,35 @@
         // After all this, the original should not change.
         dbProvider.StateDb.GetAllKeys().Count().Should().Be(originalKeyCount);
     }
+
+    [Test]
+    [NonParallelizable]
+    public void OverlayStore_read_only_lookup_uses_shared_base_nodes()
+    {
+        using TrieStore existingStore = TestTrieStoreFactory.Build(new MemDb(), No.Pruning, No.Persistence, LimboLogs.Instance);
+        PatriciaTree patriciaTree = new(existingStore, LimboLogs.Instance);
+        {
+            using IBlockCommitter _ = existingStore.BeginBlockCommit(0);
+            patriciaTree.Set(TestItem.Keccaks[0].Bytes, TestItem.Keccaks[0].BytesToArray());
+            patriciaTree.Set(TestItem.Keccaks[1].Bytes, TestItem.Keccaks[1].BytesToArray());
+            patriciaTree.Commit();
+        }
+
+        IDbProvider dbProvider = TestMemDbProvider.Init();
+        ReadOnlyDbProvider readOnlyDbProvider = dbProvider.AsReadOnly(true);
+        ITrieStore overlayStore = new OverlayTrieStore(readOnlyDbProvider.GetDb<IDb>(DbNames.State), existingStore.AsReadOnly());
+
+        long sharedHitsBefore = existingStore.SharedNodeHitCount;
+        long clonesBefore = existingStore.CloneForReadOnlyCount;
+
+        PatriciaTree overlaidTree = new(overlayStore, LimboLogs.Instance)
+        {
+            RootHash = patriciaTree.RootHash
+        };
+
+        overlaidTree.Get(TestItem.Keccaks[0].Bytes).ToArray().Should().BeEquivalentTo(TestItem.Keccaks[0].BytesToArray());
+
+        existingStore.SharedNodeHitCount.Should().BeGreaterThan(sharedHitsBefore);
+        existingStore.CloneForReadOnlyCount.Should().Be(clonesBefore);
+    }
 }