Squashed improved inmemory pruning

src/Nethermind/Nethermind.Db/IPruningConfig.cs

@@ -92,4 +92,7 @@ public interface IPruningConfig : IConfig
 
     [ConfigItem(Description = "Whether to enables available disk space check.", DefaultValue = "true")]
     bool AvailableSpaceCheckEnabled { get; set; }
+
+    [ConfigItem(Description = "[TECHNICAL] Number of past persisted keys to keep track off for possible pruning.", DefaultValue = "1000000")]
+    int TrackedPastKeyCount { get; set; }
 }

src/Nethermind/Nethermind.Db/MemDb.cs

@@ -141,6 +141,11 @@ public virtual void Set(ReadOnlySpan<byte> key, byte[]? value, WriteFlags flags
             }
 
             WritesCount++;
+            if (value == null)
+            {
+                _db.TryRemove(key, out _);
+                return;
+            }
             _db[key] = value;
         }
     }

src/Nethermind/Nethermind.Db/PruningConfig.cs

@@ -32,5 +32,6 @@ public bool Enabled
         public int FullPruningMinimumDelayHours { get; set; } = 240;
         public FullPruningCompletionBehavior FullPruningCompletionBehavior { get; set; } = FullPruningCompletionBehavior.None;
         public bool AvailableSpaceCheckEnabled { get; set; } = true;
+        public int TrackedPastKeyCount { get; set; } = 1000000;
     }
 }

src/Nethermind/Nethermind.Init/InitializeStateDb.cs

@@ -90,7 +90,8 @@ public Task Execute(CancellationToken cancellationToken)
                 persistenceStrategy = persistenceStrategy.Or(triggerPersistenceStrategy);
             }
 
-            pruningStrategy = Prune.WhenCacheReaches(pruningConfig.CacheMb.MB()); // TODO: memory hint should define this
+            pruningStrategy = Prune.WhenCacheReaches(pruningConfig.CacheMb.MB()) // TODO: memory hint should define this
+                .TrackingPastKeys(pruningConfig.TrackedPastKeyCount);
         }
         else
         {

src/Nethermind/Nethermind.Trie.Test/Pruning/TestPruningStrategy.cs

@@ -8,10 +8,13 @@ namespace Nethermind.Trie.Test.Pruning
     public class TestPruningStrategy : IPruningStrategy
     {
         private readonly bool _pruningEnabled;
-        public TestPruningStrategy(bool pruningEnabled, bool shouldPrune = false)
+        private readonly int _trackedPastKeyCount;
+
+        public TestPruningStrategy(bool pruningEnabled, bool shouldPrune = false, int trackedPastKeyCount = 0)
         {
             _pruningEnabled = pruningEnabled;
             ShouldPruneEnabled = shouldPrune;
+            _trackedPastKeyCount = trackedPastKeyCount;
         }
 
         public bool PruningEnabled => _pruningEnabled;
@@ -27,5 +30,7 @@ public bool ShouldPrune(in long currentMemory)
 
             return false;
         }
+
+        public int TrackedPastKeyCount => _trackedPastKeyCount;
     }
 }

src/Nethermind/Nethermind.Trie.Test/Pruning/TreeStoreTests.cs

@@ -35,12 +35,22 @@ public TreeStoreTests(INodeStorage.KeyScheme scheme)
             _scheme = scheme;
         }
 
-        public TrieStore CreateTrieStore(IPruningStrategy? pruningStrategy = null, IKeyValueStoreWithBatching? kvStore = null, IPersistenceStrategy? persistenceStrategy = null)
+        private TrieStore CreateTrieStore(
+            IPruningStrategy? pruningStrategy = null,
+            IKeyValueStoreWithBatching? kvStore = null,
+            IPersistenceStrategy? persistenceStrategy = null,
+            long? reorgDepthOverride = null
+        )
         {
             pruningStrategy ??= No.Pruning;
             kvStore ??= new TestMemDb();
             persistenceStrategy ??= No.Persistence;
-            return new(new NodeStorage(kvStore, _scheme, requirePath: _scheme == INodeStorage.KeyScheme.HalfPath), pruningStrategy, persistenceStrategy, _logManager);
+            return new(
+                new NodeStorage(kvStore, _scheme, requirePath: _scheme == INodeStorage.KeyScheme.HalfPath),
+                pruningStrategy,
+                persistenceStrategy,
+                _logManager,
+                reorgDepthOverride: reorgDepthOverride);
         }
 
         [SetUp]
@@ -811,5 +821,64 @@ public void After_commit_should_have_has_root()
             stateTree.Commit(0);
             trieStore.HasRoot(stateTree.RootHash).Should().BeTrue();
         }
