cleanup duplicate code in CollectGarbage

4 files changed, 50 insertions, 55 deletions

diff --git a/zenstore/blockstore.cpp b/zenstore/blockstore.cpp
index 1aebae0f3..e865ede8a 100644
--- a/zenstore/blockstore.cpp
+++ b/zenstore/blockstore.cpp
@@ -4,8 +4,6 @@
 
 #include <zenstore/blockstore.h>
 
-#include <zencore/filesystem.h>
-
 #if ZEN_WITH_TESTS
 #	include <zencore/compactbinarybuilder.h>
 #	include <zencore/testing.h>
diff --git a/zenstore/compactcas.cpp b/zenstore/compactcas.cpp
index fbdac491e..dc16305d0 100644
--- a/zenstore/compactcas.cpp
+++ b/zenstore/compactcas.cpp
@@ -341,6 +341,26 @@ CasContainerStrategy::Scrub(ScrubContext& Ctx)
 }
 
 void
+CasContainerStrategy::UpdateLocationMap(const std::unordered_map<IoHash, BlockStoreDiskLocation>& MovedChunks,
+										const std::unordered_set<IoHash, IoHash::Hasher>&		  DeletedChunks)
+{
+	for (const auto& MovedEntry : MovedChunks)
+	{
+		m_LocationMap[MovedEntry.first] = MovedEntry.second;
+		m_CasLog.Append({.Key = MovedEntry.first, .Location = MovedEntry.second});
+	}
+
+	for (const auto& ChunkHash : DeletedChunks)
+	{
+		auto	 KeyIt	   = m_LocationMap.find(ChunkHash);
+		uint64_t ChunkSize = KeyIt->second.GetSize();
+		m_CasLog.Append({.Key = ChunkHash, .Location = KeyIt->second, .Flags = CasDiskIndexEntry::kTombstone});
+		m_TotalSize.fetch_sub(ChunkSize);
+		m_LocationMap.erase(KeyIt);
+	}
+}
+
+void
 CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 {
 	namespace fs = std::filesystem;
@@ -483,20 +503,13 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 		if (KeepMap.empty())
 		{
 			std::shared_ptr<BlockStoreFile> BlockFile;
+			const auto&						DeleteMap = DeleteChunks[ChunkMapIndex];
 			{
 				RwLock::ExclusiveLockScope _i(m_LocationMapLock);
-				const auto&				   DeleteMap = DeleteChunks[ChunkMapIndex];
-				for (const auto& ChunkHash : DeleteMap)
-				{
-					auto	 KeyIt	   = m_LocationMap.find(ChunkHash);
-					uint64_t ChunkSize = KeyIt->second.GetSize();
-					m_CasLog.Append({.Key = ChunkHash, .Location = KeyIt->second, .Flags = CasDiskIndexEntry::kTombstone});
-					m_TotalSize.fetch_sub(ChunkSize);
-					m_LocationMap.erase(KeyIt);
-				}
-				DeletedChunks.insert(DeletedChunks.end(), DeleteMap.begin(), DeleteMap.end());
+				UpdateLocationMap({}, DeleteMap);
 				m_ChunkBlocks[BlockIndex].swap(BlockFile);
 			}
+			DeletedChunks.insert(DeletedChunks.end(), DeleteMap.begin(), DeleteMap.end());
 			ZEN_DEBUG("marking cas store file for delete {}, block {}", m_ContainerBaseName, std::to_string(BlockIndex));
 			std::error_code Ec;
 			BlockFile->MarkAsDeleteOnClose(Ec);
@@ -528,25 +541,13 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 				uint32_t NextBlockIndex = m_WriteBlockIndex.load(std::memory_order::memory_order_relaxed);
 				{
 					RwLock::ExclusiveLockScope _l(m_LocationMapLock);
-					if (NewBlockFile)
-					{
-						for (const auto& MovedEntry : MovedBlockChunks)
-						{
-							m_LocationMap[MovedEntry.first] = MovedEntry.second;
-							m_CasLog.Append({.Key = MovedEntry.first, .Location = MovedEntry.second});
-							MovedChunks.push_back(MovedEntry.first);
-						}
-						if (m_ChunkBlocks.size() == BlockStoreDiskLocation::MaxBlockIndex)
-						{
-							ZEN_ERROR("unable to allocate a new block in {}, count limit {} exeeded",
-									  m_ContainerBaseName,
-									  static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) + 1);
-							return;
-						}
-					}
+					UpdateLocationMap(MovedBlockChunks, {});
 					if (m_ChunkBlocks.size() == BlockStoreDiskLocation::MaxBlockIndex)
 					{
-						throw std::runtime_error(fmt::format("unable to allocate a new block in {}", m_ContainerBaseName));
+						ZEN_ERROR("unable to allocate a new block in {}, count limit {} exeeded",
+								  m_ContainerBaseName,
+								  static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) + 1);
+						return;
 					}
 					while (m_ChunkBlocks.contains(NextBlockIndex))
 					{
@@ -557,6 +558,12 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 					m_ChunkBlocks[NextBlockIndex] = NewBlockFile;
 				}
 
+				for (const auto& MovedEntry : MovedBlockChunks)
+				{
+					MovedChunks.push_back(MovedEntry.first);
+				}
+				MovedBlockChunks.clear();
+
 				std::error_code Error;
 				DiskSpace		Space = DiskSpaceInfo(m_Config.RootDirectory, Error);
 				if (Error)
@@ -591,8 +598,7 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 					NewBlockFile->Create(m_MaxBlockSize);
 				}
 				NewBlockIndex = NextBlockIndex;
