Compact algorithm that can be interrupted and still regain space at end

1 files changed, 156 insertions, 81 deletions

diff --git a/zenstore/compactcas.cpp b/zenstore/compactcas.cpp
index a2806a7fe..3541f22ee 100644
--- a/zenstore/compactcas.cpp
+++ b/zenstore/compactcas.cpp
@@ -75,6 +75,11 @@ CasContainerStrategy::InsertChunk(const void* ChunkData, size_t ChunkSize, const
 
 	const uint64_t InsertOffset = m_CurrentInsertOffset;
 	m_SmallObjectFile.Write(ChunkData, ChunkSize, InsertOffset);
+	auto VerifyChunkHash = IoHash::HashBuffer(IoBuffer(IoBuffer::Wrap, ChunkData, ChunkSize));
+	if (VerifyChunkHash != ChunkHash)
+	{
+		ZEN_ASSERT(false);
+	}
 
 	m_CurrentInsertOffset = (m_CurrentInsertOffset + ChunkSize + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
 
@@ -255,10 +260,6 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 {
 	namespace fs = std::filesystem;
 
-	// A naive garbage collection implementation that just copies evicted chunks
-	// into a new container file. We probably need to partition the container file
-	// into several parts to prevent needing to keep the entire container file during GC.
-
 	ZEN_INFO("collecting garbage from '{}'", m_Config.RootDirectory / m_ContainerBaseName);
 
 	std::vector<CasDiskLocation>			   ChunkLocations;	// Sorted by position
@@ -266,47 +267,53 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 	std::unordered_set<IoHash, IoHash::Hasher> ChunksToKeep;
 	const uint64_t							   ChunkCount = m_LocationMap.size();
 	uint64_t								   TotalSize{};
-	{
-		RwLock::ExclusiveLockScope _(m_LocationMapLock);
 
-		Flush();
+	RwLock::ExclusiveLockScope _i(m_InsertLock);
+	RwLock::ExclusiveLockScope _l(m_LocationMapLock);
 
-		ChunkLocations.reserve(m_LocationMap.size());
-		ChunkHashes.reserve(m_LocationMap.size());
+	Flush();
 
-		for (auto& Entry : m_LocationMap)
-		{
-			ChunkHashes.push_back(Entry.first);
-			TotalSize += Entry.second.GetSize();
-		}
+	if (m_LocationMap.empty())
+	{
+		ZEN_INFO("garbage collect SKIPPED, for '{}', container is empty", m_Config.RootDirectory / m_ContainerBaseName);
+		return;
+	}
 
-		GcCtx.FilterCas(ChunkHashes, [&ChunksToKeep](const IoHash& Hash, bool Keep) {
-			if (Keep)
-			{
-				ChunksToKeep.insert(Hash);
-			}
-		});
+	ChunkLocations.reserve(m_LocationMap.size());
+	ChunkHashes.reserve(m_LocationMap.size());
+
+	for (auto& Entry : m_LocationMap)
+	{
+		ChunkHashes.push_back(Entry.first);
+		TotalSize += Entry.second.GetSize();
+	}
 
-		if (m_LocationMap.empty() || ChunksToKeep.size() == m_LocationMap.size())
+	GcCtx.FilterCas(ChunkHashes, [&ChunksToKeep](const IoHash& Hash, bool Keep) {
+		if (Keep)
 		{
-			ZEN_INFO("garbage collect DONE, scanned #{} {} chunks from '{}', nothing to delete",
-						ChunkCount,
-						NiceBytes(TotalSize),
-						m_Config.RootDirectory / m_ContainerBaseName);
-			return;
+			ChunksToKeep.insert(Hash);
 		}
+	});
 
-		std::sort(begin(ChunkHashes), end(ChunkHashes), [&](IoHash Lhs, IoHash Rhs) {
-			auto LhsKeyIt = m_LocationMap.find(Lhs);
-			auto RhsKeyIt = m_LocationMap.find(Rhs);
-			return LhsKeyIt->second.GetOffset() < RhsKeyIt->second.GetOffset();
-		});
+	if (ChunksToKeep.size() == m_LocationMap.size())
+	{
+		ZEN_INFO("garbage collect DONE, scanned #{} {} chunks from '{}', nothing to delete",
+				 ChunkCount,
+				 NiceBytes(TotalSize),
+				 m_Config.RootDirectory / m_ContainerBaseName);
+		return;
+	}
 
-		for (auto Entry : ChunkHashes)
-		{
-			auto KeyIt = m_LocationMap.find(Entry);
-			ChunkLocations.push_back(KeyIt->second);
-		}
+	std::sort(begin(ChunkHashes), end(ChunkHashes), [&](IoHash Lhs, IoHash Rhs) {
+		auto LhsKeyIt = m_LocationMap.find(Lhs);
+		auto RhsKeyIt = m_LocationMap.find(Rhs);
+		return LhsKeyIt->second.GetOffset() < RhsKeyIt->second.GetOffset();
+	});
+
+	for (auto Entry : ChunkHashes)
+	{
+		auto KeyIt = m_LocationMap.find(Entry);
+		ChunkLocations.push_back(KeyIt->second);
 	}
 
 	const uint64_t NewChunkCount = ChunksToKeep.size();
@@ -323,79 +330,147 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 	if (!CollectSmallObjects)
 	{
 		ZEN_INFO("garbage collect from '{}' DISABLED, found #{} {} chunks of total #{} {}",
-					m_Config.RootDirectory / m_ContainerBaseName,
-					ChunkCount - NewChunkCount,
-					NiceBytes(TotalSize - NewTotalSize),
-					ChunkCount,
-					NiceBytes(TotalSize));
+				 m_Config.RootDirectory / m_ContainerBaseName,
+				 ChunkCount - NewChunkCount,
+				 NiceBytes(TotalSize - NewTotalSize),
+				 ChunkCount,
+				 NiceBytes(TotalSize));
 		return;
 	}
 
 	std::vector<IoHash> DeletedChunks;
 	std::vector<IoHash> MovedChunks;
-	uint64_t			WriteOffset{};
-	for (uint64_t ChunkIndex = 0; ChunkIndex < ChunkCount; ChunkIndex++)
+
+	uint64_t WriteOffset{};
+	uint64_t ChunkIndex{};
+	while (ChunkIndex < ChunkHashes.size())
 	{
-		IoHash		ChunkHash		  = ChunkHashes[ChunkIndex];
-		const auto& ChunkLocation	  = ChunkLocations[ChunkIndex];
-		uint64_t	ChunkEnd		  = ChunkLocation.GetOffset() + ChunkLocation.GetSize();
-		uint64_t	NextChunkWritePos = (ChunkEnd + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
-		bool		KeepChunk		  = ChunksToKeep.end() != ChunksToKeep.find(ChunkHash);
-		if (KeepChunk && WriteOffset == ChunkLocation.GetOffset())
-		{
-			WriteOffset = ChunkLocation.GetOffset() + ChunkLocation.GetSize();
-			continue;
-		}
+		IoHash		ChunkHash	  = ChunkHashes[ChunkIndex];
+		const auto& ChunkLocation = ChunkLocations[ChunkIndex];
+		bool		KeepChunk	  = ChunksToKeep.end() != ChunksToKeep.find(ChunkHash);
 
-		RwLock::ExclusiveLockScope _(m_LocationMapLock);
+		uint64_t NextWriteOffset = ChunkLocation.GetOffset() + ChunkLocation.GetSize();
+		NextWriteOffset			 = (NextWriteOffset + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
 
-		// Verify if we should still keep the chunk
-		std::vector<IoHash> Chunks;
-		GcCtx.FilterCas(ChunkHashes, [&KeepChunk](const IoHash& /*Hash*/, bool Keep) { KeepChunk = Keep; });
 		if (!KeepChunk)
 		{
 			DeletedChunks.push_back(ChunkHash);
 			m_CasLog.Append({.Key = ChunkHash, .Location = ChunkLocation, .Flags = CasDiskIndexEntry::kTombstone});
 			m_LocationMap.erase(ChunkHash);
-			if (m_CurrentInsertOffset == NextChunkWritePos)
+			if (m_CurrentInsertOffset == NextWriteOffset)
 			{
 				m_CurrentInsertOffset = WriteOffset;
 			}
 			m_TotalSize.fetch_sub(static_cast<uint64_t>(ChunkLocation.GetSize()));
+			ChunkIndex++;
 			continue;
 		}
 
-		// Move the chunk
-		std::vector<uint8_t> Chunk;
-		m_SmallObjectFile.Read(Chunk.data(), Chunk.size(), ChunkLocation.GetOffset());
-		CasDiskLocation NewChunkLocation(WriteOffset, ChunkLocation.GetSize());
-		m_SmallObjectFile.Write(Chunk.data(), Chunk.size(), NewChunkLocation.GetOffset());
+		uint64_t FreeChunkSize = ChunkLocation.GetOffset() - WriteOffset;
+
+		// TODO: We could keep some wiggle room here, only try to find the last keep block if there is a reasonable amount of space free
+		while (FreeChunkSize >= m_PayloadAlignment)
+		{
+			// We should move as many keep chunk at the end as we can possibly fit
+			uint64_t LastKeepChunkIndex = ChunkHashes.size() - 1;
+			while (LastKeepChunkIndex > ChunkIndex)
+			{
+				IoHash LastChunkHash = ChunkHashes[LastKeepChunkIndex];
+				bool   LastKeepChunk = ChunksToKeep.end() != ChunksToKeep.find(LastChunkHash);
+				if (LastKeepChunk)
+				{
+					break;
+				}
+				LastKeepChunkIndex--;
+			}
+
+			if (LastKeepChunkIndex == ChunkIndex)
+			{
+				break;
+			}
+
+			IoHash		LastChunkHash	  = ChunkHashes[LastKeepChunkIndex];
+			const auto& LastChunkLocation = ChunkLocations[LastKeepChunkIndex];
+			if (LastChunkLocation.GetSize() > FreeChunkSize)
+			{
+				break;
+			}
+
+			// Move the last chunk to our write location
+			std::vector<uint8_t> Chunk;
+			Chunk.resize(LastChunkLocation.GetSize());
+			m_SmallObjectFile.Read(Chunk.data(), Chunk.size(), LastChunkLocation.GetOffset());
+			auto VerifyChunkHash = IoHash::HashBuffer(IoBuffer(IoBuffer::Wrap, Chunk.data(), Chunk.size()));
+			if (VerifyChunkHash != LastChunkHash)
+			{
+				ZEN_ASSERT(false);
+			}
+			CasDiskLocation NewChunkLocation(WriteOffset, LastChunkLocation.GetSize());
+			m_SmallObjectFile.Write(Chunk.data(), Chunk.size(), NewChunkLocation.GetOffset());
+
+			CasDiskIndexEntry IndexEntry{.Key = ChunkHash, .Location = NewChunkLocation};
+			m_CasLog.Append(IndexEntry);
+			m_LocationMap[LastChunkHash] = NewChunkLocation;
+			ChunkHashes.pop_back();
+
+			uint64_t OldNextChunkWritePos = LastChunkLocation.GetOffset() + Chunk.size();
+			OldNextChunkWritePos		  = (OldNextChunkWritePos + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
+
+			if (m_CurrentInsertOffset == OldNextChunkWritePos)
+			{
+				m_CurrentInsertOffset = LastChunkLocation.GetOffset();
+			}
+
+			WriteOffset	  = (WriteOffset + NewChunkLocation.GetSize() + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
+			FreeChunkSize = ChunkLocation.GetOffset() - WriteOffset;
+			MovedChunks.push_back(LastChunkHash);
+		}
+
+		// TODO: We could keep some wiggle room here, don't move chunk if we only move it a very small amount
+		if (FreeChunkSize > m_PayloadAlignment)
+		{
+			std::vector<uint8_t> Chunk;
+			Chunk.resize(ChunkLocation.GetSize());
+			m_SmallObjectFile.Read(Chunk.data(), Chunk.size(), ChunkLocation.GetOffset());
+			auto VerifyChunkHash = IoHash::HashBuffer(IoBuffer(IoBuffer::Wrap, Chunk.data(), Chunk.size()));
+			if (VerifyChunkHash != ChunkHash)
+			{
+				ZEN_ASSERT(false);
+			}
+			CasDiskLocation NewChunkLocation(WriteOffset, ChunkLocation.GetSize());
+			m_SmallObjectFile.Write(Chunk.data(), Chunk.size(), NewChunkLocation.GetOffset());
 
-		CasDiskIndexEntry IndexEntry{.Key = ChunkHash, .Location = NewChunkLocation};
-		m_CasLog.Append(IndexEntry);
+			CasDiskIndexEntry IndexEntry{.Key = ChunkHash, .Location = NewChunkLocation};
+			m_CasLog.Append(IndexEntry);
+			m_LocationMap[ChunkHash] = NewChunkLocation;
 
-		MovedChunks.push_back(ChunkHash);
-		m_LocationMap[ChunkHash] = NewChunkLocation;
+			uint64_t ChunkEnd		   = ChunkLocation.GetOffset() + ChunkLocation.GetSize();
+			uint64_t NextChunkWritePos = (ChunkEnd + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
+			// Update insert location if this is the last chunk in the file
 
-		WriteOffset = (WriteOffset + ChunkLocation.GetSize() + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
+			if (m_CurrentInsertOffset == NextChunkWritePos)
+			{
+				uint64_t NewChunkEnd  = NewChunkLocation.GetOffset() + NewChunkLocation.GetSize();
+				m_CurrentInsertOffset = (NewChunkEnd + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
+			}
 
-		// Update insert location if this is the last chunk in the file
-		if (m_CurrentInsertOffset == NextChunkWritePos)
+			MovedChunks.push_back(ChunkHash);
+			WriteOffset = (WriteOffset + ChunkLocation.GetSize() + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
+		}
+		else
 		{
-			uint64_t NewChunkEnd  = NewChunkLocation.GetOffset() + NewChunkLocation.GetSize();
-			m_CurrentInsertOffset = (NewChunkEnd + m_PayloadAlignment - 1) & ~(m_PayloadAlignment - 1);
+			WriteOffset = NextWriteOffset;
 		}
+		ChunkIndex++;
 	}
+
+	uint64_t CurrentSize = m_SmallObjectFile.FileSize();
+	if (CurrentSize > m_CurrentInsertOffset)
 	{
-		RwLock::ExclusiveLockScope _(m_LocationMapLock);
-		uint64_t				   CurrentSize = m_SmallObjectFile.FileSize();
-		if (CurrentSize > m_CurrentInsertOffset)
-		{
-			ZEN_INFO("truncate '{}' DISABLED, from {} to {}",
-						m_Config.RootDirectory / m_ContainerBaseName,
-						NiceBytes(CurrentSize),
-						NiceBytes(m_CurrentInsertOffset));
-		}
+		ZEN_INFO("new write position '{}', from {} to {}",
+				 m_Config.RootDirectory / m_ContainerBaseName,
+				 NiceBytes(CurrentSize),
+				 NiceBytes(m_CurrentInsertOffset));
 	}
 }