+
+        public async Task Will_RemovePastKeys_OnSnapshot()
+        {
+            MemDb memDb = new();
+
+            using TrieStore fullTrieStore = CreateTrieStore(
+                kvStore: memDb,
+                pruningStrategy: new TestPruningStrategy(true, true, 100000),
+                persistenceStrategy: No.Persistence,
+                reorgDepthOverride: 2);
+
+            IScopedTrieStore trieStore = fullTrieStore.GetTrieStore(null);
+
+            for (int i = 0; i < 64; i++)
+            {
+                TrieNode node = new(NodeType.Leaf, TestItem.Keccaks[i], new byte[2]);
+                trieStore.CommitNode(i, new NodeCommitInfo(node, TreePath.Empty));
+                trieStore.FinishBlockCommit(TrieType.State, i, node);
+
+                // Pruning is done in background
+                await Task.Delay(TimeSpan.FromMilliseconds(10));
+            }
+
+            if (_scheme == INodeStorage.KeyScheme.Hash)
+            {
+                memDb.Count.Should().NotBe(1);
+            }
+            else
+            {
+                memDb.Count.Should().Be(1);
+            }
+        }
+
+        [Test]
+        public async Task Will_NotRemove_ReCommittedNode()
+        {
+            MemDb memDb = new();
+
+            using TrieStore fullTrieStore = CreateTrieStore(
+                kvStore: memDb,
+                pruningStrategy: new TestPruningStrategy(true, true, 100000),
+                persistenceStrategy: No.Persistence,
+                reorgDepthOverride: 2);
+
+            IScopedTrieStore trieStore = fullTrieStore.GetTrieStore(null);
+
+            for (int i = 0; i < 64; i++)
+            {
+                TrieNode node = new(NodeType.Leaf, TestItem.Keccaks[i % 4], new byte[2]);
+                trieStore.CommitNode(i, new NodeCommitInfo(node, TreePath.Empty));
+                node = trieStore.FindCachedOrUnknown(TreePath.Empty, node.Keccak);
+                trieStore.FinishBlockCommit(TrieType.State, i, node);
+
+                // Pruning is done in background
+                await Task.Delay(TimeSpan.FromMilliseconds(10));
+            }
+
+            memDb.Count.Should().Be(4);
+        }
     }
 }

src/Nethermind/Nethermind.Trie.Test/PruningScenariosTests.cs

@@ -76,7 +76,7 @@ public static PruningContext InMemory
             public static PruningContext InMemoryAlwaysPrune
             {
                 [DebuggerStepThrough]
-                get => new(new TestPruningStrategy(true, true), No.Persistence);
+                get => new(new TestPruningStrategy(true, true, 1000000), No.Persistence);
             }
 
             public static PruningContext SetupWithPersistenceEveryEightBlocks
@@ -746,5 +746,34 @@ public void StateRoot_reset_at_lower_level_and_accessed_at_just_the_right_time()
                 .VerifyAccountBalance(3, 101);
 
         }
+
+        [Test]
+        public void When_Reorg_OldValueIsNotRemoved()
+        {
+            Reorganization.MaxDepth = 2;
+
+            PruningContext.InMemoryAlwaysPrune
+                .SetAccountBalance(1, 100)
+                .SetAccountBalance(2, 100)
+                .Commit()
+
+                .SetAccountBalance(3, 100)
+                .SetAccountBalance(4, 100)
+                .Commit()
+
+                .SaveBranchingPoint("revert_main")
+
+                .SetAccountBalance(4, 200)
+                .Commit()
+
+                .RestoreBranchingPoint("revert_main")
+
+                .Commit()
+                .Commit()
+                .Commit()
+                .Commit()
+
+                .VerifyAccountBalance(4, 100);
+        }
     }
 }

src/Nethermind/Nethermind.Trie.Test/TinyTreePathTests.cs

@@ -0,0 +1,32 @@
+// SPDX-FileCopyrightText: 2023 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using FluentAssertions;
+using Nethermind.Core.Crypto;
+using NUnit.Framework;
+
+namespace Nethermind.Trie.Test;
+
+public class TinyTreePathTests
+{
+    [Test]
+    public void Should_ConvertFromAndToTreePath()
+    {
+        TreePath path = new TreePath(new ValueHash256("0123456789abcd00000000000000000000000000000000000000000000000000"), 14);
+
+        TinyTreePath tinyPath = new TinyTreePath(path);
+
+        tinyPath.ToTreePath().Should().Be(path);
+    }
+
+    [Test]
+    public void When_PathIsTooLong_Should_Throw()
+    {
+        TreePath path = new TreePath(new ValueHash256("0123456789000000000000000000000000000000000000000000000000000000"), 15);
+
+        Action act = () => new TinyTreePath(path);
+        act.Should().Throw<InvalidOperationException>();
+    }
+}
+

src/Nethermind/Nethermind.Trie/INodeStorage.cs

@@ -44,5 +44,6 @@ public enum KeyScheme
     public interface WriteBatch : IDisposable
     {
         void Set(Hash256? address, in TreePath path, in ValueHash256 currentNodeKeccak, byte[] toArray, WriteFlags writeFlags);
+        void Remove(Hash256? address, in TreePath path, in ValueHash256 currentNodeKeccak);
     }
 }

