Merged latest from main

1 files changed, 186 insertions, 85 deletions

diff --git a/zenserver/cache/structuredcachestore.cpp b/zenserver/cache/structuredcachestore.cpp
index 3d80bb14c..5e93ebaa9 100644
--- a/zenserver/cache/structuredcachestore.cpp
+++ b/zenserver/cache/structuredcachestore.cpp
@@ -22,8 +22,6 @@
 #include <gsl/gsl-lite.hpp>
 #include <unordered_map>
 
-#include <atlfile.h>
-
 //////////////////////////////////////////////////////////////////////////
 
 namespace zen {
@@ -131,23 +129,18 @@ ZenCacheMemoryLayer::~ZenCacheMemoryLayer()
 bool
 ZenCacheMemoryLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCacheValue& OutValue)
 {
-	CacheBucket* Bucket = nullptr;
-
-	{
-		RwLock::SharedLockScope _(m_Lock);
+	RwLock::SharedLockScope _(m_Lock);
 
-		auto it = m_Buckets.find(std::string(InBucket));
+	auto it = m_Buckets.find(std::string(InBucket));
 
-		if (it != m_Buckets.end())
-		{
-			Bucket = &it->second;
-		}
+	if (it == m_Buckets.end())
+	{
+		return false;
 	}
 
-	if (Bucket == nullptr)
-		return false;
+	CacheBucket* Bucket = Bucket = &it->second;
 
-	ZEN_ASSERT(Bucket != nullptr);
+	_.ReleaseNow();
 
 	return Bucket->Get(HashKey, OutValue);
 }
@@ -177,8 +170,6 @@ ZenCacheMemoryLayer::Put(std::string_view InBucket, const IoHash& HashKey, const
 		Bucket = &m_Buckets[std::string(InBucket)];
 	}
 
-	ZEN_ASSERT(Bucket != nullptr);
-
 	// Note that since the underlying IoBuffer is retained, the content type is also
 
 	Bucket->Put(HashKey, Value);
@@ -195,7 +186,31 @@ ZenCacheMemoryLayer::DropBucket(std::string_view Bucket)
 void
 ZenCacheMemoryLayer::Scrub(ScrubContext& Ctx)
 {
-	ZEN_UNUSED(Ctx);
+	RwLock::SharedLockScope _(m_Lock);
+
+	for (auto& Kv : m_Buckets)
+	{
+		Kv.second.Scrub(Ctx);
+	}
+}
+
+void
+ZenCacheMemoryLayer::CacheBucket::Scrub(ScrubContext& Ctx)
+{
+	std::vector<IoHash> BadHashes;
+
+	for (auto& Kv : m_cacheMap)
+	{
+		if (Kv.first != IoHash::HashBuffer(Kv.second))
+		{
+			BadHashes.push_back(Kv.first);
+		}
+	}
+
+	if (!BadHashes.empty())
+	{
+		Ctx.ReportBadChunks(BadHashes);
+	}
 }
 
 bool
@@ -203,16 +218,16 @@ ZenCacheMemoryLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutV
 {
 	RwLock::SharedLockScope _(m_bucketLock);
 
-	auto bucketIt = m_cacheMap.find(HashKey);
-
-	if (bucketIt == m_cacheMap.end())
+	if (auto bucketIt = m_cacheMap.find(HashKey); bucketIt == m_cacheMap.end())
 	{
 		return false;
 	}
+	else
+	{
+		OutValue.Value = bucketIt->second;
 
-	OutValue.Value = bucketIt->second;
-
-	return true;
+		return true;
+	}
 }
 
 void
@@ -241,8 +256,19 @@ struct DiskLocation
 
 	static uint64_t CombineOffsetAndFlags(uint64_t Offset, uint64_t Flags) { return Offset | Flags; }
 
-	inline uint64_t Offset() const { return OffsetAndFlags & kOffsetMask; }
-	inline uint64_t IsFlagSet(uint64_t Flag) const { return OffsetAndFlags & Flag; }
+	inline uint64_t		  Offset() const { return OffsetAndFlags & kOffsetMask; }
+	inline uint64_t		  IsFlagSet(uint64_t Flag) const { return OffsetAndFlags & Flag; }
+	inline ZenContentType GetContentType() const
+	{
+		ZenContentType ContentType = ZenContentType::kBinary;
+
+		if (IsFlagSet(DiskLocation::kStructured))
+		{
+			ContentType = ZenContentType::kCbObject;
+		}
+
+		return ContentType;
+	}
 };
 
 struct DiskIndexEntry
