z$: implemented basic bucket tracking

1 files changed, 65 insertions, 0 deletions

diff --git a/zenserver/cache/cachetracking.cpp b/zenserver/cache/cachetracking.cpp
index 865ac38b9..93aeb4aba 100644
--- a/zenserver/cache/cachetracking.cpp
+++ b/zenserver/cache/cachetracking.cpp
@@ -13,6 +13,9 @@ ZEN_THIRD_PARTY_INCLUDES_START
 #pragma comment(lib, "Rpcrt4.lib")	// RocksDB made me do this
 #include <fmt/format.h>
 #include <rocksdb/db.h>
+#include <tsl/robin_map.h>
+#include <tsl/robin_set.h>
+#include <gsl/gsl-lite.hpp>
 ZEN_THIRD_PARTY_INCLUDES_END
 
 namespace zen {
@@ -32,6 +35,67 @@ GetCurrentTimeStamp()
 	return Minutes;
 }
 
+struct CacheAccessSnapshot
+{
+public:
+	void TrackAccess(std::string_view BucketSegment, const IoHash& HashKey)
+	{
+		BucketTracker* Tracker = MapBucketToIndex(std::string(BucketSegment));
+
+		Tracker->Track(HashKey);
+	}
+
+private:
+	struct BucketTracker
+	{
+		RwLock				   Lock;
+		tsl::robin_set<IoHash> AccessedKeys;
+
+		void Track(const IoHash& HashKey)
+		{
+			RwLock::ExclusiveLockScope _(Lock);
+			AccessedKeys.insert(HashKey);
+		}
+	};
+
+	BucketTracker* MapBucketToIndex(const std::string& BucketName)
+	{
+		RwLock::SharedLockScope _(m_Lock);
+		if (auto It = m_BucketMapping.find(BucketName); It == m_BucketMapping.end())
+		{
+			_.ReleaseNow();
+
+			return AddNewBucket(BucketName);
+		}
+		else
+		{
+			return m_Buckets[It->second].get();
+		}
+	}
+
+	BucketTracker* AddNewBucket(const std::string& BucketName)
+	{
+		RwLock::ExclusiveLockScope _(m_Lock);
+
+		if (auto It = m_BucketMapping.find(BucketName); It == m_BucketMapping.end())
+		{
+			const uint32_t BucketIndex = gsl::narrow<uint32_t>(m_Buckets.size());
+			m_Buckets.emplace_back(std::make_unique<BucketTracker>());
+			m_BucketMapping[BucketName] = BucketIndex;
+
+			return m_Buckets[BucketIndex].get();
+		}
+		else
+		{
+			return m_Buckets[It->second].get();
+		}
+	}
+
+	RwLock										m_Lock;
+	std::vector<std::unique_ptr<BucketTracker>> m_Buckets;
+	tsl::robin_map<std::string, uint32_t>		m_BucketMapping;
+};
+
 struct ZenCacheTracker::Impl
 {
 	Impl(std::filesystem::path StateDirectory)
@@ -143,6 +207,7 @@ struct ZenCacheTracker::Impl
 
 	std::unique_ptr<rocksdb::DB>			  m_RocksDb;
 	std::vector<rocksdb::ColumnFamilyHandle*> m_RocksDbColumnHandles;
+	CacheAccessSnapshot						  m_CurrentSnapshot;
 };
 
 ZenCacheTracker::ZenCacheTracker(std::filesystem::path StateDirectory) : m_Impl(new Impl(StateDirectory))