Replace Environment.Exit with ProcessExitSource and cancelation token

NethermindEth · Apr 27, 2023 · df6dbb2 · df6dbb2
1 parent f6dd83d
commit df6dbb2
Show file tree

Hide file tree

Showing 13 changed files with 121 additions and 48 deletions.
diff --git a/src/Nethermind/Nethermind.Api/IBasicApi.cs b/src/Nethermind/Nethermind.Api/IBasicApi.cs
@@ -47,6 +47,7 @@ public interface IBasicApi
         IBetterPeerStrategy? BetterPeerStrategy { get; set; }
         ITimestamper Timestamper { get; }
         ITimerFactory TimerFactory { get; }
+        IProcessExitSource ProcessExit { get; set; }
 
         public IConsensusPlugin? GetConsensusPlugin() =>
             Plugins

diff --git a/src/Nethermind/Nethermind.Api/NethermindApi.cs b/src/Nethermind/Nethermind.Api/NethermindApi.cs
@@ -229,5 +229,6 @@ public ISealEngine SealEngine
         public IList<IPublisher> Publishers { get; } = new List<IPublisher>(); // this should be called publishers
         public CompositePruningTrigger PruningTrigger { get; } = new();
         public ISnapProvider SnapProvider { get; set; }
+        public IProcessExitSource ProcessExit { get; set; }
     }
 }
diff --git a/src/Nethermind/Nethermind.Blockchain.Test/FullPruning/FullPrunerTests.cs b/src/Nethermind/Nethermind.Blockchain.Test/FullPruning/FullPrunerTests.cs
@@ -7,6 +7,7 @@
 using System.Threading.Tasks;
 using FluentAssertions;
 using Nethermind.Blockchain.FullPruning;
+using Nethermind.Config;
 using Nethermind.Core;
 using Nethermind.Core.Crypto;
 using Nethermind.Core.Extensions;
@@ -31,8 +32,8 @@ namespace Nethermind.Blockchain.Test.FullPruning
     [Parallelizable(ParallelScope.All)]
     public class FullPrunerTests
     {
-        private int _fullPrunerMemoryBudgetMb;
-        private int _degreeOfParallelism;
+        private readonly int _fullPrunerMemoryBudgetMb;
+        private readonly int _degreeOfParallelism;
 
         public FullPrunerTests(int fullPrunerMemoryBudgetMb, int degreeOfParallelism)
         {
@@ -84,6 +85,28 @@ public async Task pruning_keeps_new_db_on_success()
             test.CopyDb.Count.Should().Be(count);
         }
 
+        [Timeout(Timeout.MaxTestTime)]
+        [TestCase(true, FullPruningCompletionBehavior.None, false)]
+        [TestCase(true, FullPruningCompletionBehavior.ShutdownOnSuccess, true)]
+        [TestCase(true, FullPruningCompletionBehavior.AlwaysShutdown, true)]
+        [TestCase(false, FullPruningCompletionBehavior.None, false)]
+        [TestCase(false, FullPruningCompletionBehavior.ShutdownOnSuccess, false)]
+        [TestCase(false, FullPruningCompletionBehavior.AlwaysShutdown, true)]
+        public async Task pruning_shuts_down_node(bool success, FullPruningCompletionBehavior behavior, bool expectedShutdown)
+        {
+            TestContext test = CreateTest(successfulPruning: success, completionBehavior: behavior);
+            await test.WaitForPruning();
+
+            if (expectedShutdown)
+            {
+                test.ProcessExitSource.Received(1).Exit(ExitCodes.Ok);
+            }
+            else
+            {
+                test.ProcessExitSource.DidNotReceiveWithAnyArgs().Exit(ExitCodes.Ok);
+            }
+        }
+
         [Test, Timeout(Timeout.MaxTestTime)]
         public async Task can_not_start_pruning_when_other_is_in_progress()
         {
@@ -136,7 +159,8 @@ public async Task should_duplicate_writes_to_batches_while_pruning()
             test.FullPruningDb[key].Should().BeEquivalentTo(key);
         }
 
-        private TestContext CreateTest(bool successfulPruning = true, bool clearPrunedDb = false) => new(successfulPruning, clearPrunedDb, _fullPrunerMemoryBudgetMb, _degreeOfParallelism);
+        private TestContext CreateTest(bool successfulPruning = true, bool clearPrunedDb = false, FullPruningCompletionBehavior completionBehavior = FullPruningCompletionBehavior.None) =>
+            new(successfulPruning, clearPrunedDb, completionBehavior, _fullPrunerMemoryBudgetMb, _degreeOfParallelism);
 
         private class TestContext
         {
@@ -151,7 +175,14 @@ private class TestContext
             public MemDb TrieDb { get; }
             public MemDb CopyDb { get; }
 
-            public TestContext(bool successfulPruning, bool clearPrunedDb = false, int fullScanMemoryBudgetMb = 0, int degreeOfParallelism = 0)
+            public IProcessExitSource ProcessExitSource { get; } = Substitute.For<IProcessExitSource>();
+
+            public TestContext(
+                bool successfulPruning,
+                bool clearPrunedDb = false,
+                FullPruningCompletionBehavior completionBehavior = FullPruningCompletionBehavior.None,
+                int fullScanMemoryBudgetMb = 0,
+                int degreeOfParallelism = 0)
             {
                 BlockTree.OnUpdateMainChain += (_, e) => _head = e.Blocks[^1].Number;
                 _clearPrunedDb = clearPrunedDb;
@@ -169,7 +200,8 @@ public TestContext(bool successfulPruning, bool clearPrunedDb = false, int fullS
                 {
                     FullPruningMaxDegreeOfParallelism = degreeOfParallelism,
                     FullPruningMemoryBudgetMb = fullScanMemoryBudgetMb,
-                }, BlockTree, StateReader, LimboLogs.Instance);
+                    FullPruningCompletionBehavior = completionBehavior
+                }, BlockTree, StateReader, ProcessExitSource, LimboLogs.Instance);
             }
 
             public async Task<bool> WaitForPruning()

