Don't prune block locations due to missing blocks a startup (#487)

* Don't prune block locations due to missing blocks a startup
This makes the behaviour consistent with FileCas - you can have an index that is not fully backed by data.
Asking for a location that is not backed by data results in getting an empty result back
Also, don't try to GC blocks that are unknown to the block store at the time of snapshot (to avoid removing data that comes in after GatherReferences in GC)

5 files changed, 113 insertions, 147 deletions

diff --git a/src/zenserver/cache/cachedisklayer.cpp b/src/zenserver/cache/cachedisklayer.cpp
index ebefa2e54..3cd725a27 100644
--- a/src/zenserver/cache/cachedisklayer.cpp
+++ b/src/zenserver/cache/cachedisklayer.cpp
@@ -604,8 +604,7 @@ ZenCacheDiskLayer::CacheBucket::OpenLog(const bool IsNew)
 
 	CreateDirectories(m_BucketDir);
 
-	std::unordered_map<uint32_t, uint64_t> BlockSizes =
-		m_BlockStore.Initialize(m_BlocksBasePath, MaxBlockSize, BlockStoreDiskLocation::MaxBlockIndex + 1);
+	m_BlockStore.Initialize(m_BlocksBasePath, MaxBlockSize, BlockStoreDiskLocation::MaxBlockIndex + 1);
 
 	if (std::filesystem::is_regular_file(IndexPath))
 	{
@@ -636,7 +635,6 @@ ZenCacheDiskLayer::CacheBucket::OpenLog(const bool IsNew)
 
 	std::vector<BlockStoreLocation> KnownLocations;
 	KnownLocations.reserve(m_Index.size());
-	std::vector<DiskIndexEntry> BadEntries;
 	for (const auto& Entry : m_Index)
 	{
 		size_t				 EntryIndex = Entry.second;
@@ -649,48 +647,10 @@ ZenCacheDiskLayer::CacheBucket::OpenLog(const bool IsNew)
 			continue;
 		}
 		const BlockStoreLocation& BlockLocation = Location.GetBlockLocation(m_PayloadAlignment);
-
-		auto BlockIt = BlockSizes.find(BlockLocation.BlockIndex);
-		if (BlockIt == BlockSizes.end())
-		{
-			ZEN_WARN("Unknown block {} for entry {} in '{}'", BlockLocation.BlockIndex, Entry.first.ToHexString(), m_BucketDir);
-		}
-		else
-		{
-			uint64_t BlockSize = BlockIt->second;
-			if (BlockLocation.Offset + BlockLocation.Size > BlockSize)
-			{
-				ZEN_WARN("Range is outside of block {} for entry {} in '{}'",
-						 BlockLocation.BlockIndex,
-						 Entry.first.ToHexString(),
-						 m_BucketDir);
-			}
-			else
-			{
-				KnownLocations.push_back(BlockLocation);
-				continue;
-			}
-		}
-
-		DiskLocation NewLocation = Payload.Location;
-		NewLocation.Flags |= DiskLocation::kTombStone;
-		BadEntries.push_back(DiskIndexEntry{.Key = Entry.first, .Location = NewLocation});
+		KnownLocations.push_back(BlockLocation);
 	}
 
-	if (!BadEntries.empty())
-	{
-		m_SlogFile.Append(BadEntries);
-		m_SlogFile.Flush();
-
-		LogEntryCount += BadEntries.size();
-
-		for (const DiskIndexEntry& BadEntry : BadEntries)
-		{
-			m_Index.erase(BadEntry.Key);
-		}
-	}
-
-	m_BlockStore.Prune(KnownLocations);
+	m_BlockStore.CreateMissingBlocks(KnownLocations);
 
 	if (IsNew || LogEntryCount > 0)
 	{
@@ -1558,11 +1518,6 @@ ZenCacheDiskLayer::CacheBucket::CollectGarbage(GcContext& GcCtx)
 		}
 		DeleteCacheKeys.push_back(ChunkHash);
 	});
-	if (DeleteCacheKeys.empty())
-	{
-		ZEN_DEBUG("garbage collect SKIPPED, for '{}', no expired cache keys found", m_BucketDir);
-		return;
-	}
 
 	auto __ = MakeGuard([&]() {
 		if (!DeletedChunks.empty())
@@ -1638,11 +1593,6 @@ ZenCacheDiskLayer::CacheBucket::CollectGarbage(GcContext& GcCtx)
 				WriteBlockTimeUs += ElapsedUs;
 				WriteBlockLongestTimeUs = std::max(ElapsedUs, WriteBlockLongestTimeUs);
 			});