@@ -267,6 +293,7 @@ struct ZenCacheDiskLayer::CacheBucket
 	void Put(const IoHash& HashKey, const ZenCacheValue& Value);
 	void Drop();
 	void Flush();
+	void Scrub(ScrubContext& Ctx);
 
 	inline bool IsOk() const { return m_Ok; }
 
@@ -277,15 +304,19 @@ private:
 	bool				  m_Ok					 = false;
 	uint64_t			  m_LargeObjectThreshold = 64 * 1024;
 
+	// These files are used to manage storage of small objects for this bucket
+
 	BasicFile					m_SobsFile;
 	TCasLogFile<DiskIndexEntry> m_SlogFile;
 
-	void BuildPath(WideStringBuilderBase& Path, const IoHash& HashKey);
-	void PutLargeObject(const IoHash& HashKey, const ZenCacheValue& Value);
-
 	RwLock												 m_IndexLock;
 	tsl::robin_map<IoHash, DiskLocation, IoHash::Hasher> m_Index;
 	uint64_t											 m_WriteCursor = 0;
+
+	void BuildPath(WideStringBuilderBase& Path, const IoHash& HashKey);
+	void PutLargeObject(const IoHash& HashKey, const ZenCacheValue& Value);
+	bool GetStandaloneCacheValue(const IoHash& HashKey, ZenCacheValue& OutValue, const DiskLocation& Loc);
+	bool GetInlineCacheValue(const DiskLocation& Loc, ZenCacheValue& OutValue);
 };
 
 ZenCacheDiskLayer::CacheBucket::CacheBucket(CasStore& Cas) : m_CasStore(Cas)
@@ -320,27 +351,24 @@ ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir)
 	std::filesystem::path SobsPath{m_BucketDir / "zen.sobs"};
 	std::filesystem::path SlogPath{m_BucketDir / "zen.slog"};
 
-	CAtlFile ManifestFile;
+	BasicFile ManifestFile;
 
 	// Try opening existing manifest file first
 
 	bool IsNew = false;
 
-	HRESULT hRes = ManifestFile.Create(ManifestPath.c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ, OPEN_EXISTING);
+	std::error_code Ec;
+	ManifestFile.Open(ManifestPath, /* IsCreate */ false, Ec);
 
-	if (SUCCEEDED(hRes))
+	if (!Ec)
 	{
-		ULONGLONG FileSize;
-		ManifestFile.GetSize(FileSize);
+		uint64_t FileSize = ManifestFile.FileSize();
 
 		if (FileSize == sizeof(Oid))
 		{
-			hRes = ManifestFile.Read(&m_BucketId, sizeof(Oid));
+			ManifestFile.Read(&m_BucketId, sizeof(Oid), 0);
 
-			if (SUCCEEDED(hRes))
-			{
-				m_Ok = true;
-			}
+			m_Ok = true;
 		}
 
 		if (!m_Ok)
@@ -353,16 +381,16 @@ ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir)
 	{
 		// No manifest file found, this is a new bucket
 
-		hRes = ManifestFile.Create(ManifestPath.c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ, CREATE_ALWAYS);
+		ManifestFile.Open(ManifestPath, /* IsCreate */ true, Ec);
 
-		if (FAILED(hRes))
+		if (Ec)
 		{
-			ThrowLastError("Failed to create bucket manifest '{}'"_format(ManifestPath));
+			throw std::system_error(Ec, "Failed to create bucket manifest '{}'"_format(ManifestPath));
 		}
 
 		m_BucketId.Generate();
 
-		hRes = ManifestFile.Write(&m_BucketId, sizeof(Oid));
+		ManifestFile.Write(&m_BucketId, sizeof(Oid), /* FileOffset */ 0);
 
 		IsNew = true;
 	}
