Refactor WriteChunk to not need callback

3 files changed, 22 insertions, 36 deletions

diff --git a/zenstore/blockstore.cpp b/zenstore/blockstore.cpp
index d293eb97d..178687b5f 100644
--- a/zenstore/blockstore.cpp
+++ b/zenstore/blockstore.cpp
@@ -193,8 +193,8 @@ BlockStore::Initialize(const std::filesystem::path&			  BlocksBasePath,
 	}
 }
 
-void
-BlockStore::WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment, WriteChunkCallback Callback)
+BlockStoreLocation
+BlockStore::WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment)
 {
 	RwLock::ExclusiveLockScope InsertLock(m_InsertLock);
 
@@ -227,15 +227,11 @@ BlockStore::WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment, Writ
 	uint64_t InsertOffset		   = m_CurrentInsertOffset;
 	m_CurrentInsertOffset		   = RoundUp(InsertOffset + Size, Alignment);
 	Ref<BlockStoreFile> WriteBlock = m_WriteBlock;
-	m_ActiveWriteBlockIndexes.push_back(WriteBlockIndex);
 	InsertLock.ReleaseNow();
 
 	WriteBlock->Write(Data, Size, InsertOffset);
 
-	Callback({.BlockIndex = WriteBlockIndex, .Offset = InsertOffset, .Size = Size});
-
-	RwLock::ExclusiveLockScope _(m_InsertLock);
-	m_ActiveWriteBlockIndexes.erase(std::find(m_ActiveWriteBlockIndexes.begin(), m_ActiveWriteBlockIndexes.end(), WriteBlockIndex));
+	return {.BlockIndex = WriteBlockIndex, .Offset = InsertOffset, .Size = Size};
 }
 
 BlockStore::ReclaimSnapshotState
@@ -243,11 +239,8 @@ BlockStore::GetReclaimSnapshotState()
 {
 	ReclaimSnapshotState	   State;
 	RwLock::ExclusiveLockScope _(m_InsertLock);
-	for (uint32_t BlockIndex : m_ActiveWriteBlockIndexes)
-	{
-		State.ExcludeBlockIndexes.insert(BlockIndex);
-	}
-	State.BlockCount = m_ChunkBlocks.size();
+	State.ExcludeBlockIndex = m_WriteBlock ? m_WriteBlockIndex.load(std::memory_order_acquire) : 0xffffffffu;
+	State.BlockCount		= m_ChunkBlocks.size();
 	_.ReleaseNow();
 	return State;
 }
