TrieStore to Channels (#7845)

Author: benaadams

src/Nethermind/Nethermind.Trie/Pruning/TrieStore.cs

@@ -3,15 +3,16 @@
 
 using System;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Channels;
 using System.Threading.Tasks;
 using Nethermind.Core;
 using Nethermind.Core.Collections;
+using Nethermind.Core.Cpu;
 using Nethermind.Core.Crypto;
 using Nethermind.Core.Extensions;
 using Nethermind.Logging;
@@ -32,7 +33,7 @@ public class TrieStore : ITrieStore, IPruningTrieStore
     private readonly Task[] _dirtyNodesTasks = [];
     private readonly ConcurrentDictionary<HashAndTinyPath, Hash256?>[] _persistedHashes = [];
     private readonly Action<TreePath, Hash256?, TrieNode> _persistedNodeRecorder;
-    private readonly Task[] _disposeTasks = new Task[Environment.ProcessorCount];
+    private readonly Task[] _disposeTasks = new Task[RuntimeInformation.PhysicalCoreCount];
 
     // This seems to attempt prevent multiple block processing at the same time and along with pruning at the same time.
     private readonly object _dirtyNodesLock = new object();
@@ -790,27 +791,33 @@ void TopLevelPersist(TrieNode tn, Hash256? address2, TreePath path)
         // However, anything that we are trying to persist here should still be in dirty cache.
         // So parallel read should go there first instead of to the database for these dataset,
         // so it should be fine for these to be non atomic.
-        using BlockingCollection<INodeStorage.WriteBatch> disposeQueue = new BlockingCollection<INodeStorage.WriteBatch>(4);
-
-        for (int index = 0; index < _disposeTasks.Length; index++)
+        Task[] disposeTasks = _disposeTasks;
+        Channel<INodeStorage.WriteBatch> disposeQueue = Channel.CreateBounded<INodeStorage.WriteBatch>(disposeTasks.Length * 2);
+        try
         {
-            _disposeTasks[index] = Task.Run(() =>
+            for (int index = 0; index < disposeTasks.Length; index++)
             {
-                while (disposeQueue.TryTake(out INodeStorage.WriteBatch disposable, Timeout.Infinite))
+                disposeTasks[index] = Task.Run(async () =>
                 {
-                    disposable.Dispose();
-                }
-            });
-        }
+                    await foreach (IDisposable disposable in disposeQueue.Reader.ReadAllAsync())
+                    {
+                        disposable.Dispose();
+                    }
+                });
+            }
 
-        using ArrayPoolList<Task> persistNodeStartingFromTasks = parallelStartNodes.Select(
-            entry => Task.Run(() => PersistNodeStartingFrom(entry.trieNode, entry.address2, entry.path, persistedNodeRecorder, writeFlags, disposeQueue)))
-            .ToPooledList(parallelStartNodes.Count);
+            using ArrayPoolList<Task> persistNodeStartingFromTasks = parallelStartNodes.Select(
+                entry => Task.Run(() => PersistNodeStartingFrom(entry.trieNode, entry.address2, entry.path, persistedNodeRecorder, writeFlags, disposeQueue)))
+                .ToPooledList(parallelStartNodes.Count);
 
-        Task.WaitAll(persistNodeStartingFromTasks.AsSpan());
+            Task.WaitAll(persistNodeStartingFromTasks.AsSpan());
+        }
+        finally
+        {
+            disposeQueue.Writer.Complete();
+        }
 
-        disposeQueue.CompleteAdding();
-        Task.WaitAll(_disposeTasks);
+        Task.WaitAll(disposeTasks);
 
         // Dispose top level last in case something goes wrong, at least the root won't be stored
         topLevelWriteBatch.Dispose();
@@ -824,28 +831,28 @@ void TopLevelPersist(TrieNode tn, Hash256? address2, TreePath path)
         LastPersistedBlockNumber = commitSet.BlockNumber;
     }
 
-    private void PersistNodeStartingFrom(TrieNode tn, Hash256 address2, TreePath path,
+    private async Task PersistNodeStartingFrom(TrieNode tn, Hash256 address2, TreePath path,
         Action<TreePath, Hash256?, TrieNode>? persistedNodeRecorder,
-        WriteFlags writeFlags, BlockingCollection<INodeStorage.WriteBatch> disposeQueue)
+        WriteFlags writeFlags, Channel<INodeStorage.WriteBatch> disposeQueue)
     {
         long persistedNodeCount = 0;
         INodeStorage.WriteBatch writeBatch = _nodeStorage.StartWriteBatch();
 
-        void DoPersist(TrieNode node, Hash256? address3, TreePath path2)
+        async ValueTask DoPersist(TrieNode node, Hash256? address3, TreePath path2)
         {
             persistedNodeRecorder?.Invoke(path2, address3, node);
             PersistNode(address3, path2, node, writeBatch, writeFlags);
 
             persistedNodeCount++;
             if (persistedNodeCount % 512 == 0)
             {
-                disposeQueue.Add(writeBatch);
+                await disposeQueue.Writer.WriteAsync(writeBatch);
                 writeBatch = _nodeStorage.StartWriteBatch();
             }
         }
 
-        tn.CallRecursively(DoPersist, address2, ref path, GetTrieStore(address2), true, _logger);
-        disposeQueue.Add(writeBatch);
+        await tn.CallRecursivelyAsync(DoPersist, address2, ref path, GetTrieStore(address2), _logger);
+        await disposeQueue.Writer.WriteAsync(writeBatch);
     }
 
     private void PersistNode(Hash256? address, in TreePath path, TrieNode currentNode, INodeStorage.WriteBatch writeBatch, WriteFlags writeFlags = WriteFlags.None)

