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

1 files changed, 27 insertions, 0 deletions

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(); });
 }