-			if (m_Index.empty())
-			{
-				ZEN_DEBUG("garbage collect SKIPPED, for '{}', container is empty", m_BucketDir);
-				return;
-			}
 
 			BlockStoreState = m_BlockStore.GetReclaimSnapshotState();
 
@@ -1769,10 +1719,6 @@ ZenCacheDiskLayer::CacheBucket::CollectGarbage(GcContext& GcCtx)
 		}
 	}
 
-	if (TotalChunkHashes.empty())
-	{
-		return;
-	}
 	TotalChunkCount = TotalChunkHashes.size();
 
 	std::vector<BlockStoreLocation> ChunkLocations;
diff --git a/src/zenserver/cache/structuredcachestore.cpp b/src/zenserver/cache/structuredcachestore.cpp
index 786053adc..207323451 100644
--- a/src/zenserver/cache/structuredcachestore.cpp
+++ b/src/zenserver/cache/structuredcachestore.cpp
@@ -1123,7 +1123,15 @@ TEST_CASE("z$.gc")
 			}
 		}
 		{
+			// GC could not remove the currently written block so size will not be zero
 			ZenCacheNamespace Zcs(Gc, *JobQueue, TempDir.Path() / "cache", true);
+			CHECK_NE(0, Zcs.StorageSize().DiskSize);
+
+			// GC again and now the stray write block should be removed
+			GcContext GcCtx(GcClock::Now() + std::chrono::minutes(2), GcClock::Now() + std::chrono::minutes(2));
+			GcCtx.CollectSmallObjects(true);
+			Gc.CollectGarbage(GcCtx);
+
 			CHECK_EQ(0, Zcs.StorageSize().DiskSize);
 		}
 	}
diff --git a/src/zenstore/blockstore.cpp b/src/zenstore/blockstore.cpp
index 12ea24752..c4453a39f 100644
--- a/src/zenstore/blockstore.cpp
+++ b/src/zenstore/blockstore.cpp
@@ -109,6 +109,10 @@ BlockStoreFile::MarkAsDeleteOnClose()
 IoBuffer
 BlockStoreFile::GetChunk(uint64_t Offset, uint64_t Size)
 {
+	if (Offset + Size > m_IoBuffer.GetSize())
+	{
+		return {};
+	}
 	return IoBuffer(m_IoBuffer, Offset, Size);
 }
 
@@ -154,7 +158,7 @@ BlockStore::~BlockStore()
 {
 }
 
-std::unordered_map<uint32_t, uint64_t>
+void
 BlockStore::Initialize(const std::filesystem::path& BlocksBasePath, uint64_t MaxBlockSize, uint64_t MaxBlockCount)
 {
 	ZEN_TRACE_CPU("BlockStore::Initialize");
@@ -217,40 +221,30 @@ BlockStore::Initialize(const std::filesystem::path& BlocksBasePath, uint64_t Max
 	{
 		CreateDirectories(m_BlocksBasePath);
 	}
-	return FoundBlocks;
 }
 
 void
-BlockStore::Prune(const std::vector<BlockStoreLocation>& KnownLocations)
+BlockStore::CreateMissingBlocks(const std::vector<BlockStoreLocation>& KnownLocations)
 {
-	ZEN_TRACE_CPU("BlockStore::Prune");
+	ZEN_TRACE_CPU("BlockStore::CreateMissingBlocks");
 
 	RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
 
-	tsl::robin_set<uint32_t> KnownBlocks;
+	tsl::robin_set<uint32_t> MissingBlocks;
 	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))
+		if (auto It = m_ChunkBlocks.find(Entry.BlockIndex); It != m_ChunkBlocks.end() && !It->second.IsNull())
 		{
-			Ref<BlockStoreFile> BlockFile = m_ChunkBlocks[BlockIndex];
-			BlocksToDelete.push_back(BlockIndex);
+			continue;
 		}
+		MissingBlocks.insert(Entry.BlockIndex);
 	}
