Feature/recover whole path

src/Nethermind/Nethermind.Consensus/Processing/BlockCachePreWarmer.cs

@@ -18,6 +18,7 @@
 using Nethermind.State;
 using Nethermind.Core.Eip2930;
 using Nethermind.Core.Collections;
+using Nethermind.Trie;
 
 namespace Nethermind.Consensus.Processing;
 
@@ -106,6 +107,9 @@ private void WarmupWithdrawals(ParallelOptions parallelOptions, IReleaseSpec spe
                             using IReadOnlyTxProcessingScope scope = env.Build(state.stateRoot);
                             scope.WorldState.WarmUp(state.block.Withdrawals[i].Address);
                         }
+                        catch (MissingTrieNodeException)
+                        {
+                        }
                         finally
                         {
                             state.envPool.Return(env);
@@ -276,14 +280,21 @@ private void WarmupAddresses(ParallelOptions parallelOptions, Block block)
                     Transaction tx = state.Block.Transactions[i];
                     Address? sender = tx.SenderAddress;
 
-                    if (sender is not null)
+                    try
                     {
-                        state.Scope.WorldState.WarmUp(sender);
+                        if (sender is not null)
+                        {
+                            state.Scope.WorldState.WarmUp(sender);
+                        }
+
+                        Address to = tx.To;
+                        if (to is not null)
+                        {
+                            state.Scope.WorldState.WarmUp(to);
+                        }
                     }
-                    Address to = tx.To;
-                    if (to is not null)
+                    catch (MissingTrieNodeException)
                     {
-                        state.Scope.WorldState.WarmUp(to);
                     }
 
                     return state;

src/Nethermind/Nethermind.Core.Test/AutoCancelTokenSourceTests.cs

@@ -0,0 +1,39 @@
+// SPDX-FileCopyrightText: 2025 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System.Threading;
+using FluentAssertions;
+using Nethermind.Core.Extensions;
+using Nethermind.Core.Utils;
+using NUnit.Framework;
+
+namespace Nethermind.Core.Test;
+
+public class AutoCancelTokenSourceTests
+{
+    [Test]
+    public void AutoCancelOnExitClosure()
+    {
+        CancellationToken TaskWithInnerCancellation(CancellationToken token)
+        {
+            using AutoCancelTokenSource cts = token.CreateChildTokenSource();
+            return cts.Token;
+        }
+
+        TaskWithInnerCancellation(default).IsCancellationRequested.Should().BeTrue();
+    }
+
+    [Test]
+    public void AutoCancelPropagateParentCancellation()
+    {
+        using CancellationTokenSource cts = new CancellationTokenSource();
+
+        using AutoCancelTokenSource acts = cts.Token.CreateChildTokenSource();
+
+        acts.Token.IsCancellationRequested.Should().BeFalse();
+
+        cts.Cancel();
+
+        acts.Token.IsCancellationRequested.Should().BeTrue();
+    }
+}

src/Nethermind/Nethermind.Core.Test/Modules/NetworkModule.cs

@@ -77,8 +77,12 @@ protected override void Load(ContainerBuilder builder)
             {
                 ILogManager logManager = ctx.Resolve<ILogManager>();
                 ctx.Resolve<IWorldStateManager>().InitializeNetwork(
-                    new GetNodeDataTrieNodeRecovery(peerPool, logManager),
-                    new SnapTrieNodeRecovery(peerPool, logManager));
+                    new PathNodeRecovery(
+                        new NodeDataRecovery(peerPool!, ctx.Resolve<INodeStorage>(), logManager),
+                        new SnapRangeRecovery(peerPool!, logManager),
+                        logManager
+                    )
+                );
             })
 
             // TODO: LastNStateRootTracker

src/Nethermind/Nethermind.Core/Extensions/CancellationTokenExtensions.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Nethermind.Core.Utils;
 
 namespace Nethermind.Core.Extensions
 {
@@ -54,5 +55,13 @@ public static CancellationTokenSource ThatCancelAfter(this CancellationTokenSour
             cts.CancelAfter(delay);
             return cts;
         }
+
+        public static AutoCancelTokenSource CreateChildTokenSource(this CancellationToken parentToken, TimeSpan delay = default)
+        {
+            CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(parentToken);
+            if (delay != TimeSpan.Zero) cts.CancelAfter(delay);
+
+            return new AutoCancelTokenSource(cts);
+        }
     }
 }

