Remove any unreferenced blocks in block store on open (#492)

* Remove any unreferenced blocks in block store on open

1 files changed, 21 insertions, 4 deletions

diff --git a/src/zenstore/blockstore.cpp b/src/zenstore/blockstore.cpp
index c4453a39f..c5f978d18 100644
--- a/src/zenstore/blockstore.cpp
+++ b/src/zenstore/blockstore.cpp
@@ -224,15 +224,22 @@ BlockStore::Initialize(const std::filesystem::path& BlocksBasePath, uint64_t Max
 }
 
 void
-BlockStore::CreateMissingBlocks(const std::vector<BlockStoreLocation>& KnownLocations)
+BlockStore::SyncExistingBlocksOnDisk(const std::vector<BlockStoreLocation>& KnownLocations)
 {
-	ZEN_TRACE_CPU("BlockStore::CreateMissingBlocks");
+	ZEN_TRACE_CPU("BlockStore::SyncExistingBlocksOnDisk");
 
 	RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
 
 	tsl::robin_set<uint32_t> MissingBlocks;
+	tsl::robin_set<uint32_t> DeleteBlocks;
+	DeleteBlocks.reserve(m_ChunkBlocks.size());
+	for (auto It : m_ChunkBlocks)
+	{
+		DeleteBlocks.insert(It.first);
+	}
 	for (const auto& Entry : KnownLocations)
 	{
+		DeleteBlocks.erase(Entry.BlockIndex);
 		if (auto It = m_ChunkBlocks.find(Entry.BlockIndex); It != m_ChunkBlocks.end() && !It->second.IsNull())
 		{
 			continue;
@@ -246,6 +253,16 @@ BlockStore::CreateMissingBlocks(const std::vector<BlockStoreLocation>& KnownLoca
 		NewBlockFile->Create(0);
 		m_ChunkBlocks[BlockIndex] = NewBlockFile;
 	}
+	for (std::uint32_t BlockIndex : DeleteBlocks)
+	{
+		std::filesystem::path BlockPath = GetBlockPath(m_BlocksBasePath, BlockIndex);
+		if (m_ChunkBlocks[BlockIndex])
+		{
+			m_TotalSize.fetch_sub(m_ChunkBlocks[BlockIndex]->FileSize(), std::memory_order::relaxed);
+			m_ChunkBlocks[BlockIndex]->MarkAsDeleteOnClose();
+		}
+		m_ChunkBlocks.erase(BlockIndex);
+	}
 }
 
 void
@@ -1193,8 +1210,8 @@ TEST_CASE("blockstore.clean.stray.blocks")
 	ThirdChunk = Store.TryGetChunk(ThirdChunkLocation);
 	CHECK(!ThirdChunk);
 
-	// Recreate a fake block for an missing chunk location
-	Store.CreateMissingBlocks({FirstChunkLocation, SecondChunkLocation, ThirdChunkLocation});
+	// Recreate a fake block for a missing chunk location
+	Store.SyncExistingBlocksOnDisk({FirstChunkLocation, SecondChunkLocation, ThirdChunkLocation});
 	// We create a fake block for the location - we should still not be able to get the chunk
 	CHECK(GetDirectoryContent(RootDirectory / "store", true, false).size() == 2);
 	ThirdChunk = Store.TryGetChunk(ThirdChunkLocation);