Added z$ memory/disk layer size.

1 files changed, 151 insertions, 5 deletions

diff --git a/zenserver/cache/structuredcachestore.cpp b/zenserver/cache/structuredcachestore.cpp
index 44226457c..ed6b065b5 100644
--- a/zenserver/cache/structuredcachestore.cpp
+++ b/zenserver/cache/structuredcachestore.cpp
@@ -168,6 +168,12 @@ ZenCacheStore::GatherReferences(GcContext& GcCtx)
 	m_DiskLayer.GatherReferences(GcCtx);
 }
 
+ZenCacheSize
+ZenCacheStore::TotalSize() const
+{
+	return {.MemorySize = m_MemLayer.TotalSize(), .DiskSize = m_DiskLayer.TotalSize()};
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 ZenCacheMemoryLayer::ZenCacheMemoryLayer()
@@ -261,6 +267,20 @@ ZenCacheMemoryLayer::GatherReferences(GcContext& GcCtx)
 	}
 }
 
+uint64_t
+ZenCacheMemoryLayer::TotalSize() const
+{
+	uint64_t				TotalSize{};
+	RwLock::SharedLockScope _(m_Lock);
+
+	for (auto& Kv : m_Buckets)
+	{
+		TotalSize += Kv.second.TotalSize();
+	}
+
+	return TotalSize;
+}
+
 void
 ZenCacheMemoryLayer::CacheBucket::Scrub(ScrubContext& Ctx)
 {
@@ -340,9 +360,12 @@ ZenCacheMemoryLayer::CacheBucket::GetCurrentTimeStamp()
 void
 ZenCacheMemoryLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue& Value)
 {
-	RwLock::ExclusiveLockScope _(m_bucketLock);
+	{
+		RwLock::ExclusiveLockScope _(m_bucketLock);
+		m_cacheMap.insert_or_assign(HashKey, BucketValue{.LastAccess = GetCurrentTimeStamp(), .Payload = Value.Value});
+	}
 
-	m_cacheMap.insert_or_assign(HashKey, BucketValue{.LastAccess = GetCurrentTimeStamp(), .Payload = Value.Value});
+	m_TotalSize.fetch_add(Value.Value.GetSize());
 }
 
 //////////////////////////////////////////////////////////////////////////
@@ -415,7 +438,8 @@ struct ZenCacheDiskLayer::CacheBucket
 	void		Scrub(ScrubContext& Ctx);
 	void		GatherReferences(GcContext& GcCtx);
 
-	inline bool IsOk() const { return m_IsOk; }
+	inline bool		IsOk() const { return m_IsOk; }
+	inline uint64_t TotalSize() const { return m_TotalSize; }
 
 private:
 	std::filesystem::path m_BucketDir;
@@ -431,6 +455,7 @@ private:
 	RwLock												 m_IndexLock;
 	tsl::robin_map<IoHash, DiskLocation, IoHash::Hasher> m_Index;
 	uint64_t											 m_WriteCursor = 0;
+	std::atomic_uint64_t								 m_TotalSize{};
 
 	void BuildPath(WideStringBuilderBase& Path, const IoHash& HashKey);
 	void PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value);
@@ -547,12 +572,14 @@ ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir, bo
 			if (Record.Key == IoHash::Zero)
 			{
 				++InvalidEntryCount;
+				m_TotalSize.fetch_sub(Record.Location.Size());
 			}
 			else
 			{
 				m_Index[Record.Key] = Record.Location;
 
 				MaxFileOffset = std::max<uint64_t>(MaxFileOffset, Record.Location.Offset() + Record.Location.Size());
+				m_TotalSize.fetch_add(Record.Location.Size());
 			}
 		});
 
@@ -686,6 +713,7 @@ ZenCacheDiskLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue&
 
 		m_SlogFile.Append({.Key = HashKey, .Location = Loc});
 		m_SobsFile.Write(Value.Value.Data(), Loc.Size(), Loc.Offset());
+		m_TotalSize.fetch_add(Loc.Size());
 	}
 }
 
@@ -751,7 +779,9 @@ ZenCacheDiskLayer::CacheBucket::Scrub(ScrubContext& Ctx)
 		{
 			// Log a tombstone and delete the in-memory index for the bad entry
 
-			m_SlogFile.Append(DiskIndexEntry{.Key = BadKey, .Location = {0, 0, 0, DiskLocation::kTombStone}});
+			const auto			It		 = m_Index.find(BadKey);
+			const DiskLocation& Location = It->second;
+			m_SlogFile.Append(DiskIndexEntry{.Key = BadKey, .Location = {Location.Offset(), Location.Size(), 0, DiskLocation::kTombStone}});
 			m_Index.erase(BadKey);
 		}
 	}
