Merge branch 'main' into zs/put-overwrite-policy

1 files changed, 177 insertions, 0 deletions

diff --git a/src/zenutil/bufferedwritefilecache.cpp b/src/zenutil/bufferedwritefilecache.cpp
new file mode 100644
index 000000000..a52850314
--- /dev/null
+++ b/src/zenutil/bufferedwritefilecache.cpp
@@ -0,0 +1,177 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include <zenutil/bufferedwritefilecache.h>
+
+#include <zencore/logging.h>
+#include <zencore/trace.h>
+
+ZEN_THIRD_PARTY_INCLUDES_START
+#include <gsl/gsl-lite.hpp>
+ZEN_THIRD_PARTY_INCLUDES_END
+
+namespace zen {
+
+BufferedWriteFileCache::BufferedWriteFileCache() : m_CacheHitCount(0), m_CacheMissCount(0), m_OpenHandleCount(0), m_DroppedHandleCount(0)
+{
+}
+
+BufferedWriteFileCache::~BufferedWriteFileCache()
+{
+	ZEN_TRACE_CPU("~BufferedWriteFileCache()");
+
+	try
+	{
+		for (TOpenHandles& OpenHandles : m_OpenFiles)
+		{
+			while (BasicFile* File = OpenHandles.Pop())
+			{
+				std::unique_ptr<BasicFile> FileToClose(File);
+				m_OpenHandleCount--;
+			}
+		}
+		m_OpenFiles.clear();
+		m_ChunkWriters.clear();
+	}
+	catch (const std::exception& Ex)
+	{
+		ZEN_ERROR("~BufferedWriteFileCache() threw exeption: {}", Ex.what());
+	}
+}
+
+std::unique_ptr<BasicFile>
+BufferedWriteFileCache::Get(uint32_t FileIndex)
+{
+	ZEN_TRACE_CPU("BufferedWriteFileCache::Get");
+
+	RwLock::ExclusiveLockScope _(m_WriterLock);
+	if (auto It = m_ChunkWriters.find(FileIndex); It != m_ChunkWriters.end())
+	{
+		const uint32_t HandleIndex = It->second;
+		TOpenHandles&  OpenHandles = m_OpenFiles[HandleIndex];
+		if (BasicFile* File = OpenHandles.Pop(); File != nullptr)
+		{
+			m_OpenHandleCount--;
+			m_CacheHitCount++;
+			return std::unique_ptr<BasicFile>(File);
+		}
+	}
+	m_CacheMissCount++;
+	return nullptr;
+}
+
+void
+BufferedWriteFileCache::Put(uint32_t FileIndex, std::unique_ptr<BasicFile>&& Writer)
+{
+	ZEN_TRACE_CPU("BufferedWriteFileCache::Put");
+
+	if (m_OpenHandleCount.load() >= MaxBufferedCount)
+	{
+		m_DroppedHandleCount++;
+		return;
+	}
+	RwLock::ExclusiveLockScope _(m_WriterLock);
+	if (auto It = m_ChunkWriters.find(FileIndex); It != m_ChunkWriters.end())
+	{
+		const uint32_t HandleIndex = It->second;
+		TOpenHandles&  OpenHandles = m_OpenFiles[HandleIndex];
+		if (OpenHandles.Push(Writer.get()))
+		{
+			Writer.release();
+			m_OpenHandleCount++;
+		}
+		else
+		{
+			m_DroppedHandleCount++;
+		}
+	}
+	else
+	{
+		const uint32_t HandleIndex = gsl::narrow<uint32_t>(m_OpenFiles.size());
+		m_OpenFiles.push_back(TOpenHandles{});
+		m_OpenFiles.back().Push(Writer.release());
+		m_ChunkWriters.insert_or_assign(FileIndex, HandleIndex);
+		m_OpenHandleCount++;
+	}
+}
+
+void
+BufferedWriteFileCache::Close(std::span<uint32_t> FileIndexes)
+{
+	ZEN_TRACE_CPU("BufferedWriteFileCache::Close");
+
+	std::vector<std::unique_ptr<BasicFile>> FilesToClose;
+	FilesToClose.reserve(FileIndexes.size());
+	{
+		RwLock::ExclusiveLockScope _(m_WriterLock);
+		for (uint32_t FileIndex : FileIndexes)
+		{
+			if (auto It = m_ChunkWriters.find(FileIndex); It != m_ChunkWriters.end())
+			{
+				const uint32_t HandleIndex = It->second;
+				TOpenHandles&  OpenHandles = m_OpenFiles[HandleIndex];
+				while (BasicFile* File = OpenHandles.Pop())
+				{
+					FilesToClose.emplace_back(std::unique_ptr<BasicFile>(File));
+					m_OpenHandleCount--;
+				}
+				m_ChunkWriters.erase(It);
+			}
+		}
+	}
+	FilesToClose.clear();
+}
+
+BufferedWriteFileCache::Local::Local(BufferedWriteFileCache& Cache) : m_Cache(Cache)
+{
+}
+
+BufferedWriteFileCache::Local::Writer*
+BufferedWriteFileCache::Local::GetWriter(uint32_t FileIndex)
+{
+	if (auto It = m_FileIndexToWriterIndex.find(FileIndex); It != m_FileIndexToWriterIndex.end())
+	{
+		return m_ChunkWriters[It->second].get();
+	}
+	std::unique_ptr<BasicFile> File = m_Cache.Get(FileIndex);
+	if (File)
+	{
+		const uint32_t WriterIndex = gsl::narrow<uint32_t>(m_ChunkWriters.size());
+		m_FileIndexToWriterIndex.insert_or_assign(FileIndex, WriterIndex);
+		m_ChunkWriters.emplace_back(std::make_unique<Writer>(Writer{.File = std::move(File)}));
+		return m_ChunkWriters.back().get();
+	}
+	return nullptr;
+}
+
+BufferedWriteFileCache::Local::Writer*
+BufferedWriteFileCache::Local::PutWriter(uint32_t FileIndex, std::unique_ptr<Writer> Writer)
+{
+	ZEN_ASSERT(!m_FileIndexToWriterIndex.contains(FileIndex));
+	const uint32_t WriterIndex = gsl::narrow<uint32_t>(m_ChunkWriters.size());
+	m_FileIndexToWriterIndex.insert_or_assign(FileIndex, WriterIndex);
+	m_ChunkWriters.emplace_back(std::move(Writer));
+	return m_ChunkWriters.back().get();
+}
+
+BufferedWriteFileCache::Local::~Local()
+{
+	ZEN_TRACE_CPU("BufferedWriteFileCache::~Local()");
+	try
+	{
+		for (auto& It : m_FileIndexToWriterIndex)
+		{
+			const uint32_t FileIndex   = It.first;
+			const uint32_t WriterIndex = It.second;
+			m_ChunkWriters[WriterIndex]->Writer.reset();
+			std::unique_ptr<BasicFile> File;
+			File.swap(m_ChunkWriters[WriterIndex]->File);
+			m_Cache.Put(FileIndex, std::move(File));
+		}
+	}
+	catch (const std::exception& Ex)
+	{
+		ZEN_ERROR("BufferedWriteFileCache::~Local() threw exeption: {}", Ex.what());
+	}
+}
+
+}  // namespace zen