@@ -407,6 +435,37 @@ ZenCacheDiskLayer::CacheBucket::BuildPath(WideStringBuilderBase& Path, const IoH
 }
 
 bool
+ZenCacheDiskLayer::CacheBucket::GetInlineCacheValue(const DiskLocation& Loc, ZenCacheValue& OutValue)
+{
+	if (!Loc.IsFlagSet(DiskLocation::kStandaloneFile))
+	{
+		OutValue.Value = IoBufferBuilder::MakeFromFileHandle(m_SobsFile.Handle(), Loc.Offset(), Loc.Size);
+		OutValue.Value.SetContentType(Loc.GetContentType());
+
+		return true;
+	}
+
+	return false;
+}
+
+bool
+ZenCacheDiskLayer::CacheBucket::GetStandaloneCacheValue(const IoHash& HashKey, ZenCacheValue& OutValue, const DiskLocation& Loc)
+{
+	WideStringBuilder<128> DataFilePath;
+	BuildPath(DataFilePath, HashKey);
+
+	if (IoBuffer Data = IoBufferBuilder::MakeFromFile(DataFilePath.c_str()))
+	{
+		OutValue.Value = Data;
+		OutValue.Value.SetContentType(Loc.GetContentType());
+
+		return true;
+	}
+
+	return false;
+}
+
+bool
 ZenCacheDiskLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutValue)
 {
 	if (!m_Ok)
@@ -420,35 +479,14 @@ ZenCacheDiskLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutVal
 	{
 		const DiskLocation& Loc = it->second;
 
-		ZenContentType ContentType = ZenContentType::kBinary;
-
-		if (Loc.IsFlagSet(DiskLocation::kStructured))
-		{
-			ContentType = ZenContentType::kCbObject;
-		}
-
-		if (!Loc.IsFlagSet(DiskLocation::kStandaloneFile))
+		if (GetInlineCacheValue(Loc, OutValue))
 		{
-			OutValue.Value = IoBufferBuilder::MakeFromFileHandle(m_SobsFile.Handle(), Loc.Offset(), Loc.Size);
-			OutValue.Value.SetContentType(ContentType);
-
 			return true;
 		}
-		else
-		{
-			_.ReleaseNow();
-
-			WideStringBuilder<128> DataFilePath;
-			BuildPath(DataFilePath, HashKey);
 
-			if (IoBuffer Data = IoBufferBuilder::MakeFromFile(DataFilePath.c_str()))
-			{
-				OutValue.Value = Data;
-				OutValue.Value.SetContentType(ContentType);
+		_.ReleaseNow();
 
-				return true;
-			}
-		}
+		return GetStandaloneCacheValue(HashKey, OutValue, Loc);
 	}
 
 	return false;
@@ -518,9 +556,58 @@ ZenCacheDiskLayer::CacheBucket::Flush()
 }
 
 void
