Fix snapshot disk space check to run before download

src/Nethermind/Nethermind.Init.Snapshot.Test/InitDatabaseSnapshotTests.cs

@@ -0,0 +1,204 @@
+// SPDX-FileCopyrightText: 2026 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Formats.Tar;
+using System.IO;
+using System.IO.Abstractions;
+using System.Net;
+using System.Threading;
+using System.Threading.Tasks;
+using Nethermind.Api;
+using Nethermind.Core.Test.IO;
+using Nethermind.Logging;
+using NSubstitute;
+using NUnit.Framework;
+using Testably.Abstractions;
+
+namespace Nethermind.Init.Snapshot.Test;
+
+public class InitDatabaseSnapshotTests
+{
+    private const int SnapshotPayloadSize = 100_000;
+
+    private TempPath _tempDir = null!;
+    private string _dbPath = null!;
+    private string _snapshotPath = null!;
+    private SnapshotConfig _snapshotConfig = null!;
+    private INethermindApi _api = null!;
+
+    [SetUp]
+    public void SetUp()
+    {
+        _tempDir = TempPath.GetTempDirectory();
+        string snapshotDirectory = Path.Combine(_tempDir.Path, "snapshot");
+        Directory.CreateDirectory(snapshotDirectory);
+        _dbPath = Path.Combine(_tempDir.Path, "db");
+        _snapshotConfig = new SnapshotConfig
+        {
+            Enabled = true,
+            DownloadUrl = "http://127.0.0.1:1/snapshot.tar",
+            SnapshotDirectory = snapshotDirectory,
+            SnapshotFileName = "snapshot.tar",
+            StripComponents = 1
+        };
+        _snapshotPath = Path.Combine(snapshotDirectory, _snapshotConfig.SnapshotFileName);
+
+        _api = Substitute.For<INethermindApi>();
+        _api.Config<ISnapshotConfig>().Returns(_snapshotConfig);
+        _api.Config<IInitConfig>().Returns(new InitConfig { BaseDbPath = _dbPath });
+        _api.LogManager.Returns(LimboLogs.Instance);
+        _api.FileSystem.Returns(new RealFileSystem());
+    }
+
+    [TearDown]
+    public void TearDown() => _tempDir.Dispose();
+
+    [Test]
+    public async Task Execute_FreeSpaceCoversExtractionButNotLegacyMultiplier_ExtractsSnapshot()
+    {
+        long snapshotSize = WriteSnapshotTar();
+        AdvanceCheckpoint(SnapshotStage.Verified);
+        long freeSpace = snapshotSize * 2;
+        Assert.That(freeSpace, Is.GreaterThanOrEqualTo(InitDatabaseSnapshot.GetRequiredSpaceForExtraction(snapshotSize)),
+            "precondition: free space must cover the extraction estimate");
+        Assert.That(freeSpace, Is.LessThan((long)(snapshotSize * 2.5)),
+            "precondition: free space must be below the legacy 2.5x requirement to prove the regression is fixed");
+        InitDatabaseSnapshot step = new(_api, DrivesWithFreeSpace(freeSpace));
+
+        await step.Execute(CancellationToken.None);
+
+        Assert.That(File.Exists(Path.Combine(_dbPath, "state.bin")), Is.True,
+            "the snapshot content should be extracted into the database directory");
+        Assert.That(File.Exists(_snapshotPath), Is.False,
+            "the snapshot archive should be deleted after a successful extraction");
+    }
+
+    [Test]
+    public void Execute_FreeSpaceBelowExtractionEstimate_ThrowsIOException()
+    {
+        long snapshotSize = WriteSnapshotTar();
+        AdvanceCheckpoint(SnapshotStage.Verified);
+        InitDatabaseSnapshot step = new(_api, DrivesWithFreeSpace(snapshotSize));
+
+        IOException exception = Assert.ThrowsAsync<IOException>(() => step.Execute(CancellationToken.None))!;
+
+        Assert.That(exception.Message, Does.Contain("Insufficient disk space"),
+            "the extraction should be rejected when free space is below the extraction estimate");
+        Assert.That(Directory.Exists(_dbPath), Is.False,
+            "nothing should be extracted when free space is insufficient");
+    }
+
+    [Test]
+    public void Execute_FreeSpaceBelowDownloadRequirement_ThrowsIOExceptionBeforeDownloading()
+    {
+        using SnapshotServer server = SnapshotServer.Start(contentLength: 1_000_000);
+        _snapshotConfig.DownloadUrl = server.Url;
+        InitDatabaseSnapshot step = new(_api, DrivesWithFreeSpace(1_000_000));
+
+        IOException exception = Assert.ThrowsAsync<IOException>(() => step.Execute(CancellationToken.None))!;
+
+        Assert.That(exception.Message, Does.Contain("Insufficient disk space"),
+            "the download should be rejected when free space cannot fit the snapshot and its extraction");
+        Assert.That(File.Exists(_snapshotPath), Is.False,
+            "no bytes should be downloaded when free space is insufficient");
+    }
+
+    [TestCase(1_000, 0, 2_500, TestName = "FreshDownload")]
+    [TestCase(1_000, 400, 2_100, TestName = "ResumedDownload")]
+    public void GetRequiredSpaceForDownload_ForGivenSizes_AddsRemainingBytesToExtractionEstimate(
+        long totalSize, long existingSize, long expected) =>
+        Assert.That(InitDatabaseSnapshot.GetRequiredSpaceForDownload(totalSize, existingSize), Is.EqualTo(expected),
+            "the pre-download requirement should be the remaining bytes plus the extraction estimate of the full snapshot");
+
+    private long WriteSnapshotTar()
+    {
+        using (FileStream fileStream = File.Create(_snapshotPath))
+        using (TarWriter tarWriter = new(fileStream))
+        {
+            tarWriter.WriteEntry(new PaxTarEntry(TarEntryType.Directory, "data"));
+            PaxTarEntry fileEntry = new(TarEntryType.RegularFile, "data/state.bin")
+            {
+                DataStream = new MemoryStream(new byte[SnapshotPayloadSize])
+            };
+            tarWriter.WriteEntry(fileEntry);
+        }
+
+        return new FileInfo(_snapshotPath).Length;
+    }
+
+    private void AdvanceCheckpoint(SnapshotStage stage) =>
+        new SnapshotCheckpoint(_snapshotConfig, LimboLogs.Instance).Advance(stage);
+
+    private static IDriveInfo[] DrivesWithFreeSpace(long freeSpace)
+    {
+        IDriveInfo drive = Substitute.For<IDriveInfo>();
+        drive.AvailableFreeSpace.Returns(freeSpace);
+        drive.RootDirectory.FullName.Returns("/");
+        return [drive];
+    }
+
+    private sealed class SnapshotServer : IDisposable
+    {
+        private readonly HttpListener _listener;
+
+        public string Url { get; }
+
+        private SnapshotServer(HttpListener listener, string url)
+        {
+            _listener = listener;
+            Url = url;
+        }
+
+        public static SnapshotServer Start(long contentLength)
+        {
+            (HttpListener listener, int port) = StartListener();
+
+            _ = Task.Run(async () =>
+            {
+                while (listener.IsListening)
+                {
+                    try
+                    {
+                        HttpListenerContext context = await listener.GetContextAsync();
+                        context.Response.ContentLength64 = contentLength;
+                        byte[] chunk = new byte[1024];
+                        await context.Response.OutputStream.WriteAsync(chunk);
+                        context.Response.OutputStream.Flush();
+                    }
+                    catch (Exception)
+                    {
+                        return;
+                    }
+                }
+            });
+
+            return new SnapshotServer(listener, $"http://127.0.0.1:{port}/snapshot.tar");
+        }
+
+        private static (HttpListener Listener, int Port) StartListener()
+        {
+            for (int attempt = 0; ; attempt++)
+            {
+                HttpListener listener = new();
+                int port = Random.Shared.Next(20000, 60000);
+                listener.Prefixes.Add($"http://127.0.0.1:{port}/");
+                try
+                {
+                    listener.Start();
+                    return (listener, port);
+                }
+                catch (HttpListenerException) when (attempt < 5)
+                {
+                    listener.Close();
+                }
+            }
+        }
+
+        public void Dispose()
+        {
+            _listener.Stop();
+            _listener.Close();
+        }
+    }
+}

