path: root/zenstore/filecas.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#include "FileCas.h"

#include <zencore/except.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/memory.h>
#include <zencore/string.h>
#include <zencore/thread.h>
#include <zencore/uid.h>

#include <spdlog/spdlog.h>
#include <gsl/gsl-lite.hpp>

#include <functional>
#include <unordered_map>

struct IUnknown;  // Workaround for "combaseapi.h(229): error C2187: syntax error: 'identifier' was unexpected here" when using /permissive-
#include <atlfile.h>
#include <filesystem>

// Used for getting My Documents for default CAS
#include <ShlObj.h>
#pragma comment(lib, "shell32.lib")

//////////////////////////////////////////////////////////////////////////

namespace zen {

using namespace fmt::literals;

WideStringBuilderBase&
FileCasStrategy::MakeShardedPath(WideStringBuilderBase& ShardedPath, const IoHash& ChunkHash, size_t& OutShard2len)
{
	ExtendableStringBuilder<96> HashString;
	ChunkHash.ToHexString(HashString);

	const char* str = HashString.c_str();

	// Shard into a path with two directory levels containing 12 bits and 8 bits
	// respectively.
	//
	// This results in a maximum of 4096 * 256 directories
	//
	// The numbers have been chosen somewhat arbitrarily but are large to scale
	// to very large chunk repositories. It may or may not make sense to make
	// this a configurable policy, and it would probably be a good idea to
	// measure performance for different policies and chunk counts

	ShardedPath.AppendAsciiRange(str, str + 3);

	ShardedPath.Append('\\');
	ShardedPath.AppendAsciiRange(str + 3, str + 5);
	OutShard2len = ShardedPath.Size();

	ShardedPath.Append('\\');
	ShardedPath.AppendAsciiRange(str + 6, str + 64);

	return ShardedPath;
}

CasStore::InsertResult
FileCasStrategy::InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash)
{
	// File-based chunks have special case handling whereby we move the file into
	// place in the file store directory, thus avoiding unnecessary copying

	IoBufferFileReference FileRef;
	if (Chunk.IsWholeFile() && Chunk.GetFileReference(/* out */ FileRef))
	{
		size_t							 Shard2len = 0;
		ExtendableWideStringBuilder<128> ShardedPath;
		ShardedPath.Append(m_Config.RootDirectory.c_str());
		ShardedPath.Append(std::filesystem::path::preferred_separator);
		MakeShardedPath(ShardedPath, ChunkHash, /* out */ Shard2len);

		auto DeletePayloadFileOnClose = [&] {
			FILE_DISPOSITION_INFO Fdi{};
			Fdi.DeleteFile = TRUE;
			BOOL Success   = SetFileInformationByHandle(FileRef.FileHandle, FileDispositionInfo, &Fdi, sizeof Fdi);

			if (!Success)
			{
				spdlog::warn("Failed to flag temporary payload file for deletion: '{}'", PathFromHandle(FileRef.FileHandle));
			}
		};

		// See if file already exists
		//
		// Future improvement: maintain Bloom filter to avoid expensive file system probes?

		{
			CAtlFile PayloadFile;

			HRESULT hRes = PayloadFile.Create(ShardedPath.c_str(), GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING);

			if (SUCCEEDED(hRes))
			{
				// If we succeeded in opening the file then we don't need to do anything else because it already exists and should contain
				// the content we were about to insert

				// We do need to ensure the file goes away on close, however

				DeletePayloadFileOnClose();

				return CasStore::InsertResult{.New = false};
			}
		}

		std::filesystem::path FullPath(ShardedPath.c_str());

		std::filesystem::path FilePath = FullPath.parent_path();
		std::wstring		  FileName = FullPath.native();

		const DWORD		  BufferSize = sizeof(FILE_RENAME_INFO) + gsl::narrow<DWORD>(FileName.size() * sizeof(WCHAR));
		FILE_RENAME_INFO* RenameInfo = reinterpret_cast<FILE_RENAME_INFO*>(Memory::Alloc(BufferSize));
		memset(RenameInfo, 0, BufferSize);

		RenameInfo->ReplaceIfExists = FALSE;
		RenameInfo->FileNameLength	= gsl::narrow<DWORD>(FileName.size());
		memcpy(RenameInfo->FileName, FileName.c_str(), FileName.size() * sizeof(WCHAR));
		RenameInfo->FileName[FileName.size()] = 0;

		// Try to move file into place

		BOOL Success = SetFileInformationByHandle(FileRef.FileHandle, FileRenameInfo, RenameInfo, BufferSize);

		if (!Success)
		{
			CAtlFile DirHandle;

			auto InternalCreateDirectoryHandle = [&] {
				return DirHandle.Create(FilePath.c_str(),
										GENERIC_READ | GENERIC_WRITE,
										FILE_SHARE_READ | FILE_SHARE_WRITE,
										OPEN_EXISTING,
										FILE_FLAG_BACKUP_SEMANTICS);
			};

			HRESULT hRes = InternalCreateDirectoryHandle();

			if (FAILED(hRes))
			{
				zen::CreateDirectories(FilePath.c_str());

				hRes = InternalCreateDirectoryHandle();
			}

			if (FAILED(hRes))
			{
				ThrowSystemException(hRes, "Failed to open shard directory '{}'"_format(FilePath));
			}

			// Retry

			Success = SetFileInformationByHandle(FileRef.FileHandle, FileRenameInfo, RenameInfo, BufferSize);
		}

		Memory::Free(RenameInfo);

		if (Success)
		{
			return CasStore::InsertResult{.New = true};
		}

		spdlog::warn("rename of CAS payload file failed, falling back to regular write for {}", ChunkHash);

		DeletePayloadFileOnClose();
	}

	return InsertChunk(Chunk.Data(), Chunk.Size(), ChunkHash);
}