-ZenCacheDiskLayer::Scrub(ScrubContext& Ctx)
+ZenCacheDiskLayer::CacheBucket::Scrub(ScrubContext& Ctx)
 {
-	ZEN_UNUSED(Ctx);
+	std::vector<DiskIndexEntry> StandaloneFiles;
+
+	std::vector<IoHash>	  BadChunks;
+	std::vector<IoBuffer> BadStandaloneChunks;
+
+	{
+		RwLock::SharedLockScope _(m_IndexLock);
+
+		for (auto& Kv : m_Index)
+		{
+			const IoHash&		Hash = Kv.first;
+			const DiskLocation& Loc	 = Kv.second;
+
+			ZenCacheValue Value;
+
+			if (!GetInlineCacheValue(Loc, Value))
+			{
+				ZEN_ASSERT(Loc.IsFlagSet(DiskLocation::kStandaloneFile));
+				StandaloneFiles.push_back({.Key = Hash, .Location = Loc});
+			}
+			else
+			{
+				if (GetStandaloneCacheValue(Hash, Value, Loc))
+				{
+					// Hash contents
+
+					const IoHash ComputedHash = HashBuffer(Value.Value);
+
+					if (ComputedHash != Hash)
+					{
+						BadChunks.push_back(Hash);
+					}
+				}
+				else
+				{
+					// Non-existent
+				}
+			}
+		}
+	}
+
+	if (Ctx.RunRecovery())
+	{
+		// Clean out bad chunks
+	}
+
+	if (!BadChunks.empty())
+	{
+		Ctx.ReportBadChunks(BadChunks);
+	}
 }
 
 void
@@ -529,35 +616,38 @@ ZenCacheDiskLayer::CacheBucket::PutLargeObject(const IoHash& HashKey, const ZenC
 	WideStringBuilder<128> DataFilePath;
 	BuildPath(DataFilePath, HashKey);
 
-	// TODO: replace this process with a more efficient implementation with proper atomic rename
-	// and also avoid creating directories if we can
-
-	std::filesystem::path ParentPath = std::filesystem::path(DataFilePath.c_str()).parent_path();
-	CreateDirectories(ParentPath);
+	TemporaryFile DataFile;
 
-	CAtlTemporaryFile DataFile;
+	std::error_code Ec;
+	DataFile.CreateTemporary(m_BucketDir.c_str(), Ec);
 
-	HRESULT hRes = DataFile.Create(m_BucketDir.c_str());
-
-	if (FAILED(hRes))
+	if (Ec)
 	{
-		ThrowSystemException(hRes, "Failed to open temporary file for put at '{}'"_format(m_BucketDir));
+		throw std::system_error(Ec, "Failed to open temporary file for put at '{}'"_format(m_BucketDir));
 	}
 
-	hRes = DataFile.Write(Value.Value.Data(), gsl::narrow<DWORD>(Value.Value.Size()));
+	DataFile.WriteAll(Value.Value, Ec);
 
-	if (FAILED(hRes))
+	if (Ec)
 	{
-		ThrowSystemException(hRes, "Failed to write payload ({} bytes) to file"_format(NiceBytes(Value.Value.Size())));
+		throw std::system_error(Ec, "Failed to write payload ({} bytes) to file"_format(NiceBytes(Value.Value.Size())));
 	}
 
-	// Move file into place (note: not fully atomic!)
+	// Move file into place (atomically)
 
-	hRes = DataFile.Close(DataFilePath.c_str());
+	DataFile.MoveTemporaryIntoPlace(DataFilePath.c_str(), Ec);
 
-	if (FAILED(hRes))
+	if (Ec)
 	{
-		ThrowSystemException(hRes, "Failed to finalize file '{}'"_format(WideToUtf8(DataFilePath)));
+		std::filesystem::path ParentPath = std::filesystem::path(DataFilePath.c_str()).parent_path();
+		CreateDirectories(ParentPath);
+
+		DataFile.MoveTemporaryIntoPlace(DataFilePath.c_str(), Ec);
+
+		if (Ec)
+		{
+			throw std::system_error(Ec, "Failed to finalize file '{}'"_format(WideToUtf8(DataFilePath)));
+		}
 	}
 
 	// Update index
@@ -729,6 +819,17 @@ ZenCacheDiskLayer::Flush()
 	}
 }
 
+void
+ZenCacheDiskLayer::Scrub(ScrubContext& Ctx)
+{
+	RwLock::SharedLockScope _(m_Lock);
+
+	for (auto& Kv : m_Buckets)
+	{
+		Kv.second.Scrub(Ctx);
+	}
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 ZenCacheTracker::ZenCacheTracker(ZenCacheStore& CacheStore)