Perf/Non blocking memory pruning (#9096)

* Move lastseen to triestore

* Slight cleanup

* Non blocking memory pruning

* Fix test

* Fix another test

* Fix pruning test

* LastSeen to LastCommit

* Add extra test

* allow disabling by setting to 0

* Debugging pruning....

* Stabilize test

* Fix build

* Remove console log

* BBetter name

* Whitespace

* Fix test

* Fix verify trie

---------

Co-authored-by: Ben {chmark} Adams <thundercat@illyriad.co.uk>
Co-authored-by: Kamil Chodoła <43241881+kamilchodola@users.noreply.github.com>

Author: 3 people

src/Nethermind/Nethermind.Blockchain.Test/FullPruning/FullPruningDiskTest.cs

@@ -75,13 +75,19 @@ protected override async Task<TestBlockchain> Build(Action<ContainerBuilder>? co
             return chain;
         }
 
-        protected override ContainerBuilder ConfigureContainer(ContainerBuilder builder, IConfigProvider configProvider) =>
+        protected override ContainerBuilder
+            ConfigureContainer(ContainerBuilder builder, IConfigProvider configProvider) =>
             // Reenable rocksdb
             base.ConfigureContainer(builder, configProvider)
                 .AddSingleton<IDbFactory, RocksDbFactory>()
                 .Intercept<IInitConfig>((initConfig) =>
                 {
                     initConfig.BaseDbPath = TempDirectory.Path;
+                })
+                .Intercept<IPruningConfig>((pruningConfig) =>
+                {
+                    // Make test faster otherwise it may potentially buffer 128 block.
+                    pruningConfig.MaxBufferedCommitCount = 1;
                 });
 
         public override void Dispose()

src/Nethermind/Nethermind.Blockchain/FullPruning/FullPruner.cs

@@ -118,16 +118,13 @@ protected virtual async Task RunFullPruning(CancellationToken cancellationToken)
         {
             IPruningContext? pruningContext = null;
 
-            // we don't want to start pruning in the middle of block processing, lets wait for new head.
-            await WaitForMainChainChange((e) =>
+            using (_trieStore.PrepareStableState(cancellationToken))
             {
                 if (_fullPruningDb.TryStartPruning(_pruningConfig.Mode.IsMemory(), out IPruningContext fromDbPruningContext))
                 {
                     pruningContext = fromDbPruningContext;
                 }
-
-                return true;
-            }, cancellationToken);
+            }
 
             if (pruningContext is null) return;
 
@@ -147,8 +144,6 @@ await WaitForMainChainChange((e) =>
 
         private async Task RunFullPruning(IPruningContext pruningContext, CancellationToken cancellationToken)
         {
-            _trieStore.PersistCache(cancellationToken);
-
             long blockToWaitFor = 0;
             await WaitForMainChainChange((e) =>
             {
@@ -168,7 +163,7 @@ await WaitForMainChainChange((e) =>
             }, cancellationToken);
 
             long stateToCopy = _blockTree.BestPersistedState.Value;
-            long blockToPruneAfter = stateToCopy + Reorganization.MaxDepth;
+            long blockToPruneAfter = stateToCopy + _pruningConfig.PruningBoundary;
 
             await WaitForMainChainChange((e) =>
             {
@@ -225,7 +220,7 @@ private void HandlePruningFinished(object? sender, PruningEventArgs e)
             }
         }
 
-        private async Task CopyTrie(IPruningContext pruning, Hash256 stateRoot, CancellationToken cancellationToken)
+        private Task CopyTrie(IPruningContext pruning, Hash256 stateRoot, CancellationToken cancellationToken)
         {
             INodeStorage.KeyScheme originalKeyScheme = _nodeStorage.Scheme;
             ICopyTreeVisitor visitor = null;
@@ -277,16 +272,12 @@ private async Task CopyTrie(IPruningContext pruning, Hash256 stateRoot, Cancella
                 {
                     visitor.Finish();
 
-                    _nodeStorage.Scheme = targetNodeStorage.Scheme;
-                    // Note: This does means that during full pruning some of the key copied will be of old key scheme.
-                    await WaitForMainChainChange((e) =>
+                    using (_trieStore.PrepareStableState(cancellationToken))
                     {
-                        // The db swap happens here. We do it within the event handler of main chain change to block
-                        // so that it does not happen during block processing.
                         pruning.Commit();
-                        return true;
-                    }, cancellationToken);
+                    }
 
+                    _nodeStorage.Scheme = targetNodeStorage.Scheme;
                     _lastPruning = DateTime.UtcNow;
                 }
             }
@@ -300,6 +291,8 @@ await WaitForMainChainChange((e) =>
             {
                 visitor?.Dispose();
             }
+
+            return Task.CompletedTask;
         }
 
         private ICopyTreeVisitor CopyTree<TContext>(

src/Nethermind/Nethermind.Core.Test/Blockchain/TestBlockchain.cs

@@ -249,7 +249,7 @@ protected virtual async Task<TestBlockchain> Build(Action<ContainerBuilder>? con
 
             Block? genesis = GetGenesisBlock(WorldStateManager.GlobalWorldState);
             BlockTree.SuggestBlock(genesis);
-            await waitGenesis;
+            waitGenesis.Wait(_cts.Token);
         }
 
         if (testConfiguration.AddBlockOnStart)

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

@@ -63,7 +63,7 @@ private async Task PrepareGenesis(CancellationToken cancellation)
 
         Block genesis = genesisLoader.Load();
         blockTree.SuggestBlock(genesis);
-        await newHeadTask;
+        newHeadTask.Wait();
     }
 
     public async Task StartNetwork(CancellationToken cancellationToken)

src/Nethermind/Nethermind.Core.Test/TestRawTrieStore.cs

@@ -62,6 +62,11 @@ public bool HasRoot(Hash256 stateRoot)
         return nodeStorage.KeyExists(null, TreePath.Empty, stateRoot);
     }
 
+    public IDisposable BeginScope(BlockHeader? baseBlock)
+    {
+        return new Reactive.AnonymousDisposable(() => { });
+    }
+
     public IScopedTrieStore GetTrieStore(Hash256? address)
     {
         return new RawScopedTrieStore(nodeStorage, address);
@@ -86,4 +91,18 @@ public event EventHandler<ReorgBoundaryReached>? ReorgBoundaryReached
     }
 
     public IReadOnlyKeyValueStore TrieNodeRlpStore => throw new Exception("Unsupported operatioon");
+
+    private Lock _scopeLock = new Lock();
+    private Lock _pruneLock = new Lock();
+    public TrieStore.StableLockScope PrepareStableState(CancellationToken cancellationToken)
+    {
+        var scopeLockScope = _scopeLock.EnterScope();
+        var pruneLockScope = _pruneLock.EnterScope();
+
+        return new TrieStore.StableLockScope
+        {
+            scopeLockScope = scopeLockScope,
+            pruneLockScope = pruneLockScope,
+        };
+    }
 }

src/Nethermind/Nethermind.Db/IPruningConfig.cs

@@ -97,4 +97,7 @@ The max number of parallel tasks that can be used by full pruning.
 
     [ConfigItem(Description = "Minimum number of block worth of unpersisted state in memory. Prevent memory pruning too often due to insufficient dirty cache memory.", DefaultValue = "8")]
     long MinUnpersistedBlockCount { get; set; }
+
+    [ConfigItem(Description = "Maximum number of block in commit buffer before blocking.", DefaultValue = "128", HiddenFromDocs = true)]
+    int MaxBufferedCommitCount { get; set; }
 }

src/Nethermind/Nethermind.Db/PruningConfig.cs

@@ -62,5 +62,6 @@ public int DirtyNodeShardBit
         public long PrunePersistedNodeMinimumTarget { get; set; } = 50.MiB();
         public long MaxUnpersistedBlockCount { get; set; } = 300; // About 1 hour on mainnet
         public long MinUnpersistedBlockCount { get; set; } = 8; // About slightly more than 1 minute
+        public int MaxBufferedCommitCount { get; set; } = 128;
     }
 }