@@ -891,6 +921,7 @@ ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, c
 	}
 
 	m_SlogFile.Append({.Key = HashKey, .Location = Loc});
+	m_TotalSize.fetch_add(Loc.Size());
 }
 
 //////////////////////////////////////////////////////////////////////////
@@ -1118,11 +1149,25 @@ ZenCacheDiskLayer::GatherReferences(GcContext& GcCtx)
 	}
 }
 
+uint64_t
+ZenCacheDiskLayer::TotalSize() const
+{
+	uint64_t				TotalSize{};
+	RwLock::SharedLockScope _(m_Lock);
+
+	for (auto& Kv : m_Buckets)
+	{
+		TotalSize += Kv.second.TotalSize();
+	}
+
+	return TotalSize;
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 #if ZEN_WITH_TESTS
 
-TEST_CASE("z$.store")
+TEST_CASE("zcache.store")
 {
 	using namespace fmt::literals;
 	using namespace std::literals;
@@ -1165,6 +1210,107 @@ TEST_CASE("z$.store")
 	}
 }
 
+TEST_CASE("zcache.size")
+{
+	using namespace fmt::literals;
+	using namespace std::literals;
+
+	const auto CreateCacheValue = [](size_t Size) -> CbObject {
+		std::vector<uint8_t> Buf;
+		Buf.resize(Size);
+
+		CbObjectWriter Writer;
+		Writer.AddBinary("Binary"sv, Buf.data(), Buf.size());
+		return Writer.Save();
+	};
+
+	SUBCASE("mem/disklayer")
+	{
+		const size_t			 Count = 16;
+		ScopedTemporaryDirectory TempDir;
+
+		ZenCacheSize CacheSize;
+
+		{
+			CasGc		  Gc;
+			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");
+
+			CbObject CacheValue = CreateCacheValue(Zcs.DiskLayerThreshold() - 256);
+
+			IoBuffer Buffer = CacheValue.GetBuffer().AsIoBuffer();
+			Buffer.SetContentType(ZenContentType::kCbObject);
+
+			for (size_t Key = 0; Key < Count; ++Key)
+			{
+				const size_t Bucket = Key % 4;
+				Zcs.Put("test_bucket-{}"_format(Bucket), IoHash::HashBuffer(&Key, sizeof(uint32_t)), {.Value = Buffer});
+			}
+
+			CacheSize = Zcs.TotalSize();
+			CHECK_EQ(CacheValue.GetSize() * Count, CacheSize.DiskSize);
+			CHECK_EQ(CacheValue.GetSize() * Count, CacheSize.MemorySize);
+		}
+
+		{
+			CasGc		  Gc;
+			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");
+
+			const ZenCacheSize SerializedSize = Zcs.TotalSize();
+			CHECK_EQ(SerializedSize.MemorySize, 0);
+			CHECK_EQ(SerializedSize.DiskSize, CacheSize.DiskSize);
+
+			for (size_t Bucket = 0; Bucket < 4; ++Bucket)
+			{
+				Zcs.DropBucket("test_bucket-{}"_format(Bucket));
+			}
+			CHECK_EQ(0, Zcs.TotalSize().DiskSize);
+		}
+	}
+
+	SUBCASE("disklayer")
+	{
+		const size_t			 Count = 16;
+		ScopedTemporaryDirectory TempDir;
+
+		ZenCacheSize CacheSize;
+
+		{
+			CasGc		  Gc;
+			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");
+
+			CbObject CacheValue = CreateCacheValue(Zcs.DiskLayerThreshold() + 64);
+
+			IoBuffer Buffer = CacheValue.GetBuffer().AsIoBuffer();
+			Buffer.SetContentType(ZenContentType::kCbObject);
+
+			for (size_t Key = 0; Key < Count; ++Key)
+			{
+				const size_t Bucket = Key % 4;
+				Zcs.Put("test_bucket-{}"_format(Bucket), IoHash::HashBuffer(&Key, sizeof(uint32_t)), {.Value = Buffer});
+			}
+
+			CacheSize = Zcs.TotalSize();
+			CHECK_EQ(CacheValue.GetSize() * Count, CacheSize.DiskSize);
+			CHECK_EQ(0, CacheSize.MemorySize);
+		}
+
+		{
+			CasGc		  Gc;
+			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");
+
+			const ZenCacheSize SerializedSize = Zcs.TotalSize();
+			CHECK_EQ(SerializedSize.MemorySize, 0);
+			CHECK_EQ(SerializedSize.DiskSize, CacheSize.DiskSize);
+
+			for (size_t Bucket = 0; Bucket < 4; ++Bucket)
+			{
+				Zcs.DropBucket("test_bucket-{}"_format(Bucket));
+			}
+			CHECK_EQ(0, Zcs.TotalSize().DiskSize);
+		}
+	}
+}
+
 #endif
 
 void