Restructured structured cache store (#314)

This change separates out the disk and memory storage strategies into separate cpp/h files to improve maintainability.

1 files changed, 317 insertions, 0 deletions

diff --git a/src/zenserver/cache/cachememorylayer.cpp b/src/zenserver/cache/cachememorylayer.cpp
new file mode 100644
index 000000000..1ef581ba5
--- /dev/null
+++ b/src/zenserver/cache/cachememorylayer.cpp
@@ -0,0 +1,317 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "cachememorylayer.h"
+
+#include <zencore/compactbinaryvalidation.h>
+#include <zencore/compress.h>
+
+//////////////////////////////////////////////////////////////////////////
+
+namespace zen {
+
+ZenCacheMemoryLayer::ZenCacheMemoryLayer()
+{
+}
+
+ZenCacheMemoryLayer::~ZenCacheMemoryLayer()
+{
+}
+
+bool
+ZenCacheMemoryLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCacheValue& OutValue)
+{
+	RwLock::SharedLockScope _(m_Lock);
+
+	auto It = m_Buckets.find(std::string(InBucket));
+
+	if (It == m_Buckets.end())
+	{
+		return false;
+	}
+
+	CacheBucket* Bucket = It->second.get();
+
+	_.ReleaseNow();
+
+	// There's a race here. Since the lock is released early to allow
+	// inserts, the bucket delete path could end up deleting the
+	// underlying data structure
+
+	return Bucket->Get(HashKey, OutValue);
+}
+
+void
+ZenCacheMemoryLayer::Put(std::string_view InBucket, const IoHash& HashKey, const ZenCacheValue& Value)
+{
+	const auto	 BucketName = std::string(InBucket);
+	CacheBucket* Bucket		= nullptr;
+
+	{
+		RwLock::SharedLockScope _(m_Lock);
+
+		if (auto It = m_Buckets.find(std::string(InBucket)); It != m_Buckets.end())
+		{
+			Bucket = It->second.get();
+		}
+	}
+
+	if (Bucket == nullptr)
+	{
+		// New bucket
+
+		RwLock::ExclusiveLockScope _(m_Lock);
+
+		if (auto It = m_Buckets.find(std::string(InBucket)); It != m_Buckets.end())
+		{
+			Bucket = It->second.get();
+		}
+		else
+		{
+			auto InsertResult = m_Buckets.emplace(BucketName, std::make_unique<CacheBucket>());
+			Bucket			  = InsertResult.first->second.get();
+		}
+	}
+
+	// Note that since the underlying IoBuffer is retained, the content type is also
+
+	Bucket->Put(HashKey, Value);
+}
+
+bool
+ZenCacheMemoryLayer::DropBucket(std::string_view InBucket)
+{
+	RwLock::ExclusiveLockScope _(m_Lock);
+
+	auto It = m_Buckets.find(std::string(InBucket));
+
+	if (It != m_Buckets.end())
+	{
+		CacheBucket& Bucket = *It->second;
+		m_DroppedBuckets.push_back(std::move(It->second));
+		m_Buckets.erase(It);
+		Bucket.Drop();
+		return true;
+	}
+	return false;
+}
+
+void
+ZenCacheMemoryLayer::Drop()
+{
+	RwLock::ExclusiveLockScope				  _(m_Lock);
+	std::vector<std::unique_ptr<CacheBucket>> Buckets;
+	Buckets.reserve(m_Buckets.size());
+	while (!m_Buckets.empty())
+	{
+		const auto&	 It		= m_Buckets.begin();
+		CacheBucket& Bucket = *It->second;
+		m_DroppedBuckets.push_back(std::move(It->second));
+		m_Buckets.erase(It->first);
+		Bucket.Drop();
+	}
+}
+
+void
+ZenCacheMemoryLayer::ScrubStorage(ScrubContext& Ctx)
+{
+	RwLock::SharedLockScope _(m_Lock);
+
+	for (auto& Kv : m_Buckets)
+	{
+		Kv.second->ScrubStorage(Ctx);
+	}
+}
+
+void
+ZenCacheMemoryLayer::GatherAccessTimes(zen::access_tracking::AccessTimes& AccessTimes)
+{
+	using namespace zen::access_tracking;
+
+	RwLock::SharedLockScope _(m_Lock);
+
+	for (auto& Kv : m_Buckets)
+	{
+		std::vector<KeyAccessTime>& Bucket = AccessTimes.Buckets[Kv.first];
+		Kv.second->GatherAccessTimes(Bucket);
+	}
+}
+
+void
+ZenCacheMemoryLayer::Reset()
+{
+	RwLock::ExclusiveLockScope _(m_Lock);
+	m_Buckets.clear();
+}
+
+uint64_t
+ZenCacheMemoryLayer::TotalSize() const
+{
+	uint64_t				TotalSize{};
+	RwLock::SharedLockScope _(m_Lock);
+
+	for (auto& Kv : m_Buckets)
+	{
+		TotalSize += Kv.second->TotalSize();
+	}
+
+	return TotalSize;
+}
+
+ZenCacheMemoryLayer::Info
+ZenCacheMemoryLayer::GetInfo() const
+{
+	ZenCacheMemoryLayer::Info Info = {.Config = m_Configuration, .TotalSize = TotalSize()};
+
+	RwLock::SharedLockScope _(m_Lock);
+	Info.BucketNames.reserve(m_Buckets.size());
+	for (auto& Kv : m_Buckets)
+	{
+		Info.BucketNames.push_back(Kv.first);
+		Info.EntryCount += Kv.second->EntryCount();
+	}
+	return Info;
+}
+
+std::optional<ZenCacheMemoryLayer::BucketInfo>
+ZenCacheMemoryLayer::GetBucketInfo(std::string_view Bucket) const
+{
+	RwLock::SharedLockScope _(m_Lock);
+
+	if (auto It = m_Buckets.find(std::string(Bucket)); It != m_Buckets.end())
+	{
+		return ZenCacheMemoryLayer::BucketInfo{.EntryCount = It->second->EntryCount(), .TotalSize = It->second->TotalSize()};
+	}
+	return {};
+}
+
+void
+ZenCacheMemoryLayer::CacheBucket::ScrubStorage(ScrubContext& Ctx)
+{
+	RwLock::SharedLockScope _(m_BucketLock);
+
+	std::vector<IoHash> BadHashes;
+
+	auto ValidateEntry = [](const IoHash& Hash, ZenContentType ContentType, IoBuffer Buffer) {
+		if (ContentType == ZenContentType::kCbObject)
+		{
+			CbValidateError Error = ValidateCompactBinary(Buffer, CbValidateMode::All);
+			return Error == CbValidateError::None;
+		}
+		if (ContentType == ZenContentType::kCompressedBinary)
+		{
+			IoHash	 RawHash;
+			uint64_t RawSize;
+			if (!CompressedBuffer::ValidateCompressedHeader(Buffer, RawHash, RawSize))
+			{
+				return false;
+			}
+			if (Hash != RawHash)
+			{
+				return false;
+			}
+		}
+		return true;
+	};
+
+	for (auto& Kv : m_CacheMap)
+	{
+		const BucketPayload& Payload = m_Payloads[Kv.second];
+		if (!ValidateEntry(Kv.first, Payload.Payload.GetContentType(), Payload.Payload))
+		{
+			BadHashes.push_back(Kv.first);
+		}
+	}
+
+	if (!BadHashes.empty())
+	{
+		Ctx.ReportBadCidChunks(BadHashes);
+	}
+}
+
+void
+ZenCacheMemoryLayer::CacheBucket::GatherAccessTimes(std::vector<zen::access_tracking::KeyAccessTime>& AccessTimes)
+{
+	RwLock::SharedLockScope _(m_BucketLock);
+	std::transform(m_CacheMap.begin(), m_CacheMap.end(), std::back_inserter(AccessTimes), [this](const auto& Kv) {
+		return access_tracking::KeyAccessTime{.Key = Kv.first, .LastAccess = m_AccessTimes[Kv.second]};
+	});
+}
+
+bool
+ZenCacheMemoryLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutValue)
+{
+	RwLock::SharedLockScope _(m_BucketLock);
+
+	if (auto It = m_CacheMap.find(HashKey); It != m_CacheMap.end())
+	{
+		uint32_t EntryIndex = It.value();
+		ZEN_ASSERT_SLOW(EntryIndex < m_Payloads.size());
+		ZEN_ASSERT_SLOW(m_AccessTimes.size() == m_Payloads.size());
+
+		const BucketPayload& Payload = m_Payloads[EntryIndex];
+		OutValue					 = {.Value = Payload.Payload, .RawSize = Payload.RawSize, .RawHash = Payload.RawHash};
+		m_AccessTimes[EntryIndex]	 = GcClock::TickCount();
+
+		return true;
+	}
+
+	return false;
+}
+
+void
+ZenCacheMemoryLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue& Value)
+{
+	size_t PayloadSize = Value.Value.GetSize();
+	{
+		GcClock::Tick			   AccessTime = GcClock::TickCount();
+		RwLock::ExclusiveLockScope _(m_BucketLock);
+		if (m_CacheMap.size() == std::numeric_limits<uint32_t>::max())
+		{
+			// No more space in our memory cache!
+			return;
+		}
+		if (auto It = m_CacheMap.find(HashKey); It != m_CacheMap.end())
+		{
+			uint32_t EntryIndex = It.value();
+			ZEN_ASSERT_SLOW(EntryIndex < m_Payloads.size());
+
+			m_TotalSize.fetch_sub(PayloadSize, std::memory_order::relaxed);
+			BucketPayload& Payload	  = m_Payloads[EntryIndex];
+			Payload.Payload			  = Value.Value;
+			Payload.RawHash			  = Value.RawHash;
+			Payload.RawSize			  = gsl::narrow<uint32_t>(Value.RawSize);
+			m_AccessTimes[EntryIndex] = AccessTime;
+		}
+		else
+		{
+			uint32_t EntryIndex = gsl::narrow<uint32_t>(m_Payloads.size());
+			m_Payloads.emplace_back(
+				BucketPayload{.Payload = Value.Value, .RawSize = gsl::narrow<uint32_t>(Value.RawSize), .RawHash = Value.RawHash});
+			m_AccessTimes.emplace_back(AccessTime);
+			m_CacheMap.insert_or_assign(HashKey, EntryIndex);
+		}
+		ZEN_ASSERT_SLOW(m_Payloads.size() == m_CacheMap.size());
+		ZEN_ASSERT_SLOW(m_AccessTimes.size() == m_Payloads.size());
+	}
+
+	m_TotalSize.fetch_add(PayloadSize, std::memory_order::relaxed);
+}
+
+void
+ZenCacheMemoryLayer::CacheBucket::Drop()
+{
+	RwLock::ExclusiveLockScope _(m_BucketLock);
+	m_CacheMap.clear();
+	m_AccessTimes.clear();
+	m_Payloads.clear();
+	m_TotalSize.store(0);
+}
+
+uint64_t
+ZenCacheMemoryLayer::CacheBucket::EntryCount() const
+{
+	RwLock::SharedLockScope _(m_BucketLock);
+	return static_cast<uint64_t>(m_CacheMap.size());
+}
+
+}  // namespace zen