z$: basic access tracking

7 files changed, 259 insertions, 51 deletions

diff --git a/zenserver/cache/cachetracking.cpp b/zenserver/cache/cachetracking.cpp
new file mode 100644
index 000000000..865ac38b9
--- /dev/null
+++ b/zenserver/cache/cachetracking.cpp
@@ -0,0 +1,196 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "cachetracking.h"
+
+#include <zencore/filesystem.h>
+#include <zencore/logging.h>
+#include <zencore/string.h>
+
+#include <zencore/testing.h>
+#include <zencore/testutils.h>
+
+ZEN_THIRD_PARTY_INCLUDES_START
+#pragma comment(lib, "Rpcrt4.lib")	// RocksDB made me do this
+#include <fmt/format.h>
+#include <rocksdb/db.h>
+ZEN_THIRD_PARTY_INCLUDES_END
+
+namespace zen {
+
+using namespace fmt::literals;
+
+namespace rocksdb = ROCKSDB_NAMESPACE;
+
+static constinit auto Epoch = std::chrono::time_point<std::chrono::system_clock>{};
+
+uint32_t
+GetCurrentTimeStamp()
+{
+	auto Duration = std::chrono::system_clock::now() - Epoch;
+	auto Minutes  = std::chrono::duration_cast<std::chrono::minutes>(Duration).count();
+
+	return Minutes;
+}
+
+struct ZenCacheTracker::Impl
+{
+	Impl(std::filesystem::path StateDirectory)
+	{
+		std::filesystem::path StatsDbPath{StateDirectory / ".zdb"};
+
+		std::string RocksdbPath = ToUtf8(StatsDbPath);
+
+		ZEN_DEBUG("opening tracker db at '{}'", RocksdbPath);
+
+		rocksdb::DB*	   Db = nullptr;
+		rocksdb::DBOptions Options;
+		Options.create_if_missing = true;
+
+		std::vector<std::string> ExistingColumnFamilies;
+		rocksdb::Status			 Status = rocksdb::DB::ListColumnFamilies(Options, RocksdbPath, &ExistingColumnFamilies);
+
+		std::vector<rocksdb::ColumnFamilyDescriptor> ColumnDescriptors;
+
+		if (Status.IsPathNotFound())
+		{
+			ColumnDescriptors.emplace_back(rocksdb::ColumnFamilyDescriptor{rocksdb::kDefaultColumnFamilyName, {}});
+		}
+		else if (Status.ok())
+		{
+			for (const std::string& Column : ExistingColumnFamilies)
+			{
+				rocksdb::ColumnFamilyDescriptor ColumnFamily;
+				ColumnFamily.name = Column;
+				ColumnDescriptors.push_back(ColumnFamily);
+			}
+		}
+		else
+		{
+			throw std::runtime_error("column family iteration failed for '{}': '{}'"_format(RocksdbPath, Status.getState()).c_str());
+		}
+
+		Status = rocksdb::DB::Open(Options, RocksdbPath, ColumnDescriptors, &m_RocksDbColumnHandles, &Db);
+
+		if (!Status.ok())
+		{
+			throw std::runtime_error("database open failed for '{}': '{}'"_format(RocksdbPath, Status.getState()).c_str());
+		}
+
+		m_RocksDb.reset(Db);
+	}
+
+	~Impl()
+	{
+		for (auto* Column : m_RocksDbColumnHandles)
+		{
+			delete Column;
+		}
+
+		m_RocksDbColumnHandles.clear();
+	}
+
+	struct KeyStruct
+	{
+		uint16_t BucketId;
+		IoHash	 HashKey;
+	};
+
+	struct ValueStruct
+	{
+		uint32_t CreateTime	 = 0;
+		uint32_t AccessTime	 = 0;
+		uint32_t AccessCount = 0;
+		uint32_t LastGcCount = 0;
+	};
+
+	void TrackAccess(std::string_view BucketSegment, const IoHash& HashKey)
+	{
+		const uint32_t		  Ts = GetCurrentTimeStamp();
+		rocksdb::WriteOptions Wo;
+		rocksdb::ReadOptions  Ro;
+
+		ZEN_UNUSED(BucketSegment);
+
+		ValueStruct Value;
+
+		rocksdb::Slice ValueSlice{(char*)&Value, sizeof Value};
+
+		KeyStruct Key;
+		Key.BucketId = 0;
+		Key.HashKey	 = HashKey;
+
+		rocksdb::Slice KeySlice{(char*)&Key, sizeof Key};
+
+		std::string		ValueString;
+		rocksdb::Status Status = m_RocksDb->Get(Ro, KeySlice, &ValueString);
+
+		if (Status.ok() && ValueString.size() == sizeof(ValueStruct))
+		{
+			memcpy(&Value, ValueString.data(), ValueString.size());
+
+			++Value.AccessCount;
+		}
+		else
+		{
+			Value.CreateTime  = Ts;
+			Value.AccessCount = 1;
+		}
+
+		Value.AccessTime = Ts;
+
+		m_RocksDb->Put(Wo, KeySlice, ValueSlice);
+	}
+
+	std::unique_ptr<rocksdb::DB>			  m_RocksDb;
+	std::vector<rocksdb::ColumnFamilyHandle*> m_RocksDbColumnHandles;
+};
+
+ZenCacheTracker::ZenCacheTracker(std::filesystem::path StateDirectory) : m_Impl(new Impl(StateDirectory))
+{
+}
+
+ZenCacheTracker::~ZenCacheTracker()
+{
+	delete m_Impl;
+}
+
+void
+ZenCacheTracker::TrackAccess(std::string_view BucketSegment, const IoHash& HashKey)
+{
+	m_Impl->TrackAccess(BucketSegment, HashKey);
+}
+
+#if ZEN_WITH_TESTS
+
+TEST_CASE("z$.tracker")
+{
+	using namespace fmt::literals;
+	using namespace std::literals;
+
+	ScopedTemporaryDirectory TempDir;
+
+	ZenCacheTracker Zcs(TempDir.Path());
+
+	for (int i = 0; i < 10000; ++i)
+	{
+		IoHash KeyHash = IoHash::HashBuffer(&i, sizeof i);
+
+		Zcs.TrackAccess("foo", KeyHash);
+	}
+
+	for (int i = 0; i < 10000; ++i)
+	{
+		IoHash KeyHash = IoHash::HashBuffer(&i, sizeof i);
+
+		Zcs.TrackAccess("foo", KeyHash);
+	}
+}
+
+#endif
+
+void
+cachetracker_forcelink()
+{
+}
+
+}  // namespace zen
diff --git a/zenserver/cache/cachetracking.h b/zenserver/cache/cachetracking.h
new file mode 100644
index 000000000..8ebe5682c
--- /dev/null
+++ b/zenserver/cache/cachetracking.h
@@ -0,0 +1,28 @@
+#pragma once
+
+#include <zencore/iohash.h>
+
+#include <filesystem>
+
+namespace zen {
+
+/**
+ */
+
+class ZenCacheTracker
+{
+public:
+	ZenCacheTracker(std::filesystem::path StateDirectory);
+	~ZenCacheTracker();
+
+	void TrackAccess(std::string_view BucketSegment, const IoHash& HashKey);
+
+private:
+	struct Impl;
+
+	Impl* m_Impl = nullptr;
+};
+
+void cachetracker_forcelink();
+
+}  // namespace zen
diff --git a/zenserver/cache/structuredcache.cpp b/zenserver/cache/structuredcache.cpp
index d71329a4b..177cdbf55 100644
--- a/zenserver/cache/structuredcache.cpp
+++ b/zenserver/cache/structuredcache.cpp
@@ -147,36 +147,6 @@ ParseCachePolicy(const HttpServerRequest::QueryParams& QueryParams)
 
 //////////////////////////////////////////////////////////////////////////
 
