Added CAS total size.

6 files changed, 130 insertions, 12 deletions

diff --git a/zenstore/CAS.cpp b/zenstore/CAS.cpp
index 86c6eb849..d2ff1514e 100644
--- a/zenstore/CAS.cpp
+++ b/zenstore/CAS.cpp
@@ -105,6 +105,7 @@ public:
 	virtual void				   Flush() override;
 	virtual void				   Scrub(ScrubContext& Ctx) override;
 	virtual void				   GarbageCollect(GcContext& GcCtx) override;
+	virtual CasStoreSize		   TotalSize() const override;
 
 private:
 	CasContainerStrategy m_TinyStrategy;
@@ -318,6 +319,16 @@ CasImpl::GarbageCollect(GcContext& GcCtx)
 	m_LargeStrategy.CollectGarbage(GcCtx);
 }
 
+CasStoreSize
+CasImpl::TotalSize() const
+{
+	const uint64_t Tiny	 = m_TinyStrategy.TotalSize();
+	const uint64_t Small = m_SmallStrategy.TotalSize();
+	const uint64_t Large = m_LargeStrategy.TotalSize();
+
+	return {.TinySize = Tiny, .SmallSize = Small, .LargeSize = Large, .TotalSize = Tiny + Small + Large};
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 CasStore*
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
diff --git a/zenstore/compactcas.h b/zenstore/compactcas.h
index a3f3121e6..9154768b3 100644
--- a/zenstore/compactcas.h
+++ b/zenstore/compactcas.h
@@ -55,6 +55,8 @@ private:
 
 struct CasDiskIndexEntry
 {
+	static const uint8_t kTombstone = 0x01;
+
 	IoHash			Key;
 	CasDiskLocation Location;
 	ZenContentType	ContentType = ZenContentType::kUnknownContentType;
@@ -79,7 +81,7 @@ struct CasContainerStrategy
 	~CasContainerStrategy();
 
 	CasStore::InsertResult InsertChunk(const void* ChunkData, size_t ChunkSize, const IoHash& ChunkHash);
-	CasStore::InsertResult InsertChunk(IoBuffer Chunk, const IoHash& chunkHash);
+	CasStore::InsertResult InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash);
 	IoBuffer			   FindChunk(const IoHash& ChunkHash);
 	bool				   HaveChunk(const IoHash& ChunkHash);
 	void				   FilterChunks(CasChunkSet& InOutChunks);
@@ -87,6 +89,7 @@ struct CasContainerStrategy
 	void				   Flush();
 	void				   Scrub(ScrubContext& Ctx);
 	void				   CollectGarbage(GcContext& GcCtx);
+	uint64_t			   TotalSize() const { return m_TotalSize; }
 
 private:
 	const CasStoreConfiguration&   m_Config;
@@ -100,8 +103,9 @@ private:
 	RwLock														m_LocationMapLock;
 	std::unordered_map<IoHash, CasDiskLocation, IoHash::Hasher> m_LocationMap;
 	RwLock														m_InsertLock;  // used to serialize inserts
-	std::atomic<uint64_t>										m_CurrentInsertOffset = 0;
-	std::atomic<uint64_t>										m_CurrentIndexOffset  = 0;
+	std::atomic_uint64_t										m_CurrentInsertOffset{};
+	std::atomic_uint64_t										m_CurrentIndexOffset{};
+	std::atomic_uint64_t										m_TotalSize{};
 
 	void MakeSnapshot();
 };
diff --git a/zenstore/filecas.cpp b/zenstore/filecas.cpp
index 0efdc96f5..3c4add82a 100644
--- a/zenstore/filecas.cpp
+++ b/zenstore/filecas.cpp
@@ -90,7 +90,7 @@ FileCasStrategy::Initialize(bool IsNewStore)
 
 	m_CasLog.Open(m_Config.RootDirectory / "cas.ulog", IsNewStore);
 
-	m_CasLog.Replay([&](const FileCasIndexEntry& Entry) { ZEN_UNUSED(Entry); });
+	m_CasLog.Replay([&](const FileCasIndexEntry& Entry) { m_TotalSize.fetch_add(Entry.Size); });
 }
 
 CasStore::InsertResult
@@ -138,6 +138,16 @@ FileCasStrategy::InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash)
 
 				// We do need to ensure the source file goes away on close, however
 
+				uint64_t FileSize = 0;
+				if (HRESULT hSizeRes = PayloadFile.GetSize(FileSize); SUCCEEDED(hSizeRes))
+				{
+					m_TotalSize.fetch_add(static_cast<int64_t>(FileSize));
+				}
+				else
+				{
+					ZEN_WARN("get file size FAILED, file cas '{}'", WideToUtf8(Name.ShardedPath));
+				}
+
 				DeletePayloadFileOnClose();
 
 				return CasStore::InsertResult{.New = false};
@@ -229,7 +239,8 @@ FileCasStrategy::InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash)
 
 		if (Success)
 		{
-			m_CasLog.Append({.Key = ChunkHash, .Size = Chunk.Size()});
+			m_TotalSize.fetch_add(static_cast<int64_t>(Chunk.Size()));
+			m_CasLog.Append({.Key = ChunkHash, .Size = static_cast<int64_t>(Chunk.Size())});
 
 			return CasStore::InsertResult{.New = true};
 		}
