fix AuRa pre-warmer issues

Author: LukaszRozmej

src/Nethermind/Nethermind.Consensus.AuRa/AuRaBlockProcessor.cs

@@ -45,7 +45,8 @@ public AuRaBlockProcessor(
             IAuRaValidator? auRaValidator,
             ITxFilter? txFilter = null,
             AuRaContractGasLimitOverride? gasLimitOverride = null,
-            ContractRewriter? contractRewriter = null)
+            ContractRewriter? contractRewriter = null,
+            IBlockCachePreWarmer? preWarmer = null)
             : base(
                 specProvider,
                 blockValidator,
@@ -55,7 +56,8 @@ public AuRaBlockProcessor(
                 receiptStorage,
                 new BlockhashStore(blockTree, specProvider, stateProvider),
                 logManager,
-                withdrawalProcessor)
+                withdrawalProcessor,
+                preWarmer: preWarmer)
         {
             _specProvider = specProvider;
             _blockTree = blockTree ?? throw new ArgumentNullException(nameof(blockTree));

src/Nethermind/Nethermind.Consensus.AuRa/InitializationSteps/InitializeBlockchainAuRa.cs

@@ -71,7 +71,7 @@ protected override async Task InitBlockchain()
         }
     }
 
-    protected override BlockProcessor CreateBlockProcessor()
+    protected override BlockProcessor CreateBlockProcessor(BlockCachePreWarmer? preWarmer)
     {
         if (_api.SpecProvider is null) throw new StepDependencyException(nameof(_api.SpecProvider));
         if (_api.ChainHeadStateProvider is null) throw new StepDependencyException(nameof(_api.ChainHeadStateProvider));
@@ -89,10 +89,10 @@ protected override BlockProcessor CreateBlockProcessor()
             _api,
             new ServiceTxFilter(_api.SpecProvider));
 
-        return NewAuraBlockProcessor(auRaTxFilter);
+        return NewAuraBlockProcessor(auRaTxFilter, preWarmer);
     }
 
-    protected virtual AuRaBlockProcessor NewAuraBlockProcessor(ITxFilter txFilter)
+    protected virtual AuRaBlockProcessor NewAuraBlockProcessor(ITxFilter txFilter, BlockCachePreWarmer? preWarmer)
     {
         IDictionary<long, IDictionary<Address, byte[]>> rewriteBytecode = _api.ChainSpec.AuRa.RewriteBytecode;
         ContractRewriter? contractRewriter = rewriteBytecode?.Count > 0 ? new ContractRewriter(rewriteBytecode) : null;
@@ -112,7 +112,8 @@ protected virtual AuRaBlockProcessor NewAuraBlockProcessor(ITxFilter txFilter)
             CreateAuRaValidator(),
             txFilter,
             GetGasLimitCalculator(),
-            contractRewriter
+            contractRewriter,
+            preWarmer: preWarmer
         );
     }
 

src/Nethermind/Nethermind.Consensus/Processing/BlockProcessor.cs

@@ -116,9 +116,9 @@ the previous head state.*/
                 }
 
                 using CancellationTokenSource cancellationTokenSource = new();
-                Task? preWarmTask = suggestedBlock.Transactions.Length < 3 ?
-                    null :
-                    _preWarmer?.PreWarmCaches(suggestedBlock, preBlockStateRoot!, cancellationTokenSource.Token);
+                Task? preWarmTask = suggestedBlock.Transactions.Length < 3
+                    ? null
+                    : _preWarmer?.PreWarmCaches(suggestedBlock, preBlockStateRoot!, cancellationTokenSource.Token);
                 (Block processedBlock, TxReceipt[] receipts) = ProcessOne(suggestedBlock, options, blockTracer);
                 cancellationTokenSource.Cancel();
                 preWarmTask?.GetAwaiter().GetResult();