@@ -351,7 +344,7 @@ BlockStore::ReclaimSpace(const ReclaimSnapshotState&			Snapshot,
 	{
 		const BlockStoreLocation& Location = ChunkLocations[Index];
 		OldTotalSize += Location.Size;
-		if (Snapshot.ExcludeBlockIndexes.contains(Location.BlockIndex))
+		if (Location.BlockIndex == Snapshot.ExcludeBlockIndex)
 		{
 			continue;
 		}
diff --git a/zenstore/compactcas.cpp b/zenstore/compactcas.cpp
index 5a0ba974b..c6115024b 100644
--- a/zenstore/compactcas.cpp
+++ b/zenstore/compactcas.cpp
@@ -250,16 +250,15 @@ CasContainerStrategy::InsertChunk(const void* ChunkData, size_t ChunkSize, const
 	// This should be a rare occasion and the current flow reduces the time we block for
 	// reads, insert and GC.
 
-	m_BlockStore.WriteChunk(ChunkData, ChunkSize, m_PayloadAlignment, [this, &ChunkHash, ChunkSize](const BlockStoreLocation& Location) {
-		BlockStoreDiskLocation	DiskLocation(Location, m_PayloadAlignment);
-		const CasDiskIndexEntry IndexEntry{.Key = ChunkHash, .Location = DiskLocation};
-		m_CasLog.Append(IndexEntry);
-		{
-			RwLock::ExclusiveLockScope _(m_LocationMapLock);
-			m_LocationMap.emplace(ChunkHash, DiskLocation);
-		}
-		m_TotalSize.fetch_add(static_cast<uint64_t>(ChunkSize), std::memory_order_seq_cst);
-	});
+	BlockStoreLocation		Location = m_BlockStore.WriteChunk(ChunkData, ChunkSize, m_PayloadAlignment);
+	BlockStoreDiskLocation	DiskLocation(Location, m_PayloadAlignment);
+	const CasDiskIndexEntry IndexEntry{.Key = ChunkHash, .Location = DiskLocation};
+	m_CasLog.Append(IndexEntry);
+	{
+		RwLock::ExclusiveLockScope _(m_LocationMapLock);
+		m_LocationMap.emplace(ChunkHash, DiskLocation);
+	}
+	m_TotalSize.fetch_add(static_cast<uint64_t>(ChunkSize), std::memory_order_seq_cst);
 
 	return CasStore::InsertResult{.New = true};
 }
diff --git a/zenstore/include/zenstore/blockstore.h b/zenstore/include/zenstore/blockstore.h
index 4a1642413..5019e257d 100644
--- a/zenstore/include/zenstore/blockstore.h
+++ b/zenstore/include/zenstore/blockstore.h
@@ -114,8 +114,8 @@ class BlockStore
 public:
 	struct ReclaimSnapshotState
 	{
-		std::unordered_set<uint32_t> ExcludeBlockIndexes;
-		size_t						 BlockCount;
+		size_t ExcludeBlockIndex;
+		size_t BlockCount;
 	};
 
 	typedef std::vector<std::pair<size_t, BlockStoreLocation>> MovedChunksArray;
@@ -123,7 +123,6 @@ public:
 
 	typedef std::function<void(const MovedChunksArray& MovedChunks, const ChunkIndexArray& RemovedChunks)> ReclaimCallback;
 	typedef std::function<uint64_t()>																	   ClaimDiskReserveCallback;
-	typedef std::function<void(const BlockStoreLocation& Location)>										   WriteChunkCallback;
 	typedef std::function<void(size_t ChunkIndex, const void* Data, uint64_t Size)>						   IterateChunksSmallSizeCallback;
 	typedef std::function<void(size_t ChunkIndex, BasicFile& BlockFile, uint64_t Offset, uint64_t Size)>   IterateChunksLargeSizeCallback;
 	typedef std::function<void(const MovedChunksArray& MovedChunks)>									   SplitCallback;
@@ -133,11 +132,7 @@ public:
 					uint64_t							   MaxBlockCount,
 					const std::vector<BlockStoreLocation>& KnownLocations);
 
-	void WriteChunk(
-		const void*		   Data,
-		uint64_t		   Size,
-		uint64_t		   Alignment,
-		WriteChunkCallback Callback = [](const BlockStoreLocation&) {});
+	BlockStoreLocation WriteChunk(const void* Data, uint64_t Size, uint64_t Alignment);
 
 	Ref<BlockStoreFile> GetChunkBlock(const BlockStoreLocation& Location);
 	void				Flush();
@@ -171,11 +166,10 @@ public:
 private:
 	std::unordered_map<uint32_t, Ref<BlockStoreFile>> m_ChunkBlocks;
 
-	RwLock				  m_InsertLock;	 // used to serialize inserts
-	Ref<BlockStoreFile>	  m_WriteBlock;
-	std::uint64_t		  m_CurrentInsertOffset = 0;
-	std::atomic_uint32_t  m_WriteBlockIndex{};
-	std::vector<uint32_t> m_ActiveWriteBlockIndexes;
+	RwLock				 m_InsertLock;	// used to serialize inserts
+	Ref<BlockStoreFile>	 m_WriteBlock;
+	std::uint64_t		 m_CurrentInsertOffset = 0;
+	std::atomic_uint32_t m_WriteBlockIndex{};
 
 	uint64_t			  m_MaxBlockSize  = 1u << 28;
 	uint64_t			  m_MaxBlockCount = BlockStoreDiskLocation::MaxBlockIndex + 1;