-class CacheAccessTracker
-{
-public:
-	CacheAccessTracker();
-	~CacheAccessTracker();
-
-	void TrackAccess(std::string_view BucketSegment, const IoHash& HashKey);
-
-private:
-	RwLock	m_Lock;
-};
-
-CacheAccessTracker::CacheAccessTracker()
-{
-}
-
-CacheAccessTracker::~CacheAccessTracker()
-{
-}
-
-void
-CacheAccessTracker::TrackAccess(std::string_view BucketSegment, const IoHash& HashKey)
-{
-	ZEN_UNUSED(BucketSegment, HashKey);
-}
-
-CacheAccessTracker g_AccessTracker;
-
-//////////////////////////////////////////////////////////////////////////
-
 HttpStructuredCacheService::HttpStructuredCacheService(ZenCacheStore&				  InCacheStore,
 													   CidStore&					  InCidStore,
 													   HttpStatsService&			  StatsService,
@@ -329,8 +299,6 @@ HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request
 	{
 		Success = true;
 
-		g_AccessTracker.TrackAccess(Ref.BucketSegment, Ref.HashKey);
-		 
 		if (!SkipData && AcceptType == ZenContentType::kCbPackage)
 		{
 			CbPackage	 Package;
diff --git a/zenserver/cache/structuredcachestore.cpp b/zenserver/cache/structuredcachestore.cpp
index 1b6c30cbd..ac9f628d3 100644
--- a/zenserver/cache/structuredcachestore.cpp
+++ b/zenserver/cache/structuredcachestore.cpp
@@ -2,22 +2,24 @@
 
 #include "structuredcachestore.h"
 
-#include <zencore/except.h>
-#include <zencore/windows.h>
+#include "cachetracking.h"
 
 #include <zencore/compactbinary.h>
 #include <zencore/compactbinarybuilder.h>
 #include <zencore/compactbinarypackage.h>
 #include <zencore/compactbinaryvalidation.h>
 #include <zencore/compress.h>
+#include <zencore/except.h>
 #include <zencore/filesystem.h>
 #include <zencore/fmtutils.h>
 #include <zencore/iobuffer.h>
 #include <zencore/logging.h>
+#include <zencore/scopeguard.h>
 #include <zencore/string.h>
 #include <zencore/testing.h>
 #include <zencore/testutils.h>
 #include <zencore/thread.h>
+#include <zencore/windows.h>
 #include <zenstore/basicfile.h>
 #include <zenstore/cas.h>
 #include <zenstore/caslog.h>
@@ -41,12 +43,14 @@ namespace zen {
 
 using namespace fmt::literals;
 
-ZenCacheStore::ZenCacheStore(CasGc& Gc, const std::filesystem::path& RootDir) : GcContributor(Gc), m_DiskLayer{RootDir}
+ZenCacheStore::ZenCacheStore(CasGc& Gc, const std::filesystem::path& RootDir) : GcContributor(Gc), m_DiskLayer(RootDir)
 {
 	ZEN_INFO("initializing structured cache at '{}'", RootDir);
 	CreateDirectories(RootDir);
 
 	m_DiskLayer.DiscoverBuckets();
+
+	m_AccessTracker.reset(new ZenCacheTracker(RootDir));
 }
 
 ZenCacheStore::~ZenCacheStore()
@@ -58,21 +62,27 @@ ZenCacheStore::Get(std::string_view InBucket, const IoHash& HashKey, ZenCacheVal
 {
 	bool Ok = m_MemLayer.Get(InBucket, HashKey, OutValue);
 
+	auto _ = MakeGuard([&] {
+		if (!Ok)
+			return;
+
+		m_AccessTracker->TrackAccess(InBucket, HashKey);
+	});
+
 	if (Ok)
 	{
 		ZEN_ASSERT(OutValue.Value.Size());
+
+		return true;
 	}
 
-	if (!Ok)
-	{
-		Ok = m_DiskLayer.Get(InBucket, HashKey, OutValue);
+	Ok = m_DiskLayer.Get(InBucket, HashKey, OutValue);
 
-		if (Ok)
-		{
-			ZEN_ASSERT(OutValue.Value.Size());
-		}
+	if (Ok)
+	{
+		ZEN_ASSERT(OutValue.Value.Size());
 
-		if (Ok && (OutValue.Value.Size() <= m_DiskLayerSizeThreshold))
+		if (OutValue.Value.Size() <= m_DiskLayerSizeThreshold)
 		{
 			m_MemLayer.Put(InBucket, HashKey, OutValue);
 		}
@@ -528,7 +538,7 @@ ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir, bo
 
 	m_SlogFile.Open(SlogPath, IsNew);
 
-	uint64_t MaxFileOffset = 0;
+	uint64_t MaxFileOffset	   = 0;
 	uint64_t InvalidEntryCount = 0;
 
 	if (RwLock::ExclusiveLockScope _(m_IndexLock); m_Index.empty())
@@ -1011,7 +1021,7 @@ ZenCacheDiskLayer::DiscoverBuckets()
 	{
 		// New bucket needs to be created
 
- 		const std::string BucketName8 = ToUtf8(BucketName);
+		const std::string BucketName8 = ToUtf8(BucketName);
 
 		if (auto It = m_Buckets.find(BucketName8); It != m_Buckets.end())
 		{
diff --git a/zenserver/cache/structuredcachestore.h b/zenserver/cache/structuredcachestore.h
index 6beecf78b..5a3191cc5 100644
--- a/zenserver/cache/structuredcachestore.h
+++ b/zenserver/cache/structuredcachestore.h
@@ -20,9 +20,10 @@ ZEN_THIRD_PARTY_INCLUDES_END
 
 namespace zen {
 
-class WideStringBuilderBase;
 class CasStore;
 class CasGc;
+class WideStringBuilderBase;
+class ZenCacheTracker;
 
 /******************************************************************************
 
@@ -142,11 +143,12 @@ public:
 	virtual void GatherReferences(GcContext& GcCtx) override;
 
 private:
-	std::filesystem::path m_RootDir;
-	ZenCacheMemoryLayer	  m_MemLayer;
-	ZenCacheDiskLayer	  m_DiskLayer;
-	uint64_t			  m_DiskLayerSizeThreshold = 1 * 1024;
-	uint64_t			  m_LastScrubTime		   = 0;
+	std::filesystem::path			 m_RootDir;
+	ZenCacheMemoryLayer				 m_MemLayer;
+	ZenCacheDiskLayer				 m_DiskLayer;
+	uint64_t						 m_DiskLayerSizeThreshold = 1 * 1024;
+	uint64_t						 m_LastScrubTime		  = 0;
+	std::unique_ptr<ZenCacheTracker> m_AccessTracker;
 };
 
 void z$_forcelink();
diff --git a/zenserver/zenserver.vcxproj b/zenserver/zenserver.vcxproj
index d954d3f8d..13589ee3b 100644
--- a/zenserver/zenserver.vcxproj
+++ b/zenserver/zenserver.vcxproj
@@ -105,6 +105,7 @@
   </ItemDefinitionGroup>
   <ItemGroup>
     <ClInclude Include="admin\admin.h" />
+    <ClInclude Include="cache\cachetracking.h" />
     <ClInclude Include="cache\structuredcache.h" />
     <ClInclude Include="cache\structuredcachestore.h" />
     <ClInclude Include="compute\apply.h" />
@@ -131,6 +132,7 @@
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="admin\admin.cpp" />
+    <ClCompile Include="cache\cachetracking.cpp" />
     <ClCompile Include="cache\structuredcache.cpp" />
     <ClCompile Include="cache\structuredcachestore.cpp" />
     <ClCompile Include="compute\apply.cpp" />
diff --git a/zenserver/zenserver.vcxproj.filters b/zenserver/zenserver.vcxproj.filters
index 04c6267ba..87591ef56 100644
--- a/zenserver/zenserver.vcxproj.filters
+++ b/zenserver/zenserver.vcxproj.filters
@@ -41,6 +41,7 @@
     <ClInclude Include="experimental\vfs.h" />
     <ClInclude Include="monitoring\httpstats.h" />
     <ClInclude Include="monitoring\httpstatus.h" />
+    <ClInclude Include="cache\cachetracking.h" />
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="zenserver.cpp" />
@@ -76,6 +77,7 @@
     <ClCompile Include="experimental\vfs.cpp" />
     <ClCompile Include="monitoring\httpstats.cpp" />
     <ClCompile Include="monitoring\httpstatus.cpp" />
+    <ClCompile Include="cache\cachetracking.cpp" />
   </ItemGroup>
   <ItemGroup>
     <Filter Include="cache">