Fix issue in payload disposing and add unit test for asking multiple times by the same payloadId

src/Nethermind/Nethermind.Core/Extensions/CancellationTokenExtensions.cs

@@ -1,31 +1,59 @@
 //  Copyright (c) 2021 Demerzel Solutions Limited
 //  This file is part of the Nethermind library.
-// 
+//
 //  The Nethermind library is free software: you can redistribute it and/or modify
 //  it under the terms of the GNU Lesser General Public License as published by
 //  the Free Software Foundation, either version 3 of the License, or
 //  (at your option) any later version.
-// 
+//
 //  The Nethermind library is distributed in the hope that it will be useful,
 //  but WITHOUT ANY WARRANTY; without even the implied warranty of
 //  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 //  GNU Lesser General Public License for more details.
-// 
+//
 //  You should have received a copy of the GNU Lesser General Public License
 //  along with the Nethermind. If not, see <http://www.gnu.org/licenses/>.
-// 
+//
 
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Nethermind.Core.Extensions
 {
     public static class CancellationTokenExtensions
     {
-        public static Task AsTask(this System.Threading.CancellationToken token)
+        /// <summary>
+        /// Converts <see cref="CancellationToken"/> to a awaitable <see cref="Task"/> when token is cancelled.
+        /// </summary>
+        /// <param name="token"></param>
+        /// <returns></returns>
+        public static Task AsTask(this CancellationToken token)
         {
             TaskCompletionSource taskCompletionSource = new();
             token.Register(() => taskCompletionSource.TrySetCanceled(), false);
             return taskCompletionSource.Task;
         }
+
+        /// <summary>
+        /// <see cref="CancellationTokenSource.Cancel()"/>s and <see cref="CancellationTokenSource.Dispose"/>s the <see cref="CancellationTokenSource"/>
+        /// and also sets it to null so multiple calls to this method are safe and <see cref="CancellationTokenSource.Cancel()"/> will be called only once.
+        /// </summary>
+        /// <param name="cancellationTokenSource"></param>
+        /// <remarks>
+        /// This method is thread-sage and uses <see cref="Interlocked.CompareExchange{T}(ref T,T,T)"/> to safely manage reference.
+        /// </remarks>
+        public static void CancelDisposeAndClear(ref CancellationTokenSource? cancellationTokenSource)
+        {
+            CancellationTokenSource? source = cancellationTokenSource;
+            if (source is not null)
+            {
+                source = Interlocked.CompareExchange(ref cancellationTokenSource, null, source);
+                if (source is not null)
+                {
+                    source.Cancel();
+                    source.Dispose();
+                }
+            }
+        }
     }
 }

src/Nethermind/Nethermind.Merge.Plugin.Test/EngineModuleTests.DelayBlockImprovementContextFactory.cs

@@ -1,33 +1,34 @@
 //  Copyright (c) 2021 Demerzel Solutions Limited
 //  This file is part of the Nethermind library.
-// 
+//
 //  The Nethermind library is free software: you can redistribute it and/or modify
 //  it under the terms of the GNU Lesser General Public License as published by
 //  the Free Software Foundation, either version 3 of the License, or
 //  (at your option) any later version.
-// 
+//
 //  The Nethermind library is distributed in the hope that it will be useful,
 //  but WITHOUT ANY WARRANTY; without even the implied warranty of
 //  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 //  GNU Lesser General Public License for more details.
-// 
+//
 //  You should have received a copy of the GNU Lesser General Public License
 //  along with the Nethermind. If not, see <http://www.gnu.org/licenses/>.
-// 
+//
 
 using System;
 using System.Threading;
 using System.Threading.Tasks;
 using Nethermind.Consensus.Producers;
 using Nethermind.Core;
+using Nethermind.Core.Extensions;
 using Nethermind.Evm.Tracing;
 using Nethermind.Merge.Plugin.BlockProduction;
 
 namespace Nethermind.Merge.Plugin.Test;
 
 public partial class EngineModuleTests
 {
-    private class DelayBlockImprovementContextFactory : IBlockImprovementContextFactory 
+    private class DelayBlockImprovementContextFactory : IBlockImprovementContextFactory
     {
         private readonly IManualBlockProductionTrigger _productionTrigger;
         private readonly TimeSpan _timeout;
@@ -40,13 +41,13 @@ public DelayBlockImprovementContextFactory(IManualBlockProductionTrigger product
             _delay = delay;
         }
 
-        public IBlockImprovementContext StartBlockImprovementContext(Block currentBestBlock, BlockHeader parentHeader, PayloadAttributes payloadAttributes) => 
+        public IBlockImprovementContext StartBlockImprovementContext(Block currentBestBlock, BlockHeader parentHeader, PayloadAttributes payloadAttributes) =>
             new DelayBlockImprovementContext(currentBestBlock, _productionTrigger, _timeout, parentHeader, payloadAttributes, _delay);
     }
 
     private class DelayBlockImprovementContext : IBlockImprovementContext
     {
-        private readonly CancellationTokenSource _cancellationTokenSource;
+        private CancellationTokenSource? _cancellationTokenSource;
 
         public DelayBlockImprovementContext(
             Block currentBestBlock,
@@ -58,13 +59,18 @@ public DelayBlockImprovementContext(
         {
             _cancellationTokenSource = new CancellationTokenSource(timeout);
             CurrentBestBlock = currentBestBlock;
-            ImprovementTask = BuildBlock(blockProductionTrigger, parentHeader, payloadAttributes, delay);
+            ImprovementTask = BuildBlock(blockProductionTrigger, parentHeader, payloadAttributes, delay, _cancellationTokenSource.Token);
         }
 
-        private async Task<Block?> BuildBlock(IManualBlockProductionTrigger blockProductionTrigger, BlockHeader parentHeader, PayloadAttributes payloadAttributes, int delay)
+        private async Task<Block?> BuildBlock(
+            IManualBlockProductionTrigger blockProductionTrigger,
+            BlockHeader parentHeader,
+            PayloadAttributes payloadAttributes,
+            int delay,
+            CancellationToken cancellationToken)
         {
-            Block? block = await blockProductionTrigger.BuildBlock(parentHeader, _cancellationTokenSource.Token, NullBlockTracer.Instance, payloadAttributes);
-            await Task.Delay(delay);
+            Block? block = await blockProductionTrigger.BuildBlock(parentHeader, cancellationToken, NullBlockTracer.Instance, payloadAttributes);
+            await Task.Delay(delay, cancellationToken);
             if (block is not null)
             {
                 CurrentBestBlock = block;
@@ -79,8 +85,7 @@ public DelayBlockImprovementContext(
 
         public void Dispose()
         {
-            _cancellationTokenSource.Cancel();
-            _cancellationTokenSource.Dispose();
+            CancellationTokenExtensions.CancelDisposeAndClear(ref _cancellationTokenSource);
         }
     }
 }