CasStore::InsertResult
FileCasStrategy::InsertChunk(const void* const ChunkData, const size_t ChunkSize, const IoHash& ChunkHash)
{
	size_t							 Shard2len = 0;
	ExtendableWideStringBuilder<128> ShardedPath;
	ShardedPath.Append(m_Config.RootDirectory.c_str());
	ShardedPath.Append(std::filesystem::path::preferred_separator);
	MakeShardedPath(ShardedPath, ChunkHash, /* out */ Shard2len);

	// See if file already exists
	//
	// Future improvement: maintain Bloom filter to avoid expensive file system probes?

	CAtlFile PayloadFile;

	HRESULT hRes = PayloadFile.Create(ShardedPath.c_str(), GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING);

	if (SUCCEEDED(hRes))
	{
		// If we succeeded in opening the file then we don't need to do anything else because it already exists and should contain the
		// content we were about to insert
		return CasStore::InsertResult{.New = false};
	}

	PayloadFile.Close();

	RwLock::ExclusiveLockScope _(LockForHash(ChunkHash));

	// For now, use double-checked locking to see if someone else was first

	hRes = PayloadFile.Create(ShardedPath.c_str(), GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING);

	if (SUCCEEDED(hRes))
	{
		// If we succeeded in opening the file then we don't need to do anything
		// else because someone else managed to create the file before we did. Just return.
		return {.New = false};
	}

	if ((hRes != HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND)) && (hRes != HRESULT_FROM_WIN32(ERROR_PATH_NOT_FOUND)))
	{
		spdlog::warn("Unexpected error code when opening shard file for read: {:#x}", uint32_t(hRes));
	}

	auto InternalCreateFile = [&] { return PayloadFile.Create(ShardedPath.c_str(), GENERIC_WRITE, FILE_SHARE_DELETE, CREATE_ALWAYS); };

	hRes = InternalCreateFile();

	if (hRes == HRESULT_FROM_WIN32(ERROR_PATH_NOT_FOUND))
	{
		// Ensure parent directories exist and retry file creation

		std::filesystem::create_directories(std::wstring_view(ShardedPath.c_str(), Shard2len));

		hRes = InternalCreateFile();
	}

	if (FAILED(hRes))
	{
		ThrowSystemException(hRes, "Failed to open shard file '{}'"_format(WideToUtf8(ShardedPath)));
	}

	size_t ChunkRemain = ChunkSize;
	auto   ChunkCursor = reinterpret_cast<const uint8_t*>(ChunkData);

	while (ChunkRemain != 0)
	{
		uint32_t ByteCount = uint32_t(std::min<size_t>(4 * 1024 * 1024ull, ChunkRemain));

		PayloadFile.Write(ChunkCursor, ByteCount);

		ChunkCursor += ByteCount;
		ChunkRemain -= ByteCount;
	}

	// We cannot rely on RAII to close the file handle since it would be closed
	// *after* the lock is released due to the initialization order.
	PayloadFile.Close();

	AtomicIncrement(m_Stats.PutCount);
	AtomicAdd(m_Stats.PutBytes, ChunkSize);

	return {.New = true};
}