src/Nethermind/Nethermind.Core/Tasks/WaitForPassingTasks.cs

@@ -0,0 +1,48 @@
+// SPDX-FileCopyrightText: 2025 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Linq;
+using System.Threading.Tasks;
+
+namespace Nethermind.Core.Tasks;
+
+public static class Wait
+{
+    /// <summary>
+    /// Wait for any of the task that passed the predicate and forward the result, or all of the task to complete.
+    /// </summary>
+    /// <param name="cond"></param>
+    /// <param name="tasks"></param>
+    /// <typeparam name="T"></typeparam>
+    /// <returns></returns>
+    public static async Task<T> ForPassingTask<T>(Func<T, bool> cond, params IEnumerable<Task<T>> tasks)
+    {
+        HashSet<Task<T>> taskSet = new HashSet<Task<T>>(tasks);
+        while (taskSet.Any())
+        {
+            Task<T> resolved = await Task.WhenAny<T>(taskSet);
+            taskSet.Remove(resolved);
+
+            T result = await resolved;
+
+            if (cond(result))
+            {
+                // Its ok, then immediately return.
+                return result;
+            }
+
+            if (!taskSet.Any())
+            {
+                // No more tasks, just return the last one.
+                return result;
+            }
+
+            // Otherwise, we try WhenAny again.
+        }
+
+        throw new UnreachableException();
+    }
+}

src/Nethermind/Nethermind.Core/Utils/AutoCancelTokenSource.cs

@@ -0,0 +1,29 @@
+// SPDX-FileCopyrightText: 2025 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Threading;
+
+namespace Nethermind.Core.Utils;
+
+/// <summary>
+/// Automatically cancel and dispose underlying cancellation token source.
+/// Make it easy to have golang style defer cancel pattern.
+/// </summary>
+public readonly struct AutoCancelTokenSource(CancellationTokenSource cancellationTokenSource) : IDisposable
+{
+    public CancellationToken Token => cancellationTokenSource.Token;
+
+    public static AutoCancelTokenSource ThatCancelAfter(TimeSpan delay)
+    {
+        CancellationTokenSource cancellationTokenSource = new();
+        cancellationTokenSource.CancelAfter(delay);
+        return new AutoCancelTokenSource(cancellationTokenSource);
+    }
+
+    public void Dispose()
+    {
+        cancellationTokenSource.Cancel();
+        cancellationTokenSource.Dispose();
+    }
+}

src/Nethermind/Nethermind.Init/PruningTrieStateFactory.cs

@@ -131,17 +131,11 @@ ILogManager logManager
 
         INodeStorage mainNodeStorage = nodeStorageFactory.WrapKeyValueStore(stateDb);
 
-        TrieStore trieStore = syncConfig.TrieHealing
-            ? new HealingTrieStore(
-                mainNodeStorage,
-                pruningStrategy,
-                persistenceStrategy,
-                logManager)
-            : new TrieStore(
-                mainNodeStorage,
-                pruningStrategy,
-                persistenceStrategy,
-                logManager);
+        TrieStore trieStore = new TrieStore(
+            mainNodeStorage,
+            pruningStrategy,
+            persistenceStrategy,
+            logManager);
 
         ITrieStore mainWorldTrieStore = trieStore;
         PreBlockCaches? preBlockCaches = null;

src/Nethermind/Nethermind.Init/Steps/InitializeNetwork.cs

@@ -123,7 +123,13 @@ private async Task Initialize(CancellationToken cancellationToken)
             _api.DisposeStack.Push((IAsyncDisposable)container);
         }
 
-        _api.WorldStateManager!.InitializeNetwork(new GetNodeDataTrieNodeRecovery(_api.SyncPeerPool!, _api.LogManager), new SnapTrieNodeRecovery(_api.SyncPeerPool!, _api.LogManager));
+        _api.WorldStateManager!.InitializeNetwork(
+            new PathNodeRecovery(
+                new NodeDataRecovery(_api.SyncPeerPool!, _api.MainNodeStorage!, _api.LogManager),
+                new SnapRangeRecovery(_api.SyncPeerPool!, _api.LogManager),
+                _api.LogManager
+            )
+        );
 
         _api.TxGossipPolicy.Policies.Add(new SyncedTxGossipPolicy(_api.SyncModeSelector));
 

src/Nethermind/Nethermind.JsonRpc/Modules/Eth/EthRpcModule.cs