diff --git a/src/Nethermind/Nethermind.Blockchain.Test/FullPruning/FullPruningDiskTest.cs b/src/Nethermind/Nethermind.Blockchain.Test/FullPruning/FullPruningDiskTest.cs
@@ -10,6 +10,7 @@
 using System.Threading.Tasks;
 using FluentAssertions;
 using Nethermind.Blockchain.FullPruning;
+using Nethermind.Config;
 using Nethermind.Core;
 using Nethermind.Core.Crypto;
 using Nethermind.Core.Extensions;
@@ -37,17 +38,18 @@ public class PruningTestBlockchain : TestBlockchain
             public IPruningTrigger PruningTrigger { get; } = Substitute.For<IPruningTrigger>();
             public FullTestPruner FullPruner { get; private set; }
             public IPruningConfig PruningConfig { get; set; } = new PruningConfig();
+            public IProcessExitSource ProcessExitSource { get; } = Substitute.For<IProcessExitSource>();
 
             public PruningTestBlockchain()
             {
                 TempDirectory = TempPath.GetTempDirectory();
             }
 
-            protected override async Task<TestBlockchain> Build(ISpecProvider specProvider = null, UInt256? initialValues = null)
+            protected override async Task<TestBlockchain> Build(ISpecProvider? specProvider = null, UInt256? initialValues = null)
             {
                 TestBlockchain chain = await base.Build(specProvider, initialValues);
                 PruningDb = (IFullPruningDb)DbProvider.StateDb;
-                FullPruner = new FullTestPruner(PruningDb, PruningTrigger, PruningConfig, BlockTree, StateReader, LogManager);
+                FullPruner = new FullTestPruner(PruningDb, PruningTrigger, PruningConfig, BlockTree, StateReader, ProcessExitSource, LogManager);
                 return chain;
             }
 
@@ -85,8 +87,9 @@ public FullTestPruner(
                     IPruningConfig pruningConfig,
                     IBlockTree blockTree,
                     IStateReader stateReader,
+                    IProcessExitSource processExitSource,
                     ILogManager logManager)