-
-	for (uint32_t BlockIndex : BlocksToDelete)
+	for (std::uint32_t BlockIndex : MissingBlocks)
 	{
-		// Clear out unused blocks
-		Ref<BlockStoreFile> BlockFile = m_ChunkBlocks[BlockIndex];
-		m_TotalSize.fetch_sub(BlockFile->FileSize(), std::memory_order::relaxed);
-		m_ChunkBlocks.erase(BlockIndex);
-		ZEN_DEBUG("marking block store file '{}' for delete, block #{}", BlockFile->GetPath(), BlockIndex);
-		BlockFile->MarkAsDeleteOnClose();
+		std::filesystem::path BlockPath = GetBlockPath(m_BlocksBasePath, BlockIndex);
+		Ref<BlockStoreFile>	  NewBlockFile(new BlockStoreFile(BlockPath));
+		NewBlockFile->Create(0);
+		m_ChunkBlocks[BlockIndex] = NewBlockFile;
 	}
 }
 
@@ -367,7 +361,10 @@ BlockStore::GetReclaimSnapshotState()
 	{
 		State.m_ActiveWriteBlocks.insert(m_WriteBlockIndex);
 	}
-	State.BlockCount = m_ChunkBlocks.size();
+	for (auto It : m_ChunkBlocks)
+	{
+		State.m_BlockIndexes.insert(It.first);
+	}
 	return State;
 }
 
@@ -421,11 +418,6 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 						 const ReclaimCallback&					ChangeCallback,
 						 const ClaimDiskReserveCallback&		DiskReserveCallback)
 {
-	if (ChunkLocations.empty())
-	{
-		return;
-	}
-
 	ZEN_TRACE_CPU("BlockStore::ReclaimSpace");
 
 	uint64_t WriteBlockTimeUs		 = 0;
@@ -460,7 +452,7 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 			NiceBytes(OldTotalSize));
 	});
 
-	size_t BlockCount = Snapshot.BlockCount;
+	size_t BlockCount = Snapshot.m_BlockIndexes.size();
 	if (BlockCount == 0)
 	{
 		ZEN_DEBUG("garbage collect for '{}' SKIPPED, no blocks to process", m_BlocksBasePath);
@@ -487,6 +479,11 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 	for (size_t Index = 0; Index < TotalChunkCount; ++Index)
 	{
 		const BlockStoreLocation& Location = ChunkLocations[Index];
+		if (!Snapshot.m_BlockIndexes.contains(Location.BlockIndex))
+		{
+			// We did not know about the block when we took the snapshot, don't touch it
+			continue;
+		}
 		OldTotalSize += Location.Size;
 		auto   BlockIndexPtr = BlockIndexToChunkMapIndex.find(Location.BlockIndex);
 		size_t ChunkMapIndex = 0;
@@ -516,7 +513,7 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 		DeleteCount++;
 	}
 
-	tsl::robin_set<uint32_t> BlocksToReWrite;
+	std::vector<uint32_t> BlocksToReWrite;
 	BlocksToReWrite.reserve(BlockIndexToChunkMapIndex.size());
 	for (const auto& Entry : BlockIndexToChunkMapIndex)
 	{
@@ -527,7 +524,33 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 		{
 			continue;
 		}
-		BlocksToReWrite.insert(BlockIndex);
+		BlocksToReWrite.push_back(BlockIndex);
+	}
+
+	{
+		// Any known block not referenced should be added as well
+		RwLock::SharedLockScope __(m_InsertLock);
+		for (std::uint32_t BlockIndex : Snapshot.m_BlockIndexes)
+		{
+			if (!m_ChunkBlocks.contains(BlockIndex))
+			{
+				continue;
+			}
+			bool WasActiveWriteBlock = Snapshot.m_ActiveWriteBlocks.contains(BlockIndex);
+			if (WasActiveWriteBlock)
+			{
+				continue;
+			}
+			if (BlockIndexToChunkMapIndex.contains(BlockIndex))
+			{
+				continue;
+			}
+			size_t ChunkMapIndex = ChunkMapIndex  = BlockKeepChunks.size();
+			BlockIndexToChunkMapIndex[BlockIndex] = ChunkMapIndex;
+			BlockKeepChunks.resize(ChunkMapIndex + 1);
+			BlockDeleteChunks.resize(ChunkMapIndex + 1);
+			BlocksToReWrite.push_back(BlockIndex);
+		}
 	}
 
 	if (DryRun)
@@ -618,6 +641,17 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 				BlockDeleteChunks[ChunkMapIndex].insert(BlockDeleteChunks[ChunkMapIndex].end(), KeepMap.begin(), KeepMap.end());
 				KeepMap.clear();
 			}
