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, 34 insertions, 0 deletions

diff --git a/src/zenstore/compactcas.cpp b/src/zenstore/compactcas.cpp
index 9982e7571..f85b19264 100644
--- a/src/zenstore/compactcas.cpp
+++ b/src/zenstore/compactcas.cpp
@@ -10,6 +10,7 @@
 #include <zencore/fmtutils.h>
 #include <zencore/iobuffer.h>
 #include <zencore/logging.h>
+#include <zencore/memory/llm.h>
 #include <zencore/scopeguard.h>
 #include <zencore/trace.h>
 #include <zencore/workthreadpool.h>
@@ -36,6 +37,14 @@ ZEN_THIRD_PARTY_INCLUDES_END
 
 namespace zen {
 
+const FLLMTag&
+GetCasContainerTag()
+{
+	static FLLMTag _("container", FLLMTag("cas"));
+
+	return _;
+}
+
 struct CasDiskIndexHeader
 {
 	static constexpr uint32_t ExpectedMagic	 = 0x75696478;	// 'uidx';
@@ -124,6 +133,8 @@ static const float IndexMaxLoadFactor = 0.7f;
 
 CasContainerStrategy::CasContainerStrategy(GcManager& Gc) : m_Log(logging::Get("containercas")), m_Gc(Gc)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 	m_LocationMap.min_load_factor(IndexMinLoadFactor);
 	m_LocationMap.max_load_factor(IndexMaxLoadFactor);
 
@@ -144,6 +155,8 @@ CasContainerStrategy::Initialize(const std::filesystem::path& RootDirectory,
 								 uint32_t					  Alignment,
 								 bool						  IsNewStore)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 	ZEN_ASSERT(IsPow2(Alignment));
 	ZEN_ASSERT(!m_IsInitialized);
 	ZEN_ASSERT(MaxBlockSize > 0);
@@ -171,6 +184,8 @@ CasContainerStrategy::InsertChunk(const void* ChunkData, size_t ChunkSize, const
 		}
 	}
 
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 	// We can end up in a situation that InsertChunk writes the same chunk data in
 	// different locations.
 	// We release the insert lock once we have the correct WriteBlock ready and we know
@@ -201,6 +216,8 @@ CasContainerStrategy::InsertChunk(const void* ChunkData, size_t ChunkSize, const
 CasStore::InsertResult
 CasContainerStrategy::InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 #if !ZEN_WITH_TESTS
 	ZEN_ASSERT(Chunk.GetContentType() == ZenContentType::kCompressedBinary);
 #endif
@@ -210,6 +227,8 @@ CasContainerStrategy::InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash)
 std::vector<CasStore::InsertResult>
 CasContainerStrategy::InsertChunks(std::span<IoBuffer> Chunks, std::span<IoHash> ChunkHashes)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 	ZEN_ASSERT(Chunks.size() == ChunkHashes.size());
 	std::vector<CasStore::InsertResult> Result(Chunks.size());
 	std::vector<size_t>					NewChunkIndexes;
@@ -308,6 +327,8 @@ CasContainerStrategy::IterateChunks(std::span<IoHash>												  ChunkHashes,
 									WorkerThreadPool*												  OptionalWorkerPool,
 									uint64_t														  LargeSizeLimit)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 	const size_t ChunkCount = ChunkHashes.size();
 	if (ChunkCount < 3)
 	{
@@ -407,6 +428,8 @@ CasContainerStrategy::IterateChunks(std::span<IoHash>												  ChunkHashes,
 void
 CasContainerStrategy::Flush()
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 	ZEN_TRACE_CPU("CasContainer::Flush");
 	m_BlockStore.Flush(/*ForceNewBlock*/ false);
 	m_CasLog.Flush();
@@ -416,6 +439,8 @@ CasContainerStrategy::Flush()
 void
 CasContainerStrategy::ScrubStorage(ScrubContext& Ctx)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
+
 	ZEN_TRACE_CPU("CasContainer::ScrubStorage");
 
 	if (Ctx.IsSkipCas())
@@ -563,6 +588,8 @@ public:
 
 	virtual void CompactStore(GcCtx& Ctx, GcCompactStoreStats& Stats, const std::function<uint64_t()>& ClaimDiskReserveCallback) override
 	{
+		ZEN_MEMSCOPE(GetCasContainerTag());
+
 		ZEN_TRACE_CPU("CasContainer::CompactStore");
 
 		auto Log = [&Ctx]() { return Ctx.Logger; };
@@ -717,6 +744,7 @@ public:
 													 GcStats&						Stats,
 													 const GetUnusedReferencesFunc& GetUnusedReferences) override
 	{
+		ZEN_MEMSCOPE(GetCasContainerTag());
 		ZEN_TRACE_CPU("CasContainer::RemoveUnreferencedData");
 
 		auto Log = [&Ctx]() { return Ctx.Logger; };
@@ -799,6 +827,7 @@ CasContainerStrategy::GetGcName(GcCtx&)
 GcReferencePruner*
 CasContainerStrategy::CreateReferencePruner(GcCtx& Ctx, GcReferenceStoreStats&)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
 	ZEN_TRACE_CPU("CasContainer::CreateReferencePruner");
 
 	auto Log = [&Ctx]() { return Ctx.Logger; };
@@ -842,6 +871,7 @@ CasContainerStrategy::CreateReferencePruner(GcCtx& Ctx, GcReferenceStoreStats&)
 void
 CasContainerStrategy::CompactIndex(RwLock::ExclusiveLockScope&)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
 	ZEN_TRACE_CPU("CasContainer::CompactIndex");
 
 	size_t								EntryCount = m_LocationMap.size();
@@ -868,6 +898,7 @@ CasContainerStrategy::StorageSize() const
 void
 CasContainerStrategy::MakeIndexSnapshot()
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
 	ZEN_TRACE_CPU("CasContainer::MakeIndexSnapshot");
 
 	uint64_t LogCount = m_CasLog.GetLogCount();
@@ -981,6 +1012,7 @@ CasContainerStrategy::MakeIndexSnapshot()
 uint64_t
 CasContainerStrategy::ReadIndexFile(const std::filesystem::path& IndexPath, uint32_t& OutVersion)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
 	ZEN_TRACE_CPU("CasContainer::ReadIndexFile");
 
 	uint64_t   EntryCount = 0;
@@ -1050,6 +1082,7 @@ CasContainerStrategy::ReadIndexFile(const std::filesystem::path& IndexPath, uint
 uint64_t
 CasContainerStrategy::ReadLog(const std::filesystem::path& LogPath, uint64_t SkipEntryCount)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
 	ZEN_TRACE_CPU("CasContainer::ReadLog");
 
 	if (!TCasLogFile<CasDiskIndexEntry>::IsValid(LogPath))
@@ -1106,6 +1139,7 @@ CasContainerStrategy::ReadLog(const std::filesystem::path& LogPath, uint64_t Ski
 void
 CasContainerStrategy::OpenContainer(bool IsNewStore)
 {
+	ZEN_MEMSCOPE(GetCasContainerTag());
 	ZEN_TRACE_CPU("CasContainer::OpenContainer");
 
 	// Add .running file and delete on clean on close to detect bad termination