-                    : base(pruningDb, pruningTrigger, pruningConfig, blockTree, stateReader, logManager)
+                    : base(pruningDb, pruningTrigger, pruningConfig, blockTree, stateReader, processExitSource, logManager)
                 {
                 }
 

diff --git a/src/Nethermind/Nethermind.Blockchain/FullPruning/FullPruner.cs b/src/Nethermind/Nethermind.Blockchain/FullPruning/FullPruner.cs
@@ -4,6 +4,7 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Nethermind.Config;
 using Nethermind.Core;
 using Nethermind.Core.Crypto;
 using Nethermind.Core.Extensions;
@@ -25,6 +26,7 @@ public class FullPruner : IDisposable
         private readonly IPruningConfig _pruningConfig;
         private readonly IBlockTree _blockTree;
         private readonly IStateReader _stateReader;
+        private readonly IProcessExitSource _processExitSource;
         private readonly ILogManager _logManager;
         private IPruningContext? _currentPruning;
         private int _waitingForBlockProcessed = 0;
@@ -41,13 +43,15 @@ public FullPruner(
             IPruningConfig pruningConfig,
             IBlockTree blockTree,
             IStateReader stateReader,
+            IProcessExitSource processExitSource,
             ILogManager logManager)
         {
             _fullPruningDb = fullPruningDb;
             _pruningTrigger = pruningTrigger;
             _pruningConfig = pruningConfig;
             _blockTree = blockTree;
             _stateReader = stateReader;
+            _processExitSource = processExitSource;
             _logManager = logManager;
             _pruningTrigger.Prune += OnPrune;
             _logger = _logManager.GetClassLogger();
@@ -161,7 +165,7 @@ private void HandlePruningFinished(object? sender, PruningEventArgs e)
                 case FullPruningCompletionBehavior.AlwaysShutdown:
                 case FullPruningCompletionBehavior.ShutdownOnSuccess when e.Success:
                     if (_logger.IsInfo) _logger.Info($"Full Pruning completed {(e.Success ? "successfully" : "unsuccessfully")}, shutting down as requested in the configuration.");
-                    Task.Run(() => Environment.Exit(0));
+                    _processExitSource.Exit(ExitCodes.Ok);
                     break;
             }
         }

diff --git a/src/Nethermind/Nethermind.Config/EnvConfigSource.cs b/src/Nethermind/Nethermind.Config/EnvConfigSource.cs
@@ -65,8 +65,6 @@ public interface IEnvironment
     {
         string GetEnvironmentVariable(string variableName);
         System.Collections.IDictionary GetEnvironmentVariables();
-
-        void Exit(int exitCode);
     }
 
     public class EnvironmentWrapper : IEnvironment
@@ -80,10 +78,5 @@ public System.Collections.IDictionary GetEnvironmentVariables()
         {
             return Environment.GetEnvironmentVariables();
         }
-
-        public void Exit(int exitCode)
-        {
-            Environment.Exit(exitCode);
-        }
     }
 }
diff --git a/src/Nethermind/Nethermind.Config/IProcessExitSource.cs b/src/Nethermind/Nethermind.Config/IProcessExitSource.cs
@@ -0,0 +1,28 @@
+// SPDX-FileCopyrightText: 2023 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System.Threading;
+using Nethermind.Core.Extensions;
+
+namespace Nethermind.Config;
+
+public interface IProcessExitSource
+{
+    public void Exit(int exitCode);
+}
+
+public class ProcessExitSource : IProcessExitSource
+{
+    private CancellationTokenSource _cancellationTokenSource = new();
+    public int ExitCode { get; set; } = ExitCodes.Ok;
+
+    public CancellationToken Token => _cancellationTokenSource!.Token;
+
+    public void Exit(int exitCode)
+    {
+        if (CancellationTokenExtensions.CancelDisposeAndClear(ref _cancellationTokenSource))
+        {
+            ExitCode = exitCode;
+        }
+    }
+}
diff --git a/src/Nethermind/Nethermind.Core/Extensions/CancellationTokenExtensions.cs b/src/Nethermind/Nethermind.Core/Extensions/CancellationTokenExtensions.cs
@@ -28,7 +28,7 @@ public static Task AsTask(this CancellationToken token)
         /// <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)