+			else if (OldBlockFile && (OldBlockFile->FileSize() == 0))
+			{
+				// Block created to accommodate missing blocks
+				ZEN_WARN("Missing block {} in {} - backing data for locations is missing, marking {} entries as deleted.",
+						 BlockIndex,
+						 m_BlocksBasePath,
+						 KeepMap.size());
+
+				BlockDeleteChunks[ChunkMapIndex].insert(BlockDeleteChunks[ChunkMapIndex].end(), KeepMap.begin(), KeepMap.end());
+				KeepMap.clear();
+			}
 
 			MovedChunksArray MovedChunks;
 			if (OldBlockFile)
@@ -1141,15 +1175,30 @@ TEST_CASE("blockstore.clean.stray.blocks")
 	BlockStoreLocation SecondChunkLocation = WriteStringAsChunk(Store, SecondChunkData, 4);
 	std::string		   ThirdChunkData =
 		"This is a much longer string that will not fit in the first block so it should be placed in the second block";
-	WriteStringAsChunk(Store, ThirdChunkData, 4);
+	BlockStoreLocation ThirdChunkLocation = WriteStringAsChunk(Store, ThirdChunkData, 4);
 
 	Store.Close();
 
-	// Not referencing the second block means that we should be deleted
 	Store.Initialize(RootDirectory / "store", 128, 1024);
-	Store.Prune({FirstChunkLocation, SecondChunkLocation});
-
+	CHECK(GetDirectoryContent(RootDirectory / "store", true, false).size() == 2);
+	IoBuffer ThirdChunk = Store.TryGetChunk(ThirdChunkLocation);
+	CHECK(ThirdChunk);
+
+	// Reclaim space should delete unreferenced block
+	Store.ReclaimSpace(Store.GetReclaimSnapshotState(), {FirstChunkLocation, SecondChunkLocation}, {0, 1}, 4, false);
+	// Block lives on as long as we reference it via ThirdChunk
+	CHECK(GetDirectoryContent(RootDirectory / "store", true, false).size() == 2);
+	ThirdChunk = {};
 	CHECK(GetDirectoryContent(RootDirectory / "store", true, false).size() == 1);
+	ThirdChunk = Store.TryGetChunk(ThirdChunkLocation);
+	CHECK(!ThirdChunk);
+
+	// Recreate a fake block for an missing chunk location
+	Store.CreateMissingBlocks({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);
+	CHECK(!ThirdChunk);
 }
 
 TEST_CASE("blockstore.flush.force.new.block")
@@ -1368,11 +1417,13 @@ TEST_CASE("blockstore.reclaim.space")
 		}
 	}
 