IoBuffer
FileCasStrategy::FindChunk(const IoHash& ChunkHash)
{
	size_t							 Shard2len = 0;
	ExtendableWideStringBuilder<128> ShardedPath;
	ShardedPath.Append(m_Config.RootDirectory.c_str());
	ShardedPath.Append(std::filesystem::path::preferred_separator);
	MakeShardedPath(ShardedPath, ChunkHash, /* out */ Shard2len);

	RwLock::SharedLockScope _(LockForHash(ChunkHash));

	auto Chunk = IoBufferBuilder::MakeFromFile(ShardedPath.c_str());

	if (Chunk)
	{
		AtomicIncrement(m_Stats.GetCount);
		AtomicAdd(m_Stats.GetBytes, Chunk.Size());
	}

	return Chunk;
}

bool
FileCasStrategy::HaveChunk(const IoHash& ChunkHash)
{
	size_t							 Shard2len = 0;
	ExtendableWideStringBuilder<128> ShardedPath;
	ShardedPath.Append(m_Config.RootDirectory.c_str());
	ShardedPath.Append(std::filesystem::path::preferred_separator);
	MakeShardedPath(ShardedPath, ChunkHash, /* out */ Shard2len);

	RwLock::SharedLockScope _(LockForHash(ChunkHash));

	std::error_code Ec;
	if (std::filesystem::exists(ShardedPath.c_str(), Ec))
	{
		return true;
	}

	return false;
}

void 
FileCasStrategy::FilterChunks(CasChunkSet& InOutChunks)
{
	std::unordered_set<IoHash> HaveSet;

	for (const IoHash& Hash : InOutChunks.GetChunkSet())
	{
		if (HaveChunk(Hash))
		{
			HaveSet.insert(Hash);
		}
	}

	for (const IoHash& Hash : HaveSet)
	{
		InOutChunks.RemoveIfPresent(Hash);
	}
}

void
FileCasStrategy::Flush()
{
	// Since we don't keep files open after writing there's nothing specific
	// to flush here.
	//
	// Depending on what semantics we want Flush() to provide, it could be
	// argued that this should just flush the volume which we are using to
	// store the CAS files on here?
}

/**
 * Straightforward file-per-chunk CAS store implementation
 */
class FileCasImpl : public CasStore
{
public:
	FileCasImpl() : m_Strategy(m_Config, m_Stats) {}
	virtual ~FileCasImpl() = default;

	void PickDefaultDirectory()
	{
		if (m_Config.RootDirectory.empty())
		{
			// Pick sensible default

			WCHAR	myDocumentsDir[MAX_PATH];
			HRESULT hRes = SHGetFolderPathW(NULL,
											CSIDL_PERSONAL /*  My Documents */,
											NULL,
											SHGFP_TYPE_CURRENT,
											/* out */ myDocumentsDir);

			if (SUCCEEDED(hRes))
			{
				wcscat_s(myDocumentsDir, L"\\zen\\DefaultCAS");

				m_Config.RootDirectory = myDocumentsDir;
			}
		}
	}

	virtual void Initialize(const CasStoreConfiguration& InConfig) override
	{
		m_Config = InConfig;

		if (m_Config.RootDirectory.empty())
		{
			PickDefaultDirectory();
		}

		// Ensure root directory exists - create if it doesn't exist already

		std::filesystem::create_directories(m_Config.RootDirectory);

		std::filesystem::path filepath = m_Config.RootDirectory;
		filepath /= ".cas_root";

		CAtlFile marker;
		HRESULT	 hRes = marker.Create(filepath.c_str(), GENERIC_READ, 0, OPEN_EXISTING);

		if (FAILED(hRes))
		{
			ExtendableStringBuilder<128> manifest;
			manifest.Append("CAS_ROOT");
			hRes = marker.Create(filepath.c_str(), GENERIC_WRITE, 0, CREATE_ALWAYS);

			if (SUCCEEDED(hRes))
				marker.Write(manifest.c_str(), (DWORD)manifest.Size());
		}
	}

	virtual CasStore::InsertResult InsertChunk(IoBuffer Chunk, const IoHash& chunkHash) override
	{
		return m_Strategy.InsertChunk(Chunk, chunkHash);
	}
	virtual IoBuffer FindChunk(const IoHash& chunkHash) override { return m_Strategy.FindChunk(chunkHash); }

private:
	FileCasStrategy m_Strategy;
};

}  // namespace zen