@@ -188,7 +188,7 @@ public ResultWrapper<byte[]> eth_getStorageAt(Address address, UInt256 positionI
         }
         catch (MissingTrieNodeException e)
         {
-            var hash = e.TrieNodeException.NodeHash;
+            var hash = e.Hash;
             return ResultWrapper<byte[]>.Fail($"missing trie node {hash} (path ) state {hash} is not available", ErrorCodes.InvalidInput);
         }
     }

src/Nethermind/Nethermind.Network.Stats/NodeStatsLight.cs

@@ -204,7 +204,7 @@ private void UpdateValue(ref decimal? currentValue, decimal newValue)
     {
         return (long?)(transferSpeedType switch
         {
-            TransferSpeedType.Latency => _averageLatency,
+            TransferSpeedType.Latency => _averageLatency ?? 10000,
             TransferSpeedType.NodeData => _averageNodesTransferSpeed,
             TransferSpeedType.Headers => _averageHeadersTransferSpeed,
             TransferSpeedType.Bodies => _averageBodiesTransferSpeed,

src/Nethermind/Nethermind.State/Healing/HealingStateTree.cs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using Nethermind.Core.Collections;
 using Nethermind.Core.Crypto;
@@ -14,15 +15,15 @@ namespace Nethermind.State.Healing;
 
 public class HealingStateTree : StateTree
 {
-    private ITrieNodeRecovery<GetTrieNodesRequest>? _recovery;
+    private IPathRecovery? _recovery;
 
     [DebuggerStepThrough]
     public HealingStateTree(ITrieStore? store, ILogManager? logManager)
         : base(store.GetTrieStore(null), logManager)
     {
     }
 
-    public void InitializeNetwork(ITrieNodeRecovery<GetTrieNodesRequest> recovery)
+    public void InitializeNetwork(IPathRecovery recovery)
     {
         _recovery = recovery;
     }
@@ -35,7 +36,8 @@ public override ReadOnlySpan<byte> Get(ReadOnlySpan<byte> rawKey, Hash256? rootH
         }
         catch (MissingTrieNodeException e)
         {
-            if (Recover(e.TrieNodeException.NodeHash, e.GetPathPart(), rootHash ?? RootHash))
+            Hash256 fullPath = new Hash256(rawKey);
+            if (Recover(e.Path, e.Hash, fullPath))
             {
                 return base.Get(rawKey, rootHash);
             }
@@ -52,7 +54,8 @@ public override void Set(ReadOnlySpan<byte> rawKey, byte[] value)
         }
         catch (MissingTrieNodeException e)
         {
-            if (Recover(e.TrieNodeException.NodeHash, e.GetPathPart(), RootHash))
+            Hash256 fullPath = new Hash256(rawKey);
+            if (Recover(e.Path, e.Hash, fullPath))
             {
                 base.Set(rawKey, value);
             }
@@ -63,27 +66,18 @@ public override void Set(ReadOnlySpan<byte> rawKey, byte[] value)
         }
     }
 
-    private bool Recover(in ValueHash256 rlpHash, ReadOnlySpan<byte> pathPart, Hash256 rootHash)
+    private bool Recover(in TreePath missingNodePath, in ValueHash256 hash, Hash256 fullPath)
     {
-        if (_recovery?.CanRecover == true)
+        if (_recovery is not null)
         {
-            GetTrieNodesRequest request = new()
+            IDictionary<TreePath, byte[]>? rlps = _recovery.Recover(RootHash, null, missingNodePath, hash, fullPath, default).GetAwaiter().GetResult();
+            if (rlps is not null)
             {
-                RootHash = rootHash,
-                AccountAndStoragePaths = new ArrayPoolList<PathGroup>(1)
+                foreach (KeyValuePair<TreePath, byte[]> kv in rlps)
                 {
-                    new()
-                    {
-                        Group = [Nibbles.EncodePath(pathPart)]
-                    }
+                    ValueHash256 nodeHash = ValueKeccak.Compute(kv.Value);
+                    TrieStore.Set(kv.Key, nodeHash, kv.Value);
                 }
-            };
-
-            byte[]? rlp = _recovery.Recover(rlpHash, request).GetAwaiter().GetResult();
-            if (rlp is not null)
-            {
-                TreePath path = TreePath.FromNibble(pathPart);
-                TrieStore.Set(path, rlpHash, rlp);
                 return true;
             }
         }