-	NewChunkLocations = ChunkLocations;
-	MovedChunkCount	  = 0;
-	DeletedChunkCount = 0;
+	// We need to take a new state since reclaim space add new block when compacting
+	BlockStore::ReclaimSnapshotState State2 = Store.GetReclaimSnapshotState();
+	NewChunkLocations						= ChunkLocations;
+	MovedChunkCount							= 0;
+	DeletedChunkCount						= 0;
 	Store.ReclaimSpace(
-		State1,
+		State2,
 		ChunkLocations,
 		{},
 		Alignment,
diff --git a/src/zenstore/compactcas.cpp b/src/zenstore/compactcas.cpp
index 98ad20b98..0468e95de 100644
--- a/src/zenstore/compactcas.cpp
+++ b/src/zenstore/compactcas.cpp
@@ -804,8 +804,7 @@ CasContainerStrategy::OpenContainer(bool IsNewStore)
 
 	CreateDirectories(BasePath);
 
-	std::unordered_map<uint32_t, uint64_t> BlockSizes =
-		m_BlockStore.Initialize(m_BlocksBasePath, m_MaxBlockSize, BlockStoreDiskLocation::MaxBlockIndex + 1);
+	m_BlockStore.Initialize(m_BlocksBasePath, m_MaxBlockSize, BlockStoreDiskLocation::MaxBlockIndex + 1);
 
 	std::filesystem::path LogPath	= GetLogPath(m_RootDirectory, m_ContainerBaseName);
 	std::filesystem::path IndexPath = GetIndexPath(m_RootDirectory, m_ContainerBaseName);
@@ -839,54 +838,14 @@ CasContainerStrategy::OpenContainer(bool IsNewStore)
 
 	std::vector<BlockStoreLocation> KnownLocations;
 	KnownLocations.reserve(m_LocationMap.size());
-	std::vector<CasDiskIndexEntry> BadEntries;
 	for (const auto& Entry : m_LocationMap)
 	{
-		const BlockStoreDiskLocation& DiskLocation = m_Locations[Entry.second];
-		uint32_t					  BlockIndex   = DiskLocation.GetBlockIndex();
-		auto						  BlockIt	   = BlockSizes.find(BlockIndex);
-		if (BlockIt == BlockSizes.end())
-		{
-			ZEN_WARN("Unknown block {} for entry {} in '{}'", BlockIndex, Entry.first.ToHexString(), BasePath);
-		}
-		else
-		{
-			BlockStoreLocation BlockLocation = DiskLocation.Get(m_PayloadAlignment);
-
-			uint64_t BlockSize = BlockIt->second;
-			if (BlockLocation.Offset + BlockLocation.Size > BlockSize)
-			{
-				ZEN_WARN("Range is outside of block {} for entry {} in '{}'",
-						 BlockLocation.BlockIndex,
-						 Entry.first.ToHexString(),
-						 BasePath);
-			}
-			else
-			{
-				KnownLocations.emplace_back(std::move(BlockLocation));
-				continue;
-			}
-			BadEntries.push_back({.Key = Entry.first, .Location = DiskLocation, .Flags = CasDiskIndexEntry::kTombstone});
-		}
-	}
-
-	if (!BadEntries.empty())
-	{
-		m_CasLog.Append(BadEntries);
-		m_CasLog.Flush();
-
-		LogEntryCount += BadEntries.size();
-
-		for (const CasDiskIndexEntry& BadEntry : BadEntries)
-		{
-			m_LocationMap.erase(BadEntry.Key);
-		}
-
-		RwLock::ExclusiveLockScope IndexLock(m_LocationMapLock);
-		CompactIndex(IndexLock);
+		const BlockStoreDiskLocation& DiskLocation	= m_Locations[Entry.second];
+		BlockStoreLocation			  BlockLocation = DiskLocation.Get(m_PayloadAlignment);
+		KnownLocations.emplace_back(std::move(BlockLocation));
 	}
 
-	m_BlockStore.Prune(KnownLocations);
+	m_BlockStore.CreateMissingBlocks(KnownLocations);
 
 	if (IsNewStore || (LogEntryCount > 0))
 	{
diff --git a/src/zenstore/include/zenstore/blockstore.h b/src/zenstore/include/zenstore/blockstore.h
index a2dea86c1..48d33edbc 100644
--- a/src/zenstore/include/zenstore/blockstore.h
+++ b/src/zenstore/include/zenstore/blockstore.h
@@ -117,7 +117,7 @@ public:
 	struct ReclaimSnapshotState
 	{
 		std::unordered_set<uint32_t> m_ActiveWriteBlocks;
-		size_t						 BlockCount;
+		std::unordered_set<uint32_t> m_BlockIndexes;
 	};
 
 	typedef std::vector<std::pair<size_t, BlockStoreLocation>> MovedChunksArray;
@@ -129,11 +129,13 @@ public:
 	typedef std::function<void(size_t ChunkIndex, BlockStoreFile& File, uint64_t Offset, uint64_t Size)>   IterateChunksLargeSizeCallback;
 	typedef std::function<void(const BlockStoreLocation& Location)>										   WriteChunkCallback;
 
-	std::unordered_map<uint32_t, uint64_t> Initialize(const std::filesystem::path& BlocksBasePath,
-													  uint64_t					   MaxBlockSize,
-													  uint64_t					   MaxBlockCount);
+	void Initialize(const std::filesystem::path& BlocksBasePath, uint64_t MaxBlockSize, uint64_t MaxBlockCount);
+
+	// Ask the store to create empty blocks for all locations that does not have a block
+	// This is to make sure we don't have locations that points to a block that will be written
+	// in the future which would result in us getting garbage data
+	void CreateMissingBlocks(const std::vector<BlockStoreLocation>& KnownLocations);
 
-	void Prune(const std::vector<BlockStoreLocation>& KnownLocations);
 	void Close();
 
 	void WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment, const WriteChunkCallback& Callback);