merge from main

1 files changed, 22 insertions, 20 deletions

diff --git a/zenstore/blockstore.cpp b/zenstore/blockstore.cpp
index d490678b5..4e61c23cf 100644
--- a/zenstore/blockstore.cpp
+++ b/zenstore/blockstore.cpp
@@ -225,23 +225,27 @@ BlockStore::WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment, Writ
 		{
 			m_WriteBlock = nullptr;
 		}
+
+		if (m_ChunkBlocks.size() == m_MaxBlockCount)
 		{
-			if (m_ChunkBlocks.size() == m_MaxBlockCount)
-			{
-				throw std::runtime_error(fmt::format("unable to allocate a new block in '{}'", m_BlocksBasePath));
-			}
-			WriteBlockIndex += IsWriting ? 1 : 0;
-			while (m_ChunkBlocks.contains(WriteBlockIndex))
-			{
-				WriteBlockIndex = (WriteBlockIndex + 1) & (m_MaxBlockCount - 1);
-			}
-			std::filesystem::path BlockPath = GetBlockPath(m_BlocksBasePath, WriteBlockIndex);
-			m_WriteBlock					= new BlockStoreFile(BlockPath);
-			m_ChunkBlocks[WriteBlockIndex]	= m_WriteBlock;
-			m_WriteBlockIndex.store(WriteBlockIndex, std::memory_order_release);
+			throw std::runtime_error(fmt::format("unable to allocate a new block in '{}'", m_BlocksBasePath));
 		}
+
+		WriteBlockIndex += IsWriting ? 1 : 0;
+		while (m_ChunkBlocks.contains(WriteBlockIndex))
+		{
+			WriteBlockIndex = (WriteBlockIndex + 1) & (m_MaxBlockCount - 1);
+		}
+
+		std::filesystem::path BlockPath = GetBlockPath(m_BlocksBasePath, WriteBlockIndex);
+
+		Ref<BlockStoreFile> NewBlockFile = new BlockStoreFile(BlockPath);
+		NewBlockFile->Create(m_MaxBlockSize);
+
+		m_ChunkBlocks[WriteBlockIndex] = NewBlockFile;
+		m_WriteBlock				   = NewBlockFile;
+		m_WriteBlockIndex.store(WriteBlockIndex, std::memory_order_release);
 		m_CurrentInsertOffset = 0;
-		m_WriteBlock->Create(m_MaxBlockSize);
 	}
 	uint64_t InsertOffset		   = m_CurrentInsertOffset;
 	m_CurrentInsertOffset		   = RoundUp(InsertOffset + Size, Alignment);
@@ -268,6 +272,10 @@ BlockStore::GetReclaimSnapshotState()
 	{
 		State.m_ActiveWriteBlocks.insert(BlockIndex);
 	}
+	if (m_WriteBlock)
+	{
+		State.m_ActiveWriteBlocks.insert(m_WriteBlockIndex);
+	}
 	State.BlockCount = m_ChunkBlocks.size();
 	return State;
 }
@@ -911,12 +919,6 @@ BlockStore::GetBlockPath(const std::filesystem::path& BlocksBasePath, const uint
 
 #if ZEN_WITH_TESTS
 
-static bool
-operator==(const BlockStoreLocation& Lhs, const BlockStoreLocation& Rhs)
-{
-	return Lhs.BlockIndex == Rhs.BlockIndex && Lhs.Offset == Rhs.Offset && Lhs.Size == Rhs.Size;
-}
-
 TEST_CASE("blockstore.blockstoredisklocation")
 {
 	BlockStoreLocation Zero = BlockStoreLocation{.BlockIndex = 0, .Offset = 0, .Size = 0};