Added CAS total size.

1 files changed, 70 insertions, 3 deletions

diff --git a/zenstore/compactcas.cpp b/zenstore/compactcas.cpp
index dbe5572b9..584db496b 100644
--- a/zenstore/compactcas.cpp
+++ b/zenstore/compactcas.cpp
@@ -21,6 +21,11 @@
 #include <functional>
 #include <gsl/gsl-lite.hpp>
 
+#if ZEN_WITH_TESTS
+#	include <algorithm>
+#	include <random>
+#endif
+
 //////////////////////////////////////////////////////////////////////////
 
 namespace zen {
@@ -62,9 +67,16 @@ CasContainerStrategy::Initialize(const std::string_view ContainerBaseName, uint6
 		RwLock::ExclusiveLockScope _(m_LocationMapLock);
 
 		m_CasLog.Replay([&](const CasDiskIndexEntry& Record) {
-			m_LocationMap[Record.Key] = Record.Location;
-
-			MaxFileOffset = std::max<uint64_t>(MaxFileOffset, Record.Location.GetOffset() + Record.Location.GetSize());
+			if (Record.Flags & CasDiskIndexEntry::kTombstone)
+			{
+				m_TotalSize.fetch_sub(Record.Location.GetSize());
+			}
+			else
+			{
+				m_TotalSize.fetch_add(Record.Location.GetSize());
+				m_LocationMap[Record.Key] = Record.Location;
+				MaxFileOffset			  = std::max<uint64_t>(MaxFileOffset, Record.Location.GetOffset() + Record.Location.GetSize());
+			}
 		});
 	}
 
@@ -103,6 +115,7 @@ CasContainerStrategy::InsertChunk(const void* ChunkData, size_t ChunkSize, const
 
 	CasDiskIndexEntry IndexEntry{.Key = ChunkHash, .Location = Location};
 
+	m_TotalSize.fetch_add(static_cast<uint64_t>(ChunkSize));
 	m_CasLog.Append(IndexEntry);
 
 	return CasStore::InsertResult{.New = true};
@@ -255,6 +268,7 @@ CasContainerStrategy::Scrub(ScrubContext& Ctx)
 	for (const CasDiskIndexEntry& Entry : BadChunks)
 	{
 		BadChunkHashes.push_back(Entry.Key);
+		m_CasLog.Append({.Key = Entry.Key, .Location = Entry.Location, .Flags = CasDiskIndexEntry::kTombstone});
 		m_LocationMap.erase(Entry.Key);
 	}
 
@@ -359,6 +373,59 @@ TEST_CASE("cas.compact.gc")
 	}
 }
 
+TEST_CASE("cas.compact.totalsize")
+{
+	std::random_device rd;
+	std::mt19937	   g(rd());
+
+	const auto CreateChunk = [&](uint64_t Size) -> IoBuffer {
+		const size_t		  Count = static_cast<size_t>(Size / sizeof(uint32_t));
+		std::vector<uint32_t> Values;
+		Values.resize(Count);
+		for (size_t Idx = 0; Idx < Count; ++Idx)
+		{
+			Values[Idx] = static_cast<uint32_t>(Idx);
+		}
+		std::shuffle(Values.begin(), Values.end(), g);
+
+		return IoBufferBuilder::MakeCloneFromMemory(Values.data(), Values.size() * sizeof(uint32_t));
+	};
+
+	ScopedTemporaryDirectory TempDir;
+
+	CasStoreConfiguration CasConfig;
+	CasConfig.RootDirectory = TempDir.Path();
+
+	CreateDirectories(CasConfig.RootDirectory);
+
+	const uint64_t kChunkSize  = 1024;
+	const int32_t  kChunkCount = 16;
+
+	{
+		CasContainerStrategy Cas(CasConfig);
+		Cas.Initialize("test", 16, true);
+
+		for (int32_t Idx = 0; Idx < kChunkCount; ++Idx)
+		{
+			IoBuffer	 Chunk		  = CreateChunk(kChunkSize);
+			const IoHash Hash		  = HashBuffer(Chunk);
+			auto		 InsertResult = Cas.InsertChunk(Chunk, Hash);
+			ZEN_ASSERT(InsertResult.New);
+		}
+
+		const uint64_t TotalSize = Cas.TotalSize();
+		CHECK_EQ(kChunkSize * kChunkCount, TotalSize);
+	}
+
+	{
+		CasContainerStrategy Cas(CasConfig);
+		Cas.Initialize("test", 16, false);
+
+		const uint64_t TotalSize = Cas.TotalSize();
+		CHECK_EQ(kChunkSize * kChunkCount, TotalSize);
+	}
+}
+
 #endif
 
 void