src/Nethermind/Nethermind.State.Test/StateReaderTests.cs

@@ -32,7 +32,7 @@ public class StateReaderTests
         private static readonly ILogManager Logger = LimboLogs.Instance;
 
         [Test]
-        public async Task Can_ask_about_balance_in_parallel()
+        public void Can_ask_about_balance_in_parallel()
         {
             IReleaseSpec spec = MainnetSpecProvider.Instance.GetSpec((ForkActivation)MainnetSpecProvider.ConstantinopleFixBlockNumber);
             IDbProvider dbProvider = TestMemDbProvider.Init();
@@ -71,11 +71,11 @@ public async Task Can_ask_about_balance_in_parallel()
             Task c = StartTask(reader, baseBlock2, 3);
             Task d = StartTask(reader, baseBlock3, 4);
 
-            await Task.WhenAll(a, b, c, d);
+            Task.WhenAll(a, b, c, d).Wait();
         }
 
         [Test]
-        public async Task Can_ask_about_storage_in_parallel()
+        public void Can_ask_about_storage_in_parallel()
         {
             StorageCell storageCell = new(_address1, UInt256.One);
             IReleaseSpec spec = MuirGlacier.Instance;
@@ -130,7 +130,7 @@ void CommitEverything()
             Task c = StartStorageTask(reader, baseBlock2, storageCell, new byte[] { 3 });
             Task d = StartStorageTask(reader, baseBlock3, storageCell, new byte[] { 4 });
 
-            await Task.WhenAll(a, b, c, d);
+            Task.WhenAll(a, b, c, d).Wait();
         }
 
         [Test]

src/Nethermind/Nethermind.State/BlockingVerifyTrie.cs

@@ -45,7 +45,7 @@ public BlockingVerifyTrie(
     public bool VerifyTrie(BlockHeader stateAtBlock, CancellationToken cancellationToken)
     {
         // This is to block processing as with halfpath old nodes will be removed
-        using IBlockCommitter? _ = _trieStore.BeginBlockCommit(stateAtBlock.Number + 1);
+        using IDisposable _ = _trieStore.BeginScope(stateAtBlock);
 
         Hash256 rootNode = stateAtBlock.StateRoot;
         TrieStats stats = _stateReader.CollectStats(rootNode, _codeDb, _logManager, cancellationToken);

src/Nethermind/Nethermind.State/WorldState.cs

@@ -246,11 +246,13 @@ public IDisposable BeginScope(BlockHeader? baseBlock)
             if (_logger.IsTrace) _logger.Trace($"Beginning WorldState scope with baseblock {baseBlock?.ToString(BlockHeader.Format.Short) ?? "null"} with stateroot {baseBlock?.StateRoot?.ToString() ?? "null"}.");
 
             StateRoot = baseBlock?.StateRoot ?? Keccak.EmptyTreeHash;
+            IDisposable trieStoreCloser = _trieStore.BeginScope(baseBlock);
 
             return new Reactive.AnonymousDisposable(() =>
             {
                 Reset();
                 StateRoot = Keccak.EmptyTreeHash;
+                trieStoreCloser.Dispose();
                 _isInScope = false;
                 if (_logger.IsTrace) _logger.Trace($"WorldState scope for baseblock {baseBlock?.ToString(BlockHeader.Format.Short) ?? "null"} closed");
             });