Merge branch 'main' of https://github.com/EpicGames/zen

1 files changed, 274 insertions, 140 deletions

diff --git a/src/zenstore/blockstore.cpp b/src/zenstore/blockstore.cpp
index 063d38707..71e306eca 100644
--- a/src/zenstore/blockstore.cpp
+++ b/src/zenstore/blockstore.cpp
@@ -15,6 +15,7 @@
 ZEN_THIRD_PARTY_INCLUDES_START
 #include <tsl/robin_map.h>
 #include <tsl/robin_set.h>
+#include <gsl/gsl-lite.hpp>
 ZEN_THIRD_PARTY_INCLUDES_END
 
 #if ZEN_WITH_TESTS
@@ -226,7 +227,17 @@ BlockStore::Initialize(const std::filesystem::path& BlocksBasePath, uint64_t Max
 }
 
 void
-BlockStore::SyncExistingBlocksOnDisk(const std::vector<BlockStoreLocation>& KnownLocations)
+BlockStore::BlockIndexSet::Add(uint32_t BlockIndex)
+{
+	if (!std::binary_search(begin(BlockIndexes), end(BlockIndexes), BlockIndex))
+	{
+		auto It = std::lower_bound(begin(BlockIndexes), end(BlockIndexes), BlockIndex);
+		BlockIndexes.insert(It, BlockIndex);
+	}
+}
+
+void
+BlockStore::SyncExistingBlocksOnDisk(const BlockIndexSet& KnownLocations)
 {
 	ZEN_TRACE_CPU("BlockStore::SyncExistingBlocksOnDisk");
 
@@ -239,14 +250,18 @@ BlockStore::SyncExistingBlocksOnDisk(const std::vector<BlockStoreLocation>& Know
 	{
 		DeleteBlocks.insert(It.first);
 	}
-	for (const auto& Entry : KnownLocations)
+
+	for (const uint32_t BlockIndex : KnownLocations.GetBlockIndices())
 	{
-		DeleteBlocks.erase(Entry.BlockIndex);
-		if (auto It = m_ChunkBlocks.find(Entry.BlockIndex); It != m_ChunkBlocks.end() && !It->second.IsNull())
+		DeleteBlocks.erase(BlockIndex);
+		if (auto It = m_ChunkBlocks.find(BlockIndex); It != m_ChunkBlocks.end() && !It->second.IsNull())
 		{
 			continue;
 		}
-		MissingBlocks.insert(Entry.BlockIndex);
+		else
+		{
+			MissingBlocks.insert(BlockIndex);
+		}
 	}
 	for (std::uint32_t BlockIndex : MissingBlocks)
 	{
@@ -267,6 +282,66 @@ BlockStore::SyncExistingBlocksOnDisk(const std::vector<BlockStoreLocation>& Know
 	}
 }
 
+BlockStore::BlockEntryCountMap
+BlockStore::GetBlocksToCompact(const BlockUsageMap& BlockUsage, uint32_t BlockUsageThresholdPercent)
+{
+	BlockEntryCountMap Result;
+	{
+		RwLock::SharedLockScope InsertLock(m_InsertLock);
+		for (const auto& It : m_ChunkBlocks)
+		{
+			uint32_t BlockIndex = It.first;
+			if ((BlockIndex == m_WriteBlockIndex.load()) && m_WriteBlock)
+			{
+				continue;
+			}
+			if (std::find(m_ActiveWriteBlocks.begin(), m_ActiveWriteBlocks.end(), BlockIndex) != m_ActiveWriteBlocks.end())
+			{
+				continue;
+			}
+
+			uint64_t UsedSize  = 0;
+			uint32_t UsedCount = 0;
+			if (auto UsageIt = BlockUsage.find(BlockIndex); UsageIt != BlockUsage.end())
+			{
+				UsedSize  = UsageIt->second.DiskUsage;
+				UsedCount = UsageIt->second.EntryCount;
+			}
+
+			uint64_t BlockSize = It.second ? It.second->FileSize() : 0u;
+			if (BlockSize == 0)
+			{
+				Result.insert_or_assign(BlockIndex, UsedCount);
+				continue;
+			}
+
+			if (BlockUsageThresholdPercent == 100)
+			{
+				if (UsedSize < BlockSize)
+				{
+					Result.insert_or_assign(BlockIndex, UsedCount);
+				}
+			}
+			else if (BlockUsageThresholdPercent == 0)
+			{
+				if (UsedSize == 0)
+				{
+					Result.insert_or_assign(BlockIndex, UsedCount);
+				}
+			}
+			else
+			{
+				const uint32_t UsedPercent = UsedSize < BlockSize ? gsl::narrow<uint32_t>((100 * UsedSize) / BlockSize) : 100u;
+				if (UsedPercent < BlockUsageThresholdPercent)
+				{
+					Result.insert_or_assign(BlockIndex, UsedCount);
+				}
+			}
+		}
+	}
+	return Result;
+}
+
 void
 BlockStore::Close()
 {
@@ -312,7 +387,7 @@ BlockStore::GetFreeBlockIndex(uint32_t ProbeIndex, RwLock::ExclusiveLockScope&,
 }
 
 void
-BlockStore::WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment, const WriteChunkCallback& Callback)
+BlockStore::WriteChunk(const void* Data, uint64_t Size, uint32_t Alignment, const WriteChunkCallback& Callback)
 {
 	ZEN_TRACE_CPU("BlockStore::WriteChunk");
 
@@ -321,12 +396,14 @@ BlockStore::WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment, cons
 	ZEN_ASSERT(Size <= m_MaxBlockSize);
 	ZEN_ASSERT(Alignment > 0u);
 
+	uint32_t ChunkSize = gsl::narrow<uint32_t>(Size);
+
 	RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
 
 	uint32_t WriteBlockIndex	 = m_WriteBlockIndex.load(std::memory_order_acquire);
 	bool	 IsWriting			 = !!m_WriteBlock;
-	uint64_t AlignedInsertOffset = RoundUp(m_CurrentInsertOffset, Alignment);
-	if (!IsWriting || (AlignedInsertOffset + Size) > m_MaxBlockSize)
+	uint32_t AlignedInsertOffset = RoundUp(m_CurrentInsertOffset, Alignment);
+	if (!IsWriting || (AlignedInsertOffset + ChunkSize) > m_MaxBlockSize)
 	{
 		if (m_WriteBlock)
 		{
@@ -351,16 +428,16 @@ BlockStore::WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment, cons
 		m_CurrentInsertOffset = 0;
 		AlignedInsertOffset	  = 0;
 	}
-	uint64_t AlignedWriteSize	   = AlignedInsertOffset - m_CurrentInsertOffset + Size;
-	m_CurrentInsertOffset		   = AlignedInsertOffset + Size;
+	uint32_t AlignedWriteSize	   = AlignedInsertOffset - m_CurrentInsertOffset + ChunkSize;
+	m_CurrentInsertOffset		   = AlignedInsertOffset + ChunkSize;
 	Ref<BlockStoreFile> WriteBlock = m_WriteBlock;
 	m_ActiveWriteBlocks.push_back(WriteBlockIndex);
 	InsertLock.ReleaseNow();
 
-	WriteBlock->Write(Data, Size, AlignedInsertOffset);
+	WriteBlock->Write(Data, ChunkSize, AlignedInsertOffset);
 	m_TotalSize.fetch_add(AlignedWriteSize, std::memory_order::relaxed);
 
-	Callback({.BlockIndex = WriteBlockIndex, .Offset = AlignedInsertOffset, .Size = Size});
+	Callback({.BlockIndex = WriteBlockIndex, .Offset = AlignedInsertOffset, .Size = ChunkSize});
 
 	{
 		RwLock::ExclusiveLockScope _(m_InsertLock);
@@ -433,7 +510,7 @@ void
 BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 						 const std::vector<BlockStoreLocation>& ChunkLocations,
 						 const ChunkIndexArray&					KeepChunkIndexes,
-						 uint64_t								PayloadAlignment,
+						 uint32_t								PayloadAlignment,
 						 bool									DryRun,
 						 const ReclaimCallback&					ChangeCallback,
 						 const ClaimDiskReserveCallback&		DiskReserveCallback)
@@ -682,9 +759,9 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 				{
 					const BlockStoreLocation ChunkLocation = ChunkLocations[ChunkIndex];
 					Chunk.resize(ChunkLocation.Size);
-					OldBlockFile->Read(Chunk.data(), Chunk.size(), ChunkLocation.Offset);
+					OldBlockFile->Read(Chunk.data(), ChunkLocation.Size, ChunkLocation.Offset);
 
-					if (!NewBlockFile || (WriteOffset + Chunk.size() > m_MaxBlockSize))
+					if (!NewBlockFile || (WriteOffset + ChunkLocation.Size > m_MaxBlockSize))
 					{
 						uint32_t NextBlockIndex = m_WriteBlockIndex.load(std::memory_order_relaxed);
 
@@ -758,10 +835,12 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 						WriteOffset	  = 0;
 					}
 
-					NewBlockFile->Write(Chunk.data(), Chunk.size(), WriteOffset);
-					MovedChunks.push_back({ChunkIndex, {.BlockIndex = NewBlockIndex, .Offset = WriteOffset, .Size = Chunk.size()}});
+					NewBlockFile->Write(Chunk.data(), ChunkLocation.Size, WriteOffset);
+					MovedChunks.push_back(
+						{ChunkIndex,
+						 {.BlockIndex = NewBlockIndex, .Offset = gsl::narrow<uint32_t>(WriteOffset), .Size = ChunkLocation.Size}});
 					uint64_t OldOffset = WriteOffset;
-					WriteOffset		   = RoundUp(WriteOffset + Chunk.size(), PayloadAlignment);
+					WriteOffset		   = RoundUp(WriteOffset + ChunkLocation.Size, PayloadAlignment);
 					m_TotalSize.fetch_add(WriteOffset - OldOffset, std::memory_order::relaxed);
 				}
 				Chunk.clear();
@@ -961,7 +1040,7 @@ BlockStore::IterateChunks(const std::vector<BlockStoreLocation>& ChunkLocations,
 
 void
 BlockStore::CompactBlocks(const BlockStoreCompactState&	  CompactState,
-						  uint64_t						  PayloadAlignment,
+						  uint32_t						  PayloadAlignment,
 						  const CompactCallback&		  ChangeCallback,
 						  const ClaimDiskReserveCallback& DiskReserveCallback)
 {
@@ -971,7 +1050,7 @@ BlockStore::CompactBlocks(const BlockStoreCompactState&	  CompactState,
 
 	Stopwatch  TotalTimer;
 	const auto _ = MakeGuard([&] {
-		ZEN_DEBUG("compact blocks for '{}' DONE after {}, deleted {} and moved {} chunks ({}) ",
+		ZEN_DEBUG("Compact blocks for '{}' DONE after {}, deleted {} and moved {} chunks ({}) ",
 				  m_BlocksBasePath,
 				  NiceTimeSpanMs(TotalTimer.GetElapsedTimeMs()),
 				  NiceBytes(DeletedSize),
@@ -983,13 +1062,14 @@ BlockStore::CompactBlocks(const BlockStoreCompactState&	  CompactState,
 	uint32_t		 NewBlockIndex = 0;
 	MovedChunksArray MovedChunks;
 
+	uint64_t AddedSize	 = 0;
 	uint64_t RemovedSize = 0;
 
 	Ref<BlockStoreFile> NewBlockFile;
 	auto				NewBlockFileGuard = MakeGuard([&]() {
 		   if (NewBlockFile)
 		   {
-			   ZEN_DEBUG("dropping incomplete cas block store file '{}'", NewBlockFile->GetPath());
+			   ZEN_DEBUG("Dropping incomplete cas block store file '{}'", NewBlockFile->GetPath());
 			   {
 				   RwLock::ExclusiveLockScope _l(m_InsertLock);
 				   if (m_ChunkBlocks[NewBlockIndex] == NewBlockFile)
@@ -1001,140 +1081,174 @@ BlockStore::CompactBlocks(const BlockStoreCompactState&	  CompactState,
 		   }
 	   });
 
+	auto ReportChanges = [&]() -> bool {
+		bool Continue = true;
+		if (!MovedChunks.empty() || RemovedSize > 0)
+		{
+			Continue = ChangeCallback(MovedChunks, RemovedSize > AddedSize ? RemovedSize - AddedSize : 0);
+			DeletedSize += RemovedSize;
+			RemovedSize = 0;
+			AddedSize	= 0;
+			MovedCount += MovedChunks.size();
+			MovedChunks.clear();
+		}
+		return Continue;
+	};
+
 	std::vector<uint32_t> RemovedBlocks;
 
-	CompactState.IterateBlocks(
-		[&](uint32_t BlockIndex, const std::vector<size_t>& KeepChunkIndexes, const std::vector<BlockStoreLocation>& ChunkLocations) {
-			Ref<BlockStoreFile> OldBlockFile;
+	CompactState.IterateBlocks([&](uint32_t								  BlockIndex,
+								   const std::vector<size_t>&			  KeepChunkIndexes,
+								   const std::vector<BlockStoreLocation>& ChunkLocations) -> bool {
+		Ref<BlockStoreFile> OldBlockFile;
+		{
+			RwLock::SharedLockScope _(m_InsertLock);
+			if ((BlockIndex == m_WriteBlockIndex.load()) && m_WriteBlock)
 			{
-				RwLock::SharedLockScope _(m_InsertLock);
-				if ((BlockIndex == m_WriteBlockIndex.load()) && m_WriteBlock)
-				{
-					// You are trying to collect the currently writing block, Report error?
-					return;
-				}
-				auto It = m_ChunkBlocks.find(BlockIndex);
-				if (It == m_ChunkBlocks.end())
-				{
-					// This block has unknown, we can't move anything. Report error?
-					return;
-				}
-				if (!It->second)
-				{
-					// This block has been removed, we can't move anything. Report error?
-					return;
-				}
-				OldBlockFile = It->second;
+				ZEN_ERROR("Compact Block was requested to rewrite the currently active write block in '{}', Block index {}",
+						  m_BlocksBasePath,
+						  BlockIndex);
+				return false;
 			}
-			ZEN_ASSERT(OldBlockFile);
-
-			uint64_t OldBlockSize = OldBlockFile->FileSize();
+			auto It = m_ChunkBlocks.find(BlockIndex);
+			if (It == m_ChunkBlocks.end())
+			{
+				ZEN_WARN("Compact Block was requested to rewrite an unknown block in '{}', Block index {}", m_BlocksBasePath, BlockIndex);
+				return true;
+			}
+			if (!It->second)
+			{
+				ZEN_WARN("Compact Block was requested to rewrite a deleted block in '{}', Block index {}", m_BlocksBasePath, BlockIndex);
+				return true;
+			}
+			OldBlockFile = It->second;
+		}
+		ZEN_ASSERT(OldBlockFile);
 
-			// TODO: Add heuristics for determining if it is worth to compact a block (if only a very small part is removed)
+		uint64_t OldBlockSize = OldBlockFile->FileSize();
 
-			std::vector<uint8_t> Chunk;
-			for (const size_t& ChunkIndex : KeepChunkIndexes)
+		std::vector<uint8_t> Chunk;
+		for (const size_t& ChunkIndex : KeepChunkIndexes)
+		{
+			const BlockStoreLocation ChunkLocation = ChunkLocations[ChunkIndex];
+			if (ChunkLocation.Offset + ChunkLocation.Size > OldBlockSize)
 			{
-				const BlockStoreLocation ChunkLocation = ChunkLocations[ChunkIndex];
-				Chunk.resize(ChunkLocation.Size);
-				OldBlockFile->Read(Chunk.data(), Chunk.size(), ChunkLocation.Offset);
+				ZEN_WARN(
+					"Compact Block skipping chunk outside of block range in '{}', Chunk start {}, Chunk size {} in Block {}, Block "
+					"size {}",
+					m_BlocksBasePath,
+					ChunkLocation.Offset,
+					ChunkLocation.Size,
+					OldBlockFile->GetPath(),
+					OldBlockSize);
+				continue;
+			}
+
+			Chunk.resize(ChunkLocation.Size);
+			OldBlockFile->Read(Chunk.data(), Chunk.size(), ChunkLocation.Offset);
 
-				if ((WriteOffset + Chunk.size()) > m_MaxBlockSize)
+			if ((WriteOffset + Chunk.size()) > m_MaxBlockSize)
+			{
+				if (NewBlockFile)
 				{
-					if (NewBlockFile)
-					{
-						NewBlockFile->Flush();
-						MovedSize += NewBlockFile->FileSize();
-						NewBlockFile = nullptr;
+					NewBlockFile->Flush();
+					MovedSize += NewBlockFile->FileSize();
+					NewBlockFile = nullptr;
 
-						ZEN_ASSERT(!MovedChunks.empty() || RemovedSize > 0);  // We should not have a new block if we haven't moved anything
+					ZEN_ASSERT(!MovedChunks.empty() || RemovedSize > 0);  // We should not have a new block if we haven't moved anything
 
-						ChangeCallback(MovedChunks, RemovedSize);
-						DeletedSize += RemovedSize;
-						RemovedSize = 0;
-						MovedCount += MovedChunks.size();
-						MovedChunks.clear();
+					if (!ReportChanges())
+					{
+						return false;
 					}
+				}
 
-					uint32_t NextBlockIndex = m_WriteBlockIndex.load(std::memory_order_relaxed);
+				uint32_t NextBlockIndex = m_WriteBlockIndex.load(std::memory_order_relaxed);
+				{
+					RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
+					std::filesystem::path	   NewBlockPath;
+					NextBlockIndex = GetFreeBlockIndex(NextBlockIndex, InsertLock, NewBlockPath);
+					if (NextBlockIndex == (uint32_t)m_MaxBlockCount)
 					{
-						RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
-						std::filesystem::path	   NewBlockPath;
-						NextBlockIndex = GetFreeBlockIndex(NextBlockIndex, InsertLock, NewBlockPath);
-						if (NextBlockIndex == (uint32_t)m_MaxBlockCount)
-						{
-							ZEN_ERROR("unable to allocate a new block in '{}', count limit {} exeeded",
-									  m_BlocksBasePath,
-									  static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) + 1);
-							return;
-						}
-
-						NewBlockFile				  = new BlockStoreFile(NewBlockPath);
-						m_ChunkBlocks[NextBlockIndex] = NewBlockFile;
+						ZEN_ERROR("unable to allocate a new block in '{}', count limit {} exeeded",
+								  m_BlocksBasePath,
+								  static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) + 1);
+						return false;
 					}
-					ZEN_ASSERT(NewBlockFile);
 
-					std::error_code Error;
-					DiskSpace		Space = DiskSpaceInfo(m_BlocksBasePath, Error);
-					if (Error)
+					NewBlockFile				  = new BlockStoreFile(NewBlockPath);
+					m_ChunkBlocks[NextBlockIndex] = NewBlockFile;
+				}
+				ZEN_ASSERT(NewBlockFile);
+
+				std::error_code Error;
+				DiskSpace		Space = DiskSpaceInfo(m_BlocksBasePath, Error);
+				if (Error)
+				{
+					ZEN_ERROR("get disk space in '{}' FAILED, reason: '{}'", m_BlocksBasePath, Error.message());
 					{
-						ZEN_ERROR("get disk space in '{}' FAILED, reason: '{}'", m_BlocksBasePath, Error.message());
-						return;
+						RwLock::ExclusiveLockScope _l(m_InsertLock);
+						ZEN_ASSERT(m_ChunkBlocks[NextBlockIndex] == NewBlockFile);
+						m_ChunkBlocks.erase(NextBlockIndex);
 					}
+					NewBlockFile->MarkAsDeleteOnClose();
+					NewBlockFile = nullptr;
+					return false;
+				}
 
-					if (Space.Free < m_MaxBlockSize)
+				if (Space.Free < m_MaxBlockSize)
+				{
+					uint64_t ReclaimedSpace = DiskReserveCallback();
+					if (Space.Free + ReclaimedSpace < m_MaxBlockSize)
 					{
-						uint64_t ReclaimedSpace = DiskReserveCallback();
-						if (Space.Free + ReclaimedSpace < m_MaxBlockSize)
-						{
-							ZEN_WARN("garbage collect for '{}' FAILED, required disk space {}, free {}",
-									 m_BlocksBasePath,
-									 m_MaxBlockSize,
-									 NiceBytes(Space.Free + ReclaimedSpace));
-							{
-								RwLock::ExclusiveLockScope _l(m_InsertLock);
-								ZEN_ASSERT(m_ChunkBlocks[NextBlockIndex] == NewBlockFile);
-								m_ChunkBlocks.erase(NextBlockIndex);
-							}
-							NewBlockFile = nullptr;
-							return;
-						}
-
-						ZEN_INFO("using gc reserve for '{}', reclaimed {}, disk free {}",
+						ZEN_WARN("garbage collect for '{}' FAILED, required disk space {}, free {}",
 								 m_BlocksBasePath,
-								 ReclaimedSpace,
+								 m_MaxBlockSize,
 								 NiceBytes(Space.Free + ReclaimedSpace));
+						{
+							RwLock::ExclusiveLockScope _l(m_InsertLock);
+							ZEN_ASSERT(m_ChunkBlocks[NextBlockIndex] == NewBlockFile);
+							m_ChunkBlocks.erase(NextBlockIndex);
+						}
+						NewBlockFile->MarkAsDeleteOnClose();
+						NewBlockFile = nullptr;
+						return false;
 					}
-					NewBlockFile->Create(m_MaxBlockSize);
-					NewBlockIndex = NextBlockIndex;
-					WriteOffset	  = 0;
-				}
 
-				NewBlockFile->Write(Chunk.data(), Chunk.size(), WriteOffset);
-				MovedChunks.push_back({ChunkIndex, {.BlockIndex = NewBlockIndex, .Offset = WriteOffset, .Size = Chunk.size()}});
-				WriteOffset = RoundUp(WriteOffset + Chunk.size(), PayloadAlignment);
+					ZEN_INFO("using gc reserve for '{}', reclaimed {}, disk free {}",
+							 m_BlocksBasePath,
+							 ReclaimedSpace,
+							 NiceBytes(Space.Free + ReclaimedSpace));
+				}
+				NewBlockFile->Create(m_MaxBlockSize);
+				NewBlockIndex = NextBlockIndex;
+				WriteOffset	  = 0;
 			}
-			Chunk.clear();
 
-			// Report what we have moved so we can purge the old block
-			if (!MovedChunks.empty() || RemovedSize > 0)
-			{
-				ChangeCallback(MovedChunks, RemovedSize);
-				DeletedSize += RemovedSize;
-				RemovedSize = 0;
-				MovedCount += MovedChunks.size();
-				MovedChunks.clear();
-			}
+			NewBlockFile->Write(Chunk.data(), ChunkLocation.Size, WriteOffset);
+			MovedChunks.push_back(
+				{ChunkIndex, {.BlockIndex = NewBlockIndex, .Offset = gsl::narrow<uint32_t>(WriteOffset), .Size = ChunkLocation.Size}});
+			WriteOffset = RoundUp(WriteOffset + ChunkLocation.Size, PayloadAlignment);
+			AddedSize += Chunk.size();
+		}
+		Chunk.clear();
+
+		if (!ReportChanges())
+		{
+			return false;
+		}
+
+		{
+			RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
+			ZEN_DEBUG("marking cas block store file '{}' for delete, block #{}", OldBlockFile->GetPath(), BlockIndex);
+			OldBlockFile->MarkAsDeleteOnClose();
+			m_ChunkBlocks.erase(BlockIndex);
+			m_TotalSize.fetch_sub(OldBlockSize);
+			RemovedSize += OldBlockSize;
+		}
+		return true;
+	});
 
-			{
-				RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
-				ZEN_DEBUG("marking cas block store file '{}' for delete, block #{}", OldBlockFile->GetPath(), BlockIndex);
-				OldBlockFile->MarkAsDeleteOnClose();
-				m_ChunkBlocks.erase(BlockIndex);
-				m_TotalSize.fetch_sub(OldBlockSize);
-				RemovedSize += OldBlockSize;
-			}
-		});
 	if (NewBlockFile)
 	{
 		NewBlockFile->Flush();
@@ -1142,14 +1256,7 @@ BlockStore::CompactBlocks(const BlockStoreCompactState&	  CompactState,
 		NewBlockFile = nullptr;
 	}
 
-	if (!MovedChunks.empty() || RemovedSize > 0)
-	{
-		ChangeCallback(MovedChunks, RemovedSize);
-		DeletedSize += RemovedSize;
-		RemovedSize = 0;
-		MovedCount += MovedChunks.size();
-		MovedChunks.clear();
-	}
+	ReportChanges();
 }
 
 const char*
@@ -1175,6 +1282,17 @@ BlockStore::GetBlockPath(const std::filesystem::path& BlocksBasePath, const uint
 	return Path.ToPath();
 }
 
+Ref<BlockStoreFile>
+BlockStore::GetBlockFile(uint32_t BlockIndex)
+{
+	RwLock::SharedLockScope _(m_InsertLock);
+	if (auto It = m_ChunkBlocks.find(BlockIndex); It != m_ChunkBlocks.end())
+	{
+		return It->second;
+	}
+	return {};
+}
+
 #if ZEN_WITH_TESTS
 
 TEST_CASE("blockstore.blockstoredisklocation")
@@ -1293,7 +1411,7 @@ TEST_CASE("blockstore.blockfile")
 }
 
 namespace blockstore::impl {
-	BlockStoreLocation WriteStringAsChunk(BlockStore& Store, std::string_view String, size_t PayloadAlignment)
+	BlockStoreLocation WriteStringAsChunk(BlockStore& Store, std::string_view String, uint32_t PayloadAlignment)
 	{
 		BlockStoreLocation Location;
 		Store.WriteChunk(String.data(), String.length(), PayloadAlignment, [&](const BlockStoreLocation& L) { Location = L; });
@@ -1392,7 +1510,12 @@ TEST_CASE("blockstore.clean.stray.blocks")
 	CHECK(!ThirdChunk);
 
 	// Recreate a fake block for a missing chunk location
-	Store.SyncExistingBlocksOnDisk({FirstChunkLocation, SecondChunkLocation, ThirdChunkLocation});
+	BlockStore::BlockIndexSet KnownBlocks;
+	KnownBlocks.Add(FirstChunkLocation.BlockIndex);
+	KnownBlocks.Add(SecondChunkLocation.BlockIndex);
+	KnownBlocks.Add(ThirdChunkLocation.BlockIndex);
+	Store.SyncExistingBlocksOnDisk(KnownBlocks);
+
 	// 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);
@@ -1760,7 +1883,10 @@ TEST_CASE("blockstore.compact.blocks")
 		Store.CompactBlocks(
 			State,
 			Alignment,
-			[&](const BlockStore::MovedChunksArray&, uint64_t) { CHECK(false); },
+			[&](const BlockStore::MovedChunksArray&, uint64_t) {
+				CHECK(false);
+				return true;
+			},
 			[]() {
 				CHECK(false);
 				return 0;
@@ -1785,6 +1911,7 @@ TEST_CASE("blockstore.compact.blocks")
 			[&](const BlockStore::MovedChunksArray& Moved, uint64_t Removed) {
 				RemovedSize += Removed;
 				CHECK(Moved.empty());
+				return true;
 			},
 			[]() { return 0; });
 		CHECK_EQ(RemovedSize, PreSize);
@@ -1810,6 +1937,7 @@ TEST_CASE("blockstore.compact.blocks")
 			[&](const BlockStore::MovedChunksArray& Moved, uint64_t Removed) {
 				RemovedSize += Removed;
 				CHECK(Moved.empty());
+				return true;
 			},
 			[]() { return 0; });
 		CHECK_EQ(Store.TotalSize() + RemovedSize, PreSize);
@@ -1830,7 +1958,10 @@ TEST_CASE("blockstore.compact.blocks")
 		Store.CompactBlocks(
 			State,
 			Alignment,
-			[&](const BlockStore::MovedChunksArray&, uint64_t) { CHECK(false); },
+			[&](const BlockStore::MovedChunksArray&, uint64_t) {
+				CHECK(false);
+				return true;
+			},
 			[]() {
 				CHECK(false);
 				return 0;
@@ -1862,6 +1993,7 @@ TEST_CASE("blockstore.compact.blocks")
 			[&](const BlockStore::MovedChunksArray& Moved, uint64_t Removed) {
 				CHECK(Moved.empty());
 				RemovedSize += Removed;
+				return true;
 			},
 			[]() {
 				CHECK(false);
@@ -1905,6 +2037,7 @@ TEST_CASE("blockstore.compact.blocks")
 					(*It) = Move.second;
 				}
 				RemovedSize += Removed;
+				return true;
 			},
 			[]() {
 				CHECK(false);
@@ -1981,6 +2114,7 @@ TEST_CASE("blockstore.compact.blocks")
 					(*It) = Move.second;
 				}
 				RemovedSize += Removed;
+				return true;
 			},
 			[]() {
 				CHECK(false);