Merge branch 'main' into zs/get-all-chunk-infoszs/get-all-chunk-infos

2 files changed, 26 insertions, 16 deletions

diff --git a/src/zenstore/include/zenstore/blockstore.h b/src/zenstore/include/zenstore/blockstore.h
index 919684e41..786780b5e 100644
--- a/src/zenstore/include/zenstore/blockstore.h
+++ b/src/zenstore/include/zenstore/blockstore.h
@@ -132,6 +132,14 @@ 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;
 
+	struct BlockUsageInfo
+	{
+		uint64_t DiskUsage;
+		uint32_t EntryCount;
+	};
+	typedef std::unordered_map<uint32_t, BlockUsageInfo> BlockUsageMap;
+	typedef std::unordered_map<uint32_t, uint32_t>		 BlockEntryCountMap;
+
 	void Initialize(const std::filesystem::path& BlocksBasePath, uint64_t MaxBlockSize, uint64_t MaxBlockCount);
 
 	struct BlockIndexSet
@@ -145,8 +153,8 @@ public:
 
 	// Ask the store to create empty blocks for all locations that does not have a block
 	// Remove any block that is not referenced
-	void				  SyncExistingBlocksOnDisk(const BlockIndexSet& KnownLocations);
-	std::vector<uint32_t> GetBlocksToCompact(const std::unordered_map<uint32_t, uint64_t>& BlockUsage, uint32_t BlockUsageThresholdPercent);
+	void			   SyncExistingBlocksOnDisk(const BlockIndexSet& KnownLocations);
+	BlockEntryCountMap GetBlocksToCompact(const BlockUsageMap& BlockUsage, uint32_t BlockUsageThresholdPercent);
 
 	void Close();
 
@@ -205,23 +213,29 @@ class BlockStoreCompactState
 public:
 	BlockStoreCompactState() = default;
 
-	void IncludeBlocks(const std::span<const uint32_t> BlockIndexes)
+	void IncludeBlocks(const BlockStore::BlockEntryCountMap& BlockEntryCountMap)
 	{
-		for (uint32_t BlockIndex : BlockIndexes)
+		size_t EntryCountTotal = 0;
+		for (auto& BlockUsageIt : BlockEntryCountMap)
 		{
-			auto It = m_BlockIndexToChunkMapIndex.find(BlockIndex);
-			if (It == m_BlockIndexToChunkMapIndex.end())
-			{
-				m_KeepChunks.emplace_back(std::vector<size_t>());
-				m_BlockIndexToChunkMapIndex.insert_or_assign(BlockIndex, m_KeepChunks.size() - 1);
-			}
+			uint32_t BlockIndex = BlockUsageIt.first;
+			ZEN_ASSERT(m_BlockIndexToChunkMapIndex.find(BlockIndex) == m_BlockIndexToChunkMapIndex.end());
+
+			m_KeepChunks.emplace_back(std::vector<size_t>());
+			m_KeepChunks.back().reserve(BlockUsageIt.second);
+			m_BlockIndexToChunkMapIndex.insert_or_assign(BlockIndex, m_KeepChunks.size() - 1);
+			EntryCountTotal += BlockUsageIt.second;
 		}
+		m_ChunkLocations.reserve(EntryCountTotal);
 	}
 
 	void IncludeBlock(uint32_t BlockIndex)
 	{
-		const uint32_t Blocks[1] = {BlockIndex};
-		IncludeBlocks(Blocks);
+		if (m_BlockIndexToChunkMapIndex.find(BlockIndex) == m_BlockIndexToChunkMapIndex.end())
+		{
+			m_KeepChunks.emplace_back(std::vector<size_t>());
+			m_BlockIndexToChunkMapIndex.insert_or_assign(BlockIndex, m_KeepChunks.size() - 1);
+		}
 	}
 
 	bool AddKeepLocation(const BlockStoreLocation& Location)
diff --git a/src/zenstore/include/zenstore/cidstore.h b/src/zenstore/include/zenstore/cidstore.h
index 319683dcb..4c9f30608 100644
--- a/src/zenstore/include/zenstore/cidstore.h
+++ b/src/zenstore/include/zenstore/cidstore.h
@@ -9,10 +9,6 @@
 #include <zenstore/hashkeyset.h>
 #include <zenutil/statsreporter.h>
 
-ZEN_THIRD_PARTY_INCLUDES_START
-#include <tsl/robin_map.h>
-ZEN_THIRD_PARTY_INCLUDES_END
-
 #include <filesystem>
 
 namespace zen {