-				MovedBlockChunks.clear();
-				WriteOffset = 0;
+				WriteOffset	  = 0;
 			}
 
 			NewBlockFile->Write(Chunk.data(), Chunk.size(), WriteOffset);
@@ -602,32 +608,19 @@ CasContainerStrategy::CollectGarbage(GcContext& GcCtx)
 		}
 		Chunk.clear();
 
-		if (NewBlockFile)
+		const auto& DeleteMap = DeleteChunks[ChunkMapIndex];
 		{
-			const auto& DeleteMap = DeleteChunks[ChunkMapIndex];
-			// Remap moved chunks to the new block file
-			{
-				RwLock::ExclusiveLockScope _l(m_LocationMapLock);
-				for (const auto& MovedEntry : MovedBlockChunks)
-				{
-					m_LocationMap[MovedEntry.first] = MovedEntry.second;
-					m_CasLog.Append({.Key = MovedEntry.first, .Location = MovedEntry.second});
-					MovedChunks.push_back(MovedEntry.first);
-				}
-
-				for (const auto& ChunkHash : DeleteMap)
-				{
-					auto	 KeyIt	   = m_LocationMap.find(ChunkHash);
-					uint64_t ChunkSize = KeyIt->second.GetSize();
-					m_CasLog.Append({.Key = ChunkHash, .Location = KeyIt->second, .Flags = CasDiskIndexEntry::kTombstone});
-					m_TotalSize.fetch_sub(ChunkSize);
-					m_LocationMap.erase(KeyIt);
-				}
-				m_ChunkBlocks[BlockIndex].reset();
-			}
-			MovedBlockChunks.clear();
-			DeletedChunks.insert(DeletedChunks.end(), DeleteMap.begin(), DeleteMap.end());
+			RwLock::ExclusiveLockScope _i(m_LocationMapLock);
+			UpdateLocationMap(MovedBlockChunks, DeleteMap);
+			m_ChunkBlocks[BlockIndex].reset();
+		}
+		for (const auto& MovedEntry : MovedBlockChunks)
+		{
+			MovedChunks.push_back(MovedEntry.first);
 		}
+		DeletedChunks.insert(DeletedChunks.end(), DeleteMap.begin(), DeleteMap.end());
+		MovedBlockChunks.clear();
+
 		ZEN_DEBUG("marking cas store file for delete {}, block index {}", m_ContainerBaseName, BlockIndex);
 		std::error_code Ec;
 		OldBlockFile->MarkAsDeleteOnClose(Ec);
diff --git a/zenstore/compactcas.h b/zenstore/compactcas.h
index 42f4b3f8f..72808358d 100644
--- a/zenstore/compactcas.h
+++ b/zenstore/compactcas.h
@@ -64,6 +64,8 @@ struct CasContainerStrategy final : public GcStorage
 private:
 	void			OpenContainer(bool IsNewStore);
 	spdlog::logger& Log() { return m_Log; }
+	void			UpdateLocationMap(const std::unordered_map<IoHash, BlockStoreDiskLocation>& MovedChunks,
+									  const std::unordered_set<IoHash, IoHash::Hasher>&			DeletedChunks);
 
 	const CasStoreConfiguration&   m_Config;
 	spdlog::logger&				   m_Log;
diff --git a/zenstore/include/zenstore/blockstore.h b/zenstore/include/zenstore/blockstore.h
index 29f520722..69ad38d87 100644
--- a/zenstore/include/zenstore/blockstore.h
+++ b/zenstore/include/zenstore/blockstore.h
@@ -2,7 +2,9 @@
 
 #pragma once
 
+#include <zencore/filesystem.h>
 #include <zencore/zencore.h>
+#include <zenstore/basicfile.h>
 
 namespace zen {