+        public static bool CancelDisposeAndClear(ref CancellationTokenSource? cancellationTokenSource)
         {
             CancellationTokenSource? source = cancellationTokenSource;
             if (source is not null)
@@ -38,8 +38,11 @@ public static void CancelDisposeAndClear(ref CancellationTokenSource? cancellati
                 {
                     source.Cancel();
                     source.Dispose();
+                    return true;
                 }
             }
+
+            return false;
         }
     }
 }
diff --git a/src/Nethermind/Nethermind.HealthChecks.Test/FreeDiskSpaceCheckerTests.cs b/src/Nethermind/Nethermind.HealthChecks.Test/FreeDiskSpaceCheckerTests.cs
@@ -8,6 +8,7 @@
 using NUnit.Framework;
 using Nethermind.Core.Extensions;
 using System.Threading.Tasks;
+using Nethermind.Config;
 using NSubstitute.ReceivedExtensions;
 
 namespace Nethermind.HealthChecks.Test
@@ -27,7 +28,12 @@ public void free_disk_check_ensure_free_on_startup_no_wait(float availableDiskSp
                 LowStorageSpaceShutdownThreshold = 1,
                 LowStorageSpaceWarningThreshold = 5
             };
-            FreeDiskSpaceChecker freeDiskSpaceChecker = new(hcConfig, LimboTraceLogger.Instance, GetDriveInfos(availableDiskSpacePercent), Core.Timers.TimerFactory.Default);
+            FreeDiskSpaceChecker freeDiskSpaceChecker = new(
+                hcConfig,
+                GetDriveInfos(availableDiskSpacePercent),
+                Core.Timers.TimerFactory.Default,
+                Substitute.For<IProcessExitSource>(),
+                LimboTraceLogger.Instance);
 
             if (exceptionExpected)
                 Assert.Throws<NotEnoughDiskSpaceException>(() => freeDiskSpaceChecker.EnsureEnoughFreeSpaceOnStart(Core.Timers.TimerFactory.Default));
@@ -47,20 +53,23 @@ public void free_disk_check_ensure_free_on_startup_wait_until_enough(float avail
                 LowStorageSpaceShutdownThreshold = 1,
                 LowStorageSpaceWarningThreshold = 5
             };
-            var drives = GetDriveInfos(availableDiskSpacePercent);
+            IDriveInfo[] drives = GetDriveInfos(availableDiskSpacePercent);
             drives[0].AvailableFreeSpace.Returns(a => _freeSpaceBytes,
                                                     a => 3 * _freeSpaceBytes);
-            FreeDiskSpaceChecker freeDiskSpaceChecker = new(hcConfig, LimboTraceLogger.Instance, drives, Core.Timers.TimerFactory.Default, ts.TotalMinutes);
+            FreeDiskSpaceChecker freeDiskSpaceChecker = new(
+                hcConfig,
+                drives,
+                Core.Timers.TimerFactory.Default,
+                Substitute.For<IProcessExitSource>(),
+                LimboTraceLogger.Instance,
+                ts.TotalMinutes);
 
             Task t = Task.Run(() => freeDiskSpaceChecker.EnsureEnoughFreeSpaceOnStart(Core.Timers.TimerFactory.Default));
             bool completed = t.Wait((int)ts.TotalMilliseconds * 2);
 
             Assert.IsTrue(completed);
 
-            if (awaitsForFreeSpace)
-                _ = drives[0].Received(3).AvailableFreeSpace;
-            else
-                _ = drives[0].Received(1).AvailableFreeSpace;
+            _ = drives[0].Received(awaitsForFreeSpace ? 3 : 1).AvailableFreeSpace;
         }
 
         private static IDriveInfo[] GetDriveInfos(float availableDiskSpacePercent)

