Share stack pool between threads (#7972)

src/Nethermind/Nethermind.Evm/EvmState.cs

@@ -2,9 +2,7 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 
 using System;
-using System.Collections.Generic;
 using System.Diagnostics;
-using System.Threading;
 using Nethermind.Core;
 using Nethermind.State;
 
@@ -16,57 +14,7 @@ namespace Nethermind.Evm
     [DebuggerDisplay("{ExecutionType} to {Env.ExecutingAccount}, G {GasAvailable} R {Refund} PC {ProgramCounter} OUT {OutputDestination}:{OutputLength}")]
     public class EvmState : IDisposable // TODO: rename to CallState
     {
-        private class StackPool
-        {
-            private readonly int _maxCallStackDepth;
-            private readonly struct StackItem(byte[] dataStack, int[] returnStack)
-            {
-                public readonly byte[] DataStack = dataStack;
-                public readonly int[] ReturnStack = returnStack;
-            }
-
-            // TODO: we have wrong call depth calculation somewhere
-            public StackPool(int maxCallStackDepth = VirtualMachine.MaxCallDepth * 2)
-            {
-                _maxCallStackDepth = maxCallStackDepth;
-            }
-
-            private readonly Stack<StackItem> _stackPool = new(32);
-
-            private int _stackPoolDepth;
-
-            /// <summary>
-            /// The word 'return' acts here once as a verb 'to return stack to the pool' and once as a part of the
-            /// compound noun 'return stack' which is a stack of subroutine return values.
-            /// </summary>
-            /// <param name="dataStack"></param>
-            /// <param name="returnStack"></param>
-            public void ReturnStacks(byte[] dataStack, int[] returnStack)
-            {
-                _stackPool.Push(new(dataStack, returnStack));
-            }
-
-            public (byte[], int[]) RentStacks()
-            {
-                if (_stackPool.TryPop(out StackItem result))
-                {
-                    return (result.DataStack, result.ReturnStack);
-                }
-
-                _stackPoolDepth++;
-                if (_stackPoolDepth > _maxCallStackDepth)
-                {
-                    EvmStack.ThrowEvmStackOverflowException();
-                }
-
-                return
-                (
-                    new byte[(EvmStack.MaxStackSize + EvmStack.RegisterLength) * 32],
-                    new int[EvmStack.ReturnStackSize]
-                );
-            }
-        }
-        private static readonly ThreadLocal<StackPool> _stackPool = new(static () => new StackPool());
+        private static readonly StackPool _stackPool = new();
 
         public byte[]? DataStack;
 
@@ -81,7 +29,7 @@ public void ReturnStacks(byte[] dataStack, int[] returnStack)
         public int ReturnStackHead = 0;
         private bool _canRestore = true;
         /// <summary>
-        /// Contructor for a top level <see cref="EvmState"/>.
+        /// Constructor for a top level <see cref="EvmState"/>.
         /// </summary>
         public EvmState(
             long gasAvailable,
@@ -101,7 +49,7 @@ public EvmState(
         {
         }
         /// <summary>
-        /// Contructor for a top level <see cref="EvmState"/>.
+        /// Constructor for a top level <see cref="EvmState"/>.
         /// </summary>
         public EvmState(
             long gasAvailable,
@@ -120,7 +68,7 @@ public EvmState(
         {
         }
         /// <summary>
-        /// Contructor for a frame <see cref="EvmState"/> beneath top level.
+        /// Constructor for a frame <see cref="EvmState"/> beneath top level.
         /// </summary>
         internal EvmState(
             long gasAvailable,
@@ -215,7 +163,7 @@ public void Dispose()
             if (DataStack is not null)
             {
                 // Only Dispose once
-                _stackPool.Value.ReturnStacks(DataStack, ReturnStack!);
+                _stackPool.ReturnStacks(DataStack, ReturnStack!);
                 DataStack = null;
                 ReturnStack = null;
             }
@@ -228,14 +176,14 @@ public void InitStacks()
         {
             if (DataStack is null)
             {
-                (DataStack, ReturnStack) = _stackPool.Value.RentStacks();
+                (DataStack, ReturnStack) = _stackPool.RentStacks();
             }
         }
 
         public void CommitToParent(EvmState parentState)
         {
             parentState.Refund += Refund;
-            _canRestore = false; // we can't restore if we commited
+            _canRestore = false; // we can't restore if we committed
         }
 
         private void Restore()

src/Nethermind/Nethermind.Evm/StackPool.cs

@@ -0,0 +1,48 @@
+// SPDX-FileCopyrightText: 2022 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System.Collections.Concurrent;
+
+namespace Nethermind.Evm;
+
+internal class StackPool
+{
+    // Also have parallel prewarming and Rpc calls
+    private const int MaxStacksPooled = VirtualMachine.MaxCallDepth * 2;
+    private readonly struct StackItem(byte[] dataStack, int[] returnStack)
+    {
+        public readonly byte[] DataStack = dataStack;
+        public readonly int[] ReturnStack = returnStack;
+    }
+
+    private readonly ConcurrentQueue<StackItem> _stackPool = new();
+
+    /// <summary>
+    /// The word 'return' acts here once as a verb 'to return stack to the pool' and once as a part of the
+    /// compound noun 'return stack' which is a stack of subroutine return values.
+    /// </summary>
+    /// <param name="dataStack"></param>
+    /// <param name="returnStack"></param>
+    public void ReturnStacks(byte[] dataStack, int[] returnStack)
+    {
+        if (_stackPool.Count <= MaxStacksPooled)
+        {
+            _stackPool.Enqueue(new(dataStack, returnStack));
+        }
+    }
+
+    public (byte[], int[]) RentStacks()
+    {
+        if (_stackPool.TryDequeue(out StackItem result))
+        {
+            return (result.DataStack, result.ReturnStack);
+        }
+
+        return
+        (
+            new byte[(EvmStack.MaxStackSize + EvmStack.RegisterLength) * 32],
+            new int[EvmStack.ReturnStackSize]
+        );
+    }
+}
+