@@ -159,9 +159,12 @@ the previous head state.*/
         {
             _logger.Trace($"Encountered exception {ex} while processing blocks.");
             RestoreBranch(previousBranchStateRoot);
-            _preWarmer?.ClearCaches();
             throw;
         }
+        finally
+        {
+            _preWarmer?.ClearCaches();
+        }
     }
 
     public event EventHandler<BlocksProcessingEventArgs>? BlocksProcessing;

src/Nethermind/Nethermind.Init/Steps/InitializeBlockchain.cs

@@ -104,7 +104,11 @@ protected virtual Task InitBlockchain()
 
             setApi.TxPoolInfoProvider = new TxPoolInfoProvider(chainHeadInfoProvider.AccountStateProvider, txPool);
             setApi.GasPriceOracle = new GasPriceOracle(getApi.BlockTree!, getApi.SpecProvider, _api.LogManager, blocksConfig.MinGasPrice);
-            IBlockProcessor mainBlockProcessor = setApi.MainBlockProcessor = CreateBlockProcessor();
+            PreBlockCaches? preBlockCaches = (_api.WorldState as IPreBlockCaches)?.Caches;
+            BlockCachePreWarmer? preWarmer = blocksConfig.PreWarmStateOnBlockProcessing
+                ? new(new(_api.WorldStateManager!, _api.BlockTree!, _api.SpecProvider, _api.LogManager, preBlockCaches), _api.SpecProvider, _api.LogManager, preBlockCaches)
+                : null;
+            IBlockProcessor mainBlockProcessor = setApi.MainBlockProcessor = CreateBlockProcessor(preWarmer);
 
             BlockchainProcessor blockchainProcessor = new(
                 getApi.BlockTree,
@@ -212,7 +216,7 @@ protected virtual TxPool.TxPool CreateTxPool() =>
             _api.LogManager);
 
         // TODO: remove from here - move to consensus?
-        protected virtual BlockProcessor CreateBlockProcessor()
+        protected virtual BlockProcessor CreateBlockProcessor(BlockCachePreWarmer? preWarmer)
         {
             if (_api.DbProvider is null) throw new StepDependencyException(nameof(_api.DbProvider));
             if (_api.RewardCalculatorSource is null) throw new StepDependencyException(nameof(_api.RewardCalculatorSource));
@@ -221,13 +225,7 @@ protected virtual BlockProcessor CreateBlockProcessor()
             if (_api.WorldStateManager is null) throw new StepDependencyException(nameof(_api.WorldStateManager));
             if (_api.SpecProvider is null) throw new StepDependencyException(nameof(_api.SpecProvider));
 
-            IBlocksConfig blocksConfig = _api.Config<IBlocksConfig>();
             IWorldState worldState = _api.WorldState!;
-
-            PreBlockCaches? preBlockCaches = (worldState as IPreBlockCaches)?.Caches;
-            BlockCachePreWarmer? preWarmer = blocksConfig.PreWarmStateOnBlockProcessing
-                ? new(new(_api.WorldStateManager, _api.BlockTree, _api.SpecProvider, _api.LogManager, preBlockCaches), _api.SpecProvider, _api.LogManager, preBlockCaches)
-                : null;
             return new BlockProcessor(
                 _api.SpecProvider,
                 _api.BlockValidator,

src/Nethermind/Nethermind.Merge.AuRa/AuRaMergeBlockProcessor.cs

@@ -33,7 +33,8 @@ public AuRaMergeBlockProcessor(
         IAuRaValidator? validator,
         ITxFilter? txFilter = null,
         AuRaContractGasLimitOverride? gasLimitOverride = null,
-        ContractRewriter? contractRewriter = null
+        ContractRewriter? contractRewriter = null,
+        IBlockCachePreWarmer? preWarmer = null
     ) : base(
             specProvider,
             blockValidator,
@@ -47,7 +48,8 @@ public AuRaMergeBlockProcessor(
             validator,
             txFilter,
             gasLimitOverride,
-            contractRewriter
+            contractRewriter,
+            preWarmer
         )
     { }
 

src/Nethermind/Nethermind.Merge.AuRa/InitializationSteps/InitializeBlockchainAuRaMerge.cs

@@ -24,7 +24,7 @@ public InitializeBlockchainAuRaMerge(AuRaNethermindApi api) : base(api)
             _api = api;
         }
 
-        protected override AuRaBlockProcessor NewAuraBlockProcessor(ITxFilter txFilter)
+        protected override AuRaBlockProcessor NewAuraBlockProcessor(ITxFilter txFilter, BlockCachePreWarmer? preWarmer)
         {
             IDictionary<long, IDictionary<Address, byte[]>> rewriteBytecode = _api.ChainSpec.AuRa.RewriteBytecode;
             ContractRewriter? contractRewriter = rewriteBytecode?.Count > 0 ? new ContractRewriter(rewriteBytecode) : null;
@@ -47,7 +47,8 @@ protected override AuRaBlockProcessor NewAuraBlockProcessor(ITxFilter txFilter)
                 CreateAuRaValidator(),
                 txFilter,
                 GetGasLimitCalculator(),
-                contractRewriter
+                contractRewriter,
+                preWarmer: preWarmer
             );
         }
 

src/Nethermind/Nethermind.Optimism/InitializeBlockchainOptimism.cs

@@ -6,9 +6,11 @@
 using Nethermind.Blockchain.Blocks;
 using Nethermind.Blockchain.Services;
 using Nethermind.Config;
+using Nethermind.Consensus.AuRa.Withdrawals;
 using Nethermind.Consensus.Processing;
 using Nethermind.Consensus.Producers;
 using Nethermind.Consensus.Validators;
+using Nethermind.Consensus.Withdrawals;
 using Nethermind.Evm;
 using Nethermind.Evm.TransactionProcessing;
 using Nethermind.Init.Steps;
@@ -82,7 +84,7 @@ protected override IBlockValidator CreateBlockValidator()
         return new InvalidBlockInterceptor(blockValidator, _api.InvalidChainTracker, _api.LogManager);
     }
 
