added support for dynamic LLM tags (#245)

* added FLLMTag which can be used to register memory tags outside of core
* changed `UE_MEMSCOPE` -> `ZEN_MEMSCOPE` for consistency
* instrumented some subsystems with dynamic tags

3 files changed, 62 insertions, 0 deletions

diff --git a/src/zenstore/cache/cachedisklayer.cpp b/src/zenstore/cache/cachedisklayer.cpp
index 3046ab87c..a4f9fe78b 100644
--- a/src/zenstore/cache/cachedisklayer.cpp
+++ b/src/zenstore/cache/cachedisklayer.cpp
@@ -22,8 +22,18 @@
 
 //////////////////////////////////////////////////////////////////////////
 
+#include <zencore/memory/llm.h>
+
 namespace zen {
 
+const FLLMTag&
+GetCacheDiskTag()
+{
+	static FLLMTag _("disk", FLLMTag("cache"));
+
+	return _;
+}
+
 namespace {
 
 #pragma pack(push)
@@ -3894,6 +3904,8 @@ ZenCacheDiskLayer::DiscoverBuckets()
 	{
 		WorkLatch.AddCount(1);
 		Pool.ScheduleWork([this, &WorkLatch, &SyncLock, BucketPath]() {
+			ZEN_MEMSCOPE(GetCacheDiskTag());
+
 			auto			  _			 = MakeGuard([&]() { WorkLatch.CountDown(); });
 			const std::string BucketName = PathToUtf8(BucketPath.stem());
 			try
@@ -3980,6 +3992,7 @@ ZenCacheDiskLayer::Drop()
 void
 ZenCacheDiskLayer::Flush()
 {
+	ZEN_MEMSCOPE(GetCacheDiskTag());
 	ZEN_TRACE_CPU("Z$::Flush");
 
 	std::vector<CacheBucket*> Buckets;
@@ -4010,6 +4023,8 @@ ZenCacheDiskLayer::Flush()
 			{
 				WorkLatch.AddCount(1);
 				Pool.ScheduleWork([&WorkLatch, Bucket]() {
+					ZEN_MEMSCOPE(GetCacheDiskTag());
+
 					auto _ = MakeGuard([&]() { WorkLatch.CountDown(); });
 					try
 					{
diff --git a/src/zenstore/cache/cacherpc.cpp b/src/zenstore/cache/cacherpc.cpp
index 2a7721fe2..54c2ff7d0 100644
--- a/src/zenstore/cache/cacherpc.cpp
+++ b/src/zenstore/cache/cacherpc.cpp
@@ -17,10 +17,28 @@
 #include <zenutil/cache/cacherequests.h>
 #include <zenutil/packageformat.h>
 
+#include <zencore/memory/llm.h>
+
 //////////////////////////////////////////////////////////////////////////
 
 namespace zen {
 
+const FLLMTag&
+GetCacheTag()
+{
+	static FLLMTag CacheTag("cache");
+
+	return CacheTag;
+}
+
+const FLLMTag&
+GetCacheRpcTag()
+{
+	static FLLMTag CacheRpcTag("rpc", GetCacheTag());
+
+	return CacheRpcTag;
+}
+
 using namespace std::literals;
 
 std::optional<std::string>
@@ -165,6 +183,8 @@ CacheRpcHandler::HandleRpcRequest(const CacheRequestContext& Context,
 {
 	ZEN_TRACE_CPU("Z$::HandleRpcRequest");
 
+	ZEN_MEMSCOPE(GetCacheRpcTag());
+
 	m_CacheStats.RpcRequests.fetch_add(1);
 
 	CbPackage	 Package;
diff --git a/src/zenstore/cache/structuredcachestore.cpp b/src/zenstore/cache/structuredcachestore.cpp
index 512f1d7f2..c14ea73a8 100644
--- a/src/zenstore/cache/structuredcachestore.cpp
+++ b/src/zenstore/cache/structuredcachestore.cpp
@@ -42,8 +42,20 @@ ZEN_THIRD_PARTY_INCLUDES_END
 #	include <unordered_map>
 #endif
 
+#include <zencore/memory/llm.h>
+
+//////////////////////////////////////////////////////////////////////////
+
 namespace zen {
 
+const FLLMTag&
+GetCacheStoreTag()
+{
+	static FLLMTag _("store", FLLMTag("cache"));
+
+	return _;
+}
+
 bool
 IsKnownBadBucketName(std::string_view Bucket)
 {
@@ -122,6 +134,7 @@ ZenCacheNamespace::ZenCacheNamespace(GcManager& Gc, JobQueue& JobQueue, const st
 , m_Configuration(Config)
 , m_DiskLayer(m_Gc, m_JobQueue, m_RootDir, m_Configuration.DiskLayerConfig)
 {
+	ZEN_MEMSCOPE(GetCacheStoreTag());
 	ZEN_INFO("initializing structured cache at '{}'", m_RootDir);
 	CreateDirectories(m_RootDir);
 
@@ -403,6 +416,8 @@ ZenCacheStore::ZenCacheStore(GcManager&					  Gc,
 , m_Configuration(Configuration)
 , m_ExitLogging(false)
 {
+	ZEN_MEMSCOPE(GetCacheStoreTag());
+
 	SetLoggingConfig(m_Configuration.Logging);
 	CreateDirectories(m_BasePath);
 
@@ -459,6 +474,8 @@ ZenCacheStore::LogWorker()
 {
 	SetCurrentThreadName("ZenCacheStore::LogWorker");
 
+	ZEN_MEMSCOPE(GetCacheStoreTag());
+
 	LoggerRef ZCacheLog(logging::Get("z$"));
 
 	auto Log = [&ZCacheLog]() -> LoggerRef { return ZCacheLog; };
@@ -539,6 +556,7 @@ ZenCacheStore::LogWorker()
 ZenCacheStore::PutBatch::PutBatch(ZenCacheStore& CacheStore, std::string_view InNamespace, std::vector<bool>& OutResult)
 : m_CacheStore(CacheStore)
 {
+	ZEN_MEMSCOPE(GetCacheStoreTag());
 	if (m_Store = m_CacheStore.GetNamespace(InNamespace); m_Store)
 	{
 		m_NamespaceBatchHandle = m_Store->BeginPutBatch(OutResult);
@@ -551,6 +569,7 @@ ZenCacheStore::PutBatch::~PutBatch()
 	{
 		if (m_Store)
 		{
+			ZEN_MEMSCOPE(GetCacheStoreTag());
 			ZEN_ASSERT(m_NamespaceBatchHandle);
 			m_Store->EndPutBatch(m_NamespaceBatchHandle);
 		}
@@ -565,6 +584,7 @@ ZenCacheStore::GetBatch::GetBatch(ZenCacheStore& CacheStore, std::string_view In
 : m_CacheStore(CacheStore)
 , Results(OutResult)
 {
+	ZEN_MEMSCOPE(GetCacheStoreTag());
 	if (m_Store = m_CacheStore.GetNamespace(InNamespace); m_Store)
 	{
 		m_NamespaceBatchHandle = m_Store->BeginGetBatch(OutResult);
@@ -577,6 +597,7 @@ ZenCacheStore::GetBatch::~GetBatch()
 	{
 		if (m_Store)
 		{
+			ZEN_MEMSCOPE(GetCacheStoreTag());
 			ZEN_ASSERT(m_NamespaceBatchHandle);
 			m_Store->EndGetBatch(m_NamespaceBatchHandle);
 
@@ -613,6 +634,7 @@ ZenCacheStore::Get(const CacheRequestContext& Context,
 		return false;
 	}
 
+	ZEN_MEMSCOPE(GetCacheStoreTag());
 	ZEN_TRACE_CPU("Z$::Get");
 
 	metrics::RequestStats::Scope OpScope(m_GetOps, 0);
@@ -673,6 +695,8 @@ ZenCacheStore::Get(const CacheRequestContext& Context,
 		m_RejectedReadCount++;
 		return;
 	}
+
+	ZEN_MEMSCOPE(GetCacheStoreTag());
 	ZEN_TRACE_CPU("Z$::Get");
 
 	metrics::RequestStats::Scope OpScope(m_GetOps, 0);
@@ -709,6 +733,7 @@ ZenCacheStore::Put(const CacheRequestContext& Context,
 		return;
 	}
 
+	ZEN_MEMSCOPE(GetCacheStoreTag());
 	ZEN_TRACE_CPU("Z$::Put");
 
 	metrics::RequestStats::Scope $(m_PutOps, Value.Value.GetSize());
@@ -776,6 +801,8 @@ ZenCacheStore::DropNamespace(std::string_view InNamespace)
 void
 ZenCacheStore::Flush()
 {
+	ZEN_MEMSCOPE(GetCacheStoreTag());
+
 	ZEN_INFO("flushing cache store at '{}'", m_BasePath);
 	IterateNamespaces([&](std::string_view, ZenCacheNamespace& Store) { Store.Flush(); });
 }