From 924169d85b2af36a95e1844140dd01573cfb113e Mon Sep 17 00:00:00 2001
From: Stefan Boberg <stefan@boberg.org>
Date: Wed, 3 Nov 2021 22:05:29 +0100
Subject: z$: basic access tracking

---
 zenserver/cache/cachetracking.cpp | 196 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 196 insertions(+)
 create mode 100644 zenserver/cache/cachetracking.cpp

(limited to 'zenserver/cache/cachetracking.cpp')

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
-- 
cgit v1.2.3