3 files changed, 53 insertions, 33 deletions

diff --git a/src/flatfile.cpp b/src/flatfile.cpp
index 535f4eda91..d9fd4041b7 100644
--- a/src/flatfile.cpp
+++ b/src/flatfile.cpp
@@ -7,6 +7,7 @@
 #include <flatfile.h>
 #include <logging.h>
 #include <tinyformat.h>
+#include <util/system.h>
 
 FlatFileSeq::FlatFileSeq(fs::path dir, const char* prefix, size_t chunk_size) :
     m_dir(std::move(dir)),
@@ -45,3 +46,29 @@ FILE* FlatFileSeq::Open(const CDiskBlockPos& pos, bool fReadOnly)
     }
     return file;
 }
+
+size_t FlatFileSeq::Allocate(const CDiskBlockPos& pos, size_t add_size, bool& out_of_space)
+{
+    out_of_space = false;
+
+    unsigned int n_old_chunks = (pos.nPos + m_chunk_size - 1) / m_chunk_size;
+    unsigned int n_new_chunks = (pos.nPos + add_size + m_chunk_size - 1) / m_chunk_size;
+    if (n_new_chunks > n_old_chunks) {
+        size_t old_size = pos.nPos;
+        size_t new_size = n_new_chunks * m_chunk_size;
+        size_t inc_size = new_size - old_size;
+
+        if (CheckDiskSpace(m_dir, inc_size)) {
+            FILE *file = Open(pos);
+            if (file) {
+                LogPrintf("Pre-allocating up to position 0x%x in %s%05u.dat\n", new_size, m_prefix, pos.nFile);
+                AllocateFileRange(file, pos.nPos, inc_size);
+                fclose(file);
+                return inc_size;
+            }
+        } else {
+            out_of_space = true;
+        }
+    }
+    return 0;
+}
diff --git a/src/flatfile.h b/src/flatfile.h
index b1eea93140..1b5d83e8e6 100644
--- a/src/flatfile.h
+++ b/src/flatfile.h
@@ -34,6 +34,17 @@ public:
 
     /** Open a handle to the file at the given position. */
     FILE* Open(const CDiskBlockPos& pos, bool fReadOnly = false);
+
+    /**
+     * Allocate additional space in a file after the given starting position. The amount allocated
+     * will be the minimum multiple of the sequence chunk size greater than add_size.
+     *
+     * @param[in] pos The starting position that bytes will be allocated after.
+     * @param[in] add_size The minimum number of bytes to be allocated.
+     * @param[out] out_of_space Whether the allocation failed due to insufficient disk space.
+     * @return The number of bytes successfully allocated.
+     */
+    size_t Allocate(const CDiskBlockPos& pos, size_t add_size, bool& out_of_space);
 };
 
 #endif // BITCOIN_FLATFILE_H
diff --git a/src/validation.cpp b/src/validation.cpp
index 4135063505..aca555ee0a 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -3014,21 +3014,13 @@ static bool FindBlockPos(CDiskBlockPos &pos, unsigned int nAddSize, unsigned int
         vinfoBlockFile[nFile].nSize += nAddSize;
 
     if (!fKnown) {
-        unsigned int nOldChunks = (pos.nPos + BLOCKFILE_CHUNK_SIZE - 1) / BLOCKFILE_CHUNK_SIZE;
-        unsigned int nNewChunks = (vinfoBlockFile[nFile].nSize + BLOCKFILE_CHUNK_SIZE - 1) / BLOCKFILE_CHUNK_SIZE;
-        if (nNewChunks > nOldChunks) {
-            if (fPruneMode)
-                fCheckForPruning = true;
-            if (CheckDiskSpace(GetBlocksDir(), nNewChunks * BLOCKFILE_CHUNK_SIZE - pos.nPos)) {
-                FILE *file = OpenBlockFile(pos);
-                if (file) {
-                    LogPrintf("Pre-allocating up to position 0x%x in blk%05u.dat\n", nNewChunks * BLOCKFILE_CHUNK_SIZE, pos.nFile);
-                    AllocateFileRange(file, pos.nPos, nNewChunks * BLOCKFILE_CHUNK_SIZE - pos.nPos);
-                    fclose(file);
-                }
-            }
-            else
-                return AbortNode("Disk space is low!", _("Error: Disk space is low!"));
+        bool out_of_space;
+        size_t bytes_allocated = BlockFileSeq().Allocate(pos, nAddSize, out_of_space);
+        if (out_of_space) {
+            return AbortNode("Disk space is low!", _("Error: Disk space is low!"));
+        }
+        if (bytes_allocated != 0 && fPruneMode) {
+            fCheckForPruning = true;
         }
     }
 
@@ -3042,27 +3034,17 @@ static bool FindUndoPos(CValidationState &state, int nFile, CDiskBlockPos &pos,
 
     LOCK(cs_LastBlockFile);
 
-    unsigned int nNewSize;
     pos.nPos = vinfoBlockFile[nFile].nUndoSize;
-    nNewSize = vinfoBlockFile[nFile].nUndoSize += nAddSize;
+    vinfoBlockFile[nFile].nUndoSize += nAddSize;
     setDirtyFileInfo.insert(nFile);
 
-    unsigned int nOldChunks = (pos.nPos + UNDOFILE_CHUNK_SIZE - 1) / UNDOFILE_CHUNK_SIZE;
-    unsigned int nNewChunks = (nNewSize + UNDOFILE_CHUNK_SIZE - 1) / UNDOFILE_CHUNK_SIZE;
-    if (nNewChunks > nOldChunks) {
-        if (fPruneMode)
-            fCheckForPruning = true;
-        if (CheckDiskSpace(GetBlocksDir(), nNewChunks * UNDOFILE_CHUNK_SIZE - pos.nPos)) {
-            FILE *file = OpenUndoFile(pos);
-            if (file) {
-                LogPrintf("Pre-allocating up to position 0x%x in rev%05u.dat\n", nNewChunks * UNDOFILE_CHUNK_SIZE, pos.nFile);
-                AllocateFileRange(file, pos.nPos, nNewChunks * UNDOFILE_CHUNK_SIZE - pos.nPos);
-                fclose(file);
-            }
-        }
-        else {
-            return AbortNode(state, "Disk space is low!", _("Error: Disk space is low!"));
-        }
+    bool out_of_space;
+    size_t bytes_allocated = UndoFileSeq().Allocate(pos, nAddSize, out_of_space);
+    if (out_of_space) {
+        return AbortNode(state, "Disk space is low!", _("Error: Disk space is low!"));
+    }
+    if (bytes_allocated != 0 && fPruneMode) {
+        fCheckForPruning = true;
     }
 
     return true;