Validate that entries points inside valid blocks at startup (#280)

* Separate initialization of block store from pruning of unknown blocks
* Validate that entries points inside valid blocks

1 files changed, 46 insertions, 33 deletions

diff --git a/src/zenstore/blockstore.cpp b/src/zenstore/blockstore.cpp
index e19712c40..05bc69fcb 100644
--- a/src/zenstore/blockstore.cpp
+++ b/src/zenstore/blockstore.cpp
@@ -118,28 +118,19 @@ BlockStoreFile::StreamByteRange(uint64_t FileOffset, uint64_t Size, std::functio
 
 constexpr uint64_t ScrubSmallChunkWindowSize = 4 * 1024 * 1024;
 
-void
-BlockStore::Initialize(const std::filesystem::path&			  BlocksBasePath,
-					   uint64_t								  MaxBlockSize,
-					   uint64_t								  MaxBlockCount,
-					   const std::vector<BlockStoreLocation>& KnownLocations)
+std::unordered_map<uint32_t, uint64_t>
+BlockStore::Initialize(const std::filesystem::path& BlocksBasePath, uint64_t MaxBlockSize, uint64_t MaxBlockCount)
 {
 	ZEN_ASSERT(MaxBlockSize > 0);
 	ZEN_ASSERT(MaxBlockCount > 0);
 	ZEN_ASSERT(IsPow2(MaxBlockCount));
 
+	std::unordered_map<uint32_t, uint64_t> FoundBlocks;
+
 	m_TotalSize		 = 0;
 	m_BlocksBasePath = BlocksBasePath;
 	m_MaxBlockSize	 = MaxBlockSize;
 
-	m_ChunkBlocks.clear();
-
-	std::unordered_set<uint32_t> KnownBlocks;
-	for (const auto& Entry : KnownLocations)
-	{
-		KnownBlocks.insert(Entry.BlockIndex);
-	}
-
 	if (std::filesystem::is_directory(m_BlocksBasePath))
 	{
 		std::vector<std::filesystem::path> FoldersToScan;
@@ -168,23 +159,11 @@ BlockStore::Initialize(const std::filesystem::path&			  BlocksBasePath,
 					{
 						continue;
 					}
-					if (!KnownBlocks.contains(BlockIndex))
-					{
-						// Log removing unreferenced block
-						// Clear out unused blocks
-						ZEN_DEBUG("removing unused block at '{}'", Path);
-						std::error_code Ec;
-						std::filesystem::remove(Path, Ec);
-						if (Ec)
-						{
-							ZEN_WARN("Failed to delete file '{}' reason: '{}'", Path, Ec.message());
-						}
-						continue;
-					}
 					Ref<BlockStoreFile> BlockFile{new BlockStoreFile(Path)};
 					BlockFile->Open();
 					m_TotalSize.fetch_add(BlockFile->FileSize(), std::memory_order::relaxed);
 					m_ChunkBlocks[BlockIndex] = BlockFile;
+					FoundBlocks[BlockIndex]	  = BlockFile->FileSize();
 				}
 			}
 			++FolderOffset;
@@ -194,6 +173,39 @@ BlockStore::Initialize(const std::filesystem::path&			  BlocksBasePath,
 	{
 		CreateDirectories(m_BlocksBasePath);
 	}
+	return FoundBlocks;
+}
+
+void
+BlockStore::Prune(const std::vector<BlockStoreLocation>& KnownLocations)
+{
+	RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
+
+	std::unordered_set<uint32_t> KnownBlocks;
+	for (const auto& Entry : KnownLocations)
+	{
+		KnownBlocks.insert(Entry.BlockIndex);
+	}
+	std::vector<uint32_t> BlocksToDelete;
+	for (auto It = m_ChunkBlocks.begin(); It != m_ChunkBlocks.end(); ++It)
+	{
+		uint32_t BlockIndex = It->first;
+		if (!KnownBlocks.contains(BlockIndex))
+		{
+			Ref<BlockStoreFile> BlockFile = m_ChunkBlocks[BlockIndex];
+			m_TotalSize.fetch_add(BlockFile->FileSize(), std::memory_order::relaxed);
+			BlocksToDelete.push_back(BlockIndex);
+		}
+	}
+
+	for (uint32_t BlockIndex : BlocksToDelete)
+	{
+		// Clear out unused blocks
+		Ref<BlockStoreFile> BlockFile = m_ChunkBlocks[BlockIndex];
+		m_ChunkBlocks.erase(BlockIndex);
+		ZEN_DEBUG("marking block store file '{}' for delete, block #{}", BlockFile->GetPath(), BlockIndex);
+		BlockFile->MarkAsDeleteOnClose();
+	}
 }
 
 void
@@ -939,7 +951,7 @@ TEST_CASE("blockstore.chunks")
 	auto					 RootDirectory = TempDir.Path();
 
 	BlockStore Store;
-	Store.Initialize(RootDirectory, 128, 1024, {});
+	Store.Initialize(RootDirectory, 128, 1024);
 	IoBuffer BadChunk = Store.TryGetChunk({.BlockIndex = 0, .Offset = 0, .Size = 512});
 	CHECK(!BadChunk);
 
@@ -969,7 +981,7 @@ TEST_CASE("blockstore.clean.stray.blocks")
 	auto					 RootDirectory = TempDir.Path();
 
 	BlockStore Store;
-	Store.Initialize(RootDirectory / "store", 128, 1024, {});
+	Store.Initialize(RootDirectory / "store", 128, 1024);
 
 	std::string		   FirstChunkData	   = "This is the data of the first chunk that we will write";
 	BlockStoreLocation FirstChunkLocation  = WriteStringAsChunk(Store, FirstChunkData, 4);
@@ -982,7 +994,8 @@ TEST_CASE("blockstore.clean.stray.blocks")
 	Store.Close();
 
 	// Not referencing the second block means that we should be deleted
-	Store.Initialize(RootDirectory / "store", 128, 1024, {FirstChunkLocation, SecondChunkLocation});
+	Store.Initialize(RootDirectory / "store", 128, 1024);
+	Store.Prune({FirstChunkLocation, SecondChunkLocation});
 
 	CHECK(GetDirectoryContent(RootDirectory / "store", true, false).size() == 1);
 }
@@ -995,7 +1008,7 @@ TEST_CASE("blockstore.flush.forces.new.block")
 	auto					 RootDirectory = TempDir.Path();
 
 	BlockStore Store;
-	Store.Initialize(RootDirectory / "store", 128, 1024, {});
+	Store.Initialize(RootDirectory / "store", 128, 1024);
 
 	std::string FirstChunkData = "This is the data of the first chunk that we will write";
 	WriteStringAsChunk(Store, FirstChunkData, 4);
@@ -1018,7 +1031,7 @@ TEST_CASE("blockstore.iterate.chunks")
 	auto					 RootDirectory = TempDir.Path();
 
 	BlockStore Store;
-	Store.Initialize(RootDirectory / "store", ScrubSmallChunkWindowSize * 2, 1024, {});
+	Store.Initialize(RootDirectory / "store", ScrubSmallChunkWindowSize * 2, 1024);
 	IoBuffer BadChunk = Store.TryGetChunk({.BlockIndex = 0, .Offset = 0, .Size = 512});
 	CHECK(!BadChunk);
 
@@ -1114,7 +1127,7 @@ TEST_CASE("blockstore.reclaim.space")
 	auto					 RootDirectory = TempDir.Path();
 
 	BlockStore Store;
-	Store.Initialize(RootDirectory / "store", 512, 1024, {});
+	Store.Initialize(RootDirectory / "store", 512, 1024);
 
 	constexpr size_t				ChunkCount = 200;
 	constexpr size_t				Alignment  = 8;
@@ -1231,7 +1244,7 @@ TEST_CASE("blockstore.thread.read.write")
 	auto					 RootDirectory = TempDir.Path();
 
 	BlockStore Store;
-	Store.Initialize(RootDirectory / "store", 1088, 1024, {});
+	Store.Initialize(RootDirectory / "store", 1088, 1024);
 
 	constexpr size_t	  ChunkCount = 1000;
 	constexpr size_t	  Alignment	 = 8;