src/Nethermind/Nethermind.Trie/NodeStorage.cs

@@ -220,5 +220,11 @@ public void Set(Hash256? address, in TreePath path, in ValueHash256 keccak, byte
 
             _writeBatch.Set(_nodeStorage.GetExpectedPath(stackalloc byte[StoragePathLength], address, path, keccak), toArray, writeFlags);
         }
+
+        public void Remove(Hash256? address, in TreePath path, in ValueHash256 keccak)
+        {
+            // Only delete half path key. DO NOT delete hash based key.
+            _writeBatch.Remove(GetHalfPathNodeStoragePathSpan(stackalloc byte[StoragePathLength], address, path, keccak));
+        }
     }
 }

src/Nethermind/Nethermind.Trie/Pruning/ISnapshotStrategy.cs

@@ -7,5 +7,6 @@ public interface IPruningStrategy
     {
         bool PruningEnabled { get; }
         bool ShouldPrune(in long currentMemory);
+        int TrackedPastKeyCount { get; }
     }
 }

src/Nethermind/Nethermind.Trie/Pruning/MemoryLimit.cs

@@ -21,5 +21,7 @@ public bool ShouldPrune(in long currentMemory)
         {
             return PruningEnabled && currentMemory >= _memoryLimit;
         }
+
+        public int TrackedPastKeyCount => 0;
     }
 }

src/Nethermind/Nethermind.Trie/Pruning/NoPruning.cs

@@ -15,5 +15,7 @@ public bool ShouldPrune(in long currentMemory)
         {
             return false;
         }
+
+        public int TrackedPastKeyCount => 0;
     }
 }

src/Nethermind/Nethermind.Trie/Pruning/Prune.cs

@@ -7,5 +7,10 @@ public static class Prune
     {
         public static IPruningStrategy WhenCacheReaches(long sizeInBytes)
             => new MemoryLimit(sizeInBytes);
+
+        public static IPruningStrategy TrackingPastKeys(this IPruningStrategy baseStrategy, int trackedPastKeyCount)
+            => trackedPastKeyCount <= 0
+                ? baseStrategy
+                : new TrackedPastKeyCountStrategy(baseStrategy, trackedPastKeyCount);
     }
 }

src/Nethermind/Nethermind.Trie/Pruning/TinyTreePath.cs

@@ -0,0 +1,43 @@
+// SPDX-FileCopyrightText: 2023 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using Nethermind.Core.Crypto;
+
+namespace Nethermind.Trie;
+
+/// <summary>
+/// Like TreePath, but tiny. Fit in 8 byte, like a long. Can only represent 14 nibble.
+/// </summary>
+public struct TinyTreePath
+{
+    public const int MaxNibbleLength = 14;
+
+    long _data;
+
+    // Eh.. readonly?
+    private Span<byte> AsSpan => MemoryMarshal.AsBytes(MemoryMarshal.CreateSpan(ref _data, 1));
+
+    public TinyTreePath(in TreePath path)
+    {
+        if (path.Length > MaxNibbleLength) throw new InvalidOperationException("Unable to represent more than 14 nibble");
+        Span<byte> pathSpan = path.Path.BytesAsSpan;
+        Span<byte> selfSpan = AsSpan;
+        pathSpan[..7].CopyTo(selfSpan);
+        selfSpan[7] = (byte)path.Length;
+    }
+
+    public int Length => AsSpan[7];
+
+    public TreePath ToTreePath()
+    {
+        ValueHash256 rawPath = Keccak.Zero;
+        Span<byte> pathSpan = rawPath.BytesAsSpan;
+        Span<byte> selfSpan = AsSpan;
+        selfSpan[..7].CopyTo(pathSpan);
+
+        return new TreePath(rawPath, selfSpan[7]);
+    }
+}

src/Nethermind/Nethermind.Trie/Pruning/TrackedPastKeyCountStrategy.cs

@@ -0,0 +1,24 @@
+// SPDX-FileCopyrightText: 2023 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+namespace Nethermind.Trie.Pruning;
+
+public class TrackedPastKeyCountStrategy : IPruningStrategy
+{
+    private IPruningStrategy _baseStrategy;
+    private readonly int _trackedPastKeyCount;
+    public bool PruningEnabled => _baseStrategy.PruningEnabled;
+
+    public TrackedPastKeyCountStrategy(IPruningStrategy baseStrategy, int trackedPastKeyCount)
+    {
+        _baseStrategy = baseStrategy;
+        _trackedPastKeyCount = trackedPastKeyCount;
+    }
+
+    public bool ShouldPrune(in long currentMemory)
+    {
+        return _baseStrategy.ShouldPrune(in currentMemory);
+    }
+
+    public int TrackedPastKeyCount => _trackedPastKeyCount;
+}