Skip to content

Commit

Permalink
Traverse trie branches in same stack frame (#6981)
Browse files Browse the repository at this point in the history
  • Loading branch information
benaadams authored May 6, 2024
1 parent a6521e9 commit 4e98f26
Show file tree
Hide file tree
Showing 4 changed files with 167 additions and 112 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ namespace Nethermind.Synchronization.Trie;
/// <summary>
/// Trie store that can recover from network using eth63-eth66 protocol and GetNodeData.
/// </summary>
public class HealingTrieStore : TrieStore
public sealed class HealingTrieStore : TrieStore
{
private ITrieNodeRecovery<IReadOnlyList<Hash256>>? _recovery;

Expand Down
213 changes: 134 additions & 79 deletions src/Nethermind/Nethermind.Trie/PatriciaTree.cs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Runtime.CompilerServices;
using System.Runtime.ExceptionServices;
using System.Threading;
using System.Threading.Tasks;
using Nethermind.Core;
Expand Down Expand Up @@ -558,7 +559,9 @@ private ref readonly CappedArray<byte> Run(
TreePath startingPath = TreePath.Empty;
TrieNode startNode = TrieStore.FindCachedOrUnknown(startingPath, startRootHash);
ResolveNode(startNode, in traverseContext, in startingPath);
return ref TraverseNode(startNode, in traverseContext, ref startingPath);
return ref startNode.IsBranch ?
ref TraverseBranches(startNode, ref startingPath, traverseContext) :
ref TraverseNode(startNode, traverseContext, ref startingPath);
}
else
{
Expand All @@ -580,7 +583,9 @@ private ref readonly CappedArray<byte> Run(
TreePath startingPath = TreePath.Empty;
ResolveNode(RootRef, in traverseContext, in startingPath);
if (_logger.IsTrace) TraceNode(in traverseContext);
return ref TraverseNode(RootRef, in traverseContext, ref startingPath);
return ref RootRef.IsBranch ?
ref TraverseBranches(RootRef, ref startingPath, traverseContext) :
ref TraverseNode(RootRef, traverseContext, ref startingPath);
}
}

Expand All @@ -607,14 +612,29 @@ void TraceNode(in TraverseContext traverseContext)

private void ResolveNode(TrieNode node, in TraverseContext traverseContext, in TreePath path)
{
if (node.NodeType != NodeType.Unknown) return;

try
{
node.ResolveNode(TrieStore, path);
node.ResolveUnknownNode(TrieStore, path);
}
catch (RlpException rlpException)
{
ThrowDecodingError(node, in traverseContext, rlpException);
}
catch (TrieNodeException e)
{
ThrowMissingTrieNodeException(e, in traverseContext);
}

[DoesNotReturn]
[StackTraceHidden]
static void ThrowDecodingError(TrieNode node, in TraverseContext traverseContext, RlpException rlpException)
{
var exception = new TrieNodeException($"Error when decoding node {node.Keccak}", node.Keccak ?? Keccak.Zero, rlpException);
exception = (TrieNodeException)ExceptionDispatchInfo.SetCurrentStackTrace(exception);
ThrowMissingTrieNodeException(exception, in traverseContext);
}
}

private ref readonly CappedArray<byte> TraverseNode(TrieNode node, scoped in TraverseContext traverseContext, scoped ref TreePath path)
Expand All @@ -628,16 +648,12 @@ private ref readonly CappedArray<byte> TraverseNode(TrieNode node, scoped in Tra

switch (node.NodeType)
{
case NodeType.Branch:
return ref TraverseBranch(node, in traverseContext, ref path);
case NodeType.Extension:
return ref TraverseExtension(node, in traverseContext, ref path);
case NodeType.Leaf:
return ref TraverseLeaf(node, in traverseContext, ref path);
case NodeType.Unknown:
return ref TraverseUnknown(node);
default:
return ref ThrowNotSupported(node);
return ref TraverseInvalid(node);
}

[MethodImpl(MethodImplOptions.NoInlining)]
Expand All @@ -646,6 +662,28 @@ void Trace(TrieNode node, in TraverseContext traverseContext)
_logger.Trace($"Traversing {node} to {(traverseContext.IsReadValue ? "READ" : traverseContext.IsDelete ? "DELETE" : "UPDATE")}");
}

[DoesNotReturn]
[StackTraceHidden]
static ref readonly CappedArray<byte> TraverseInvalid(TrieNode node)
{
switch (node.NodeType)
{
case NodeType.Branch:
return ref TraverseBranch(node);
case NodeType.Unknown:
return ref TraverseUnknown(node);
default:
return ref ThrowNotSupported(node);
}
}

[DoesNotReturn]
[StackTraceHidden]
static ref readonly CappedArray<byte> TraverseBranch(TrieNode node)
{
throw new InvalidOperationException($"Branch node {node.Keccak} should already be handled");
}

[DoesNotReturn]
[StackTraceHidden]
static ref readonly CappedArray<byte> TraverseUnknown(TrieNode node)
Expand Down Expand Up @@ -887,80 +925,39 @@ L L - - - - - - - - - - - - - - */
RootRef = nextNode;
}

private ref readonly CappedArray<byte> TraverseBranch(TrieNode node, scoped in TraverseContext traverseContext, scoped ref TreePath path)
private ref readonly CappedArray<byte> TraverseBranches(TrieNode node, scoped ref TreePath path, TraverseContext traverseContext)
{
if (traverseContext.RemainingUpdatePathLength == 0)
while (true)
{
/* all these cases when the path ends on the branch assume a trie with values in the branches
which is not possible within the Ethereum protocol which has keys of the same length (64) */

if (traverseContext.IsNodeRead)
{
return ref node.FullRlp;
}
if (traverseContext.IsReadValue)
if (traverseContext.RemainingUpdatePathLength == 0)
{
return ref node.ValueRef;
return ref ResolveBranchNode(node, in traverseContext);
}

if (traverseContext.IsDelete)
{
if (node.Value.IsNull)
{
return ref CappedArray<byte>.Null;
}

ConnectNodes(null, in traverseContext);
}
else if (Bytes.AreEqual(traverseContext.UpdateValue, node.Value))
{
return ref traverseContext.UpdateValue;
}
else
int childIdx = traverseContext.UpdatePath[traverseContext.CurrentIndex];
path.AppendMut(childIdx);
TrieNode childNode = node.GetChildWithChildPath(TrieStore, ref path, childIdx);
if (traverseContext.IsUpdate)
{
TrieNode withUpdatedValue = node.CloneWithChangedValue(in traverseContext.UpdateValue);
ConnectNodes(withUpdatedValue, in traverseContext);
PushToNodeStack(node, traverseContext.CurrentIndex, childIdx);
}

return ref traverseContext.UpdateValue;
}

int childIdx = traverseContext.UpdatePath[traverseContext.CurrentIndex];
path.AppendMut(childIdx);
TrieNode childNode = node.GetChildWithChildPath(TrieStore, ref path, childIdx);
if (traverseContext.IsUpdate)
{
PushToNodeStack(new StackedNode(node, traverseContext.CurrentIndex, childIdx));
}

if (childNode is null)
{
if (traverseContext.IsRead)
if (childNode is null)
{
return ref CappedArray<byte>.Null;
return ref ResolveCurrent(in traverseContext);
}

if (traverseContext.IsDelete)
{
if (traverseContext.IgnoreMissingDelete)
{
return ref CappedArray<byte>.Null;
}
ResolveNode(childNode, in traverseContext, in path);

ThrowMissingLeafException(in traverseContext);
traverseContext = traverseContext.WithNewIndex(traverseContext.CurrentIndex + 1);
if (!childNode.IsBranch)
{
return ref TraverseNode(childNode, in traverseContext, ref path);
}

int currentIndex = traverseContext.CurrentIndex + 1;
byte[] leafPath = traverseContext.UpdatePath[
currentIndex..].ToArray();
TrieNode leaf = TrieNodeFactory.CreateLeaf(leafPath, in traverseContext.UpdateValue);
ConnectNodes(leaf, in traverseContext);

return ref traverseContext.UpdateValue;
// Traverse next branch
node = childNode;
}

ResolveNode(childNode, in traverseContext, in path);
return ref TraverseNext(childNode, in traverseContext, ref path, 1);
}

private ref readonly CappedArray<byte> TraverseLeaf(TrieNode node, scoped in TraverseContext traverseContext, scoped ref TreePath path)
Expand Down Expand Up @@ -1044,7 +1041,7 @@ private ref readonly CappedArray<byte> TraverseLeaf(TrieNode node, scoped in Tra
{
ReadOnlySpan<byte> extensionPath = longerPath[..extensionLength];
TrieNode extension = TrieNodeFactory.CreateExtension(extensionPath.ToArray());
PushToNodeStack(new StackedNode(extension, traverseContext.CurrentIndex, 0));
PushToNodeStack(extension, traverseContext.CurrentIndex, 0);
}

TrieNode branch = TrieNodeFactory.CreateBranch();
Expand All @@ -1063,7 +1060,7 @@ private ref readonly CappedArray<byte> TraverseLeaf(TrieNode node, scoped in Tra
TrieNode withUpdatedKeyAndValue = node.CloneWithChangedKeyAndValue(
leafPath.ToArray(), longerPathValue);

PushToNodeStack(new StackedNode(branch, traverseContext.CurrentIndex, longerPath[extensionLength]));
PushToNodeStack(branch, traverseContext.CurrentIndex, longerPath[extensionLength]);
ConnectNodes(withUpdatedKeyAndValue, in traverseContext);

return ref traverseContext.UpdateValue;
Expand All @@ -1084,7 +1081,7 @@ private ref readonly CappedArray<byte> TraverseExtension(TrieNode node, scoped i
{
if (traverseContext.IsUpdate)
{
PushToNodeStack(new StackedNode(node, traverseContext.CurrentIndex, 0));
PushToNodeStack(node, traverseContext.CurrentIndex, 0);
}

node.AppendChildPath(ref path, 0);
Expand All @@ -1095,7 +1092,10 @@ private ref readonly CappedArray<byte> TraverseExtension(TrieNode node, scoped i
}

ResolveNode(next, in traverseContext, in path);
return ref TraverseNext(next, in traverseContext, ref path, extensionLength);
TraverseContext newContext = traverseContext.WithNewIndex(traverseContext.CurrentIndex + extensionLength);
return ref next.IsBranch ?
ref TraverseBranches(next, ref path, newContext) :
ref TraverseNode(next, newContext, ref path);
}

if (traverseContext.IsRead)
Expand All @@ -1118,7 +1118,7 @@ private ref readonly CappedArray<byte> TraverseExtension(TrieNode node, scoped i
{
byte[] extensionPath = node.Key.Slice(0, extensionLength);
node = node.CloneWithChangedKey(extensionPath);
PushToNodeStack(new StackedNode(node, traverseContext.CurrentIndex, 0));
PushToNodeStack(node, traverseContext.CurrentIndex, 0);
}

// The node from extension become a branch
Expand Down Expand Up @@ -1157,12 +1157,66 @@ TrieNode secondExtension
return ref traverseContext.UpdateValue;
}

private ref readonly CappedArray<byte> TraverseNext(TrieNode next, scoped in TraverseContext traverseContext, scoped ref TreePath path, int extensionLength)
private ref readonly CappedArray<byte> ResolveCurrent(scoped in TraverseContext traverseContext)
{
if (traverseContext.IsRead)
{
return ref CappedArray<byte>.Null;
}

if (traverseContext.IsDelete)
{
if (traverseContext.IgnoreMissingDelete)
{
return ref CappedArray<byte>.Null;
}

ThrowMissingLeafException(in traverseContext);
}

int currentIndex = traverseContext.CurrentIndex + 1;
byte[] leafPath = traverseContext.UpdatePath[
currentIndex..].ToArray();
TrieNode leaf = TrieNodeFactory.CreateLeaf(leafPath, in traverseContext.UpdateValue);
ConnectNodes(leaf, in traverseContext);

return ref traverseContext.UpdateValue;
}

private ref readonly CappedArray<byte> ResolveBranchNode(TrieNode node, scoped in TraverseContext traverseContext)
{
// Move large struct creation out of flow so doesn't force additional stack space
// in calling method even if not used
TraverseContext newContext = traverseContext.WithNewIndex(traverseContext.CurrentIndex + extensionLength);
return ref TraverseNode(next, in newContext, ref path);
// all these cases when the path ends on the branch assume a trie with values in the branches
// which is not possible within the Ethereum protocol which has keys of the same length (64)

if (traverseContext.IsNodeRead)
{
return ref node.FullRlp;
}
if (traverseContext.IsReadValue)
{
return ref node.ValueRef;
}

if (traverseContext.IsDelete)
{
if (node.Value.IsNull)
{
return ref CappedArray<byte>.Null;
}

ConnectNodes(null, in traverseContext);
}
else if (Bytes.AreEqual(traverseContext.UpdateValue, node.Value))
{
return ref traverseContext.UpdateValue;
}
else
{
TrieNode withUpdatedValue = node.CloneWithChangedValue(in traverseContext.UpdateValue);
ConnectNodes(withUpdatedValue, in traverseContext);
}

return ref traverseContext.UpdateValue;
}

private readonly ref struct TraverseContext
Expand Down Expand Up @@ -1357,11 +1411,12 @@ bool IsNodeStackEmpty()

void ClearNodeStack() => _nodeStack?.Clear();

void PushToNodeStack(in StackedNode value)
[MethodImpl(MethodImplOptions.NoInlining)]
void PushToNodeStack(TrieNode node, int pathLength, int pathIndex)
{
// Allocated the _nodeStack if first push
_nodeStack ??= new();
_nodeStack.Push(value);
_nodeStack.Push(new StackedNode(node, pathLength, pathIndex));
}

StackedNode PopFromNodeStack()
Expand Down
2 changes: 1 addition & 1 deletion src/Nethermind/Nethermind.Trie/Pruning/ScopedTrieStore.cs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

namespace Nethermind.Trie.Pruning;

public class ScopedTrieStore : IScopedTrieStore
public sealed class ScopedTrieStore : IScopedTrieStore
{
private ITrieStore _trieStoreImplementation;
private readonly Hash256? _address;
Expand Down
Loading

0 comments on commit 4e98f26

Please sign in to comment.