src/Nethermind/Nethermind.Trie/TrieNode.cs

@@ -7,6 +7,7 @@
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
+using System.Threading.Tasks;
 using Nethermind.Core;
 using Nethermind.Core.Buffers;
 using Nethermind.Core.Crypto;
@@ -896,12 +897,12 @@ public void CallRecursively(
 
             if (!IsLeaf)
             {
-                if (_data is not null)
+                object?[] data = _data;
+                if (data is not null)
                 {
-                    for (int i = 0; i < _data.Length; i++)
+                    for (int i = 0; i < data.Length; i++)
                     {
-                        object o = _data[i];
-                        if (o is TrieNode child)
+                        if (data[i] is TrieNode child)
                         {
                             if (logger.IsTrace) logger.Trace($"Persist recursively on child {i} {child} of {this}");
                             int previousLength = AppendChildPath(ref currentPath, i);
@@ -938,6 +939,101 @@ public void CallRecursively(
             action(this, storageAddress, currentPath);
         }
 
+        public ValueTask CallRecursivelyAsync(
+            Func<TrieNode, Hash256?, TreePath, ValueTask> action,
+            Hash256? storageAddress,
+            ref TreePath currentPath,
+            ITrieNodeResolver resolver,
+            ILogger logger)
+        {
+            if (IsPersisted)
+            {
+                if (logger.IsTrace) logger.Trace($"Skipping {this} - already persisted");
+                return default;
+            }
+
+            if (currentPath.Length >= Int32.MaxValue)
+            {
+                return action(this, storageAddress, currentPath);
+            }
+
+            if (!IsLeaf)
+            {
+                if (_data is null)
+                {
+                    return action(this, storageAddress, currentPath);
+                }
+
+                return CallRecursivelyNotLeafAsync(
+                    action,
+                    storageAddress,
+                    currentPath,
+                    resolver,
+                    logger);
+            }
+            else
+            {
+                return CallRecursivelyLeafAsync(
+                    action,
+                    storageAddress,
+                    currentPath,
+                    resolver,
+                    logger);
+            }
+        }
+
+        private async ValueTask CallRecursivelyNotLeafAsync(
+            Func<TrieNode, Hash256?, TreePath, ValueTask> action,
+            Hash256? storageAddress,
+            TreePath currentPath,
+            ITrieNodeResolver resolver,
+            ILogger logger)
+        {
+            object?[] data = _data;
+            for (int i = 0; i < data.Length; i++)
+            {
+                if (data[i] is TrieNode child)
+                {
+                    if (logger.IsTrace) logger.Trace($"Persist recursively on child {i} {child} of {this}");
+                    int previousLength = AppendChildPath(ref currentPath, i);
+                    await child.CallRecursivelyAsync(action, storageAddress, ref currentPath, resolver, logger);
+                    currentPath.TruncateMut(previousLength);
+                }
+            }
+
+            await action(this, storageAddress, currentPath);
+        }
+
+        private async ValueTask CallRecursivelyLeafAsync(
+            Func<TrieNode, Hash256?, TreePath, ValueTask> action,
+            Hash256? storageAddress,
+            TreePath currentPath,
+            ITrieNodeResolver resolver,
+            ILogger logger)
+        {
+            TrieNode? storageRoot = _storageRoot;
+            if (storageRoot is not null || TryResolveStorageRoot(resolver, ref currentPath, out storageRoot))
+            {
+                if (logger.IsTrace) logger.Trace($"Persist recursively on storage root {_storageRoot} of {this}");
+                Hash256 storagePathAddr;
+                using (currentPath.ScopedAppend(Key))
+                {
+                    if (currentPath.Length != 64) throw new Exception("unexpected storage path length. Total nibble count should add up to 64.");
+                    storagePathAddr = currentPath.Path.ToCommitment();
+                }
+
+                TreePath emptyPath = TreePath.Empty;
+                await storageRoot!.CallRecursivelyAsync(
+                    action,
+                    storagePathAddr,
+                    ref emptyPath,
+                    resolver.GetStorageTrieNodeResolver(storagePathAddr),
+                    logger);
+            }
+
+            await action(this, storageAddress, currentPath);
+        }
+
         /// <summary>
         /// Imagine a branch like this:
         ///        B