From cbb9ed149517cf781bf21bff4650d7d01bd6d567 Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Fri, 30 Aug 2024 11:26:42 +0200
Subject: meta info store (#75)

- Feature: Added option `--gc-cache-attachment-store` which caches referenced attachments in cache records on disk for faster GC - default is `false`
- Feature: Added option `--gc-projectstore-attachment-store` which caches referenced attachments in project store oplogs on disk for faster GC - default is `false`
---
 src/zenstore/cache/cachedisklayer.cpp | 129 +++++++++++++++++++++++++++++-----
 1 file changed, 113 insertions(+), 16 deletions(-)

(limited to 'src/zenstore/cache/cachedisklayer.cpp')

diff --git a/src/zenstore/cache/cachedisklayer.cpp b/src/zenstore/cache/cachedisklayer.cpp
index f85d05dce..940f78c30 100644
--- a/src/zenstore/cache/cachedisklayer.cpp
+++ b/src/zenstore/cache/cachedisklayer.cpp
@@ -14,6 +14,7 @@
 #include <zencore/trace.h>
 #include <zencore/workthreadpool.h>
 #include <zencore/xxhash.h>
+#include <zenutil/referencemetadata.h>
 #include <zenutil/workerpools.h>
 
 #include <future>
@@ -109,6 +110,8 @@ namespace {
 
 	static_assert(sizeof(BucketMetaHeader) == 32);
 
+	static constexpr uint32_t BlockMetaDataExpectedMagic = 0x61'74'6d'62;  // 'bmta';
+
 #pragma pack(pop)
 
 	//////////////////////////////////////////////////////////////////////////
@@ -3409,15 +3412,59 @@ ZenCacheDiskLayer::CacheBucket::RemoveExpiredData(GcCtx& Ctx, GcStats& Stats)
 	return new DiskBucketStoreCompactor(*this, std::move(ExpiredStandaloneKeys));
 }
 
+bool
+ZenCacheDiskLayer::CacheBucket::ReadAttachmentsFromMetaData(uint32_t					 BlockIndex,
+															std::span<const IoHash>		 InlineKeys,
+															std::span<const std::size_t> ChunkIndexes,
+															std::vector<IoHash>&		 OutReferences) const
+{
+	ZEN_TRACE_CPU("Z$::Bucket::ReadAttachmentsFromMetaData");
+	IoBuffer MetaDataPayload = m_BlockStore.GetMetaData(BlockIndex);
+	if (MetaDataPayload)
+	{
+		std::unordered_set<IoHash, IoHash::Hasher> WantedKeys;
+		WantedKeys.reserve(ChunkIndexes.size());
+		for (const size_t ChunkIndex : ChunkIndexes)
+		{
+			WantedKeys.insert(InlineKeys[ChunkIndex]);
+		}
+		ZEN_TRACE_CPU("Z$::Bucket::GetAttachmentsFromMetaData");
+		return GetAttachmentsFromMetaData<IoHash, IoHash>(
+			MetaDataPayload,
+			BlockMetaDataExpectedMagic,
+			[&](std::span<const IoHash> Keys, std::span<const uint32_t> AttachmentCounts, std::span<const IoHash> Attachments) {
+				OutReferences.reserve(OutReferences.capacity() + Attachments.size());
+				auto AttachmentStart = Attachments.begin();
+				for (uint32_t Index = 0; Index < Keys.size(); Index++)
+				{
+					uint32_t AttachmentCount = AttachmentCounts[Index];
+					if (AttachmentCount > 0)
+					{
+						if (WantedKeys.contains(Keys[Index]))
+						{
+							OutReferences.insert(OutReferences.end(), AttachmentStart, AttachmentStart + AttachmentCount);
+						}
+						AttachmentStart += AttachmentCount;
+					}
+				}
+			});
+	}
+	return false;
+}
+
 bool
 ZenCacheDiskLayer::CacheBucket::GetReferencesLocked(GcCtx& Ctx, std::vector<IoHash>& OutReferences)
 {
-	auto GetAttachments = [&](MemoryView Data) {
+	ZEN_TRACE_CPU("Z$::Bucket::GetReferencesLocked");
+
+	auto GetAttachments = [&](MemoryView Data) -> bool {
 		if (ValidateCompactBinary(Data, CbValidateMode::Default) == CbValidateError::None)
 		{
 			CbObjectView Obj(Data.GetData());
 			Obj.IterateAttachments([&](CbFieldView Field) { OutReferences.emplace_back(Field.AsAttachment()); });
+			return true;
 		}
+		return false;
 	};
 
 	std::vector<std::pair<IoHash, DiskLocation>> StandaloneKeys;
@@ -3471,26 +3518,77 @@ ZenCacheDiskLayer::CacheBucket::GetReferencesLocked(GcCtx& Ctx, std::vector<IoHa
 		{
 			ZEN_ASSERT(!ChunkIndexes.empty());
 
-			bool Continue = m_BlockStore.IterateBlock(
-				InlineLocations,
-				ChunkIndexes,
-				[&](size_t ChunkIndex, const void* Data, uint64_t Size) {
-					ZEN_UNUSED(ChunkIndex);
-					GetAttachments(MemoryView(Data, Size));
-					return !Ctx.IsCancelledFlag.load();
-				},
-				[&](size_t ChunkIndex, BlockStoreFile& File, uint64_t Offset, uint64_t Size) {
-					ZEN_UNUSED(ChunkIndex);
-					GetAttachments(File.GetChunk(Offset, Size).GetView());
-					return !Ctx.IsCancelledFlag.load();
-				});
+			uint32_t BlockIndex = InlineLocations[ChunkIndexes[0]].BlockIndex;
+			if (!m_Configuration.StoreAttachmentMetaData ||
+				!ReadAttachmentsFromMetaData(BlockIndex, InlineKeys, ChunkIndexes, OutReferences))
+			{
+				std::vector<IoHash>	  Keys;
+				std::vector<uint32_t> AttachmentCounts;
+				size_t				  PrecachedReferencesStart	   = OutReferences.size();
+				size_t				  NextPrecachedReferencesStart = PrecachedReferencesStart;
 
-			if (!Continue && Ctx.IsCancelledFlag.load())
+				bool WriteMetaData = m_Configuration.StoreAttachmentMetaData && !m_BlockStore.IsWriting(BlockIndex);
+				if (WriteMetaData)
+				{
+					Keys.reserve(InlineLocations.size());
+				}
+
+				auto CaptureAttachments = [&](size_t ChunkIndex, MemoryView Data) {
+					if (GetAttachments(Data))
+					{
+						size_t AttachmentCount = OutReferences.size() - NextPrecachedReferencesStart;
+						if (WriteMetaData && AttachmentCount > 0)
+						{
+							Keys.push_back(InlineKeys[ChunkIndex]);
+							AttachmentCounts.push_back(gsl::narrow<uint32_t>(AttachmentCount));
+							NextPrecachedReferencesStart += AttachmentCount;
+						}
+					}
+				};
+
+				bool Continue = m_BlockStore.IterateBlock(
+					InlineLocations,
+					ChunkIndexes,
+					[&](size_t ChunkIndex, const void* Data, uint64_t Size) {
+						ZEN_UNUSED(ChunkIndex);
+						CaptureAttachments(ChunkIndex, MemoryView(Data, Size));
+						return !Ctx.IsCancelledFlag.load();
+					},
+					[&](size_t ChunkIndex, BlockStoreFile& File, uint64_t Offset, uint64_t Size) {
+						ZEN_UNUSED(ChunkIndex);
+						CaptureAttachments(ChunkIndex, File.GetChunk(Offset, Size).GetView());
+						return !Ctx.IsCancelledFlag.load();
+					});
+
+				if (Continue)
+				{
+					if (WriteMetaData)
+					{
+						ZEN_ASSERT(Keys.size() == AttachmentCounts.size());
+						IoBuffer MetaDataPayload =
+							BuildReferenceMetaData<IoHash>(
+								BlockMetaDataExpectedMagic,
+								Keys,
+								AttachmentCounts,
+								std::span<const IoHash>(OutReferences)
+									.subspan(PrecachedReferencesStart, OutReferences.size() - PrecachedReferencesStart))
+								.Flatten()
+								.AsIoBuffer();
+						m_BlockStore.SetMetaData(BlockIndex, MetaDataPayload);
+					}
+				}
+				else
+				{
+					return false;
+				}
+			}
+			if (Ctx.IsCancelledFlag.load())
 			{
 				return false;
 			}
 		}
 	}
+
 	for (const auto& It : StandaloneKeys)
 	{
 		if (Ctx.IsCancelledFlag.load())
@@ -3551,7 +3649,6 @@ public:
 		RwLock::SharedLockScope IndexLock(m_CacheBucket.m_IndexLock);
 		bool					Continue = m_CacheBucket.GetReferencesLocked(Ctx, m_References);
 		IndexLock.ReleaseNow();
-
 		if (!Continue)
 		{
 			m_CacheBucket.m_IndexLock.WithExclusiveLock([&]() { m_CacheBucket.m_TrackedReferences.reset(); });
-- 
cgit v1.2.3