src/Nethermind/Nethermind.Init.Snapshot.Test/Nethermind.Init.Snapshot.Test.csproj

@@ -0,0 +1,14 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <Import Project="../tests.props" />
+
+  <ItemGroup>
+    <PackageReference Include="NSubstitute" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\Nethermind.Core.Test\Nethermind.Core.Test.csproj" />
+    <ProjectReference Include="..\Nethermind.Init.Snapshot\Nethermind.Init.Snapshot.csproj" />
+  </ItemGroup>
+
+</Project>

src/Nethermind/Nethermind.Init.Snapshot/InitDatabaseSnapshot.cs

@@ -26,6 +26,7 @@ public class InitDatabaseSnapshot(
     INethermindApi api,
     [KeyFilter(nameof(IInitConfig.BaseDbPath))] IDriveInfo[] drives) : IStep
 {
+    private const double ExtractionSpaceMultiplier = 1.5;
     private const int ExtractionRestartDelaySeconds = 5;
     private const int InitialRetryDelaySeconds = 5;
     private const int MaxRetryDelaySeconds = 300;
@@ -109,6 +110,8 @@ private async Task DownloadWithRetryAsync(
         if (checkpoint.Read() >= SnapshotStage.Downloaded)
             return;
 
+        await CheckDiskSpaceBeforeDownloadAsync(downloader, url, destinationPath, cancellationToken).ConfigureAwait(false);
+
         TimeSpan retryDelay = TimeSpan.FromSeconds(InitialRetryDelaySeconds);
         long lastSize = GetFileSize(destinationPath);
 
@@ -191,26 +194,52 @@ private async Task ExtractAsync(
         if (checkpoint.Read() >= SnapshotStage.Extracted)
             return;
 
-        CheckDiskSpace(snapshotPath);
+        CheckDiskSpace(GetRequiredSpaceForExtraction(GetFileSize(snapshotPath)), "extract");
 
         SnapshotExtractor extractor = new(api.LogManager);
         await extractor.ExtractAsync(snapshotPath, dbPath, stripComponents, cancellationToken).ConfigureAwait(false);
         checkpoint.Advance(SnapshotStage.Extracted);
     }
 
-    private void CheckDiskSpace(string snapshotPath)
+    private async Task CheckDiskSpaceBeforeDownloadAsync(
+        SnapshotDownloader downloader, string url, string destinationPath, CancellationToken cancellationToken)
     {
-        if (drives.Length == 0)
+        long existingSize = GetFileSize(destinationPath);
+        long? totalSize;
+        try
+        {
+            totalSize = await downloader.GetTotalSizeAsync(url, existingSize, cancellationToken).ConfigureAwait(false);
+        }
+        catch (Exception e) when (e is IOException or HttpRequestException)
+        {
+            if (_logger.IsWarn)
+                _logger.Warn($"Could not determine the snapshot size upfront. Skipping the pre-download disk space check. Error: {e.Message}");
+            return;
+        }
+
+        if (totalSize is null)
+        {
+            if (_logger.IsWarn)
+                _logger.Warn("The server did not report the snapshot size. Skipping the pre-download disk space check.");
             return;
+        }
 
-        long snapshotSize = api.FileSystem.FileInfo.New(snapshotPath).Length;
-        long required = (long)(snapshotSize * 2.5);
+        CheckDiskSpace(GetRequiredSpaceForDownload(totalSize.Value, existingSize), "download and extract");
+    }
+
+    internal static long GetRequiredSpaceForDownload(long totalSize, long existingSize) =>
+        totalSize - existingSize + GetRequiredSpaceForExtraction(totalSize);
 
+    internal static long GetRequiredSpaceForExtraction(long snapshotSize) =>
+        (long)(snapshotSize * ExtractionSpaceMultiplier);
+
+    private void CheckDiskSpace(long required, string operation)
+    {
         foreach (IDriveInfo drive in drives)
         {
             if (drive.AvailableFreeSpace < required)
                 throw new IOException(
-                    $"Insufficient disk space on '{drive.RootDirectory.FullName}' to extract snapshot: " +
+                    $"Insufficient disk space on '{drive.RootDirectory.FullName}' to {operation} the snapshot: " +
                     $"need at least {required} bytes, {drive.AvailableFreeSpace} available.");
         }
     }

src/Nethermind/Nethermind.Init.Snapshot/Nethermind.Init.Snapshot.csproj

@@ -12,6 +12,10 @@
   <ItemGroup>
     <ProjectReference Include="..\Nethermind.Api\Nethermind.Api.csproj" />
     <ProjectReference Include="..\Nethermind.Init\Nethermind.Init.csproj" />
+
+    <AssemblyAttribute Include="System.Runtime.CompilerServices.InternalsVisibleTo">
+      <_Parameter1>Nethermind.Init.Snapshot.Test</_Parameter1>
+    </AssemblyAttribute>
   </ItemGroup>
 
 </Project>

src/Nethermind/Nethermind.Init.Snapshot/SnapshotDownloader.cs

@@ -83,6 +83,16 @@ public async Task DownloadAsync(string url, string destinationPath, Cancellation
             _logger.Info($"Snapshot downloaded to {destinationPath}.");
     }
 
+    public async Task<long?> GetTotalSizeAsync(string url, long existingSize, CancellationToken cancellationToken)
+    {
+        using HttpResponseMessage response = await SendWithRangeAsync(_httpClient, url, existingSize, cancellationToken).ConfigureAwait(false);
+
+        if (response.StatusCode == HttpStatusCode.RequestedRangeNotSatisfiable)
+            return existingSize;
+
+        return ResolveCopyStrategy(response.StatusCode, existingSize, response.Content.Headers.ContentLength).totalSize;
+    }
+
     public void Dispose() => _httpClient.Dispose();
 
     private static (FileMode fileMode, long bytesToSkip, long? totalSize) ResolveCopyStrategy(

src/Nethermind/Nethermind.slnx

@@ -76,6 +76,7 @@
     <Project Path="Nethermind.Facade.Test/Nethermind.Facade.Test.csproj" />
     <Project Path="Nethermind.HealthChecks.Test/Nethermind.HealthChecks.Test.csproj" />
     <Project Path="Nethermind.History.Test/Nethermind.History.Test.csproj" />
+    <Project Path="Nethermind.Init.Snapshot.Test/Nethermind.Init.Snapshot.Test.csproj" />
     <Project Path="Nethermind.IntegrationTests/Nethermind.IntegrationTests.csproj" />
     <Project Path="Nethermind.JsonRpc.Test/Nethermind.JsonRpc.Test.csproj" />
     <Project Path="Nethermind.KeyStore.Test/Nethermind.KeyStore.Test.csproj" />