Added z$ memory/disk layer size.

3 files changed, 178 insertions, 18 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
diff --git a/zenserver/cache/structuredcachestore.h b/zenserver/cache/structuredcachestore.h
index 8e1260b52..d3c2227ac 100644
--- a/zenserver/cache/structuredcachestore.h
+++ b/zenserver/cache/structuredcachestore.h
@@ -16,6 +16,7 @@ ZEN_THIRD_PARTY_INCLUDES_START
 #include <tsl/robin_map.h>
 ZEN_THIRD_PARTY_INCLUDES_END
 
+#include <atomic>
 #include <compare>
 #include <filesystem>
 #include <unordered_map>
@@ -64,11 +65,12 @@ public:
 	ZenCacheMemoryLayer();
 	~ZenCacheMemoryLayer();
 
-	bool Get(std::string_view Bucket, const IoHash& HashKey, ZenCacheValue& OutValue);
-	void Put(std::string_view Bucket, const IoHash& HashKey, const ZenCacheValue& Value);
-	bool DropBucket(std::string_view Bucket);
-	void Scrub(ScrubContext& Ctx);
-	void GatherReferences(GcContext& GcCtx);
+	bool	 Get(std::string_view Bucket, const IoHash& HashKey, ZenCacheValue& OutValue);
+	void	 Put(std::string_view Bucket, const IoHash& HashKey, const ZenCacheValue& Value);
+	bool	 DropBucket(std::string_view Bucket);
+	void	 Scrub(ScrubContext& Ctx);
+	void	 GatherReferences(GcContext& GcCtx);
+	uint64_t TotalSize() const;
 
 	struct Configuration
 	{
@@ -91,16 +93,18 @@ private:
 		RwLock								m_bucketLock;
 		tsl::robin_map<IoHash, BucketValue> m_cacheMap;
 
-		bool Get(const IoHash& HashKey, ZenCacheValue& OutValue);
-		void Put(const IoHash& HashKey, const ZenCacheValue& Value);
-		void Scrub(ScrubContext& Ctx);
-		void GatherReferences(GcContext& GcCtx);
+		bool			Get(const IoHash& HashKey, ZenCacheValue& OutValue);
+		void			Put(const IoHash& HashKey, const ZenCacheValue& Value);
+		void			Scrub(ScrubContext& Ctx);
+		void			GatherReferences(GcContext& GcCtx);
+		inline uint64_t TotalSize() const { return m_TotalSize; }
 
 	private:
-		uint64_t GetCurrentTimeStamp();
+		uint64_t			 GetCurrentTimeStamp();
+		std::atomic_uint64_t m_TotalSize{};
 	};
 
-	RwLock										 m_Lock;
+	mutable RwLock								 m_Lock;
 	std::unordered_map<std::string, CacheBucket> m_Buckets;
 	Configuration								 m_Configuration;
 
@@ -121,7 +125,8 @@ public:
 	void Scrub(ScrubContext& Ctx);
 	void GatherReferences(GcContext& GcCtx);
 
-	void DiscoverBuckets();
+	void	 DiscoverBuckets();
+	uint64_t TotalSize() const;
 
 private:
 	/** A cache bucket manages a single directory containing
@@ -130,13 +135,19 @@ private:
 	struct CacheBucket;
 
 	std::filesystem::path						 m_RootDir;
-	RwLock										 m_Lock;
+	mutable RwLock								 m_Lock;
 	std::unordered_map<std::string, CacheBucket> m_Buckets;	 // TODO: make this case insensitive
 
 	ZenCacheDiskLayer(const ZenCacheDiskLayer&) = delete;
 	ZenCacheDiskLayer& operator=(const ZenCacheDiskLayer&) = delete;
 };
 
+struct ZenCacheSize
+{
+	uint64_t MemorySize{};
+	uint64_t DiskSize{};
+};
+
 class ZenCacheStore : public GcContributor
 {
 public:
@@ -149,6 +160,8 @@ public:
 	void		 Flush();
 	void		 Scrub(ScrubContext& Ctx);
 	virtual void GatherReferences(GcContext& GcCtx) override;
+	ZenCacheSize TotalSize() const;
+	uint64_t	 DiskLayerThreshold() const { return m_DiskLayerSizeThreshold; }
 
 private:
 	std::filesystem::path			 m_RootDir;
diff --git a/zenserver/zenserver.cpp b/zenserver/zenserver.cpp
index b1bcb545f..47c79526b 100644
--- a/zenserver/zenserver.cpp
+++ b/zenserver/zenserver.cpp
@@ -924,6 +924,7 @@ test_main(int argc, char** argv)
 	zen::zencore_forcelinktests();
 	zen::zenhttp_forcelinktests();
 	zen::zenstore_forcelinktests();
+	zen::z$_forcelink();
 
 	zen::logging::InitializeLogging();