diff --git a/src/Nethermind/Nethermind.HealthChecks/FreeDiskSpaceChecker.cs b/src/Nethermind/Nethermind.HealthChecks/FreeDiskSpaceChecker.cs
@@ -18,19 +18,21 @@ public class FreeDiskSpaceChecker : IHostedService, IAsyncDisposable
         private readonly IHealthChecksConfig _healthChecksConfig;
         private readonly ILogger _logger;
         private readonly IDriveInfo[] _drives;
+        private readonly IProcessExitSource _processExitSource;
         private readonly ITimer _timer;
         private readonly double _checkPeriodMinutes;
-        public static readonly int DefaultCheckPeriodMinutes = 1;
 
-        public FreeDiskSpaceChecker(IHealthChecksConfig healthChecksConfig, ILogger logger, IDriveInfo[] drives, ITimerFactory timerFactory) :
-            this(healthChecksConfig, logger, drives, timerFactory, DefaultCheckPeriodMinutes)
-        { }
-
-        public FreeDiskSpaceChecker(IHealthChecksConfig healthChecksConfig, ILogger logger, IDriveInfo[] drives, ITimerFactory timerFactory, double checkPeriodMinutes)
+        public FreeDiskSpaceChecker(IHealthChecksConfig healthChecksConfig,
+            IDriveInfo[] drives,
+            ITimerFactory timerFactory,
+            IProcessExitSource processExitSource,
+            ILogger logger,
+            double checkPeriodMinutes = 1)
         {
             _healthChecksConfig = healthChecksConfig;
             _logger = logger;
             _drives = drives;
+            _processExitSource = processExitSource;
             _checkPeriodMinutes = checkPeriodMinutes;
             _timer = timerFactory.CreateTimer(TimeSpan.FromMinutes(_checkPeriodMinutes));
             _timer.Elapsed += CheckDiskSpace;
@@ -45,7 +47,7 @@ private void CheckDiskSpace(object sender, EventArgs e)
                 if (freeSpacePercent < _healthChecksConfig.LowStorageSpaceShutdownThreshold)
                 {
                     if (_logger.IsError) _logger.Error($"Free disk space in '{drive.RootDirectory.FullName}' is below {_healthChecksConfig.LowStorageSpaceShutdownThreshold:0.00}% - shutting down...");
-                    Environment.Exit(ExitCodes.LowDiskSpace);
+                    _processExitSource.Exit(ExitCodes.LowDiskSpace);
                 }
 
                 if (freeSpacePercent < _healthChecksConfig.LowStorageSpaceWarningThreshold)

diff --git a/src/Nethermind/Nethermind.HealthChecks/HealthChecksPlugin.cs b/src/Nethermind/Nethermind.HealthChecks/HealthChecksPlugin.cs
@@ -53,7 +53,12 @@ public async ValueTask DisposeAsync()
         public bool MustInitialize => true;
 
         public FreeDiskSpaceChecker FreeDiskSpaceChecker => LazyInitializer.EnsureInitialized(ref _freeDiskSpaceChecker,
-            () => new FreeDiskSpaceChecker(_healthChecksConfig, _logger, _api.FileSystem.GetDriveInfos(_initConfig.BaseDbPath), _api.TimerFactory));
+            () => new FreeDiskSpaceChecker(
+                _healthChecksConfig,
+                _api.FileSystem.GetDriveInfos(_initConfig.BaseDbPath),
+                _api.TimerFactory,
+                _api.ProcessExit,
+                _logger));
 
         public Task Init(INethermindApi api)
         {

diff --git a/src/Nethermind/Nethermind.Init/Steps/InitializeBlockchain.cs b/src/Nethermind/Nethermind.Init/Steps/InitializeBlockchain.cs
@@ -314,7 +314,7 @@ private static void InitializeFullPruning(
                         api.PruningTrigger.Add(pruningTrigger);
                     }
 
-                    FullPruner pruner = new(fullPruningDb, api.PruningTrigger, pruningConfig, api.BlockTree!, stateReader, api.LogManager);
+                    FullPruner pruner = new(fullPruningDb, api.PruningTrigger, pruningConfig, api.BlockTree!, stateReader, api.ProcessExit, api.LogManager);
                     api.DisposeStack.Push(pruner);
                 }
             }