-    protected override BlockProcessor CreateBlockProcessor()
+    protected override BlockProcessor CreateBlockProcessor(BlockCachePreWarmer? preWarmer)
     {
         if (_api.DbProvider is null) throw new StepDependencyException(nameof(_api.DbProvider));
         if (_api.RewardCalculatorSource is null) throw new StepDependencyException(nameof(_api.RewardCalculatorSource));
@@ -105,7 +107,9 @@ protected override BlockProcessor CreateBlockProcessor()
             new BlockhashStore(_api.BlockTree, _api.SpecProvider, _api.WorldState),
             _api.LogManager,
             _api.SpecHelper,
-            contractRewriter);
+            contractRewriter,
+            new BlockProductionWithdrawalProcessor(new NullWithdrawalProcessor()),
+            preWarmer: preWarmer);
     }
 
     protected override IUnclesValidator CreateUnclesValidator() => Always.Valid;

src/Nethermind/Nethermind.Optimism/OptimismBlockProcessor.cs

@@ -34,9 +34,11 @@ public OptimismBlockProcessor(
         ILogManager? logManager,
         IOPConfigHelper opConfigHelper,
         Create2DeployerContractRewriter contractRewriter,
-        IWithdrawalProcessor? withdrawalProcessor = null)
+        IWithdrawalProcessor? withdrawalProcessor = null,
+        IBlockCachePreWarmer? preWarmer = null)
         : base(specProvider, blockValidator, rewardCalculator, blockTransactionsExecutor,
-            stateProvider, receiptStorage, blockhashStore, logManager, withdrawalProcessor, ReceiptsRootCalculator.Instance)
+            stateProvider, receiptStorage, blockhashStore, logManager, withdrawalProcessor,
+            ReceiptsRootCalculator.Instance, preWarmer)
     {
         ArgumentNullException.ThrowIfNull(stateProvider);
         _contractRewriter = contractRewriter;