@@ -272,6 +283,8 @@ FileCasStrategy::InsertChunk(const void* const ChunkData, const size_t ChunkSize
 	{
 		// If we succeeded in opening the file then we don't need to do anything else because it already exists and should contain the
 		// content we were about to insert
+
+		m_TotalSize.fetch_add(static_cast<int64_t>(ChunkSize));
 		return CasStore::InsertResult{.New = false};
 	}
 
@@ -287,6 +300,8 @@ FileCasStrategy::InsertChunk(const void* const ChunkData, const size_t ChunkSize
 	{
 		// If we succeeded in opening the file then we don't need to do anything
 		// else because someone else managed to create the file before we did. Just return.
+
+		m_TotalSize.fetch_add(static_cast<int64_t>(ChunkSize));
 		return {.New = false};
 	}
 
@@ -330,7 +345,8 @@ FileCasStrategy::InsertChunk(const void* const ChunkData, const size_t ChunkSize
 	// *after* the lock is released due to the initialization order
 	PayloadFile.Close();
 
-	m_CasLog.Append({.Key = ChunkHash, .Size = ChunkSize});
+	m_TotalSize.fetch_add(static_cast<int64_t>(ChunkSize));
+	m_CasLog.Append({.Key = ChunkHash, .Size = static_cast<int64_t>(ChunkSize)});
 
 	return {.New = true};
 }
@@ -370,13 +386,21 @@ FileCasStrategy::DeleteChunk(const IoHash& ChunkHash, std::error_code& Ec)
 {
 	ShardingHelper Name(m_Config.RootDirectory.c_str(), ChunkHash);
 
-	ZEN_DEBUG("deleting CAS payload file '{}'", WideToUtf8(Name.ShardedPath));
+	uint64_t FileSize = static_cast<uint64_t>(std::filesystem::file_size(Name.ShardedPath.c_str(), Ec));
+	if (Ec)
+	{
+		ZEN_WARN("get file size FAILED, file cas '{}'", WideToUtf8(Name.ShardedPath));
+		FileSize = 0;
+	}
+
+	ZEN_DEBUG("deleting CAS payload file '{}' {}", WideToUtf8(Name.ShardedPath), NiceBytes(FileSize));
 
 	std::filesystem::remove(Name.ShardedPath.c_str(), Ec);
 
 	if (!Ec)
 	{
-		m_CasLog.Append({.Key = ChunkHash, .Size = ~(0ull)});
+		m_TotalSize.fetch_sub(FileSize);
+		m_CasLog.Append({.Key = ChunkHash, .Size = -static_cast<int64_t>(FileSize)});
 	}
 }
 
diff --git a/zenstore/filecas.h b/zenstore/filecas.h
index ec2ca3f31..06f35785b 100644
--- a/zenstore/filecas.h
+++ b/zenstore/filecas.h
@@ -12,6 +12,7 @@
 #include <zenstore/caslog.h>
 #include <zenstore/gc.h>
 
+#include <atomic>
 #include <functional>
 
 namespace spdlog {
@@ -39,6 +40,7 @@ struct FileCasStrategy : public GcStorage
 	void				   Flush();
 	virtual void		   CollectGarbage(GcContext& GcCtx) override;
 	void				   Scrub(ScrubContext& Ctx);
+	uint64_t			   TotalSize() const { return static_cast<uint64_t>(m_TotalSize); }
 
 private:
 	const CasStoreConfiguration& m_Config;
@@ -46,13 +48,14 @@ private:
 	RwLock						 m_ShardLocks[256];	 // TODO: these should be spaced out so they don't share cache lines
 	spdlog::logger&				 m_Log;
 	spdlog::logger&				 Log() { return m_Log; }
+	std::atomic_int64_t			 m_TotalSize{};
 	bool						 m_IsInitialized = false;
 
 	struct FileCasIndexEntry
 	{
 		IoHash	 Key;
 		uint32_t Pad  = 0;
-		uint64_t Size = 0;
+		int64_t	 Size = 0;
 	};
 
 	static_assert(sizeof(FileCasIndexEntry) == 32);
diff --git a/zenstore/include/zenstore/CAS.h b/zenstore/include/zenstore/CAS.h
index 5b508baa0..bba1bb721 100644
--- a/zenstore/include/zenstore/CAS.h
+++ b/zenstore/include/zenstore/CAS.h
@@ -91,6 +91,14 @@ private:
 	CasChunkSet			  m_BadCid;
 };
 
+struct CasStoreSize
+{
+	uint64_t TinySize{};
+	uint64_t SmallSize{};
+	uint64_t LargeSize{};
+	uint64_t TotalSize{};
+};
+
 /** Content Addressable Storage interface
 
   */
@@ -114,6 +122,7 @@ public:
 	virtual void		 Flush()											 = 0;
 	virtual void		 Scrub(ScrubContext& Ctx)							 = 0;
 	virtual void		 GarbageCollect(GcContext& GcCtx)					 = 0;
+	virtual CasStoreSize TotalSize() const									 = 0;
 
 protected:
 	CasStoreConfiguration m_Config;